Update app.py

app.py

@@ -32,33 +32,35 @@ 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, k=5):
-        query_vec = self.vectorizer.transform([query]).toarray()[0]
-        scores = np.dot(self.doc_vectors, query_vec)
-        indeces = np.argpartition(scores, -k)[-k:]
-        try:
-            res_docs = [self.docs[i] for i in indeces]
-        except IndexError:
-            res_docs = self.docs
-        return res_docs
+    def search(self, query: str, k: int) -> list:
+        # get vector representation of text query
+        query_vector = self.encoder.encode([query])
+        # perform search via faiss FlatIP index
+        distances, indeces = self.index.search(query_vector, k*10)
+        # get docs by index
+        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]
+        return sorted(results, key=lambda x: x['score'], reverse=True)[:k]
+
+
+if __name__ == "__main__":
+    # get instance of DocumentSearch class
+    surfer = DocumentSearch()
 
 
 if __name__ == "__main__":