n

app.py

@@ -45,8 +45,24 @@ model = AutoModelForCausalLM.from_pretrained(model_name, device_map="balanced",
 
 # Step 4: Define the Retrieval Function
 def retrieve_documents(query, top_k=3):
-    query_embedding = np.mean([embeddings[i] for i in range(len(metadata)) if query.lower() in metadata[i].lower()], axis=0)
-    distances, indices = index.search(np.array([query_embedding]), top_k)
+    # Find embeddings matching the query
+    matched_embeddings = [embeddings[i] for i in range(len(metadata)) if query.lower() in metadata[i].lower()]
+    
+    # If no matches found, set a default query embedding
+    if matched_embeddings:
+        query_embedding = np.mean(matched_embeddings, axis=0)
+    else:
+        # Fallback: use the mean of all embeddings as a default embedding
+        query_embedding = np.mean(embeddings, axis=0)
+        print("No exact matches found for query. Using default query embedding.")
+    
+    # Reshape query_embedding to match FAISS expected shape (1, d)
+    query_embedding = query_embedding.reshape(1, -1)
+
+    # Perform the search
+    distances, indices = index.search(query_embedding, top_k)
+    
+    # Retrieve document metadata based on indices
     retrieved_docs = [metadata[idx] for idx in indices[0]]
     return retrieved_docs