import streamlit as st
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

device=torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
# Initialize the model and tokenizer
@st.cache_resource
def load_model(device):
    tokenizer = AutoTokenizer.from_pretrained('tirthadagr8/Japanese_to_english_gpt2CasualLM_GemmaTokenizer')
    model = AutoModelForCausalLM.from_pretrained('tirthadagr8/Japanese_to_english_gpt2CasualLM_GemmaTokenizer')
    model.to(device)  # Ensure the model runs on GPU if available
    return tokenizer, model

tokenizer, model = load_model(device)

# Streamlit App UI
st.title("Japanese to English Translation")
st.subheader("Using a Hugging Face GPT-2 model")

# Input text box
src_text = st.text_area("Enter Japanese text for translation:")
# print(src_text)
if st.button("Translate"):
    if src_text.strip():
        with st.spinner("Translating..."):
            # Prepare the input for the model
            prompt = f"Translate the following Japanese sentence to English:\n\nJapanese:{src_text}\nEnglish:"
            input_ids = tokenizer.encode(prompt, return_tensors='pt')[:,:-1].to(device)
            # print(tokenizer.batch_decode(input_ids))
            # Generate translation
            output_ids = model.generate(input_ids, max_length=128)
            translation = tokenizer.batch_decode(output_ids[:, input_ids.size(-1):-1])[0]

        st.success("Translation completed!")
        st.write("**English Translation:**")
        st.write(translation)
    else:
        st.warning("Please enter some text to translate.")

# Footer
st.markdown("---")
st.caption("Model by tirthadagr8 | Powered by Hugging Face Transformers and Streamlit")