import gradio as gr
from transformers import pipeline

# Load the translation pipeline
pipe = pipeline("translation", model="facebook/nllb-200-distilled-600M")

# Define the translation function
def translate_text(text, source_language, target_language, max_length=200):
    # Define source and target languages (e.g., 'eng_Latn' -> 'fra_Latn')
    language_pair = f"{source_language}-{target_language}"
    response = pipe(
        text,
        src_lang=source_language,
        tgt_lang=target_language,
        max_length=max_length
    )
    return response[0]['translation_text']

# Create the Gradio app interface
with gr.Blocks() as app:
    gr.Markdown("## Translation App")
    gr.Markdown(
        "Enter text in the source language, and the model will translate it to the target language. "
        "This app uses the `facebook/nllb-200-distilled-600M` model."
    )

    gr.Markdown("### To choose the languages: [https://github.com/facebookresearch/flores/blob/main/flores200/README.md#languages-in-flores-200](https://github.com/facebookresearch/flores/blob/main/flores200/README.md#languages-in-flores-200)")
    
    with gr.Row():
        input_text = gr.Textbox(
            label="Input Text",
            placeholder="Type the text to translate...",
            lines=4
        )
        output_translation = gr.Textbox(label="Translated Text", lines=4)

    with gr.Row():
        source_language = gr.Textbox(
            label="Source Language Code",
            placeholder="e.g., eng_Latn (for English)"
        )
        target_language = gr.Textbox(
            label="Target Language Code",
            placeholder="e.g., fra_Latn (for French)"
        )

    max_length = gr.Slider(
        label="Max Length",
        minimum=50,
        maximum=500,
        step=50,
        value=200
    )

    submit_button = gr.Button("Translate")
    
    submit_button.click(
        fn=translate_text,
        inputs=[input_text, source_language, target_language, max_length],
        outputs=output_translation
    )

# Launch the app
if __name__ == "__main__":
    app.launch()