app.py

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

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

tokenizer, model = load_model()

# 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:", "あなたとは遊びたくない")

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').cuda()

            # Generate translation
            output_ids = model.generate(input_ids, max_length=128)
            translation = tokenizer.batch_decode(output_ids[:, input_ids.size(-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")