Update app.py
Browse files
app.py
CHANGED
@@ -1,6 +1,6 @@
|
|
1 |
import gradio as gr
|
2 |
import spaces
|
3 |
-
from transformers import T5Tokenizer, T5ForConditionalGeneration, AutoTokenizer, AutoModelForSeq2SeqLM
|
4 |
|
5 |
langs = {"English": "en", "Romanian": "ro", "German": "de", "French": "fr", "Spanish": "es", "Italian": "it"}
|
6 |
options = list(langs.keys())
|
@@ -26,6 +26,9 @@ def translate_text(input_text, sselected_language, tselected_language, model_nam
|
|
26 |
model_name_full = "facebook/nllb-200-distilled-1.3B"
|
27 |
tokenizer = AutoTokenizer.from_pretrained(model_name_full)
|
28 |
model = AutoModelForSeq2SeqLM.from_pretrained(model_name_full)
|
|
|
|
|
|
|
29 |
else:
|
30 |
tokenizer = T5Tokenizer.from_pretrained(model_name)
|
31 |
model = T5ForConditionalGeneration.from_pretrained(model_name, device_map="auto")
|
|
|
1 |
import gradio as gr
|
2 |
import spaces
|
3 |
+
from transformers import T5Tokenizer, T5ForConditionalGeneration, AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
|
4 |
|
5 |
langs = {"English": "en", "Romanian": "ro", "German": "de", "French": "fr", "Spanish": "es", "Italian": "it"}
|
6 |
options = list(langs.keys())
|
|
|
26 |
model_name_full = "facebook/nllb-200-distilled-1.3B"
|
27 |
tokenizer = AutoTokenizer.from_pretrained(model_name_full)
|
28 |
model = AutoModelForSeq2SeqLM.from_pretrained(model_name_full)
|
29 |
+
translator = pipeline('translation', model=model, tokenizer=tokenizer, src_lang='eng_Latn', tgt_lang='ron_Latn')
|
30 |
+
translated_text = translator(text, max_length<=360)
|
31 |
+
return translated_text
|
32 |
else:
|
33 |
tokenizer = T5Tokenizer.from_pretrained(model_name)
|
34 |
model = T5ForConditionalGeneration.from_pretrained(model_name, device_map="auto")
|