Update app.py

app.py

@@ -2,23 +2,42 @@ import streamlit as st
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
 
-# Load Phi-2 model and tokenizer
-MODEL_NAME = "microsoft/phi-2"
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
+# Load the Mistral-7B-Instruct model
+tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-Instruct-v0.1")
+model = AutoModelForCausalLM.from_pretrained("mistralai/Mistral-7B-Instruct-v0.1", torch_dtype=torch.float16, device_map="auto")
+
+# Initialize conversation history if not present
+if "conversation" not in st.session_state:
+    st.session_state.conversation = []
 
 def get_response(user_input):
-    input_ids = tokenizer(user_input, return_tensors="pt").input_ids
+    """Generate a thoughtful response that includes a follow-up question."""
+    history = "\n".join(st.session_state.conversation[-5:])  # Keep only the last 5 turns
+    prompt = (
+        f"You are a knowledgeable study coach. Engage the student in conversation. "
+        f"Ask open-ended questions to deepen understanding. Provide feedback and encourage explanations.\n\n"
+        f"Previous conversation:\n{history}\n\n"
+        f"Student: {user_input}\n"
+        f"Coach: "
+    )
+    input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(model.device)
     with torch.no_grad():
-        output = model.generate(input_ids, max_length=200)
+        output = model.generate(input_ids, max_length=300)
     response = tokenizer.decode(output[:, input_ids.shape[-1]:][0], skip_special_tokens=True)
     return response
 
 # Streamlit UI
-st.title("Study Buddy Chatbot 📚")
-st.write("Ask a question or type a topic, and I'll help you learn!")
+st.title("📚 Study Buddy Chatbot")
+st.write("Ask a question or type a topic, and I'll help you learn interactively!")
 
 user_input = st.text_input("Type your question or topic:")
 if user_input:
     response = get_response(user_input)
-    st.write("🤖 Chatbot:", response)
+    st.session_state.conversation.append(f"Student: {user_input}")
+    st.session_state.conversation.append(f"Coach: {response}")
+    st.write("🤖 Coach:", response)
+
+# Display conversation history
+st.subheader("Conversation History")
+for chat in st.session_state.conversation[-10:]:
+    st.write(chat)