# -*- coding: utf-8 -*-
"""Untitled4.ipynb

Automatically generated by Colab.

Original file is located at
    https://colab.research.google.com/drive/1tC3Zqo0dc6hrck0oxp5luNsOywKjdSGc
"""

!pip install transformers gradio

from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
import gradio as gr

LANGUAGES = {
    'en-fr': 'Helsinki-NLP/opus-mt-en-fr',
    'en-es': 'Helsinki-NLP/opus-mt-en-es',
    'en-de': 'Helsinki-NLP/opus-mt-en-de',
    'fr-en': 'Helsinki-NLP/opus-mt-fr-en',
    'es-en': 'Helsinki-NLP/opus-mt-es-en'
}

def load_model(lang_pair):
    model_name = LANGUAGES.get(lang_pair)
    if model_name:
        tokenizer = AutoTokenizer.from_pretrained(model_name)
        model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
        return tokenizer, model
    else:
        return None, None

def translate_text(text, lang_pair='en-fr'):
    tokenizer, model = load_model(lang_pair)
    if tokenizer and model:
        inputs = tokenizer.encode(text, return_tensors="pt", max_length=512, truncation=True)
        outputs = model.generate(inputs, max_length=512, num_beams=4, early_stopping=True)
        return tokenizer.decode(outputs[0], skip_special_tokens=True)
    else:
        return "Language pair not supported"

# Gradio interface for language pair selection and translation
def translate_interface(input_text, lang_pair='en-fr'):
    return translate_text(input_text, lang_pair)

# Define the Gradio interface
interface = gr.Interface(
    fn=translate_interface,
    inputs=[gr.Textbox(label="Input Text"), gr.Dropdown(choices=list(LANGUAGES.keys()), label="Select Language Pair")],
    outputs="text",
    title="Multilingual Translator",
    description="Translate text between various languages using Hugging Face models."
)

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