fix

app.py

@@ -1,10 +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():
     model_name_dict = {
         "nllb-distilled-600M": "facebook/nllb-200-distilled-600M",
@@ -23,8 +24,11 @@ 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_lang, target_lang, input_text, model_dict):
+def translate_text(source, target, text):
     model_name = "nllb-distilled-600M"
 
     if model_name in model_dict:
@@ -32,58 +36,57 @@ def translate_text(source_lang, target_lang, input_text, model_dict):
         tokenizer = model_dict[model_name]["tokenizer"]
 
         start_time = time.time()
-        source_code = flores_codes[source_lang]
-        target_code = flores_codes[target_lang]
+        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,
         )
-        translated_output = translator(input_text, max_length=400)
+        output = translator(text, max_length=400)
 
         end_time = time.time()
 
-        translated_result = {
+        output = output[0]["translation_text"]
+        result = {
             "inference_time": end_time - start_time,
-            "source": source_lang,
-            "target": target_lang,
-            "result": translated_output[0]["translation_text"],
+            "source": source,
+            "target": target,
+            "result": output,
         }
-        return translated_result
+        return result
     else:
         raise KeyError(f"Model '{model_name}' not found in model_dict")
 
 if __name__ == "__main__":
     print("\tInitializing models")
 
-    # Load models and tokenizers
-    model_dict = load_models()
-
     lang_codes = list(flores_codes.keys())
     inputs = [
-        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"),
+        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"),
     ]
 
     outputs = gr.outputs.JSON()
 
     title = "The Master Betters Translator"
 
-    app_description = (
-        "This is a beta version of The Master Betters Translator that utilizes pre-trained language models for translation."
+    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}"
     )
-    examples = [["English", "Nepali", "Hello, how are you?"]]
+    examples = [["English", "Nepali", "Hi. nice to meet you"]]
 
     gr.Interface(
-        fn=translate_text,
-        inputs=inputs,
-        outputs=outputs,
+        translate_text,
+        inputs,
+        outputs,
         title=title,
-        description=app_description,
+        description=description,
         examples=examples,
         examples_per_page=50,
     ).launch()