Spaces:
Sleeping
Sleeping
import gradio as gr | |
import spaces | |
from transformers import pipeline | |
import os | |
import torch | |
print(torch.cuda.is_available()) | |
device = 'cuda' if torch.cuda.is_available() else 'cpu' | |
summarizer = pipeline("summarization", model="ShynBui/Bartpho_spelling_correction", device=device) | |
# @spaces.GPU | |
def generate(prompt): | |
return summarizer(prompt, max_new_tokens = 512)[0]['summary_text'] | |
gr.Interface( | |
fn=generate, | |
inputs=gr.Text(), | |
outputs=gr.Text(), | |
examples= eval(os.environ['DES_EXAMPLE']) | |
).launch() |