rapacious commited on
Commit
a26d575
·
verified ·
1 Parent(s): da84d97

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +27 -0
app.py ADDED
@@ -0,0 +1,27 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
3
+ import torch
4
+
5
+ # Load model and tokenizer
6
+ model_name = "VietAI/envit5-translation"
7
+ tokenizer = AutoTokenizer.from_pretrained(model_name)
8
+ model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
9
+
10
+ def translate(text, source_lang):
11
+ """Translate text based on the source language."""
12
+ input_text = f"{source_lang}: {text}"
13
+ inputs = tokenizer(input_text, return_tensors="pt", padding=True).input_ids.to('cpu')
14
+ outputs = model.generate(inputs, max_length=512)
15
+ return tokenizer.decode(outputs[0], skip_special_tokens=True)
16
+
17
+ # Create UI
18
+ demo = gr.Interface(
19
+ fn=translate,
20
+ inputs=[gr.Textbox(label="Input Text"), gr.Radio(["vi", "en"], label="Source Language")],
21
+ outputs=gr.Textbox(label="Translated Text"),
22
+ title="VietAI Translation",
23
+ description="Translate between Vietnamese and English using envit5-translation model."
24
+ )
25
+
26
+ # Launch app
27
+ demo.launch()