Spaces:
Runtime error
Runtime error
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", " "))) | |