USC-GPT / reranker.py
bhulston's picture
Update reranker.py
26ed99e
raw
history blame
585 Bytes
from sentence_transformers import CrossEncoder
import numpy as np
# Let's use a reranker to get better results from our semantic search
def reranker(query, matches):
pairs = []
for match in matches:
pairs.append([query, match["metadata"]["text"]])
model = CrossEncoder('cross-encoder/ms-marco-TinyBERT-L-2-v2', max_length = 512)
print("Pairs variable:", pairs)
scores = model.predict(pairs)
top_indices = np.argsort(scores)[::-5]
top_results = ["Class: " + matches[index]["metadata"]["text"] for index in top_indices]
return top_results