import gradio as gr
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, pipeline

# Load the translation model and tokenizer from Hugging Face
model_name = "Helsinki-NLP/opus-mt-en-fr"  # Example model for English to French
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSeq2SeqLM.from_pretrained(model_name)

# Define the translation function
def translate_text(text, target_language):
    model_name = f"Helsinki-NLP/opus-mt-en-{target_language}"
    translator = pipeline("translation", model=model_name)
    translation = translator(text)[0]['translation_text']
    return translation

# Define available languages
languages = {
    "english": "English",
    "spanish": "Spanish",
    "french": "French",
    "german": "German",
    "chinese": "Chinese (Mandarin)",
    "japanese": "Japanese",
    "korean": "Korean",
    "italian": "Italian",
    "portuguese": "Portuguese",
    "russian": "Russian",
    "hindi": "Hindi",
    "arabic": "Arabic",
    "dutch": "Dutch",
    "turkish": "Turkish",
    "greek": "Greek",
    "urdu": "Urdu"
}


# Set up the Gradio interface with a submit button and side-by-side layout
with gr.Blocks() as iface:
    gr.Markdown("# Text Translator")
    gr.Markdown("Translate text into multiple languages using Hugging Face models.")

    # Create a row for input and output
    with gr.Row():
        # Input components on the left
        with gr.Column(scale=1):  # This column is smaller
            text_input = gr.Textbox(label="Enter text to translate")
            language_dropdown = gr.Dropdown(list(languages.keys()), label="Target Language", type="value")

        # Output components on the right
        with gr.Column(scale=1):  # This column is also smaller
            translation_output = gr.Textbox(label="Translation", interactive=False)

    # Button to submit the translation
    submit_button = gr.Button("Translate")

    # When button is clicked, trigger the translation
    submit_button.click(fn=translate_text, inputs=[text_input, language_dropdown], outputs=translation_output)

iface.launch()