bug fix

app.py

@@ -1,9 +1,10 @@
-import time
+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():
     model_name_dict = {
         "nllb-distilled-600M": "facebook/nllb-200-distilled-600M",
@@ -22,37 +23,38 @@ 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, model_dict):
+def translate_text(source, target, text):
     model_name = "nllb-distilled-600M"
 
-    if model_name in model_dict:
-        model_info = model_dict[model_name]
-        model = model_info["model"]
-        tokenizer = model_info["tokenizer"]
+    if model_name in model_dict and model_dict[model_name]["model"] is not None:
+        model = model_dict[model_name]["model"]
+        tokenizer = model_dict[model_name]["tokenizer"]
 
         start_time = time.time()
-
-        source_code = flores_codes[source]
-        target_code = flores_codes[target]
+        source = flores_codes[source]
+        target = flores_codes[target]
 
         translator = pipeline(
             "translation",
             model=model,
             tokenizer=tokenizer,
-            src_lang=source_code,
-            tgt_lang=target_code,
+            src_lang=source,
+            tgt_lang=target,
         )
         output = translator(text, max_length=400)
 
         end_time = time.time()
 
-        output_text = output[0]["translation_text"]
+        output = output[0]["translation_text"]
         result = {
             "inference_time": end_time - start_time,
-            "source": source_code,
-            "target": target_code,
-            "result": output_text,
+            "source": source,
+            "target": target,
+            "result": output,
         }
         return result
     else:
@@ -60,7 +62,6 @@ def translate_text(source, target, text, model_dict):
 
 if __name__ == "__main__":
     print("\tInitializing models")
-    model_dict = load_models()
 
     lang_codes = list(flores_codes.keys())
     inputs = [
@@ -68,13 +69,16 @@ 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 = (
-        "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."
+        f"{desc}"
     )
-    examples = [["English", "Nepali", "Hello, how are you"]]
+    examples = [["English", "Nepali", "Hi. nice to meet you"]]
 
     gr.Interface(
         translate_text,