Update app.py

app.py

@@ -1,12 +1,12 @@
 import streamlit as st
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import T5ForConditionalGeneration, T5Tokenizer
 import torch
 import random
 
-# Load pre-trained DialoGPT-small model and tokenizer
-model_name = "microsoft/DialoGPT-small"
-model = AutoModelForCausalLM.from_pretrained(model_name)
-tokenizer = AutoTokenizer.from_pretrained(model_name)
+# Load pre-trained T5 model and tokenizer
+model_name = "t5-small"  # You can use "t5-base" or "t5-large" for better quality but slower response
+model = T5ForConditionalGeneration.from_pretrained(model_name)
+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")
@@ -20,10 +20,10 @@ if 'conversation' not in st.session_state:
 
 # Define multiple system prompts to control bot's behavior
 system_prompts = [
-    "You are a friendly and professional assistant. You respond in a polite and helpful manner.",
-    "You are a casual chatbot that likes to engage in fun and interesting conversations, but always stay respectful.",
-    "You are a helpful assistant. Your goal is to provide clear and precise answers to any questions.",
-    "You are a compassionate and empathetic listener, always responding with kindness and understanding."
+    "You are a helpful assistant. Respond in a polite, friendly, and informative manner.",
+    "You are a conversational chatbot. Provide friendly, engaging, and empathetic responses.",
+    "You are an informative assistant. Respond clearly and concisely to any questions asked.",
+    "You are a fun, casual chatbot. Keep the conversation light-hearted and interesting."
 ]
 
 # Select a random system prompt to start the conversation
@@ -35,28 +35,25 @@ def generate_response(input_text):
     if len(st.session_state['history']) == 0:
         system_prompt = get_system_prompt()
         st.session_state['conversation'].append(f"System: {system_prompt}")
-        system_input_ids = tokenizer.encode(system_prompt + tokenizer.eos_token, return_tensors='pt').to(device)
-        st.session_state['history'] = system_input_ids[0].tolist()  # Save only the system prompt to history
+        system_input = f"conversation: {system_prompt} "
+        st.session_state['history'].append(system_input)
 
-    # Encode the new user input, add end of string token
-    new_user_input_ids = tokenizer.encode(input_text + tokenizer.eos_token, return_tensors='pt').to(device)
+    # Prepare the user input by appending it to the history
+    user_input = f"conversation: {input_text} "
 
-    # Append the new user input to the history
-    if len(st.session_state['history']) > 0:
-        history_tensor = torch.tensor(st.session_state['history']).unsqueeze(0).to(device)
-        bot_input_ids = torch.cat([history_tensor, new_user_input_ids], dim=-1)
-    else:
-        bot_input_ids = new_user_input_ids
+    # Concatenate history (system prompt + user input)
+    full_input = "".join(st.session_state['history']) + user_input
 
-    # Generate a response from the model
-    chat_history_ids = model.generate(bot_input_ids, max_length=1000, pad_token_id=tokenizer.eos_token_id, top_k=50, top_p=0.95, temperature=0.7)
+    # Tokenize input text and generate response from the model
+    input_ids = tokenizer.encode(full_input, return_tensors="pt").to(device)
+    outputs = model.generate(input_ids, max_length=1000, num_beams=5, top_p=0.95, temperature=0.7, pad_token_id=tokenizer.eos_token_id)
 
-    # Decode the model's output and add it to the history
-    chat_history_ids = chat_history_ids[:, bot_input_ids.shape[-1]:]  # only take the latest generated tokens
-    bot_output = tokenizer.decode(chat_history_ids[0], skip_special_tokens=True)
+    # Decode the model's output
+    bot_output = tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-    # Update session state history with the new tokens (flattened)
-    st.session_state['history'] = chat_history_ids[0].tolist()
+    # Update the history with the new user input and the model's output
+    st.session_state['history'].append(user_input)
+    st.session_state['history'].append(f"bot: {bot_output} ")
 
     # Add both user input and bot response to the conversation history for display
     st.session_state['conversation'].append(f"You: {input_text}")
@@ -65,7 +62,7 @@ def generate_response(input_text):
     return bot_output
 
 # Streamlit Interface
-st.title("Chat with DialoGPT")
+st.title("Chat with T5")
 
 # Display the conversation history
 if st.session_state['conversation']: