Update app.py

app.py

@@ -1,6 +1,37 @@
+import streamlit as st
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+import torch
+
+st.title("📚 Study Buddy Chatbot")
+st.write("Ask a question or type a topic, and I'll help you learn interactively!")
+
+# Initialize session state for conversation history
+if "conversation" not in st.session_state:
+    st.session_state.conversation = []
+
+# Load model with better caching and memory management
+@st.cache_resource
+def load_model():
+    MODEL_NAME = "HuggingFaceH4/zephyr-7b-alpha"
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+    model = AutoModelForCausalLM.from_pretrained(
+        MODEL_NAME, 
+        torch_dtype=torch.float16, 
+        device_map="auto",
+        low_cpu_mem_usage=True
+    )
+    return tokenizer, model
+
+# Only load model when needed
+if "model_loaded" not in st.session_state:
+    with st.spinner("Loading AI model (this may take a minute)..."):
+        tokenizer, model = load_model()
+        st.session_state.model_loaded = True
+
 def get_response(user_input):
-    """Generate a thoughtful response that includes a follow-up question."""
-    history = "\n".join(st.session_state.conversation[-5:])  # Keep only the last 5 turns
+    # Format conversation history for context
+    history = "\n".join(st.session_state.conversation[-6:])  # Last 6 exchanges
+    
     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"
@@ -8,19 +39,42 @@ def get_response(user_input):
         f"Student: {user_input}\n"
         f"Coach: "
     )
-
-    # Tokenize input with padding and attention mask
-    inputs = tokenizer(prompt, return_tensors="pt", padding=True, truncation=True)
-    input_ids = inputs.input_ids.to(model.device)
-    attention_mask = inputs.attention_mask.to(model.device)
-
+    
+    # Better generation parameters
+    input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(model.device)
     with torch.no_grad():
         output = model.generate(
-            input_ids,
-            attention_mask=attention_mask,
-            max_length=300,
-            pad_token_id=tokenizer.eos_token_id  # Ensures correct token handling
+            input_ids, 
+            max_new_tokens=250,
+            temperature=0.7,
+            top_p=0.9,
+            do_sample=True,
+            repetition_penalty=1.2
         )
     
-    response = tokenizer.decode(output[:, input_ids.shape[-1]:][0], skip_special_tokens=True)
-    return response
+    response = tokenizer.decode(output[0, input_ids.shape[1]:], skip_special_tokens=True)
+    return response
+
+# User interface
+user_input = st.text_input("Type your question or topic:")
+
+if user_input:
+    with st.spinner("Thinking..."):
+        response = get_response(user_input)
+    
+    # Add to conversation history
+    st.session_state.conversation.append(f"Student: {user_input}")
+    st.session_state.conversation.append(f"Coach: {response}")
+
+# Display conversation in a better format
+st.subheader("Conversation History")
+for i, message in enumerate(st.session_state.conversation[-10:]):
+    if i % 2 == 0:  # Student messages
+        st.markdown(f"**You**: {message.replace('Student: ', '')}")
+    else:  # Coach messages
+        st.markdown(f"**Coach**: {message.replace('Coach: ', '')}")
+
+# Add a clear conversation button
+if st.button("Clear Conversation"):
+    st.session_state.conversation = []
+    st.experimental_rerun()