import streamlit as st
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

# Load model and tokenizer
tokenizer = AutoTokenizer.from_pretrained("facebook/mbart-large-50")
model = AutoModelForSeq2SeqLM.from_pretrained("facebook/mbart-large-50")

def translate_to_japanese(english_text):
    # Add language tokens for translation (from English > to Japanese)
    input_text = f">en< {english_text} </s>"

    # Tokenize input text
    inputs = tokenizer(input_text, return_tensors="pt", padding=True, truncation=True)

    # Generate translation (the model's output should be a token sequence for the translation)
    outputs = model.generate(inputs['input_ids'], max_length=128, num_beams=4, early_stopping=True)

    # Decode the generated token IDs into a string
    translated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)

    return translated_text

# Streamlit UI
st.title("English to Japanese Translator")
st.write("This app translates English sentences to Japanese using the mBART model.")

# Create an input box for the user to enter an English sentence
english_sentence = st.text_area("Enter English sentence:", "")

# Translate and display the result
if english_sentence:
    with st.spinner("Translating..."):
        translation = translate_to_japanese(english_sentence)
    st.subheader("Translation (English to Japanese):")
    st.write(translation)