diff --git a/examples/benchmarking/using_beir_with_a_custom_sbert_model.py b/examples/benchmarking/using_beir_with_a_custom_sbert_model.py new file mode 100644 index 00000000..5547fa05 --- /dev/null +++ b/examples/benchmarking/using_beir_with_a_custom_sbert_model.py @@ -0,0 +1,235 @@ +# -*- coding: utf-8 -*- +"""Using BEIR with a Custom SBERT Model + +Automatically generated by Colaboratory. + +Original file is located at + https://colab.research.google.com/drive/1552z6RVGVaLgIlLy4AOXXGaweQgiEA4r + +## Downloading Sentence Transformers +""" + +!pip install sentence_transformers + +"""Downloading the Semantic Textua Similarity dataset""" + +!wget http://ixa2.si.ehu.es/stswiki/images/4/48/Stsbenchmark.tar.gz + +"""Taking a look at the files""" + +# Commented out IPython magic to ensure Python compatibility. +# %%bash +# +# tar -xvf /content/Stsbenchmark.tar.gz + +!cat stsbenchmark/readme.txt + +"""Looking into the file contents""" + +!head -n 10 stsbenchmark/sts-dev.csv + +"""What is the average length of the train/dev/test files?""" + +# Commented out IPython magic to ensure Python compatibility. +# %%bash +# dev_len=$(wc -l return np.asarray(self.model.encode(queries, batch_size=batch_size, **kwargs)) + def encode_queries(self, queries: List[str], batch_size: int = 16, **kwargs) -> np.ndarray: + return np.asarray(self.model.encode(queries, batch_size=batch_size, **kwargs).cpu()) + + # Write your own encoding corpus function (Returns: Document embeddings as numpy array) + # For eg ==> sentences = [(doc["title"] + " " + doc["text"]).strip() for doc in corpus] + # ==> return np.asarray(self.model.encode(sentences, batch_size=batch_size, **kwargs)) + def encode_corpus(self, corpus: List[Dict[str, str]], batch_size: int = 8, **kwargs) -> np.ndarray: + sentences = [(doc["title"] + " " + doc["text"]).strip() for doc in corpus] + return np.asarray(self.model.encode(sentences, batch_size=batch_size, **kwargs).cpu()) + +"""Finally evaluate the results on Scifact dataset using our custom-model""" + +from beir.retrieval.evaluation import EvaluateRetrieval +#from beir.retrieval import models +from beir.retrieval.search.dense import DenseRetrievalExactSearch as DRES + +model = DRES(DistilBertModel(model_path=model_save_path), batch_size=64) +retriever = EvaluateRetrieval(model, score_function="cos_sim") + +#### Retrieve dense results (format of results is identical to qrels) +results = retriever.retrieve(corpus, queries) + +ndcg, _map, recall, precision = retriever.evaluate(qrels, results, retriever.k_values) + +print(ndcg, _map, recall, precision) \ No newline at end of file