import streamlit as st
from transformers import MBartForConditionalGeneration, MBart50TokenizerFast

# Define model name
model_checkpoint = "aryaumesh/english-to-telugu"

# Load tokenizer and model
tokenizer = MBart50TokenizerFast.from_pretrained(model_checkpoint)
model = MBartForConditionalGeneration.from_pretrained(model_checkpoint)

# Streamlit UI
st.title("English to Telugu Translator")
st.write("Enter English text to translate into Telugu.")

# User input
input_text = st.text_input("Enter English text:")

# When the button is clicked
if st.button("Translate"):
    if input_text:
        # Tokenize and generate output
        inputs = tokenizer(input_text, return_tensors="pt", padding=True, truncation=True)
        outputs = model.generate(**inputs)
        translation = tokenizer.decode(outputs[0], skip_special_tokens=True)
        
        # Display the result
        st.success(f"Translation: {translation}")
    else:
        st.warning("Please enter some text to translate.")