Update app.py

app.py

@@ -32,7 +32,7 @@ class DocumentSearch:
         # loading faiss index
         self.index = faiss.read_index(DocumentSearch.idx_path)
         # loading sbert cross_encoder
-        self.cross_encoder = CrossEncoder(DocumentSearch.cross_enc_path)
+        # self.cross_encoder = CrossEncoder(DocumentSearch.cross_enc_path)
 
     def search(self, query: str, k: int) -> list:
         # get vector representation of text query
@@ -43,24 +43,21 @@ class DocumentSearch:
         res_docs = [self.docs[i] for i in indeces[0]]
         # get scores by index
         dists = [dist for dist in distances[0]]
-
+        
+        return[{'doc': doc[0], 'url': doc[1], 'score': dist} for doc, dist in zip(res_docs, dists)][:k]
+        ##### OLD VERSION WITH CROSS-ENCODER #####
         # get answers by index
-        answers = [self.docs[i] for i in indeces[0]]
+        #answers = [self.docs[i] for i in indeces[0]]
         # prepare inputs for cross encoder
-        model_inputs = [[query, pairs[0]] for pairs in answers]
-        urls = [pairs[1] for pairs in answers]
+        # model_inputs = [[query, pairs[0]] for pairs in answers]
+        # urls = [pairs[1] for pairs in answers]
         # get similarity score between query and documents
-        scores = self.cross_encoder.predict(model_inputs, batch_size=1)
+        # scores = self.cross_encoder.predict(model_inputs, batch_size=1)
         # compose results into list of dicts
-        results = [{'doc': doc[1], 'url': url, 'score': score} for doc, url, score in zip(model_inputs, urls, scores)]
+        # results = [{'doc': doc[1], 'url': url, 'score': score} for doc, url, score in zip(model_inputs, urls, scores)]
         
         # return results sorted by similarity scores
-        return sorted(results, key=lambda x: x['score'], reverse=True)[:k]
-
-
-if __name__ == "__main__":
-    # get instance of DocumentSearch class
-    surfer = DocumentSearch()
+        # return sorted(results, key=lambda x: x['score'], reverse=True)[:k]
 
 
 if __name__ == "__main__":
@@ -89,7 +86,7 @@ if __name__ == "__main__":
         # set start time
         stt = time.time()
         # retrieve top 5 documents
-        results = surfer.search(query, k=1)
+        results = surfer.search(query, k=10)
         # set endtime
         ent = time.time()
         # measure resulting time
@@ -114,4 +111,4 @@ if __name__ == "__main__":
     else:
         st.markdown("Typical queries looks like this: _**\"What is flu?\"**_,\
                     _**\"How to cure breast cancer?\"**_,\
-                    _**\"I have headache, what should I do?\"**_")
+                    _**\"I have headache, what should I do?\"**_")