import gradio as gr
from transformers import MarianMTModel, MarianTokenizer

# Load the MarianMT model and tokenizer for English to Urdu translation
def load_model():
    model_name = 'Helsinki-NLP/opus-mt-en-ur'  # English to Urdu pre-trained model
    model = MarianMTModel.from_pretrained(model_name)
    tokenizer = MarianTokenizer.from_pretrained(model_name)
    return model, tokenizer

# Define the translation function
def translate(text, model, tokenizer):
    # Prepare the text for translation
    translated = tokenizer(text, return_tensors="pt", padding=True, truncation=True)
    # Generate translation
    translated = model.generate(**translated)
    # Decode the translated text
    translated_text = tokenizer.decode(translated[0], skip_special_tokens=True)
    return translated_text

# Load model and tokenizer
model, tokenizer = load_model()

# Define Gradio interface function
def gradio_interface(text):
    return translate(text, model, tokenizer)

# Set up Gradio interface for the translation web app
interface = gr.Interface(fn=gradio_interface,
                         inputs="text",  # User input type (text box)
                         outputs="text", # Output type (translated text)
                         live=True,  # Updates live as user types
                         title="English to Urdu Translation Web App", 
                         description="Translate English text to Urdu using the MarianMT model.")

# Launch the Gradio interface
interface.launch()