from transformers import pipeline
import gradio as gr

# Load the translation pipeline using the NLLB-200 model
pipe = pipeline("translation", model="facebook/nllb-200-3.3B")

# Function to translate text
def translate_text(text, source_lang, target_lang):
    # Translate the input text
    translation = pipe(text, src_lang=source_lang, tgt_lang=target_lang)
    return translation[0]['translation_text']

# Set up Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("# Text Translation using NLLB-200 (facebook/nllb-200-3.3B)")

    # Input for user's text
    text_input = gr.Textbox(label="Enter Text", placeholder="Type the text you want to translate")

    # Input for source language
    source_lang = gr.Textbox(label="Source Language", placeholder="e.g., 'eng' for English", value="eng")

    # Input for target language
    target_lang = gr.Textbox(label="Target Language", placeholder="e.g., 'fra' for French", value="fra")

    # Output for displaying the translation
    output_text = gr.Textbox(label="Translated Text")

    # Button to trigger translation
    translate_button = gr.Button("Translate")

    # Link button click with translation function
    translate_button.click(fn=translate_text, inputs=[text_input, source_lang, target_lang], outputs=output_text)

# Launch the Gradio app
demo.launch()