-
Notifications
You must be signed in to change notification settings - Fork 192
Evaluate your custom model
Nandan Thakur edited this page Jun 29, 2022
·
1 revision
Mention your dual-encoder model in a class and have two functions: 1. encode_queries
and 2. encode_corpus
.
from beir.retrieval.search.dense import DenseRetrievalExactSearch as DRES
class YourCustomDEModel:
def __init__(self, model_path=None, **kwargs)
self.model = None # ---> HERE Load your custom model
# Write your own encoding query function (Returns: Query embeddings as numpy array)
def encode_queries(self, queries: List[str], batch_size: int, **kwargs) -> np.ndarray:
pass
# Write your own encoding corpus function (Returns: Document embeddings as numpy array)
def encode_corpus(self, corpus: List[Dict[str, str]], batch_size: int, **kwargs) -> np.ndarray:
pass
custom_model = DRES(YourCustomDEModel(model_path="your-custom-model-path"))
Mention your cross-encoder model in a class and have a single function: predict
from beir.reranking import Rerank
class YourCustomCEModel:
def __init__(self, model_path=None, **kwargs)
self.model = None # ---> HERE Load your custom model
# Write your own score function, which takes in query-document text pairs and returns the similarity scores
def predict(self, sentences: List[Tuple[str,str]], batch_size: int, **kwags) -> List[float]:
pass # return only the list of float scores
reranker = Rerank(YourCustomCEModel(model_path="your-custom-model-path"), batch_size=128)