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

# Define available languages with their corresponding model suffixes
languages = {
    "english": "en",
    "spanish": "es",
    "french": "fr",
    "german": "de",
    "chinese": "zh",
    "japanese": "ja",
    "korean": "ko",
    "italian": "it",
    "portuguese": "pt",
    "russian": "ru",
    "hindi": "hi",
    "arabic": "ar",
    "dutch": "nl",
    "turkish": "tr",
    "greek": "el",
    "urdu": "ur"
}

# Define the translation function
def translate_text(text, target_language):
    # Get the correct model name based on the target language
    target_language_code = languages.get(target_language.lower())
    
    if target_language_code:
        model_name = f"Helsinki-NLP/opus-mt-en-{target_language_code}"
        tokenizer = AutoTokenizer.from_pretrained(model_name)
        model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
        translator = pipeline("translation", model=model, tokenizer=tokenizer)
        translation = translator(text)[0]['translation_text']
        return translation
    else:
        return "Error: Language not supported or incorrect."

# 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()