Update app.py

app.py

@@ -41,56 +41,36 @@
 #     else:
 #         st.write(f"**Bot:** {message['content']}")
 import streamlit as st
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import pipeline
 
-st.title("🤖 Improved Chatbot")
+st.title("🤖 Conversational Chatbot")
 
-# Initialize model and tokenizer
 @st.cache_resource
-def load_model():
-    model_name = "microsoft/DialoGPT-medium"
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
-    model = AutoModelForCausalLM.from_pretrained(model_name)
-    return model, tokenizer
+def load_chatbot():
+    return pipeline("conversational", model="facebook/blenderbot-400M-distill")
 
-model, tokenizer = load_model()
+chatbot = load_chatbot()
 
-# Initialize chat history
-if "history" not in st.session_state:
-    st.session_state.history = []
+if "conversation" not in st.session_state:
+    st.session_state.conversation = []
 
-# Display chat history
-for message in st.session_state.history:
-    with st.chat_message(message["role"]):
-        st.markdown(message["content"])
+# Display history
+for msg in st.session_state.conversation:
+    with st.chat_message(msg["role"]):
+        st.markdown(msg["content"])
 
-# User input
-if prompt := st.chat_input("Type your message..."):
-    # Add user message to history
-    st.session_state.history.append({"role": "user", "content": prompt})
-    
-    # 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"
-    )
+if prompt := st.chat_input("Say something"):
+    # Add user message
+    st.session_state.conversation.append({"role": "user", "content": prompt})
     
     # Generate response
     with st.spinner("Thinking..."):
-        output = model.generate(
-            input_ids,
-            max_length=1000,
-            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})
+        result = chatbot(str(st.session_state.conversation))
+    
+    # Extract bot response
+    response = result.generated_responses[-1]
+    
+    # Add to conversation
+    st.session_state.conversation.append({"role": "assistant", "content": response})
     
     st.rerun()