File size: 522 Bytes
f17eb5d
 
3ae2dd1
f47dd3c
7f6b51d
f17eb5d
afd14d3
7f6b51d
17dad0e
0f17896
f17eb5d
16ab1a8
7f6b51d
 
f17eb5d
7f6b51d
 
 
 
a34b13a
7f6b51d
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
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()