Update app.py

app.py

@@ -1,9 +1,11 @@
 import streamlit as st
-from transformers import T5ForTranslation, T5Tokenizer
+from transformers import T5Tokenizer, T5ForConditionalGeneration
+model_name = 't5-base'
+tokenizer = T5Tokenizer.from_pretrained(model_name)
+model = T5ForConditionalGeneration.from_pretrained(model_name)
 
-# Load the pre-trained model and tokenizer
-model = T5ForTranslation.from_pretrained("Helsinki-NLP/opus-mt-en-ro")
-tokenizer = T5Tokenizer.from_pretrained("Helsinki-NLP/opus-mt-en-ro")
+# model = T5ForTranslation.from_pretrained("Helsinki-NLP/opus-mt-en-ro")
+# tokenizer = T5Tokenizer.from_pretrained("Helsinki-NLP/opus-mt-en-ro")
 
 # Create the app layout
 st.title("Text Translation App")
@@ -14,13 +16,15 @@ translated_text = st.text("")
 # Handle the submit button click
 if submit_button:
     # Encode the input text
-    encoded = tokenizer(text_input, return_tensors="pt")
+    # encoded = tokenizer(text_input, return_tensors="pt")
+    input_ids = tokenizer.encode(input_text, return_tensors='pt')
 
     # Perform translation
-    translated = model.generate(**encoded)
+    output_ids = model.generate(input_ids)
+    translated = tokenizer.decode(output_ids[0], skip_special_tokens=True)
 
     # Decode the translated text
-    translated_text = tokenizer.batch_decode(translated, skip_special_tokens=True)
+    translated_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
 
     # Display the translated text
     st.write("Translated Text:", translated_text[0])