Update app.py

app.py

@@ -1,6 +1,7 @@
 import streamlit as st
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
+import random
 
 # Load pre-trained DialoGPT-small model and tokenizer
 model_name = "microsoft/DialoGPT-small"
@@ -11,24 +12,40 @@ tokenizer = AutoTokenizer.from_pretrained(model_name)
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model.to(device)
 
-# Initialize chat history
+# Initialize chat history and conversation context
 if 'history' not in st.session_state:
     st.session_state['history'] = []
 if 'conversation' not in st.session_state:
     st.session_state['conversation'] = []
 
+# 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."
+]
+
+# Select a random system prompt to start the conversation
+def get_system_prompt():
+    return random.choice(system_prompts)
+
 def generate_response(input_text):
+    # If it's the first interaction, add the system prompt to the conversation history
+    if not st.session_state['history']:
+        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()
+
     # 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)
 
     # If there is conversation history, append the new input to it
     if st.session_state['history']:
-        # Convert history to a 2D tensor (batch_size x seq_len)
         history_tensor = torch.tensor(st.session_state['history']).unsqueeze(0).to(device)
-        # Concatenate history with the new input
         bot_input_ids = torch.cat([history_tensor, new_user_input_ids], dim=-1)
     else:
-        # If no history, just use the new user input
         bot_input_ids = new_user_input_ids
 
     # Generate a response from the model