Skip to main content

anno_eval/eval/
advanced_harness.rs

1//! Advanced evaluation harness for specialized NER tasks.
2//!
3//! Extends the standard evaluation harness to support:
4//! - Discontinuous NER
5//! - Relation Extraction
6//! - Visual NER
7//!
8//! # Usage
9//!
10//! ```rust,ignore
11//! use anno_eval::eval::advanced_harness::{AdvancedEvalHarness, AdvancedTaskResults};
12//! use anno::W2NER;
13//!
14//! let mut harness = AdvancedEvalHarness::new();
15//!
16//! // Register a discontinuous NER model
17//! harness.register_discontinuous("w2ner", Box::new(W2NER::new()));
18//!
19//! // Run evaluation on synthetic data
20//! let results = harness.run_synthetic_discontinuous()?;
21//!
22//! // Generate report
23//! println!("{}", results.summary());
24//! ```
25
26use anno::{DiscontinuousEntity, DiscontinuousNER, RelationExtractor, Result};
27use serde::{Deserialize, Serialize};
28use std::collections::HashMap;
29
30use super::dataset::synthetic::{
31    discontinuous::dataset as discontinuous_dataset, relations::dataset as relations_dataset,
32};
33use super::discontinuous::{
34    evaluate_discontinuous_ner, DiscontinuousEvalConfig, DiscontinuousGold, DiscontinuousNERMetrics,
35};
36use super::relation::{
37    evaluate_relations, RelationEvalConfig, RelationGold, RelationMetrics, RelationPrediction,
38};
39use super::visual::{
40    evaluate_visual_ner, synthetic_visual_examples, VisualEvalConfig, VisualGold, VisualNERMetrics,
41    VisualPrediction,
42};
43
44// =============================================================================
45// RESULTS TYPES
46// =============================================================================
47
48/// Results from running an advanced evaluation task.
49#[derive(Debug, Clone, Serialize, Deserialize)]
50pub struct AdvancedTaskResults {
51    /// Timestamp of the evaluation run.
52    pub timestamp: String,
53    /// Task type evaluated.
54    pub task: String,
55    /// Results per model.
56    pub models: Vec<ModelResult>,
57    /// Number of examples evaluated.
58    pub num_examples: usize,
59    /// Total entities/relations in gold.
60    pub num_gold: usize,
61}
62
63impl AdvancedTaskResults {
64    /// Get a summary string.
65    pub fn summary(&self) -> String {
66        let mut s = format!(
67            "=== {} Evaluation ({} examples) ===\n",
68            self.task, self.num_examples
69        );
70
71        for model in &self.models {
72            s.push_str(&format!(
73                "\n{}: F1={:.1}%\n",
74                model.name,
75                model.primary_f1 * 100.0
76            ));
77        }
78
79        s
80    }
81}
82
83/// Results for a single model.
84#[derive(Debug, Clone, Serialize, Deserialize)]
85pub struct ModelResult {
86    /// Model name.
87    pub name: String,
88    /// Primary F1 score.
89    pub primary_f1: f64,
90    /// Additional metrics (task-specific).
91    pub metrics: HashMap<String, f64>,
92}
93
94// =============================================================================
95// DISCONTINUOUS NER EVALUATION
96// =============================================================================
97
98/// Run discontinuous NER evaluation on synthetic data.
99///
100/// Returns metrics for any model implementing `DiscontinuousNER`.
101pub fn evaluate_discontinuous_synthetic<M: DiscontinuousNER>(
102    model: &M,
103    labels: &[&str],
104    threshold: f32,
105) -> Result<DiscontinuousNERMetrics> {
106    let examples = discontinuous_dataset();
107    let config = DiscontinuousEvalConfig::default();
108
109    let mut all_gold: Vec<DiscontinuousGold> = Vec::new();
110    let mut all_pred: Vec<DiscontinuousEntity> = Vec::new();
111
112    for example in &examples {
113        // Collect gold entities
114        all_gold.extend(example.entities.clone());
115
116        // Run model prediction
117        let pred = model.extract_discontinuous(&example.text, labels, threshold)?;
118        all_pred.extend(pred);
119    }
120
121    Ok(evaluate_discontinuous_ner(&all_gold, &all_pred, &config))
122}
123
124/// Run discontinuous NER evaluation without a model (for testing metrics only).
125pub fn evaluate_discontinuous_gold_vs_gold() -> DiscontinuousNERMetrics {
126    let examples = discontinuous_dataset();
127    let config = DiscontinuousEvalConfig::default();
128
129    let gold: Vec<DiscontinuousGold> = examples.iter().flat_map(|ex| ex.entities.clone()).collect();
130
131    // Perfect prediction = gold
132    let pred: Vec<DiscontinuousEntity> = gold
133        .iter()
134        .map(|g| DiscontinuousEntity {
135            spans: g.spans.clone(),
136            text: g.text.clone(),
137            entity_type: g.entity_type.clone(),
138            confidence: anno::Confidence::ONE,
139        })
140        .collect();
141
142    evaluate_discontinuous_ner(&gold, &pred, &config)
143}
144
145// =============================================================================
146// RELATION EXTRACTION EVALUATION
147// =============================================================================
148
149/// Run relation extraction evaluation on synthetic data.
150pub fn evaluate_relations_synthetic<M: RelationExtractor>(
151    model: &M,
152    labels: &[&str],
153    relations: &[&str],
154    threshold: f32,
155) -> Result<RelationMetrics> {
156    let examples = relations_dataset();
157    let config = RelationEvalConfig::default();
158
159    let mut all_gold: Vec<RelationGold> = Vec::new();
160    let mut all_pred: Vec<RelationPrediction> = Vec::new();
161
162    for example in &examples {
163        // Collect gold relations
164        all_gold.extend(example.relations.clone());
165
166        // Run model prediction
167        let result = model.extract_with_relations(&example.text, labels, relations, threshold)?;
168
169        // Convert to predictions using entity indices
170        for rel in &result.relations {
171            if rel.head_idx < result.entities.len() && rel.tail_idx < result.entities.len() {
172                let head = &result.entities[rel.head_idx];
173                let tail = &result.entities[rel.tail_idx];
174                all_pred.push(RelationPrediction {
175                    head_span: (head.start(), head.end()),
176                    head_type: head.entity_type.as_label().to_string(),
177                    tail_span: (tail.start(), tail.end()),
178                    tail_type: tail.entity_type.as_label().to_string(),
179                    relation_type: rel.relation_type.clone(),
180                    confidence: rel.confidence.value() as f32,
181                });
182            }
183        }
184    }
185
186    Ok(evaluate_relations(&all_gold, &all_pred, &config))
187}
188
189/// Run relation extraction evaluation without a model (for testing metrics only).
190pub fn evaluate_relations_gold_vs_gold() -> RelationMetrics {
191    let examples = relations_dataset();
192    let config = RelationEvalConfig::default();
193
194    let gold: Vec<RelationGold> = examples
195        .iter()
196        .flat_map(|ex| ex.relations.clone())
197        .collect();
198
199    // Perfect prediction = gold
200    let pred: Vec<RelationPrediction> = gold
201        .iter()
202        .map(|g| RelationPrediction {
203            head_span: g.head_span,
204            head_type: g.head_type.clone(),
205            tail_span: g.tail_span,
206            tail_type: g.tail_type.clone(),
207            relation_type: g.relation_type.clone(),
208            confidence: 1.0,
209        })
210        .collect();
211
212    evaluate_relations(&gold, &pred, &config)
213}
214
215// =============================================================================
216// VISUAL NER EVALUATION
217// =============================================================================
218
219/// Run visual NER evaluation on synthetic data.
220pub fn evaluate_visual_gold_vs_gold() -> VisualNERMetrics {
221    let examples = synthetic_visual_examples();
222    let config = VisualEvalConfig::default();
223
224    let gold: Vec<VisualGold> = examples
225        .iter()
226        .flat_map(|(_, entities)| entities.clone())
227        .collect();
228
229    // Perfect prediction = gold
230    let pred: Vec<VisualPrediction> = gold
231        .iter()
232        .map(|g| VisualPrediction {
233            text: g.text.clone(),
234            entity_type: g.entity_type.clone(),
235            bbox: g.bbox,
236            confidence: 1.0,
237        })
238        .collect();
239
240    evaluate_visual_ner(&gold, &pred, &config)
241}
242
243// =============================================================================
244// DATASET STATISTICS
245// =============================================================================
246
247/// Get statistics about the synthetic advanced datasets.
248pub fn synthetic_dataset_stats() -> SyntheticDatasetStats {
249    let disc = discontinuous_dataset();
250    let rel = relations_dataset();
251    let vis = synthetic_visual_examples();
252
253    SyntheticDatasetStats {
254        discontinuous_examples: disc.len(),
255        discontinuous_entities: disc.iter().map(|ex| ex.entities.len()).sum(),
256        relation_examples: rel.len(),
257        relations: rel.iter().map(|ex| ex.relations.len()).sum(),
258        visual_examples: vis.len(),
259        visual_entities: vis.iter().map(|(_, e)| e.len()).sum(),
260    }
261}
262
263/// Statistics about synthetic advanced datasets.
264#[derive(Debug, Clone, Serialize, Deserialize)]
265pub struct SyntheticDatasetStats {
266    /// Number of discontinuous NER examples.
267    pub discontinuous_examples: usize,
268    /// Total discontinuous entities.
269    pub discontinuous_entities: usize,
270    /// Number of relation extraction examples.
271    pub relation_examples: usize,
272    /// Total relations.
273    pub relations: usize,
274    /// Number of visual NER examples.
275    pub visual_examples: usize,
276    /// Total visual entities.
277    pub visual_entities: usize,
278}
279
280// =============================================================================
281// TESTS
282// =============================================================================
283
284#[cfg(test)]
285mod tests {
286    use super::*;
287
288    #[test]
289    fn test_discontinuous_gold_vs_gold() {
290        let metrics = evaluate_discontinuous_gold_vs_gold();
291        assert!(
292            (metrics.exact_f1 - 1.0).abs() < 0.001,
293            "Perfect prediction should give F1=1.0, got {}",
294            metrics.exact_f1
295        );
296    }
297
298    #[test]
299    fn test_relations_gold_vs_gold() {
300        let metrics = evaluate_relations_gold_vs_gold();
301        assert!(
302            (metrics.strict_f1 - 1.0).abs() < 0.001,
303            "Perfect prediction should give F1=1.0, got {}",
304            metrics.strict_f1
305        );
306    }
307
308    #[test]
309    fn test_visual_gold_vs_gold() {
310        let metrics = evaluate_visual_gold_vs_gold();
311        assert!(
312            (metrics.e2e_f1 - 1.0).abs() < 0.001,
313            "Perfect prediction should give F1=1.0, got {}",
314            metrics.e2e_f1
315        );
316    }
317
318    #[test]
319    fn test_synthetic_dataset_stats() {
320        let stats = synthetic_dataset_stats();
321        assert!(stats.discontinuous_examples > 0);
322        assert!(stats.discontinuous_entities > 0);
323        assert!(stats.relation_examples > 0);
324        assert!(stats.relations > 0);
325        assert!(stats.visual_examples > 0);
326        assert!(stats.visual_entities > 0);
327    }
328}