File size: 585 Bytes
c755297 a45a30d 56eebbb c755297 fe4ae60 c755297 26ed99e c755297 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 |
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 |