File size: 665 Bytes
c755297
 
 
 
 
 
 
 
4cf78ca
a45a30d
56eebbb
c755297
 
 
 
d6bf0f6
4cf78ca
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)

    scores = model.predict(pairs)
    top_indices = np.argsort(scores)[::-4]
    top_results = ["Class: " + matches[index]["metadata"]["text"] + "Class Url: " + matches[index]["metadata"]["url"] + "Class Time: " + str(matches[index]["metadata"]["time"]) for index in top_indices]
    return top_results