Update app.py

app.py

@@ -44,26 +44,35 @@ essay_chunks.extend(staar_chunks)
 # load an embedding model 
 model = SentenceTransformer('all-MiniLM-L6-v2')
 
-chunk_embeddings = model.encode(cleaned_chunks, convert_to_tensor=True)
+def create_embeddings(text_chunks):
+    chunk_embeddings = model.encode(text_chunks, convert_to_tensor=True)
+    return chunk_embeddings
 
-def pull_relevant_info(query, top_k=3): 
-    query_embedding = model.encode(query, convert_to_tensor=True)
+essay_embeddings = create_embeddings(essay_chunks)
+
+def pull_relevant_info(message, chunk_embeddings, text_chunks): 
+    query_embedding = model.encode(message, convert_to_tensor=True)
     query_embedding = query_embedding / query_embedding.norm()
 
     norm_chunk_embeddings = chunk_embeddings / chunk_embeddings.norm(dim=1, keepdim=True)
 
     similarities = torch.matmul(norm_chunk_embeddings, query_embedding)
 
-    top_indices = torch.topk(similarities, k=top_k).indices.cpu().numpy()
+    top_indices = torch.topk(similarities, k=3).indices
+
+    top_chunks = []
+
+    for i in top_indices:
+        relevant_info = text_chunks[i]
+        top_chunks.append(relevant_info)
 
-    relevant_info = "\n\n".join([cleaned_chunks[i] for i in top_indices])
-    return relevant_info
+    return top_chunks
 
 client = InferenceClient("microsoft/phi-4")
 
 def respond(message, history): 
     
-    info = pull_relevant_info(message, top_k=3)
+    info = pull_relevant_info(message, essay_embeddings, essay_chunks)
     system_message = (f"You are a helpful and kind teacher named Ms. Honey.  You respond clearly in no more than three complete sentences.  If a user asks you to write something for them, you refuse and remind them they are capable of writing the piece themselves. Use the following information to help answer the user's question:\n\n{info}\n\n")
     messages = [{"role": "system", "content": system_message}]