code optimization

app.py

@@ -1,11 +1,9 @@
-import os
-import torch
-import gradio as gr
 import time
+import gradio as gr
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
 from flores200_codes import flores_codes
 
-
+# Load models and tokenizers once during initialization
 def load_models():
     model_name_dict = {
         "nllb-distilled-600M": "facebook/nllb-200-distilled-600M",
@@ -24,38 +22,37 @@ def load_models():
 
     return model_dict
 
-# Load models and tokenizers once during initialization
-model_dict = load_models()
-
 # Translate text using preloaded models and tokenizers
-def translate_text(source, target, text):
+def translate_text(source, target, text, model_dict):
     model_name = "nllb-distilled-600M"
 
     if model_name in model_dict:
-        model = model_dict[model_name]["model"]
-        tokenizer = model_dict[model_name]["tokenizer"]
+        model_info = model_dict[model_name]
+        model = model_info["model"]
+        tokenizer = model_info["tokenizer"]
 
         start_time = time.time()
-        source = flores_codes[source]
-        target = flores_codes[target]
+
+        source_code = flores_codes[source]
+        target_code = flores_codes[target]
 
         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)
 
         end_time = time.time()
 
-        output = output[0]["translation_text"]
+        output_text = output[0]["translation_text"]
         result = {
             "inference_time": end_time - start_time,
-            "source": source,
-            "target": target,
-            "result": output,
+            "source": source_code,
+            "target": target_code,
+            "result": output_text,
         }
         return result
     else:
@@ -63,6 +60,7 @@ def translate_text(source, target, text):
 
 if __name__ == "__main__":
     print("\tInitializing models")
+    model_dict = load_models()
 
     lang_codes = list(flores_codes.keys())
     inputs = [
@@ -70,16 +68,13 @@ if __name__ == "__main__":
         gr.inputs.Dropdown(lang_codes, default="Nepali", label="Target"),
         gr.inputs.Textbox(lines=5, label="Input text"),
     ]
-
     outputs = gr.outputs.JSON()
 
     title = "The Master Betters Translator"
-
-    desc = "This is a beta version of The Master Betters Translator that utilizes pre-trained language models for translation. 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}"
+        "This is a beta version of The Master Betters Translator that utilizes pre-trained language models for translation. To use this app you need to have chosen the source and target language with your input text to get the output."
     )
-    examples = [["English", "Nepali", "Hi. nice to meet you"]]
+    examples = [["English", "Nepali", "Hello, how are you"]]
 
     gr.Interface(
         translate_text,