import gradio as gr
import spaces
from transformers import pipeline
import os
import torch


print(torch.cuda.is_available())
summarizer = pipeline("summarization", model="ShynBui/Bartpho_spelling_correction", device=0)

@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()