Update app.py

app.py

@@ -1,19 +1,9 @@
 import streamlit as st
-from transformers import T5ForConditionalGeneration, T5Tokenizer
-import torch
-
-# Load pre-trained model and tokenizer from the "KhantKyaw/T5-small_new_chatbot"
-model_name = "KhantKyaw/T5-small_new_chatbot"  # Use the fine-tuned model
-model = T5ForConditionalGeneration.from_pretrained(model_name)
+import transformers
+from transformers import T5Tokenizer, T5ForConditionalGeneration
+model_name = "t5-small"
 tokenizer = T5Tokenizer.from_pretrained(model_name)
-
-# Set device to GPU if available for faster inference, otherwise fallback to CPU
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-model.to(device)
-
-# Streamlit Interface
-st.title("Mental Health Chatbot with T5")
-
+model = T5ForConditionalGeneration.from_pretrained(model_name)
 def generate_response(input_text):
     input_ids = tokenizer.encode(input_text, return_tensors='pt')
     outputs = model.generate(input_ids,
@@ -24,8 +14,8 @@ def generate_response(input_text):
         outputs[0], skip_special_tokens=True)
     return generated_text
 
+
 prompt = st.chat_input(placeholder="Say Something!",key=None, max_chars=None, disabled=False, on_submit=None, args=None, kwargs=None)
 if prompt:
    with st.chat_message(name="AI",avatar=None):
     st.write(generate_response(prompt))
-