File size: 546 Bytes
c755297 a45a30d c755297 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 |
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)
scores = model.predict(pairs)
top_indices = np.argsort(scores)[::-5]
top_results = ["Class: " + docs[index]["metadata"]["text"] for index in top_indices]
return top_results |