| 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: " + docs[index]["metadata"]["text"] for index in top_indices] | |
| return top_results |