optimize code

app.py

@@ -1,13 +1,11 @@
-import os
 import torch
 import gradio as gr
 import time
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
 from flores200_codes import flores_codes
 
-
+# Load models and tokenizers once during initialization
 def load_models():
-    # build model and tokenizer
     model_name_dict = {
         "nllb-distilled-600M": "facebook/nllb-200-distilled-600M",
     }
@@ -15,76 +13,74 @@ def load_models():
     model_dict = {}
 
     for call_name, real_name in model_name_dict.items():
-        print("\tLoading model: %s" % call_name)
+        print("\tLoading model:", call_name)
         model = AutoModelForSeq2SeqLM.from_pretrained(real_name)
         tokenizer = AutoTokenizer.from_pretrained(real_name)
-        model_dict[call_name + "_model"] = model
-        model_dict[call_name + "_tokenizer"] = tokenizer
+        model_dict[call_name] = {
+            "model": model,
+            "tokenizer": tokenizer,
+        }
 
     return model_dict
 
-
-def translation(source, target, text):
-    if len(model_dict) == 2:
-        model_name = "nllb-distilled-600M"
+# Translate text using preloaded models and tokenizers
+def translate_text(source_lang, target_lang, input_text, model_dict):
+    model_name = "nllb-distilled-600M"
 
     start_time = time.time()
-    source = flores_codes[source]
-    target = flores_codes[target]
+    source_code = flores_codes[source_lang]
+    target_code = flores_codes[target_lang]
 
-    model = model_dict[model_name + "_model"]
-    tokenizer = model_dict[model_name + "_tokenizer"]
+    model = model_dict[model_name]["model"]
+    tokenizer = model_dict[model_name]["tokenizer"]
 
     translator = pipeline(
         "translation",
         model=model,
         tokenizer=tokenizer,
-        src_lang=source,
-        tgt_lang=target,
+        src_lang=source_code,
+        tgt_lang=target_code,
     )
-    output = translator(text, max_length=400)
+    translated_output = translator(input_text, max_length=400)
 
     end_time = time.time()
 
-    output = output[0]["translation_text"]
-    result = {
+    translated_result = {
         "inference_time": end_time - start_time,
-        "source": source,
-        "target": target,
-        "result": output,
+        "source": source_lang,
+        "target": target_lang,
+        "result": translated_output[0]["translation_text"],
     }
-    return result
-
+    return translated_result
 
 if __name__ == "__main__":
-    global model_dict
+    print("\tInitializing models")
 
+    # Load models and tokenizers
     model_dict = load_models()
 
-    # define gradio demo
     lang_codes = list(flores_codes.keys())
     inputs = [
-        gr.inputs.Dropdown(lang_codes, default="English", label="Source"),
-        gr.inputs.Dropdown(lang_codes, default="Nepali", label="Target"),
-        gr.inputs.Textbox(lines=5, label="Input text"),
+        gr.inputs.Dropdown(lang_codes, default="English", label="Source Language"),
+        gr.inputs.Dropdown(lang_codes, default="Nepali", label="Target Language"),
+        gr.inputs.Textbox(lines=5, label="Input Text"),
     ]
 
     outputs = gr.outputs.JSON()
 
-    title = "The Master Betters Translator"
+    title = "Masterful Translator"
 
-    desc = "This is the beta version of the master betters translator, which used the pre-trained model of facebook's no language left behind and fine-tuned with custom datasets. To use this app you need to have chosen the source and target language with your input text to get the output."
-    description = (
-        f"{desc}"
+    app_description = (
+        "This is a beta version of the Masterful Translator that utilizes pre-trained language models for translation."
     )
-    examples = [["English", "Nepali", "Hi. nice to meet you"]]
+    examples = [["English", "Nepali", "Hello, how are you?"]]
 
     gr.Interface(
-        translation,
+        translate_text,
         inputs,
         outputs,
         title=title,
-        description=description,
+        description=app_description,
         examples=examples,
         examples_per_page=50,
     ).launch()