Update app.py

app.py

@@ -1,33 +1,56 @@
-# Import the logging utility from the Hugging Face transformers library
-from transformers.utils import logging
-
-# Set the logging verbosity to 'error' level
-# This means only error messages will be printed, suppressing warnings, info, and debug messages
-logging.set_verbosity_error()
+import gradio as gr
 from transformers import pipeline
 import torch
-import gradio as gr
-
-translator = pipeline(task="translation",
-                      model="facebook/nllb-200-distilled-600M",
-                      torch_dtype=torch.bfloat16)
 
+# Load the translation model
+translator = pipeline(
+    task="translation",
+    model="facebook/nllb-200-distilled-600M",
+    torch_dtype=torch.bfloat16
+)
 
-
-def translation_func(text, source, target):
-    text_translated = translator(text, src_lang=source, tgt_lang= target)
-    return text_translated
-
-iface = gr.Interface(
-    fn=translation_func,
-    inputs = [
-        gr.Textbox(label = "Enter your text"),
-        gr.Dropdown( ["eng_Latn", "fra_Latn", "arz_Arab"], label = "Source Language" ),
-        gr.Dropdown(["eng_Latn", "fra_Latn", "arz_Arab"], label = "Target Language"),
-
+# Define the supported languages (this is a subset of possible languages, expand as needed)
+languages = {
+    "English": "eng_Latn",
+    "French": "fra_Latn",
+    "Arabic": "arb_Arab",
+    "Spanish": "spa_Latn",
+    "German": "deu_Latn",
+    "Chinese (Simplified)": "zho_Hans",
+    "Hindi": "hin_Deva"
+}
+
+# Function to translate text based on selected languages
+def translate(text, src_lang, tgt_lang):
+    if src_lang not in languages or tgt_lang not in languages:
+        return "Invalid language selection"
+
+    # Step 1: Define source and target languages using the input language codes
+    source_language = languages[src_lang]  # Look up the source language
+    target_language = languages[tgt_lang]  # Look up the target language
+
+    # Step 2: Call the translator function with the given text, source, and target languages
+    result = translator(text, src_lang=source_language, tgt_lang=target_language)
+
+    # Step 3: Extract the translated text from the result
+    translation = result[0]['translation_text']
+
+    # Step 4: Return the translated text
+    return translation
+
+# Gradio Interface
+demo = gr.Interface(
+    fn=translate,  # Translation function
+    inputs=[
+        gr.Textbox(label="Input Text", placeholder="Enter text to translate"),  # Text input
+        gr.Dropdown(choices=list(languages.keys()), label="Source Language", value="English"),  # Source language
+        gr.Dropdown(choices=list(languages.keys()), label="Target Language", value="French")  # Target language
     ],
-    outputs=gr.Textbox(label="Translation")
+    outputs=gr.Textbox(label="Translated Text"),  # Output
+    title="Translation using NLLB-200",  # Title of the interface
+    description="Select source and target languages, then translate the text."  # Description
 )
 
-iface.launch()
-    
+# Launch the Gradio interface
+demo.launch()
+