app.py

import gradio as gr
from transformers import pipeline
import torch

# Load the translation model
translator = pipeline(
    task="translation",
    model="facebook/nllb-200-distilled-600M",
    torch_dtype=torch.bfloat16
)

# 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="Translated Text"),  # Output
    title="Translation using NLLB-200",  # Title of the interface
    description="Select source and target languages, then translate the text."  # Description
)

# Launch the Gradio interface
demo.launch()