anno_eval/eval/
discontinuous.rs

1//! Discontinuous NER evaluation metrics.
2//!
3//! # Overview
4//!
5//! Discontinuous named entities span non-contiguous text regions, common in:
6//! - Coordination structures: "New York and Los Angeles airports"
7//! - Biomedical text: "left and right ventricle"
8//! - Legal documents: "paragraphs 2(a), 3(b), and 4(c)"
9//!
10//! # Evaluation Metrics
11//!
12//! This module implements metrics from the W2NER paper (arXiv:2112.10070):
13//!
14//! | Metric | Description |
15//! |--------|-------------|
16//! | Exact F1 | All spans must match exactly |
17//! | Entity Boundary F1 (EBF) | Head and tail tokens correct |
18//! | Partial Span F1 | Overlap-based matching for each segment |
19//!
20//! # Research Alignment
21//!
22//! This module’s “boundary” metric is a lenient strategy that checks whether the
23//! head/tail tokens of predicted entities are correctly identified (see W2NER,
24//! arXiv:2112.10070).
25//!
26//! Benchmark datasets for discontinuous NER include CADEC and ShARe/CLEF tasks.
27
28use anno::DiscontinuousEntity;
29use serde::{Deserialize, Serialize};
30use std::collections::HashMap;
31
32/// Ground truth discontinuous entity.
33#[derive(Debug, Clone, Serialize, Deserialize)]
34pub struct DiscontinuousGold {
35    /// Non-contiguous spans (start, end) pairs
36    pub spans: Vec<(usize, usize)>,
37    /// Entity type label
38    pub entity_type: String,
39    /// Original text (concatenated from spans)
40    pub text: String,
41}
42
43impl DiscontinuousGold {
44    /// Create a new discontinuous gold entity.
45    pub fn new(
46        spans: Vec<(usize, usize)>,
47        entity_type: impl Into<String>,
48        text: impl Into<String>,
49    ) -> Self {
50        Self {
51            spans,
52            entity_type: entity_type.into(),
53            text: text.into(),
54        }
55    }
56
57    /// Create a contiguous gold entity (single span).
58    pub fn contiguous(
59        start: usize,
60        end: usize,
61        entity_type: impl Into<String>,
62        text: impl Into<String>,
63    ) -> Self {
64        Self {
65            spans: vec![(start, end)],
66            entity_type: entity_type.into(),
67            text: text.into(),
68        }
69    }
70
71    /// Check if this entity is contiguous (single span).
72    pub fn is_contiguous(&self) -> bool {
73        self.spans.len() == 1
74    }
75
76    /// Get the bounding range (min start, max end).
77    pub fn bounding_range(&self) -> Option<(usize, usize)> {
78        let min_start = self.spans.iter().map(|(s, _)| *s).min()?;
79        let max_end = self.spans.iter().map(|(_, e)| *e).max()?;
80        Some((min_start, max_end))
81    }
82
83    /// Total character length across all spans.
84    pub fn total_length(&self) -> usize {
85        self.spans.iter().map(|(s, e)| e - s).sum()
86    }
87}
88
89/// Metrics for discontinuous NER evaluation.
90#[derive(Debug, Clone, Serialize, Deserialize)]
91pub struct DiscontinuousNERMetrics {
92    /// Exact match F1: all spans must match exactly
93    pub exact_f1: f64,
94    /// Exact match precision
95    pub exact_precision: f64,
96    /// Exact match recall
97    pub exact_recall: f64,
98    /// Entity Boundary F1: head/tail tokens correct
99    pub entity_boundary_f1: f64,
100    /// Entity Boundary precision
101    pub entity_boundary_precision: f64,
102    /// Entity Boundary recall
103    pub entity_boundary_recall: f64,
104    /// Partial span F1: overlap-based matching
105    pub partial_span_f1: f64,
106    /// Partial span precision
107    pub partial_span_precision: f64,
108    /// Partial span recall
109    pub partial_span_recall: f64,
110    /// Total predicted entities
111    pub num_predicted: usize,
112    /// Total gold entities
113    pub num_gold: usize,
114    /// Exact matches
115    pub exact_matches: usize,
116    /// Boundary matches
117    pub boundary_matches: usize,
118    /// Per-type breakdown
119    pub per_type: HashMap<String, TypeMetrics>,
120}
121
122/// Per-type metrics for discontinuous NER.
123#[derive(Debug, Clone, Default, Serialize, Deserialize)]
124pub struct TypeMetrics {
125    /// Exact F1 for this type
126    pub exact_f1: f64,
127    /// Entity boundary F1 for this type
128    pub boundary_f1: f64,
129    /// Number of gold entities of this type
130    pub gold_count: usize,
131    /// Number of predicted entities of this type
132    pub pred_count: usize,
133    /// Number of exact matches
134    pub exact_matches: usize,
135}
136
137/// Configuration for discontinuous NER evaluation.
138#[derive(Debug, Clone)]
139pub struct DiscontinuousEvalConfig {
140    /// Overlap threshold for partial matching (default: 0.5)
141    pub overlap_threshold: f64,
142    /// Whether to require type match (default: true)
143    pub require_type_match: bool,
144}
145
146impl Default for DiscontinuousEvalConfig {
147    fn default() -> Self {
148        Self {
149            overlap_threshold: 0.5,
150            require_type_match: true,
151        }
152    }
153}
154
155/// Evaluate discontinuous NER predictions against gold standard.
156///
157/// # Arguments
158/// * `gold` - Ground truth discontinuous entities
159/// * `pred` - Predicted discontinuous entities
160/// * `config` - Evaluation configuration
161///
162/// # Returns
163/// Comprehensive metrics for discontinuous NER
164///
165/// # Example
166///
167/// ```rust
168/// use anno_eval::eval::discontinuous::{evaluate_discontinuous_ner, DiscontinuousGold, DiscontinuousEvalConfig};
169/// use anno::DiscontinuousEntity;
170///
171/// let gold = vec![
172///     DiscontinuousGold::new(
173///         vec![(0, 8), (25, 33)],
174///         "location",
175///         "New York airports"
176///     ),
177/// ];
178///
179/// let pred = vec![
180///     DiscontinuousEntity {
181///         spans: vec![(0, 8), (25, 33)],
182///         text: "New York airports".to_string(),
183///         entity_type: "location".to_string(),
184///         confidence: anno::Confidence::new(0.9),
185///     },
186/// ];
187///
188/// let metrics = evaluate_discontinuous_ner(&gold, &pred, &DiscontinuousEvalConfig::default());
189/// assert!((metrics.exact_f1 - 1.0).abs() < 0.001);
190/// ```
191pub fn evaluate_discontinuous_ner(
192    gold: &[DiscontinuousGold],
193    pred: &[DiscontinuousEntity],
194    config: &DiscontinuousEvalConfig,
195) -> DiscontinuousNERMetrics {
196    if gold.is_empty() && pred.is_empty() {
197        return DiscontinuousNERMetrics {
198            exact_f1: 1.0,
199            exact_precision: 1.0,
200            exact_recall: 1.0,
201            entity_boundary_f1: 1.0,
202            entity_boundary_precision: 1.0,
203            entity_boundary_recall: 1.0,
204            partial_span_f1: 1.0,
205            partial_span_precision: 1.0,
206            partial_span_recall: 1.0,
207            num_predicted: 0,
208            num_gold: 0,
209            exact_matches: 0,
210            boundary_matches: 0,
211            per_type: HashMap::new(),
212        };
213    }
214
215    // Track matches
216    let mut gold_matched_exact = vec![false; gold.len()];
217    let mut gold_matched_boundary = vec![false; gold.len()];
218    let mut pred_matched_exact = vec![false; pred.len()];
219    let mut pred_matched_boundary = vec![false; pred.len()];
220
221    // Per-type tracking
222    let mut type_stats: HashMap<String, (usize, usize, usize, usize)> = HashMap::new(); // (gold, pred, exact, boundary)
223
224    // Count gold per type
225    for g in gold {
226        let entry = type_stats.entry(g.entity_type.clone()).or_default();
227        entry.0 += 1;
228    }
229
230    // Count pred per type
231    for p in pred {
232        let entry = type_stats.entry(p.entity_type.clone()).or_default();
233        entry.1 += 1;
234    }
235
236    // Exact matching: spans must match exactly
237    for (pi, p) in pred.iter().enumerate() {
238        for (gi, g) in gold.iter().enumerate() {
239            if gold_matched_exact[gi] {
240                continue;
241            }
242
243            // Type match check
244            if config.require_type_match && p.entity_type != g.entity_type {
245                continue;
246            }
247
248            // Exact span match
249            if spans_match_exactly(&p.spans, &g.spans) {
250                gold_matched_exact[gi] = true;
251                pred_matched_exact[pi] = true;
252
253                let entry = type_stats.entry(g.entity_type.clone()).or_default();
254                entry.2 += 1;
255                break;
256            }
257        }
258    }
259
260    // Boundary matching: head and tail tokens correct
261    for (pi, p) in pred.iter().enumerate() {
262        for (gi, g) in gold.iter().enumerate() {
263            if gold_matched_boundary[gi] {
264                continue;
265            }
266
267            if config.require_type_match && p.entity_type != g.entity_type {
268                continue;
269            }
270
271            // Boundary match: first span start and last span end
272            if boundaries_match(&p.spans, &g.spans) {
273                gold_matched_boundary[gi] = true;
274                pred_matched_boundary[pi] = true;
275
276                let entry = type_stats.entry(g.entity_type.clone()).or_default();
277                entry.3 += 1;
278                break;
279            }
280        }
281    }
282
283    // Calculate partial span overlap scores
284    let mut partial_precision_sum = 0.0;
285    let mut partial_recall_sum = 0.0;
286
287    for p in pred {
288        let best_overlap = gold
289            .iter()
290            .filter(|g| !config.require_type_match || p.entity_type == g.entity_type)
291            .map(|g| calculate_multi_span_overlap(&p.spans, &g.spans))
292            .max_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal))
293            .unwrap_or(0.0);
294        partial_precision_sum += best_overlap;
295    }
296
297    for g in gold {
298        let best_overlap = pred
299            .iter()
300            .filter(|p| !config.require_type_match || p.entity_type == g.entity_type)
301            .map(|p| calculate_multi_span_overlap(&p.spans, &g.spans))
302            .max_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal))
303            .unwrap_or(0.0);
304        partial_recall_sum += best_overlap;
305    }
306
307    // Calculate metrics
308    let exact_matches = pred_matched_exact.iter().filter(|&&m| m).count();
309    let boundary_matches = pred_matched_boundary.iter().filter(|&&m| m).count();
310
311    let exact_precision = if !pred.is_empty() {
312        exact_matches as f64 / pred.len() as f64
313    } else {
314        0.0
315    };
316    let exact_recall = if !gold.is_empty() {
317        exact_matches as f64 / gold.len() as f64
318    } else {
319        0.0
320    };
321    let exact_f1 = f1_score(exact_precision, exact_recall);
322
323    let boundary_precision = if !pred.is_empty() {
324        boundary_matches as f64 / pred.len() as f64
325    } else {
326        0.0
327    };
328    let boundary_recall = if !gold.is_empty() {
329        boundary_matches as f64 / gold.len() as f64
330    } else {
331        0.0
332    };
333    let entity_boundary_f1 = f1_score(boundary_precision, boundary_recall);
334
335    let partial_span_precision = if !pred.is_empty() {
336        partial_precision_sum / pred.len() as f64
337    } else {
338        0.0
339    };
340    let partial_span_recall = if !gold.is_empty() {
341        partial_recall_sum / gold.len() as f64
342    } else {
343        0.0
344    };
345    let partial_span_f1 = f1_score(partial_span_precision, partial_span_recall);
346
347    // Build per-type metrics
348    let per_type: HashMap<String, TypeMetrics> = type_stats
349        .into_iter()
350        .map(|(t, (gold_count, pred_count, exact, boundary))| {
351            let exact_p = if pred_count > 0 {
352                exact as f64 / pred_count as f64
353            } else {
354                0.0
355            };
356            let exact_r = if gold_count > 0 {
357                exact as f64 / gold_count as f64
358            } else {
359                0.0
360            };
361            let boundary_p = if pred_count > 0 {
362                boundary as f64 / pred_count as f64
363            } else {
364                0.0
365            };
366            let boundary_r = if gold_count > 0 {
367                boundary as f64 / gold_count as f64
368            } else {
369                0.0
370            };
371
372            (
373                t,
374                TypeMetrics {
375                    exact_f1: f1_score(exact_p, exact_r),
376                    boundary_f1: f1_score(boundary_p, boundary_r),
377                    gold_count,
378                    pred_count,
379                    exact_matches: exact,
380                },
381            )
382        })
383        .collect();
384
385    DiscontinuousNERMetrics {
386        exact_f1,
387        exact_precision,
388        exact_recall,
389        entity_boundary_f1,
390        entity_boundary_precision: boundary_precision,
391        entity_boundary_recall: boundary_recall,
392        partial_span_f1,
393        partial_span_precision,
394        partial_span_recall,
395        num_predicted: pred.len(),
396        num_gold: gold.len(),
397        exact_matches,
398        boundary_matches,
399        per_type,
400    }
401}
402
403/// Check if two span sets match exactly.
404fn spans_match_exactly(a: &[(usize, usize)], b: &[(usize, usize)]) -> bool {
405    if a.len() != b.len() {
406        return false;
407    }
408
409    let mut a_sorted: Vec<_> = a.to_vec();
410    let mut b_sorted: Vec<_> = b.to_vec();
411    a_sorted.sort();
412    b_sorted.sort();
413
414    a_sorted == b_sorted
415}
416
417/// Check if boundaries (first start, last end) match.
418fn boundaries_match(a: &[(usize, usize)], b: &[(usize, usize)]) -> bool {
419    match (a.is_empty(), b.is_empty()) {
420        (true, true) => true,
421        (true, false) | (false, true) => false,
422        (false, false) => {
423            // Safe: we've verified both are non-empty
424            let (Some(a_min), Some(a_max)) = (
425                a.iter().map(|(s, _)| *s).min(),
426                a.iter().map(|(_, e)| *e).max(),
427            ) else {
428                return false;
429            };
430            let (Some(b_min), Some(b_max)) = (
431                b.iter().map(|(s, _)| *s).min(),
432                b.iter().map(|(_, e)| *e).max(),
433            ) else {
434                return false;
435            };
436            a_min == b_min && a_max == b_max
437        }
438    }
439}
440
441/// Calculate overlap between two sets of spans.
442///
443/// Uses Intersection over Union (IoU) across all spans.
444fn calculate_multi_span_overlap(a: &[(usize, usize)], b: &[(usize, usize)]) -> f64 {
445    let a_chars: std::collections::HashSet<usize> = a.iter().flat_map(|(s, e)| *s..*e).collect();
446    let b_chars: std::collections::HashSet<usize> = b.iter().flat_map(|(s, e)| *s..*e).collect();
447
448    let intersection = a_chars.intersection(&b_chars).count();
449    let union = a_chars.union(&b_chars).count();
450
451    if union == 0 {
452        return 1.0; // Both empty
453    }
454
455    intersection as f64 / union as f64
456}
457
458/// Calculate F1 score from precision and recall.
459fn f1_score(precision: f64, recall: f64) -> f64 {
460    if precision + recall > 0.0 {
461        2.0 * precision * recall / (precision + recall)
462    } else {
463        0.0
464    }
465}
466
467#[cfg(test)]
468mod tests {
469    use super::*;
470
471    #[test]
472    fn test_exact_match() {
473        let gold = vec![DiscontinuousGold::new(
474            vec![(0, 5), (10, 15)],
475            "LOC",
476            "test",
477        )];
478        let pred = vec![DiscontinuousEntity {
479            spans: vec![(0, 5), (10, 15)],
480            text: "test".to_string(),
481            entity_type: "LOC".to_string(),
482            confidence: anno::Confidence::new(0.9),
483        }];
484
485        let metrics = evaluate_discontinuous_ner(&gold, &pred, &DiscontinuousEvalConfig::default());
486        assert!((metrics.exact_f1 - 1.0).abs() < 0.001);
487    }
488
489    #[test]
490    fn test_boundary_match() {
491        let gold = vec![DiscontinuousGold::new(
492            vec![(0, 5), (10, 15)],
493            "LOC",
494            "test",
495        )];
496        // Same boundaries but different internal span structure
497        let pred = vec![DiscontinuousEntity {
498            spans: vec![(0, 3), (3, 5), (10, 15)],
499            text: "test".to_string(),
500            entity_type: "LOC".to_string(),
501            confidence: anno::Confidence::new(0.9),
502        }];
503
504        let metrics = evaluate_discontinuous_ner(&gold, &pred, &DiscontinuousEvalConfig::default());
505        // Not exact match
506        assert!(metrics.exact_f1 < 1.0);
507        // But boundary match
508        assert!((metrics.entity_boundary_f1 - 1.0).abs() < 0.001);
509    }
510
511    #[test]
512    fn test_contiguous_entity() {
513        let gold = DiscontinuousGold::contiguous(0, 10, "PER", "John Smith");
514        assert!(gold.is_contiguous());
515        assert_eq!(gold.total_length(), 10);
516    }
517
518    #[test]
519    fn test_bounding_range() {
520        let gold = DiscontinuousGold::new(vec![(0, 5), (20, 30)], "LOC", "test");
521        assert_eq!(gold.bounding_range(), Some((0, 30)));
522    }
523
524    #[test]
525    fn test_empty_inputs() {
526        let metrics = evaluate_discontinuous_ner(&[], &[], &DiscontinuousEvalConfig::default());
527        assert!((metrics.exact_f1 - 1.0).abs() < 0.001);
528    }
529
530    #[test]
531    fn test_type_mismatch() {
532        let gold = vec![DiscontinuousGold::new(vec![(0, 5)], "PER", "John")];
533        let pred = vec![DiscontinuousEntity {
534            spans: vec![(0, 5)],
535            text: "John".to_string(),
536            entity_type: "ORG".to_string(), // Wrong type
537            confidence: anno::Confidence::new(0.9),
538        }];
539
540        let config = DiscontinuousEvalConfig {
541            require_type_match: true,
542            ..Default::default()
543        };
544        let metrics = evaluate_discontinuous_ner(&gold, &pred, &config);
545        assert!(metrics.exact_f1 < 0.001);
546    }
547
548    #[test]
549    fn test_partial_overlap() {
550        let gold = vec![DiscontinuousGold::new(vec![(0, 10)], "LOC", "test")];
551        let pred = vec![DiscontinuousEntity {
552            spans: vec![(5, 15)], // 50% overlap
553            text: "test".to_string(),
554            entity_type: "LOC".to_string(),
555            confidence: anno::Confidence::new(0.9),
556        }];
557
558        let metrics = evaluate_discontinuous_ner(&gold, &pred, &DiscontinuousEvalConfig::default());
559        // Should have partial overlap score
560        assert!(metrics.partial_span_f1 > 0.0);
561        assert!(metrics.partial_span_f1 < 1.0);
562    }
563
564    #[test]
565    fn test_multi_span_overlap() {
566        let a = vec![(0, 10), (20, 30)];
567        let b = vec![(5, 25)];
568
569        let overlap = calculate_multi_span_overlap(&a, &b);
570        // a covers: 0-10, 20-30 = 20 chars
571        // b covers: 5-25 = 20 chars
572        // intersection: 5-10, 20-25 = 10 chars
573        // union: 0-30 = 30 chars (but with gap, so actually 25 chars)
574        assert!(overlap > 0.0);
575        assert!(overlap < 1.0);
576    }
577}
anno_eval/eval/discontinuous.rs

anno_eval/eval/
discontinuous.rs
