Update app.py

app.py

@@ -113,7 +113,10 @@ def get_top_chunks(query, chunk_embeddings, text_chunks):
 
   # Normalize all chunk embeddings to unit length for consistent comparison
   # chunk_embeddings_normalized = chunk_embeddings / chunk_embeddings.norm(dim=1, keepdim=True)
-    
+  if chunk_embeddings.ndim == 1:
+      chunk_embeddings_normalized = chunk_embeddings / chunk_embeddings.norm()
+  else:
+      chunk_embeddings_normalized = chunk_embeddings / chunk_embeddings.norm(dim=1, keepdim=True)
   # Calculate cosine similarity between query and all chunks using matrix multiplication
   similarities = torch.matmul(chunk_embeddings_normalized, query_embedding_normalized)