Update app.py

app.py

@@ -24,12 +24,20 @@ translation_task_names = {
     'Japanese to English': 'translation_ja_to_en',
 }
 
+model_names = {
+    'T5-Base': 't5-base',
+    'T5-Small': 't5-small',
+    'T5-Large': 't5-large',
+    'Opus-zh-en': 'Helsinki-NLP/opus-mt-zh-en',
+    'Opus-ru-en': 'Helsinki-NLP/opus-mt-ru-en'
+}
+
 # Create a dictionary to store loaded models
 loaded_models = {}
 
 # Simple translation function
-def translate_text(task_choice, text_input, load_in_8bit, device):
-    model_key = (task_choice, load_in_8bit)  # Create a tuple to represent the unique combination of task and 8bit loading
+def translate_text(model_choice, task_choice, text_input, load_in_8bit, device):
+    model_key = (model_choice, task_choice, load_in_8bit)  # Create a tuple to represent the unique combination of task and 8bit loading
 
     # Check if the model is already loaded
     if model_key in loaded_models:
@@ -38,6 +46,7 @@ def translate_text(task_choice, text_input, load_in_8bit, device):
         model_kwargs = {"load_in_8bit": load_in_8bit} if load_in_8bit else {}
         dtype = torch.float16 if load_in_8bit else torch.float32  # Set dtype based on the value of load_in_8bit
         translator = pipeline(task=translation_task_names[task_choice],
+                            model=model_names[model_choice],  # Use selected model
                             device=device,  # Use selected device
                             model_kwargs=model_kwargs, 
                             torch_dtype=dtype,  # Set the floating point
@@ -49,14 +58,15 @@ def translate_text(task_choice, text_input, load_in_8bit, device):
     translation = translator(text_input)[0]['translation_text']
     return str(translation).strip()
 
-def launch(task_choice, text_input, load_in_8bit, device):
-    return translate_text(task_choice, text_input, load_in_8bit, device)
+def launch(model_choice, task_choice, text_input, load_in_8bit, device):
+    return translate_text(model_choice, task_choice, text_input, load_in_8bit, device)
 
+model_dropdown = gr.Dropdown(choices=list(model_names.keys()), label='Select Model')
 task_dropdown = gr.Dropdown(choices=list(translation_task_names.keys()), label='Select Translation Task')
 text_input = gr.Textbox(label="Input Text")  # Single line text input
 load_in_8bit = gr.Checkbox(label="Load model in 8bit")
 device = gr.Radio(['cpu', 'cuda'], label='Select device', default='cpu')
 
-iface = gr.Interface(launch, inputs=[task_dropdown, text_input, load_in_8bit, device], 
+iface = gr.Interface(launch, inputs=[model_dropdown, task_dropdown, text_input, load_in_8bit, device], 
                      outputs=gr.outputs.Textbox(type="text", label="Translation"))
 iface.launch()