Update app.py

app.py

@@ -1,5 +1,5 @@
 import gradio as gr
-import numpy as np
+import torch
 from huggingface_hub import hf_hub_download
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
 
@@ -8,8 +8,6 @@ REPO_ID_MARIANNMT_en = "mbarnig/MarianNMT-tatoeba-en-lb"
 REPO_ID_MARIANNMT_lb = "mbarnig/MarianNMT-tatoeba-lb-en"
 REPO_ID_T5MT5 = "mbarnig/T5-mt5-tatoeba-en-lb"
 
-translator = pipeline("translation", model=REPO_ID_NLLB)
-
 my_title = "🇬🇧 Mir iwwersetzen vun an op Lëtzebuergesch ! 🇫🇷"
 my_description = "English-Luxembourgish machine translation (MT) demo based on 3 open-source transformer models: Facebook-NLLB, Microsoft-MarianNMT & Google-T5/mt5."
 my_article = "<h3>User guide</h3><p>1. Press the submit button to translate an english text with the default values. 2. Compare the result with the luxembourgish example. 3. Select a model and a translation direction and enter your own text. Have fun !</p><p>Go to <a href='https://www.web3.lu/'>Internet with a Brain</a> to read my french publication <a href='https://www.web3.lu/'>Das Küsschen und die Sonne stritten sich ...</a> about the history of machine translation in Luxembourg from 1975 until today.</p>"
@@ -36,8 +34,31 @@ my_inputs = [
 
 my_output = gr.Textbox(lines=5, label="Translation")
 
-def iwwersetz(source_text):
-    translation = translator(source_text)
+def customization(model, direc):
+    if model == "NLLB":
+        translator = pipeline("translation", model=REPO_ID_NLLB)   
+    elif model == "MarianNMT":
+        if direc = "en -> lb":
+            translator = pipeline("translation", model=REPO_ID_MARIANNMT_en)        
+        elif direc = "lb -> en":
+            translator = pipeline("translation", model=REPO_ID_MARIANNMT_lb)        
+        else:
+            print("Please select a Translation Direction !")    
+    elif model == "T5/mt5":
+        translator = pipeline("translation", model=REPO_ID_T5MT5) 
+    else:
+       print("Please select a Translation Model !")
+    return translator
+
+def iwwersetz(source_text, model, direc):
+    translator = customization(model, direc)
+    if model == "NLLB":
+        if direc == "en -> lb":
+            translation = translator("en", "lb", source_text)
+        else:
+            translation = translator("lb", "en", source_text)        
+    else:
+        translation = translator(source_text)
     return translation
     
 demo=gr.Interface(