import streamlit as st
from transformers import MarianMTModel, MarianTokenizer, pipeline

# Load translation models for Urdu-English and English-Urdu
model_names = {
    "en-ur": "Helsinki-NLP/opus-mt-en-ur",
    "ur-en": "Helsinki-NLP/opus-mt-ur-en"
}

models = {}

# Load models once at the start
for key, model_name in model_names.items():
    tokenizer = MarianTokenizer.from_pretrained(model_name)
    model = MarianMTModel.from_pretrained(model_name)
    models[key] = pipeline("translation", model=model, tokenizer=tokenizer)

# Streamlit UI
st.title("Urdu-English Translator")

# User section on the left
st.sidebar.title("User Section")
input_lang = st.sidebar.selectbox("Select Input Language", ["English", "Urdu"])
output_lang = st.sidebar.selectbox("Select Output Language", ["Urdu", "English"])

# Text input
input_text = st.text_area("Enter Text to Translate",placeholder =  "Type here...")

# Translation logic
if st.button("Translate"):
    model_key = f"{input_lang.lower()[0:2]}-{output_lang.lower()[0:2]}"
    if model_key in models:
        translator = models[model_key]
        translation = translator(input_text)
        st.write("Translated Text:")
        st.write(translation[0]['translation_text'])
    else:
        st.write("Translation for the selected languages is not available yet.")