Skip to content
This repository was archived by the owner on Sep 9, 2026. It is now read-only.

Commit 8a4224d

Browse files
guenthermiguenthermi
authored
feat: add max_rel_per_label to support recall for labeled data (#826)
1 parent 7a5b0bf commit 8a4224d

4 files changed

Lines changed: 137 additions & 19 deletions

File tree

‎docarray/array/mixins/evaluation.py‎

Lines changed: 36 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,10 @@
11
import warnings
2-
from typing import Optional, Union, TYPE_CHECKING, Callable, List, Dict, Tuple
2+
from typing import Optional, Union, TYPE_CHECKING, Callable, List, Dict, Tuple, Any
33

44
from functools import wraps
55

66
import numpy as np
7-
from collections import defaultdict
7+
from collections import defaultdict, Counter
88

99
from docarray.score import NamedScore
1010

@@ -79,6 +79,7 @@ def evaluate(
7979
metric_names: Optional[List[str]] = None,
8080
strict: bool = True,
8181
label_tag: str = 'label',
82+
num_relevant_documents_per_label: Optional[Dict[Any, int]] = None,
8283
**kwargs,
8384
) -> Dict[str, float]:
8485
"""
@@ -109,6 +110,10 @@ def evaluate(
109110
aligned: on the length, and on the semantic of length. These are preventing
110111
you to evaluate on irrelevant matches accidentally.
111112
:param label_tag: Specifies the tag which contains the labels.
113+
:param num_relevant_documents_per_label: Some metrics, e.g., recall@k, require
114+
the number of relevant documents. To apply those to a labeled dataset, one
115+
can provide a dictionary which maps labels to the total number of documents
116+
with this label.
112117
:param kwargs: Additional keyword arguments to be passed to the metric
113118
functions.
114119
:return: A dictionary which stores for each metric name the average evaluation
@@ -161,7 +166,22 @@ def evaluate(
161166
results = defaultdict(list)
162167
caller_max_rel = kwargs.pop('max_rel', None)
163168
for d, gd in zip(self, ground_truth):
164-
max_rel = caller_max_rel or len(gd.matches)
169+
if caller_max_rel:
170+
max_rel = caller_max_rel
171+
elif ground_truth_type == 'labels':
172+
if num_relevant_documents_per_label:
173+
max_rel = num_relevant_documents_per_label.get(
174+
d.tags[label_tag], None
175+
)
176+
if max_rel is None:
177+
raise ValueError(
178+
'`num_relevant_documents_per_label` misses the label '
179+
+ str(d.tags[label_tag])
180+
)
181+
else:
182+
max_rel = None
183+
else:
184+
max_rel = len(gd.matches)
165185
if strict and hash_fn(d) != hash_fn(gd):
166186
raise ValueError(
167187
f'Document {d} from the left-hand side and '
@@ -174,7 +194,7 @@ def evaluate(
174194
f'Document {d!r} or {gd!r} has no matches, please check your Document'
175195
)
176196

177-
targets = gd.matches[:max_rel]
197+
targets = gd.matches
178198

179199
if ground_truth_type == 'matches':
180200
desired = {hash_fn(m) for m in targets}
@@ -438,12 +458,24 @@ def fuse_matches(global_matches: DocumentArray, local_matches: DocumentArray):
438458
new_matches.append(m)
439459
query_data[doc.id, 'matches'] = new_matches
440460

461+
if ground_truth and label_tag in ground_truth[0].tags:
462+
num_relevant_documents_per_label = dict(
463+
Counter([d.tags[label_tag] for d in ground_truth])
464+
)
465+
elif not ground_truth and label_tag in query_data[0].tags:
466+
num_relevant_documents_per_label = dict(
467+
Counter([d.tags[label_tag] for d in query_data])
468+
)
469+
else:
470+
num_relevant_documents_per_label = None
471+
441472
metrics_resp = query_data.evaluate(
442473
ground_truth=ground_truth,
443474
metrics=metrics,
444475
metric_names=metric_names,
445476
strict=strict,
446477
label_tag=label_tag,
478+
num_relevant_documents_per_label=num_relevant_documents_per_label,
447479
**kwargs,
448480
)
449481

‎docarray/math/evaluation.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -98,6 +98,8 @@ def recall_at_k(
9898
"""
9999
_check_k(k)
100100
binary_relevance = np.array(binary_relevance[:k]) != 0
101+
if max_rel is None:
102+
raise ValueError('The metric recall_at_k requires a max_rel parameter')
101103
if np.sum(binary_relevance) > max_rel:
102104
raise ValueError(f'Number of relevant Documents retrieved > {max_rel}')
103105
return np.sum(binary_relevance) / max_rel

‎docs/fundamentals/documentarray/evaluation.md‎

Lines changed: 60 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -69,10 +69,10 @@ da_prediction['@m'].summary()
6969
To evaluate the matches against a ground truth array, you simply provide a DocumentArray to the evaluate function like `da_groundtruth` in the call below:
7070

7171
```python
72-
da_predict.evaluate(ground_truth=da_groundtruth, metrics=['...'], **kwargs)
72+
da_prediction.evaluate(ground_truth=da_groundtruth, metrics=['...'], **kwargs)
7373
```
7474

75-
Thereby, `da_groundtruth` should contain the same documents as in `da_prediction` where each `matches` attribute contains exactly those documents which are relevant to the respective root document.
75+
Thereby, `da_groundtruth` should contain the same Documents as in `da_prediction` where each `matches` attribute contains exactly those Documents which are relevant to the respective root Document.
7676
The `metrics` argument determines the metric you want to use for your evaluation, e.g., `precision_at_k`.
7777

7878
In the code cell below, we evaluate the array `da_prediction` with the noisy matches against the original one `da_original`:
@@ -111,7 +111,8 @@ for d in da_prediction:
111111
Note that the evaluation against a ground truth DocumentArray only works if both DocumentArrays have the same length and their nested structure is the same.
112112
It makes no sense to evaluate with a completely different DocumentArray.
113113

114-
While evaluating, Document pairs are recognized as correct if they share the same identifier. By default, it simply uses {attr}`~docarray.Document.id`. One can customize this behavior by specifying `hash_fn`.
114+
While evaluating, Document pairs are recognized as correct if they share the same identifier. By default, it simply uses {attr}`~docarray.Document.id`.
115+
You can customize this behavior by specifying `hash_fn`.
115116

116117
Let's see an example by creating two DocumentArrays with some matches with identical texts.
117118

@@ -157,8 +158,8 @@ It is correct as we define the evaluation as checking if the first two character
157158

158159
## Evaluation via labels
159160

160-
Alternatively, you can add labels to your documents to evaluate them.
161-
In this case, a match is considered relevant to its root document if it has the same label:
161+
Alternatively, you can add labels to your Documents to evaluate them.
162+
In this case, a match is considered relevant to its root Document if it has the same label:
162163

163164
```python
164165
import numpy as np
@@ -198,7 +199,7 @@ Some of those metrics accept additional arguments as `kwargs` which you can simp
198199
```{danger}
199200
These metric scores might change if the `limit` argument of the match function is set differently.
200201
201-
**Note:** Not all of these metrics can be applied to a Top-K result, i.e., `ndcg_at_k` and `r_precision` are calculated correctly only if the limit is set equal or higher than the number of documents in the `DocumentArray` provided to the match function.
202+
**Note:** Not all of these metrics can be applied to a Top-K result, i.e., `ndcg_at_k` and `r_precision` are calculated correctly only if the limit is set equal or higher than the number of Documents in the `DocumentArray` provided to the match function.
202203
```
203204

204205
You can evaluate multiple metric functions at once, as you can see below:
@@ -215,13 +216,57 @@ da_prediction.evaluate(
215216

216217
In this case, the keyword argument `k` is passed to all metric functions, even though it does not fulfill any specific function for the calculation of the reciprocal rank.
217218

219+
### The max_rel parameter
220+
221+
Some metric functions shown in the table above require a `max_rel` parameter.
222+
This parameter should be set to the number of relevant Documents in the Document collection.
223+
Without the knowledge of this number, metrics like `recall_at_k` and `f1_score_at_k` cannot be calculated.
224+
225+
In the `evaluate` function, you can provide a keyword argument `max_rel`, which is then used for all queries.
226+
In the example below, we can use the datasets `da_prediction` and `da_original` from the beginning, where each query has nine relevant Documents.
227+
Therefore, we set `max_rel=9`.
228+
229+
```python
230+
da_prediction.evaluate(ground_truth=da_original, metrics=['recall_at_k'], max_rel=9)
231+
```
232+
233+
```text
234+
{'recall_at_k': 1.0}
235+
```
236+
237+
Since all relevant Documents are in the matches, the recall is one.
238+
However, this only makes sense if the number of relevant Documents is equal for each query.
239+
If you provide a `ground_truth` parameter to the `evaluate` function, `max_rel` is set to the number of matches of the query Document.
240+
241+
```python
242+
da_prediction.evaluate(ground_truth=da_original, metrics=['recall_at_k'])
243+
```
244+
```text
245+
{'recall_at_k': 1.0}
246+
```
247+
248+
For labeled datasets, this is not possible.
249+
Here, you can set the `num_relevant_documents_per_label` parameter of `evaluate`.
250+
It accepts a dictionary that contains the number of relevant Documents for each label.
251+
In this way, the function can set `max_rel` to the correct value for each query Document.
252+
253+
```python
254+
example_da.evaluate(
255+
metrics=['recall_at_k'], num_relevant_documents_per_label={0: 5, 1: 5}
256+
)
257+
```
258+
259+
```text
260+
{'recall_at_k': 1.0}
261+
```
262+
218263
### Custom metrics
219264

220265
If the pre-defined metrics do not fit your use-case, you can define a custom metric function.
221266
It should take as input a list of binary relevance judgements of a query (`1` and `0` values).
222267
The evaluate function already calculates this binary list from the `matches` attribute so that each number represents the relevancy of a match.
223268

224-
Let's write a custom metric function, which counts the number of relevant documents per query:
269+
Let's write a custom metric function, which counts the number of relevant Documents per query:
225270

226271
```python
227272
def count_relevant(binary_relevance):
@@ -282,20 +327,22 @@ print(result)
282327
{'reciprocal_rank': 0.7583333333333333}
283328
```
284329

330+
For metric functions which require a `max_rel` parameter, the `embed_and_evaluate` function (described later in this section) automatically constructs the dictionary for `num_relevant_documents_per_label` based on the `index_data` argument.
331+
285332
### Batch-wise matching
286333

287-
The ``embed_and_evaluate`` function is especially useful, when you need to evaluate the queries on a very large document collection (`example_index` in the code snippet above), which is too large to store the embeddings of all documents in main-memory.
288-
In this case, ``embed_and_evaluate`` matches the queries to batches of the document collection.
334+
The ``embed_and_evaluate`` function is especially useful, when you need to evaluate the queries on a very large Document collection (`example_index` in the code snippet above), which is too large to store the embeddings of all Documents in main-memory.
335+
In this case, ``embed_and_evaluate`` matches the queries to batches of the Document collection.
289336
After the batch is processed all embeddings are deleted.
290337
By default, the batch size for the matching (`match_batch_size`) is set to `100_000`.
291338
If you want to reduce the memory footprint, you can set it to a lower value.
292339

293340
### Sampling Queries
294341

295-
If you want to evaluate a large dataset, it might be useful to sample query documents.
342+
If you want to evaluate a large dataset, it might be useful to sample query Documents.
296343
Since the metric values returned by the `embed_and_evaluate` are mean values, sampling should not change the result significantly if the sample is large enough.
297-
By default, sampling is applied for `DocumentArray` objects with more than 1,000 documents.
298-
However, it is only applied on the `DocumentArray` itself and not on the document provided in `index_data`.
344+
By default, sampling is applied for `DocumentArray` objects with more than 1,000 Documents.
345+
However, it is only applied on the `DocumentArray` itself and not on the Documents provided in `index_data`.
299346
If you want to change the number of samples, you can ajust the `query_sample_size` argument.
300347
In the following code block an evaluation is done with 100 samples:
301348

@@ -323,7 +370,7 @@ da.embed_and_evaluate(
323370
{'precision_at_k': 0.13649999999999998}
324371
```
325372

326-
Please note that in this way only documents which are actually evaluated obtain an `.evaluations` attribute.
373+
Please note that in this way only Documents which are actually evaluated obtain an `.evaluations` attribute.
327374

328375
To test how close it is to the exact result, we execute the function again with `query_sample_size` set to 1,000:
329376

‎tests/unit/array/mixins/oldproto/test_eval_class.py‎

Lines changed: 39 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -192,10 +192,43 @@ def test_eval_mixin_one_of_n_labeled(metric_fn, metric_score, label_tag):
192192
da = DocumentArray([Document(text=str(i), tags={label_tag: i}) for i in range(3)])
193193
for d in da:
194194
d.matches = da
195-
r = da.evaluate([metric_fn], label_tag=label_tag)[metric_fn]
195+
r = da.evaluate([metric_fn], label_tag=label_tag, max_rel=3)[metric_fn]
196196
assert abs(r - metric_score) < 0.001
197197

198198

199+
@pytest.mark.parametrize('label_tag', ['label', 'custom_tag'])
200+
@pytest.mark.parametrize(
201+
'metric_fn, metric_score',
202+
[
203+
('recall_at_k', 1.0),
204+
('f1_score_at_k', 0.5),
205+
],
206+
)
207+
def test_num_relevant_documents_per_label(metric_fn, metric_score, label_tag):
208+
da = DocumentArray([Document(text=str(i), tags={label_tag: i}) for i in range(3)])
209+
num_relevant_documents_per_label = {i: 1 for i in range(3)}
210+
for d in da:
211+
d.matches = da
212+
r = da.evaluate(
213+
[metric_fn],
214+
label_tag=label_tag,
215+
num_relevant_documents_per_label=num_relevant_documents_per_label,
216+
)[metric_fn]
217+
assert abs(r - metric_score) < 0.001
218+
219+
220+
def test_missing_max_rel_should_raise():
221+
da = DocumentArray([Document(text=str(i), tags={'label': i}) for i in range(3)])
222+
num_relevant_documents_per_label = {i: 1 for i in range(2)}
223+
for d in da:
224+
d.matches = da
225+
with pytest.raises(ValueError):
226+
da.evaluate(
227+
['recall_at_k'],
228+
num_relevant_documents_per_label=num_relevant_documents_per_label,
229+
)
230+
231+
199232
@pytest.mark.parametrize(
200233
'storage, config',
201234
[
@@ -540,7 +573,11 @@ def test_embed_and_evaluate_two_das(storage, config, sample_size, start_storage)
540573
(False, {'precision_at_k': 1.0 / 3, 'reciprocal_rank': 1.0}, 'label'),
541574
(
542575
True,
543-
{'precision_at_k': 1.0 / 3, 'reciprocal_rank': 11.0 / 18.0},
576+
{
577+
'precision_at_k': 1.0 / 3,
578+
'reciprocal_rank': 11.0 / 18.0,
579+
'recall_at_k': 1.0,
580+
},
544581
'custom_tag',
545582
),
546583
],

0 commit comments

Comments
 (0)