Update app.py

app.py

@@ -57,12 +57,20 @@ corpus_embeddings.shape
 
 def find_sentences(query):
     query_embedding = model.encode(query)
-    hits = util.semantic_search(query_embedding, corpus_embeddings, top_k=5)
+    hits = util.semantic_search(query_embedding, corpus_embeddings, top_k=1)
     hit = hits[0][0]
-    
     corpus_id = hit['corpus_id']
-
-    return corpus[corpus_id]
+    saved = corpus[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
+    return saved