Update app.py

app.py

@@ -1,85 +1,59 @@
 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
 
-
 def load_models():
-    # build model and tokenizer
-    model_name_dict = {'nllb-distilled-600M': 'facebook/nllb-200-distilled-600M',
-                  #'nllb-1.3B': 'facebook/nllb-200-1.3B',
-                  #'nllb-distilled-1.3B': 'facebook/nllb-200-distilled-1.3B',
-                  #'nllb-3.3B': 'facebook/nllb-200-3.3B',
-                  }
-
+    model_name_dict = {'nllb-distilled-600M': 'facebook/nllb-200-distilled-600M'}
     model_dict = {}
-
     for call_name, real_name in model_name_dict.items():
-        print('\tLoading model: %s' % call_name)
+        print(f'\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
+        model_dict[call_name + '_tokenizer'] = tokenizer
     return model_dict
 
+global model_dict
+model_dict = load_models()
 
-def translation(source, target, text):
+def translate_text(source_lang, target_lang, input_text):
     if len(model_dict) == 2:
         model_name = 'nllb-distilled-600M'
-
     start_time = time.time()
-    source = flores_codes[source]
-    target = flores_codes[target]
+    source = flores_codes.get(source_lang)
+    target = flores_codes.get(target_lang)
+    
+    if not source or not target:
+        return {"error": "Invalid source or target language code"}
 
     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)
-    output = translator(text, max_length=400)
-
+    output = translator(input_text, max_length=400)
     end_time = time.time()
-
-    output = output[0]['translation_text']
-    result = {'inference_time': end_time - start_time,
-              'source': source,
-              'target': target,
-              'result': output}
+    output_text = output[0]['translation_text']
+    result = {
+        'inference_time': end_time - start_time,
+        'source': source_lang,
+        'target': target_lang,
+        'result': output_text
+    }
     return result
 
-
-if __name__ == '__main__':
-    print('\tinit models')
-
-    global model_dict
-
-    model_dict = load_models()
-    
-    # define gradio demo
-    lang_codes = list(flores_codes.keys())
-    #inputs = [gr.inputs.Radio(['nllb-distilled-600M', 'nllb-1.3B', 'nllb-distilled-1.3B'], label='NLLB Model'),
-    inputs = [gr.inputs.Dropdown(lang_codes, default='English', label='Source'),
-              gr.inputs.Dropdown(lang_codes, default='Korean', label='Target'),
-              gr.inputs.Textbox(lines=5, label="Input text"),
-              ]
-
-    outputs = gr.outputs.JSON()
-
-    title = "NLLB distilled 600M demo"
-
-    demo_status = "Demo is running on CPU"
-    description = f"Details: https://github.com/facebookresearch/fairseq/tree/nllb. {demo_status}"
-    examples = [
-    ['English', 'Korean', 'Hi. nice to meet you']
-    ]
-
-    gr.Interface(translation,
-                 inputs,
-                 outputs,
-                 title=title,
-                 description=description,
-                 ).launch() 
-
-
+# Define Gradio Interface
+iface = gr.Interface(
+    fn=translate_text,
+    inputs=[
+        gr.inputs.Textbox(lines=1, placeholder="Source language code", label="Source Language Code"),
+        gr.inputs.Textbox(lines=1, placeholder="Target language code", label="Target Language Code"),
+        gr.inputs.Textbox(lines=5, placeholder="Enter text to translate", label="Input Text"),
+    ],
+    outputs=gr.outputs.JSON(),
+    title="Translation API",
+    description="Translation API using NLLB model."
+)
+
+# Launch as API only
+iface.launch(share=True, enable_queue=True, show_error=True, server_name="0.0.0.0", server_port=7860)