ShynBui's picture
Update app.py
17dad0e verified
raw
history blame contribute delete
522 Bytes
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()