tirthadagr8's picture
Update app.py
95c14fd verified
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")