greg funcs

app.py

@@ -41,6 +41,9 @@ from sentence_transformers import SentenceTransformer, CrossEncoder, util
 from torch import tensor as torch_tensor
 from datasets import load_dataset
 
+
+from greg_funcs import greg_search
+
 """# import models"""
 
 bi_encoder = SentenceTransformer(
@@ -63,7 +66,6 @@ mycorpus_embeddings = torch_tensor(dataset_embed_pd.values)
 
 def search(query, top_k=20, top_n=1):
     question_embedding = bi_encoder.encode(query, convert_to_tensor=True)
-    question_embedding = question_embedding  # .cuda()
     hits = util.semantic_search(
         question_embedding, mycorpus_embeddings, top_k=top_k)
     hits = hits[0]  # Get the hits for the first query
@@ -84,6 +86,7 @@ def search(query, top_k=20, top_n=1):
 
 
 def get_text(qry):
+    # predictions = greg_search(qry)
     predictions = search(qry)
     prediction_text = []
     for hit in predictions:

greg_funcs.py

@@ -0,0 +1,47 @@
+
+
+from sentence_transformers import SentenceTransformer, CrossEncoder, util
+from torch import tensor as torch_tensor
+from datasets import load_dataset
+
+
+
+"""# import models"""
+
+bi_encoder = SentenceTransformer('multi-qa-MiniLM-L6-cos-v1')
+bi_encoder.max_seq_length = 256     #Truncate long passages to 256 tokens
+
+#The bi-encoder will retrieve top_k documents. We use a cross-encoder, to re-rank the results list to improve the quality
+cross_encoder = CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2')
+
+
+
+"""# import datasets"""
+
+dataset = load_dataset("gfhayworth/hack_policy", split='train')
+mypassages = list(dataset.to_pandas()['psg'])
+
+dataset_embed = load_dataset("gfhayworth/hack_policy_embed", split='train')
+dataset_embed_pd = dataset_embed.to_pandas()
+mycorpus_embeddings = torch_tensor(dataset_embed_pd.values)
+
+def greg_search(query, passages = mypassages, doc_embedding = mycorpus_embeddings, top_k=20, top_n = 1):
+    question_embedding = bi_encoder.encode(query, convert_to_tensor=True)
+    question_embedding = question_embedding #.cuda()
+    hits = util.semantic_search(question_embedding, doc_embedding, top_k=top_k)
+    hits = hits[0]  # Get the hits for the first query
+
+    ##### Re-Ranking #####
+    cross_inp = [[query, passages[hit['corpus_id']]] for hit in hits]
+    cross_scores = cross_encoder.predict(cross_inp)
+
+    # Sort results by the cross-encoder scores
+    for idx in range(len(cross_scores)):
+        hits[idx]['cross-score'] = cross_scores[idx]
+
+    hits = sorted(hits, key=lambda x: x['cross-score'], reverse=True)
+    predictions = hits[:top_n]
+    return predictions
+    # for hit in hits[0:3]:
+    #         print("\t{:.3f}\t{}".format(hit['cross-score'], mypassages[hit['corpus_id']].replace("\n", " ")))
+