Update app.py

app.py

@@ -28,16 +28,22 @@ def model_to_cuda(model):
         print("CUDA not available! Using CPU.")
     return model
 
-def eurollm(model_id, sl, tl, input_text):
-    model_id = "utter-project/EuroLLM-1.7B"
-    tokenizer = AutoTokenizer.from_pretrained(model_id)
-    model = AutoModelForCausalLM.from_pretrained(model_id)  
-    prompt = f"{sl}: {input_text}. {tl}:"
+def eurollm(model_name, sl, tl, input_text):
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model = AutoModelForCausalLM.from_pretrained(model_name)  
+    prompt = f"{sl}: {input_text} {tl}:"
     inputs = tokenizer(prompt, return_tensors="pt")
     outputs = model.generate(**inputs, max_new_tokens=512)
     output = tokenizer.decode(outputs[0], skip_special_tokens=True) 
-    print(output.rsplit(f'{tl}:')[-1])
-    return output
+    result = output.rsplit(f'{tl}:')[-1].strip())
+    return result
+
+def nllb(model_name, sl, tl, input_text):
+    tokenizer = AutoTokenizer.from_pretrained(model_name, src_lang=sl)
+    model = AutoModelForSeq2SeqLM.from_pretrained(model_name, device_map="auto")
+    translator = pipeline('translation', model=model, tokenizer=tokenizer, src_lang=sl, tgt_lang=tl)
+    translated_text = translator(input_text, max_length=512)
+    return translated_text[0]['translation_text']
 
 @spaces.GPU
 def translate_text(input_text, sselected_language, tselected_language, model_name):
@@ -61,12 +67,10 @@ def translate_text(input_text, sselected_language, tselected_language, model_nam
         translated_text = eurollm(model_name, sselected_language, tselected_language, input_text)
         return translated_text, message_text
         
-    if 'nllb' in model_name:
-        tokenizer = AutoTokenizer.from_pretrained(model_name, src_lang=languagecodes.nllb_language_codes[sselected_language])
-        model = AutoModelForSeq2SeqLM.from_pretrained(model_name, device_map="auto")
-        translator = pipeline('translation', model=model, tokenizer=tokenizer, src_lang=languagecodes.nllb_language_codes[sselected_language], tgt_lang=languagecodes.nllb_language_codes[tselected_language])
-        translated_text = translator(input_text, max_length=512)
-        return translated_text[0]['translation_text'], message_text
+    if 'nllb' in model_name.lower():
+        nnlbsl, nnlbtl = languagecodes.nllb_language_codes[sselected_language], languagecodes.nllb_language_codes[tselected_language]
+        translated_text = nllb(model_name, nnlbsl, nnlbtl, input_text)
+        return translated_text, message_text
     
     if model_name.startswith('facebook/mbart-large'):
         from transformers import MBartForConditionalGeneration, MBart50TokenizerFast
@@ -113,7 +117,7 @@ def swap_languages(src_lang, tgt_lang):
 
 def create_interface():
     with gr.Blocks() as interface:
-        gr.Markdown("## Machine Text Translation")
+        gr.Markdown("### Machine Text Translation - maximum 512 tokens")
 
         with gr.Row():
             input_text = gr.Textbox(label="Enter text to translate:", placeholder="Type your text here...")