import gradio as gr
from transformers import pipeline

# Create a translation pipeline
model_name = "facebook/nllb-200-1.3B"
translator = pipeline("translation", model=model_name)

# Define available languages
languages = {
    "English": "eng_Latn",
    "French": "fra_Latn",
    "Spanish": "spa_Latn",
    "German": "deu_Latn",
    "Chinese": "zho_Hans",
    "Arabic": "ara_Arab",
    "Russian": "rus_Cyrl",
    "Hindi": "hin_Deva",
    "Japanese": "jpn_Jpan"
}

def translate(text, source_lang, target_lang):
    if not text:
        return ""
    
    source_code = languages.get(source_lang)
    target_code = languages.get(target_lang)
    
    # NLLB requires specific format for translation
    translation = translator(
        text, 
        src_lang=source_code, 
        tgt_lang=target_code,
        max_length=400
    )
    
    return translation[0]["translation_text"]

# Create the Gradio interface
demo = gr.Interface(
    fn=translate,
    inputs=[
        gr.Textbox(label="Input Text", lines=5),
        gr.Dropdown(list(languages.keys()), label="Source Language", value="English"),
        gr.Dropdown(list(languages.keys()), label="Target Language", value="French")
    ],
    outputs=gr.Textbox(label="Translated Text", lines=5),
    title="NLLB-200 Multilingual Translation",
    description="Translate text between multiple languages using Facebook's NLLB-200 model."
)

if __name__ == "__main__":
    demo.launch()