Spaces:
Running
Running
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 | |
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") | |