from transformers import MBartForConditionalGeneration, MBart50TokenizerFast
import gradio as gr

# Load the tokenizer and model
model_checkpoint = "aryaumesh/english-to-telugu"
tokenizer = MBart50TokenizerFast.from_pretrained(model_checkpoint)
model = MBartForConditionalGeneration.from_pretrained(model_checkpoint)

# Function to translate text
def translate_to_telugu(text):
    inputs = tokenizer(text, return_tensors="pt")
    outputs = model.generate(**inputs)
    translation = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return translation

# Create a Gradio interface
iface = gr.Interface(
    fn=translate_to_telugu,
    inputs=gr.Textbox(lines=2, placeholder="Enter text in English...", label='English Text'),  # Text input
    outputs=gr.Textbox(label="Translation to Telugu"),  # Text output
    title="English to Telugu Translator",
    description="Translate text from English to Telugu."
)

# Launch the Gradio interface
iface.launch()