Update app.py

app.py

@@ -2,33 +2,34 @@ 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():
+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.cuda()  # Ensure the model runs on GPU if available
+    model.to(device)  # Ensure the model runs on GPU if available
     return tokenizer, model
 
-tokenizer, model = load_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:", "あなたとは遊びたくない")
-
+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').cuda()
-
+            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):])[0]
+            translation = tokenizer.batch_decode(output_ids[:, input_ids.size(-1):-1])[0]
 
         st.success("Translation completed!")
         st.write("**English Translation:**")