Update app.py

app.py

@@ -1,21 +1,3 @@
-import streamlit as st
-from transformers import AutoModelForCausalLM, AutoTokenizer
-import torch
-
-# Load the Zephyr-7B-Alpha model (fully open and optimized for instruction-following)
-MODEL_NAME = "HuggingFaceH4/zephyr-7b-alpha"
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-model = AutoModelForCausalLM.from_pretrained(
-    MODEL_NAME, 
-    torch_dtype=torch.float16, 
-    device_map="cpu",  # Forces CPU usage
-    low_cpu_mem_usage=True  # Helps reduce memory spikes
-)
-
-# Initialize conversation history if not present
-if "conversation" not in st.session_state:
-    st.session_state.conversation = []
-
 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
@@ -26,24 +8,19 @@ def get_response(user_input):
         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=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 interactively!")
 
-user_input = st.text_input("Type your question or topic:")
-if user_input:
-    response = get_response(user_input)
-    st.session_state.conversation.append(f"Student: {user_input}")
-    st.session_state.conversation.append(f"Coach: {response}")
-    st.write("🤖 Coach:", response)
+    # 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)
 
-# Display conversation history
-st.subheader("Conversation History")
-for chat in st.session_state.conversation[-10:]:
-    st.write(chat)
+    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
+        )
+    
+    response = tokenizer.decode(output[:, input_ids.shape[-1]:][0], skip_special_tokens=True)
+    return response