File size: 602 Bytes
c755297 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 |
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):
docs = matches.matches
print("matches are:", docs)
pairs = []
for match in docs:
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 |