Update app.py

app.py

@@ -2,10 +2,10 @@ import streamlit as st
 from transformers import GPT2LMHeadModel, GPT2Tokenizer
 import torch
 
-# Check if PyTorch is installed correctly and configure device
+# Configure device (CPU/GPU)
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-# Load pre-trained GPT-2 (DialoGPT is fine-tuned from GPT-2 for conversations)
+# Load the model and tokenizer
 @st.cache_resource
 def load_model():
     model = GPT2LMHeadModel.from_pretrained("microsoft/DialoGPT-medium").to(device)
@@ -14,45 +14,49 @@ def load_model():
 
 model, tokenizer = load_model()
 
-# Function to generate responses from the model
-def get_response(input_text):
-    # Encode the input and append the end-of-sentence token
-    inputs = tokenizer.encode(input_text + tokenizer.eos_token, return_tensors="pt").to(device)
-    
-    # Generate the response
+# Function to generate chatbot response
+def get_response(conversation_history, user_input):
+    input_text = " ".join(conversation_history) + " " + user_input + tokenizer.eos_token
+    inputs = tokenizer.encode(input_text, return_tensors="pt").to(device)
+
     with torch.no_grad():
-        outputs = model.generate(inputs, max_length=150, num_return_sequences=1,
-                                 no_repeat_ngram_size=3, top_k=50, top_p=0.95, temperature=0.7)
-    
-    # Decode the response
-    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return response
+        outputs = model.generate(
+            inputs, 
+            max_length=75, 
+            num_return_sequences=1, 
+            no_repeat_ngram_size=2, 
+            top_k=40, 
+            top_p=0.8, 
+            temperature=0.7
+        )
+    return tokenizer.decode(outputs[0], skip_special_tokens=True)
 
 # Streamlit app interface
-st.title("Chat with ChatGPT-like Bot")
-
-# Display instructions
-st.write("### Chat with the DialoGPT model. Type your message below and press Enter!")
+st.title("Chatbot with Hugging Face Model")
+st.write("### Chat with the chatbot powered by DialoGPT. Type your message below!")
 
-# Maintain a conversation history
+# Initialize conversation history
 if "messages" not in st.session_state:
     st.session_state.messages = []
 
-# Display the conversation so far
+# Display chat history
 for message in st.session_state.messages:
-    st.markdown(f"**{message['role']}**: {message['content']}")
+    st.markdown(f"{message['role']}: {message['content']}")
 
-# Text input for user to interact with the chatbot
-user_input = st.text_input("You: ", "")
+# User input
+user_input = st.text_input("You: ", key="user_input")
 
-# When the user submits a message, generate a response
+# Generate response if user submits a message
 if user_input:
-    # Add user message to conversation history
+    # Add user input to conversation history
     st.session_state.messages.append({"role": "User", "content": user_input})
-    
-    # Generate and display the chatbot's response
-    response = get_response(user_input)
+
+    # Prepare context for the chatbot
+    history = [msg["content"] for msg in st.session_state.messages[-3:] if msg["role"] == "User"]
+
+    # Generate chatbot response
+    response = get_response(history, user_input)
     st.session_state.messages.append({"role": "Chatbot", "content": response})
 
-    # Keep the conversation going
-    st.text_input("You: ", "", key="input")
+    # Clear input box for new input
+    st.text_input("You: ", key="user_input", value="")