TiberiuCristianLeon commited on
Commit
fe93b05
·
verified ·
1 Parent(s): 1f1ba03

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +4 -1
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")