Update app.py

app.py

@@ -41,47 +41,56 @@
 #     else:
 #         st.write(f"**Bot:** {message['content']}")
 import streamlit as st
-from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoModelForCausalLM, AutoTokenizer
 
-# Set up the app title
-st.title("🤖 Simple Chatbot")
+st.title("🤖 Improved Chatbot")
 
-# Initialize the Hugging Face pipeline
+# Initialize model and tokenizer
 @st.cache_resource
-def load_chatbot():
-    # Choose a model from Hugging Face Hub
-    model_name = "microsoft/DialoGPT-medium"  # Try other models like "gpt2", "facebook/blenderbot-400M-distill"
+def load_model():
+    model_name = "microsoft/DialoGPT-medium"
     tokenizer = AutoTokenizer.from_pretrained(model_name)
     model = AutoModelForCausalLM.from_pretrained(model_name)
-    return pipeline("text-generation", model=model, tokenizer=tokenizer)
+    return model, tokenizer
 
-chatbot = load_chatbot()
+model, tokenizer = load_model()
 
-# Initialize session state for chat history
-if "messages" not in st.session_state:
-    st.session_state.messages = []
+# Initialize chat history
+if "history" not in st.session_state:
+    st.session_state.history = []
 
 # Display chat history
-for message in st.session_state.messages:
+for message in st.session_state.history:
     with st.chat_message(message["role"]):
         st.markdown(message["content"])
 
-# Chat input
+# User input
 if prompt := st.chat_input("Type your message..."):
     # Add user message to history
-    st.session_state.messages.append({"role": "user", "content": prompt})
+    st.session_state.history.append({"role": "user", "content": prompt})
     
-    # Generate bot response
+    # Prepare context for the model
+    input_ids = tokenizer.encode(
+        "\n".join([f"{msg['role']}: {msg['content']}" for msg in st.session_state.history[-5:]]) + "\nassistant:",
+        return_tensors="pt"
+    )
+    
+    # Generate response
     with st.spinner("Thinking..."):
-        bot_response = chatbot(
-            prompt,
+        output = model.generate(
+            input_ids,
             max_length=1000,
-            num_return_sequences=1,
-            pad_token_id=chatbot.tokenizer.eos_token_id
-        )[0]['generated_text']
-    
-    # Add bot response to history
-    st.session_state.messages.append({"role": "assistant", "content": bot_response})
+            pad_token_id=tokenizer.eos_token_id,
+            no_repeat_ngram_size=3,
+            do_sample=True,
+            top_k=50,
+            top_p=0.95,
+            temperature=0.7
+        )
+        
+        response = tokenizer.decode(output[0], skip_special_tokens=True).split("assistant:")[-1].strip()
+
+    # Add assistant response to history
+    st.session_state.history.append({"role": "assistant", "content": response})
     
-    # Rerun to display new messages
-    st.rerun()
+    st.rerun()