Update app.py

app.py

@@ -24,8 +24,15 @@ cleaned_chunks = preprocess_text(info_text)
 model = SentenceTransformer('all-MiniLM-L6-v2')
 
 def create_embeddings(text_chunks):
-    chunk_embeddings = model.encode(text_chunks, convert_to_tensor=True) # Replace ... with the text_chunks list
-    return chunk_embeddings
+    embeddings = model.encode(text_chunks, convert_to_tensor=True)
+
+    if embeddings.dim() == 1:
+        embeddings = embeddings.unsqueeze(0)
+
+    return embeddings
+
+    # chunk_embeddings = model.encode(text_chunks, convert_to_tensor=True) # Replace ... with the text_chunks list
+    # return chunk_embeddings
 
 chunk_embeddings = create_embeddings(cleaned_chunks) # Complete this line
 
@@ -35,7 +42,12 @@ def get_top_chunks(query, chunk_embeddings, text_chunks):
     query_embedding_normalized = query_embedding / query_embedding.norm()
     chunk_embeddings_normalized = chunk_embeddings / chunk_embeddings.norm(dim=1, keepdim=True)
     similarities = torch.matmul(chunk_embeddings_normalized, query_embedding_normalized) # Complete this line
-    top_indices = torch.topk(similarities, k=3).indices
+    
+    # top_indices = torch.topk(similarities, k=3).indices
+
+    k = min(3, len(text_chunks))
+    top_indices = torch.topk(similarities, k=k).indices
+    
     top_chunks = []
     
     for i in top_indices: