Update app.py

app.py

@@ -58,27 +58,23 @@ corpus_embeddings.shape
 def find_sentences(query):
     query_embedding = model.encode(query)
     hits = util.semantic_search(query_embedding, corpus_embeddings, top_k=5)
-    hits = hits[0]
-
-    for hit in hits:
-        corpus_id = hit['corpus_id']
-
-        # Find source document based on sentence index
-        count = 0
-        for idx, c in enumerate(sentence_count):
-            count+=c
-            if (corpus_id > count-1):
-                continue
-            else:
-                doc = corpus[idx]
-                #return doc
-
-                break
-    return corpus[corpus_id]
-
-
-
+    hits = hits[0][0]
 
+    
+    corpus_id = hit['corpus_id']
+    # print(corpus[corpus_id], "(Score: {:.4f})".format(hit['score']))
+    # Find source document based on sentence index
+    count = 0
+    for idx, c in enumerate(sentence_count):
+        count+=c
+        if (corpus_id > count-1):
+            continue
+        else:
+            doc = corpus[idx]
+            return doc
+            #print(f"Document: {doc}, {count}")
+            break
+    return corpus[corpus_id]