wfhw1ddlmZddlZddlZddlmZddlmZmZm Z ddl Z ddl m Z ddlmZddlmZerddlmZdd lmZdd l mZdd lmZej2eZGd d e Zy)) annotationsN) defaultdict) TYPE_CHECKINGAnyCallable)NanoBEIREvaluator)#SparseInformationRetrievalEvaluator)append_to_last_row)Tensor)SimilarityFunction)DatasetNameType) SparseEncoderceZdZdZeZddgdggdgddgddddddejd dddf dfd Zdfd Z fd Z d dfd Z dfd Z dfd Z xZS)SparseNanoBEIREvaluatora$ This evaluator extends :class:`~sentence_transformers.evaluation.NanoBEIREvaluator` but is specifically designed for sparse encoder models. This class evaluates the performance of a SparseEncoder Model on the NanoBEIR collection of Information Retrieval datasets. The collection is a set of datasets based on the BEIR collection, but with a significantly smaller size, so it can be used for quickly evaluating the retrieval performance of a model before committing to a full evaluation. The datasets are available on Hugging Face in the `NanoBEIR collection `_. This evaluator will return the same metrics as the InformationRetrievalEvaluator (i.e., MRR, nDCG, Recall@k), for each dataset and on average. Args: dataset_names (List[str]): The names of the datasets to evaluate on. Defaults to all datasets. mrr_at_k (List[int]): A list of integers representing the values of k for MRR calculation. Defaults to [10]. ndcg_at_k (List[int]): A list of integers representing the values of k for NDCG calculation. Defaults to [10]. accuracy_at_k (List[int]): A list of integers representing the values of k for accuracy calculation. Defaults to [1, 3, 5, 10]. precision_recall_at_k (List[int]): A list of integers representing the values of k for precision and recall calculation. Defaults to [1, 3, 5, 10]. map_at_k (List[int]): A list of integers representing the values of k for MAP calculation. Defaults to [100]. show_progress_bar (bool): Whether to show a progress bar during evaluation. Defaults to False. batch_size (int): The batch size for evaluation. Defaults to 32. write_csv (bool): Whether to write the evaluation results to a CSV file. Defaults to True. max_active_dims (Optional[int], optional): The maximum number of active dimensions to use. `None` uses the model's current `max_active_dims`. Defaults to None. score_functions (Dict[str, Callable[[Tensor, Tensor], Tensor]]): A dictionary mapping score function names to score functions. Defaults to {SimilarityFunction.COSINE.value: cos_sim, SimilarityFunction.DOT_PRODUCT.value: dot_score}. main_score_function (Union[str, SimilarityFunction], optional): The main score function to use for evaluation. Defaults to None. aggregate_fn (Callable[[list[float]], float]): The function to aggregate the scores. Defaults to np.mean. aggregate_key (str): The key to use for the aggregated score. Defaults to "mean". query_prompts (str | dict[str, str], optional): The prompts to add to the queries. If a string, will add the same prompt to all queries. If a dict, expects that all datasets in dataset_names are keys. corpus_prompts (str | dict[str, str], optional): The prompts to add to the corpus. If a string, will add the same prompt to all corpus. If a dict, expects that all datasets in dataset_names are keys. write_predictions (bool): Whether to write the predictions to a JSONL file. Defaults to False. This can be useful for downstream evaluation as it can be used as input to the :class:`~sentence_transformers.sparse_encoder.evaluation.ReciprocalRankFusionEvaluator` that accept precomputed predictions. Example: :: import logging from sentence_transformers import SparseEncoder from sentence_transformers.sparse_encoder.evaluation import SparseNanoBEIREvaluator logging.basicConfig(format="%(message)s", level=logging.INFO) # Load a model model = SparseEncoder("naver/splade-cocondenser-ensembledistil") datasets = ["QuoraRetrieval", "MSMARCO"] evaluator = SparseNanoBEIREvaluator( dataset_names=datasets, show_progress_bar=True, batch_size=32, ) # Run evaluation results = evaluator(model) ''' Evaluating NanoQuoraRetrieval Information Retrieval Evaluation of the model on the NanoQuoraRetrieval dataset: Queries: 50 Corpus: 5046 Score-Function: dot Accuracy@1: 92.00% Accuracy@3: 96.00% Accuracy@5: 98.00% Accuracy@10: 100.00% Precision@1: 92.00% Precision@3: 40.00% Precision@5: 24.80% Precision@10: 13.20% Recall@1: 79.73% Recall@3: 92.53% Recall@5: 94.93% Recall@10: 98.27% MRR@10: 0.9439 NDCG@10: 0.9339 MAP@100: 0.9070 Model Query Sparsity: Active Dimensions: 59.4, Sparsity Ratio: 0.9981 Model Corpus Sparsity: Active Dimensions: 61.9, Sparsity Ratio: 0.9980 Information Retrieval Evaluation of the model on the NanoMSMARCO dataset: Queries: 50 Corpus: 5043 Score-Function: dot Accuracy@1: 48.00% Accuracy@3: 74.00% Accuracy@5: 76.00% Accuracy@10: 86.00% Precision@1: 48.00% Precision@3: 24.67% Precision@5: 15.20% Precision@10: 8.60% Recall@1: 48.00% Recall@3: 74.00% Recall@5: 76.00% Recall@10: 86.00% MRR@10: 0.6191 NDCG@10: 0.6780 MAP@100: 0.6277 Model Query Sparsity: Active Dimensions: 45.4, Sparsity Ratio: 0.9985 Model Corpus Sparsity: Active Dimensions: 122.6, Sparsity Ratio: 0.9960 Average Queries: 50.0 Average Corpus: 5044.5 Aggregated for Score Function: dot Accuracy@1: 70.00% Accuracy@3: 85.00% Accuracy@5: 87.00% Accuracy@10: 93.00% Precision@1: 70.00% Recall@1: 63.87% Precision@3: 32.33% Recall@3: 83.27% Precision@5: 20.00% Recall@5: 85.47% Precision@10: 10.90% Recall@10: 92.13% MRR@10: 0.7815 NDCG@10: 0.8060 Model Query Sparsity: Active Dimensions: 52.4, Sparsity Ratio: 0.9983 Model Corpus Sparsity: Active Dimensions: 92.2, Sparsity Ratio: 0.9970 ''' # Print the results print(f"Primary metric: {evaluator.primary_metric}") # => Primary metric: NanoBEIR_mean_dot_ndcg@10 print(f"Primary metric value: {results[evaluator.primary_metric]:.4f}") # => Primary metric value: 0.8060 N )rdF Tmeanc| |_tt|_t|||||||||| | | | |||||j#|xj d|jz c_yy)N) dataset_namesmrr_at_k ndcg_at_k accuracy_at_kprecision_recall_at_kmap_at_kshow_progress_bar batch_size write_csvscore_functionsmain_score_function aggregate_fn aggregate_key query_promptscorpus_promptswrite_predictions_)max_active_dimsrlistsparsity_statssuper__init__name)selfrrrrrrrr r!r*r"r#r$r%r&r'r( __class__s /home/chris/cleankitchens-env/lib/python3.12/site-packages/sentence_transformers/sparse_encoder/evaluation/SparseNanoBEIREvaluator.pyr.z SparseNanoBEIREvaluator.__init__s( /)$/ ''"7/!+ 3%'')/!  $    + II1T1123 3I ,cbt||}|j|d|jz }|S)Nr))r-_get_human_readable_namer*)r0 dataset_namehuman_readable_namer1s r2r5z0SparseNanoBEIREvaluator._get_human_readable_names>#g>|L    + Qt';';&<#= = ""r3c^t|||jjgdy)N)query_active_dimsquery_sparsity_ratiocorpus_active_dimscorpus_sparsity_ratio)r-_append_csv_headers csv_headersextend)r0similarity_fn_namesr1s r2r=z+SparseNanoBEIREvaluator._append_csv_headerss( #$78  h r3c tt|_tt|_t ||g||||d|}|j D]}|jdjt|j|jdjt|j|jjD]#\} } |j| j| %|jjD]x\} } tdt| |j| jddDt|j| jddz |j| <z|j|j!|j|j"|jD cic]4\} } | j%|j"s"| |j"k7s2| | 6} } } |j'|| ||t(j+d|jdd d |jd d t(j+d |jdd d |jdd |X|j,rLt/t0j2j5||j6|jj9|Scc} } w)N) output_pathepochstepsquerycorpusc3,K|] \}}||zyw)N).0vallengths r2 z3SparseNanoBEIREvaluator.__call__..s+!,ff +sr)rz)Model Query Sparsity: Active Dimensions: r9z.1fz, Sparsity Ratio: r:z.4fz*Model Corpus Sparsity: Active Dimensions: r;r<)rr+r,lengthsr-__call__ evaluatorsappendlenqueriesrFitemssumzipsplitupdateprefix_name_to_metricsr/ startswith store_metrics_in_model_card_dataloggerinfor!r ospathjoincsv_filevalues) r0modelrBrCrDargskwargsper_dataset_results evaluatorkeyvalueaggregated_resultsr1s r2rNz SparseNanoBEIREvaluator.__call__s*$/"4( #g.  GK *%u OU  7I LL ! ( (Y->->)? @ LL " ) )#i.>.>*? @'66<<> 7 U##C(//6 7 7 --335 5JC'*+03E4<< RUWXHY;Z0[+(DLL3!234(5D   $ 5 ""4#>#>t?R?RTXT]T]#^_)<)B)B)D %3W[W`W`HafimqmvmvfvCJ   --e5GPUV 78K8KL_8`ad7eewx|yLyLMcydehxi j   89L9LMa9bcf8ggyz~{N{NOf{ghkzl m   "t~~  [$--8##**,  #"! s#K+5K+K+c f|j|d<|jddt| |fi|S)Nr* truncate_dim)r*popr- _load_dataset)r0r6ir_evaluator_kwargsr1s r2rmz%SparseNanoBEIREvaluator._load_datasets=261E1E-.5w$\I5HIIr3cZt|}|j|j|d<|S)Nr*)r-get_config_dictr*)r0 config_dictr1s r2rpz'SparseNanoBEIREvaluator.get_config_dicts3g-/    +-1-A-AK) *r3)"rzlist[DatasetNameType] | Noner list[int]rrrrrrrrrrrrrboolr intr!rsr*z int | Noner"z4dict[str, Callable[[Tensor, Tensor], Tensor]] | Noner#zstr | SimilarityFunction | Noner$zCallable[[list[float]], float]r%strr&str | dict[str, str] | Noner'rvr(rs)r6r returnru)Nrx) rbrrBz str | NonerCrtrDrtrwzdict[str, float])r6r rwr )rwzdict[str, Any])__name__ __module__ __qualname____doc__r information_retrieval_classnprr.r5r=rNrmrp __classcell__)r1s@r2rrsb@D#F7;!d "t#0+8"e"'&*PT?C79ww#596:"'%)43)4)4 )4 ! )4 ) )4)4 )4)4)4$)4N)4=)45)4)4 3!)4"4#)4$ %)4V#  ce##"##1;##KN##\_## ##JJ+J ,Jr3r) __future__rloggingr] collectionsrtypingrrrnumpyr~2sentence_transformers.evaluation.NanoBEIREvaluatorrSsentence_transformers.sparse_encoder.evaluation.SparseInformationRetrievalEvaluatorr sentence_transformers.utilr torchr sentence_transformers.evaluationr r $sentence_transformers.sparse_encoderr getLoggerryr[rrHr3r2rsW" #//P:CRB   8 $l/lr3