from transformers import MarianMTModel, MarianTokenizer

def load_model(src_lang, tgt_lang):
    model_name = f'Helsinki-NLP/opus-mt-{src_lang}-{tgt_lang}'
    tokenizer = MarianTokenizer.from_pretrained(model_name)
    model = MarianMTModel.from_pretrained(model_name)
    return model, tokenizer

# Load the English to Urdu model
src_lang = 'en'
tgt_lang = 'ur'
model, tokenizer = load_model(src_lang, tgt_lang)

def translate(text, src_lang, tgt_lang):
    model, tokenizer = load_model(src_lang, tgt_lang)
    inputs = tokenizer.encode(text, return_tensors="pt", padding=True)
    translated = model.generate(inputs)
    translated_text = tokenizer.decode(translated[0], skip_special_tokens=True)
    return translated_text

import gradio as gr

# Define a list of supported language pairs
languages = {
    'Urdu': 'ur',
    'French': 'fr',
    'Spanish': 'es',
    'German': 'de',
    'Chinese': 'zh',
    'Italian': 'it',
    'Russian': 'ru',
    'Japanese': 'ja',
    'Arabic': 'ar',
    'Hindi': 'hi',
    # Add more languages as needed
}

def translate_ui(text, target_language):
    tgt_lang = languages[target_language]
    return translate(text, 'en', tgt_lang)


# Create Gradio interface
iface = gr.Interface(
    fn=translate_ui,
    inputs=[
        gr.Textbox(lines=2, placeholder="Enter text here...", label="Input Text"),
        gr.Dropdown(choices=list(languages.keys()), label="Target Language")
    ],
    outputs=gr.Textbox(label="Translated Text"),
    title="English to Other Languages Translator",
    description="Translate English text to various languages including Urdu."
)

# Launch the interface
iface.launch()