anno_eval/eval/
harness.rs

1//! Evaluation harness for comprehensive NER benchmarking.
2//!
3//! Integrates datasets, metrics, and backends into a unified evaluation framework.
4//!
5//! # Philosophy
6//!
7//! Following burntsushi's approach: real-world evaluation with clear, honest metrics.
8//! No cherry-picking success cases - measure what actually matters.
9//!
10//! # Usage
11//!
12//! ```rust,ignore
13//! use anno_eval::eval::harness::{EvalHarness, EvalConfig, BackendRegistry};
14//! use anno::{RegexNER, HeuristicNER, StackedNER};
15//!
16//! // Create harness with default config
17//! let mut harness = EvalHarness::new(EvalConfig::default())?;
18//!
19//! // Register backends
20//! harness.register("pattern", Box::new(RegexNER::new()));
21//! harness.register("heuristic", Box::new(HeuristicNER::new()));
22//! harness.register("stacked", Box::new(StackedNER::new()));
23//!
24//! // Run on synthetic data
25//! let results = harness.run_synthetic()?;
26//!
27//! // Or run on real datasets (requires `eval` feature)
28//! #[cfg(feature = "eval")]
29//! let results = harness.run_real_datasets(&[DatasetId::WikiGold, DatasetId::Wnut17])?;
30//!
31//! // Generate HTML report
32//! let html = harness.to_html_report(&results)?;
33//! std::fs::write("eval_report.html", html)?;
34//! ```
35
36use crate::eval::datasets::GoldEntity;
37use crate::eval::loader::{DatasetId, DatasetLoader};
38use crate::eval::synthetic::{
39    all_datasets, datasets_by_difficulty, datasets_by_domain, AnnotatedExample, Difficulty, Domain,
40};
41use crate::eval::types::MetricWithVariance;
42use crate::eval::{evaluate_ner_model, TypeMetrics};
43use anno::{Error, Model, Result};
44use serde::{Deserialize, Serialize};
45use std::collections::HashMap;
46use std::time::Instant;
47
48// =============================================================================
49// Configuration
50// =============================================================================
51
52/// Configuration for evaluation runs.
53#[derive(Debug, Clone, Serialize, Deserialize)]
54pub struct EvalConfig {
55    /// Maximum examples per dataset (0 = no limit)
56    pub max_examples_per_dataset: usize,
57    /// Include per-difficulty breakdown
58    pub breakdown_by_difficulty: bool,
59    /// Include per-domain breakdown
60    pub breakdown_by_domain: bool,
61    /// Include per-entity-type metrics
62    pub breakdown_by_type: bool,
63    /// Run warmup iteration before timing
64    pub warmup: bool,
65    /// Number of warmup iterations
66    pub warmup_iterations: usize,
67    /// Minimum confidence threshold (filter predictions below this)
68    pub min_confidence: Option<f64>,
69    /// Cache directory for downloaded datasets
70    pub cache_dir: Option<String>,
71    /// Use type mapping to normalize domain-specific entity types to standard NER types.
72    ///
73    /// When enabled:
74    /// - MIT Movie "Actor" -> "Person"
75    /// - MIT Restaurant "Restaurant_Name" -> "Organization"
76    /// - Biomedical "Disease" -> "Other" (or custom)
77    pub normalize_types: bool,
78}
79
80impl Default for EvalConfig {
81    fn default() -> Self {
82        Self {
83            max_examples_per_dataset: 0, // No limit
84            breakdown_by_difficulty: true,
85            breakdown_by_domain: true,
86            breakdown_by_type: true,
87            warmup: true,
88            warmup_iterations: 1,
89            min_confidence: None,
90            cache_dir: None,
91            normalize_types: false, // Preserve original types by default
92        }
93    }
94}
95
96impl EvalConfig {
97    /// Quick evaluation (limited examples, no breakdowns).
98    pub fn quick() -> Self {
99        Self {
100            max_examples_per_dataset: 100,
101            breakdown_by_difficulty: false,
102            breakdown_by_domain: false,
103            breakdown_by_type: true,
104            warmup: false,
105            warmup_iterations: 0,
106            min_confidence: None,
107            cache_dir: None,
108            normalize_types: false,
109        }
110    }
111
112    /// Full evaluation (all examples, all breakdowns).
113    pub fn full() -> Self {
114        Self {
115            max_examples_per_dataset: 0,
116            breakdown_by_difficulty: true,
117            breakdown_by_domain: true,
118            breakdown_by_type: true,
119            warmup: true,
120            warmup_iterations: 2,
121            min_confidence: None,
122            cache_dir: None,
123            normalize_types: true, // Full eval normalizes types
124        }
125    }
126
127    /// CI-aware configuration.
128    ///
129    /// Respects environment variables:
130    /// - `ANNO_MAX_EXAMPLES`: Max examples per dataset (default: 50 in CI, 200 otherwise; set to 0 for unlimited)
131    /// - `CI` or `GITHUB_ACTIONS`: Detects CI environment
132    ///
133    /// # Example
134    ///
135    /// ```bash
136    /// # Limit to 20 examples per dataset
137    /// ANNO_MAX_EXAMPLES=20 cargo test -p anno-eval --features eval
138    /// ```
139    pub fn ci_aware() -> Self {
140        let in_ci = std::env::var("CI").is_ok() || std::env::var("GITHUB_ACTIONS").is_ok();
141
142        // Sensible defaults:
143        // - CI: moderate (50 examples for stability while staying fast)
144        // - local: larger (100 examples for more stable results)
145        //
146        // Opt-out by setting `ANNO_MAX_EXAMPLES=0` (unlimited).
147        let default_max = if in_ci { 50 } else { 100 };
148        let max_examples = std::env::var("ANNO_MAX_EXAMPLES")
149            .ok()
150            .and_then(|s| s.parse().ok())
151            .unwrap_or(default_max);
152
153        Self {
154            max_examples_per_dataset: max_examples,
155            breakdown_by_difficulty: !in_ci,
156            breakdown_by_domain: !in_ci,
157            breakdown_by_type: true,
158            warmup: !in_ci,
159            warmup_iterations: if in_ci { 0 } else { 1 },
160            min_confidence: None,
161            cache_dir: None,
162            normalize_types: false,
163        }
164    }
165
166    /// With type normalization enabled.
167    ///
168    /// When enabled, domain-specific entity types are mapped to standard NER types:
169    /// - MIT Movie "Actor" -> "Person"
170    /// - MIT Restaurant "Restaurant_Name" -> "Organization"
171    /// - Biomedical "Disease" -> "Other"
172    #[must_use]
173    pub fn with_type_normalization(mut self) -> Self {
174        self.normalize_types = true;
175        self
176    }
177}
178
179// =============================================================================
180// Backend Registry
181// =============================================================================
182
183/// Registry for NER backends to evaluate.
184pub struct BackendRegistry {
185    backends: Vec<(String, String, Box<dyn Model>)>, // (name, description, model)
186}
187
188impl BackendRegistry {
189    /// Create empty registry.
190    #[must_use]
191    pub fn new() -> Self {
192        Self {
193            backends: Vec::new(),
194        }
195    }
196
197    /// Register a backend for evaluation.
198    pub fn register(
199        &mut self,
200        name: impl Into<String>,
201        description: impl Into<String>,
202        model: Box<dyn Model>,
203    ) {
204        self.backends.push((name.into(), description.into(), model));
205    }
206
207    /// Get number of registered backends.
208    pub fn len(&self) -> usize {
209        self.backends.len()
210    }
211
212    /// Check if registry is empty.
213    pub fn is_empty(&self) -> bool {
214        self.backends.is_empty()
215    }
216
217    /// Iterate over backends.
218    pub fn iter(&self) -> impl Iterator<Item = (&str, &str, &dyn Model)> {
219        self.backends
220            .iter()
221            .map(|(name, desc, model)| (name.as_str(), desc.as_str(), model.as_ref()))
222    }
223
224    /// Register default zero-dependency backends.
225    pub fn register_defaults(&mut self) {
226        use anno::{HeuristicNER, RegexNER, StackedNER};
227
228        self.register(
229            "RegexNER",
230            "Regex patterns (DATE/MONEY/EMAIL/etc.)",
231            Box::new(RegexNER::new()),
232        );
233        self.register(
234            "HeuristicNER",
235            "Heuristics (PER/ORG/LOC)",
236            Box::new(HeuristicNER::new()),
237        );
238        self.register(
239            "StackedNER",
240            "Pattern + Statistical combined",
241            Box::new(StackedNER::new()),
242        );
243    }
244
245    /// Register ONNX backends if feature enabled.
246    #[cfg(feature = "onnx")]
247    pub fn register_onnx(&mut self) {
248        use crate::{BertNEROnnx, GLiNEROnnx, DEFAULT_BERT_ONNX_MODEL, DEFAULT_GLINER_MODEL};
249
250        // GLiNER v2.1 (zero-shot NER)
251        match GLiNEROnnx::new(DEFAULT_GLINER_MODEL) {
252            Ok(gliner) => {
253                self.register("GLiNER", "Zero-shot NER via ONNX", Box::new(gliner));
254            }
255            Err(e) => {
256                log::warn!("Failed to load GLiNER ONNX: {}", e);
257            }
258        }
259
260        // BERT NER (fine-tuned)
261        match BertNEROnnx::new(DEFAULT_BERT_ONNX_MODEL) {
262            Ok(bert) => {
263                self.register("BertNEROnnx", "BERT NER via ONNX", Box::new(bert));
264            }
265            Err(e) => {
266                log::warn!("Failed to load BERT ONNX: {}", e);
267            }
268        }
269    }
270
271    /// Register Candle backends if feature enabled.
272    #[cfg(feature = "candle")]
273    pub fn register_candle(&mut self) {
274        use crate::{CandleNER, DEFAULT_CANDLE_MODEL};
275
276        match CandleNER::from_pretrained(DEFAULT_CANDLE_MODEL) {
277            Ok(candle) => {
278                self.register(
279                    "CandleNER",
280                    "Pure Rust BERT NER via Candle",
281                    Box::new(candle),
282                );
283            }
284            Err(e) => {
285                log::warn!("Failed to load Candle NER: {}", e);
286            }
287        }
288    }
289
290    /// Register GLiNER multi-task model.
291    ///
292    /// Loads `knowledgator/gliner-multitask-large-v0.5` (Stepanov & Shtopko 2024,
293    /// arXiv:2406.12925), which extends GLiNER v1 to multiple tasks via task-conditioned
294    /// label prompts:
295    /// - Zero-shot NER with arbitrary entity types
296    /// - Text classification (single/multi-label)
297    /// - Hierarchical structure extraction
298    ///
299    /// Note: this is NOT the fastino-ai GLiNER2 architecture (Zaratiana et al. 2025,
300    /// arXiv:2507.18546). `fastino/gliner2-*` models will not load through this backend
301    /// (different special tokens, different head structure). See issue #17.
302    ///
303    /// Requires either `onnx` or `candle` feature.
304    #[cfg(any(feature = "onnx", feature = "candle"))]
305    pub fn register_gliner_multitask(&mut self, model_id: &str) {
306        use crate::backends::gliner_multitask::GLiNERMultitask;
307
308        match GLiNERMultitask::from_pretrained(model_id) {
309            Ok(model) => {
310                self.register(
311                    "GLiNERMultitask",
312                    "Multi-task zero-shot NER, classification, structure",
313                    Box::new(model),
314                );
315            }
316            Err(e) => {
317                log::warn!("Failed to load GLiNER multi-task from {}: {}", model_id, e);
318            }
319        }
320    }
321
322    /// Register GLiNER multi-task with the default model
323    /// (`onnx-community/gliner-multitask-large-v0.5`).
324    #[cfg(any(feature = "onnx", feature = "candle"))]
325    pub fn register_gliner_multitask_default(&mut self) {
326        self.register_gliner_multitask(anno::DEFAULT_GLINER_MULTITASK_MODEL);
327    }
328
329    /// Register a custom stacked combination of backends.
330    ///
331    /// # Example
332    ///
333    /// ```rust,ignore
334    /// use anno_eval::eval::harness::BackendRegistry;
335    /// use anno::backends::stacked::ConflictStrategy;
336    ///
337    /// let mut registry = BackendRegistry::new();
338    /// registry.register_stack(
339    ///     "custom_stack",
340    ///     &["RegexNER", "HeuristicNER"],
341    ///     ConflictStrategy::HighestConf,
342    /// );
343    /// ```
344    pub fn register_stack(
345        &mut self,
346        name: impl Into<String>,
347        layer_names: &[&str],
348        strategy: anno::backends::stacked::ConflictStrategy,
349    ) {
350        use anno::backends::stacked::StackedNERBuilder;
351        use anno::{HeuristicNER, RegexNER};
352
353        let name = name.into();
354        let mut builder = StackedNERBuilder::default().strategy(strategy);
355
356        for layer_name in layer_names {
357            match *layer_name {
358                "RegexNER" | "pattern" => {
359                    builder = builder.layer(RegexNER::new());
360                }
361                "HeuristicNER" | "heuristic" => {
362                    builder = builder.layer(HeuristicNER::new());
363                }
364                _ => {
365                    eprintln!(
366                        "Warning: Unknown layer '{}' in stack '{}'",
367                        layer_name, name
368                    );
369                }
370            }
371        }
372
373        let description = format!("Stack: {} ({:?})", layer_names.join(" -> "), strategy);
374
375        self.register(name, description, Box::new(builder.build()));
376    }
377
378    /// Register all possible combinations of base backends.
379    ///
380    /// This creates:
381    /// - Individual backends (RegexNER, HeuristicNER)
382    /// - Two-layer stacks (Pattern->Heuristic, Heuristic->Pattern)
383    /// - Different conflict strategies
384    pub fn register_all_combinations(&mut self) {
385        use anno::backends::stacked::ConflictStrategy;
386
387        // Already registered as part of defaults
388        self.register_defaults();
389
390        // Additional ordering: HeuristicNER first, then RegexNER
391        self.register_stack(
392            "Heuristic->Pattern",
393            &["HeuristicNER", "RegexNER"],
394            ConflictStrategy::HighestConf,
395        );
396
397        // Different conflict strategies for default stack
398        self.register_stack(
399            "Stack_LongestSpan",
400            &["RegexNER", "HeuristicNER"],
401            ConflictStrategy::LongestSpan,
402        );
403        self.register_stack(
404            "Stack_Priority",
405            &["RegexNER", "HeuristicNER"],
406            ConflictStrategy::Priority,
407        );
408        self.register_stack(
409            "Stack_Union",
410            &["RegexNER", "HeuristicNER"],
411            ConflictStrategy::Union,
412        );
413    }
414}
415
416impl Default for BackendRegistry {
417    fn default() -> Self {
418        let mut registry = Self::new();
419        registry.register_defaults();
420
421        #[cfg(feature = "onnx")]
422        registry.register_onnx();
423
424        #[cfg(feature = "candle")]
425        registry.register_candle();
426
427        registry
428    }
429}
430
431// =============================================================================
432// Results Structures
433// =============================================================================
434
435/// Results for a single backend on a single dataset/split.
436#[derive(Debug, Clone, Serialize, Deserialize)]
437pub struct BackendDatasetResult {
438    /// Backend name
439    pub backend_name: String,
440    /// Dataset/split name
441    pub dataset_name: String,
442    /// Number of examples evaluated
443    pub num_examples: usize,
444    /// Number of gold entities
445    pub num_gold_entities: usize,
446    /// Precision
447    pub precision: f64,
448    /// Recall
449    pub recall: f64,
450    /// F1 score
451    pub f1: f64,
452    /// Macro F1 (average across entity types)
453    pub macro_f1: Option<f64>,
454    /// Entities found by model
455    pub found: usize,
456    /// Entities expected (gold)
457    pub expected: usize,
458    /// Per-entity-type metrics
459    pub per_type: HashMap<String, TypeMetrics>,
460    /// Evaluation duration
461    pub duration_ms: f64,
462    /// Tokens per second
463    pub tokens_per_second: f64,
464}
465
466/// Aggregate results across multiple datasets for a backend.
467#[derive(Debug, Clone, Serialize, Deserialize)]
468pub struct BackendAggregateResult {
469    /// Backend name
470    pub backend_name: String,
471    /// Backend description
472    pub description: String,
473    /// F1 with variance across datasets
474    pub f1: MetricWithVariance,
475    /// Precision with variance
476    pub precision: MetricWithVariance,
477    /// Recall with variance
478    pub recall: MetricWithVariance,
479    /// Total examples evaluated
480    pub total_examples: usize,
481    /// Total entities found
482    pub total_found: usize,
483    /// Total entities expected
484    pub total_expected: usize,
485    /// Total duration
486    pub total_duration_ms: f64,
487    /// Per-dataset results
488    pub per_dataset: Vec<BackendDatasetResult>,
489}
490
491/// Full evaluation results.
492#[derive(Debug, Clone, Serialize, Deserialize)]
493pub struct EvalResults {
494    /// Timestamp
495    pub timestamp: String,
496    /// Configuration used
497    pub config: EvalConfig,
498    /// Results per backend
499    pub backends: Vec<BackendAggregateResult>,
500    /// Breakdown by difficulty (if enabled)
501    pub by_difficulty: Option<HashMap<String, Vec<BackendDatasetResult>>>,
502    /// Breakdown by domain (if enabled)
503    pub by_domain: Option<HashMap<String, Vec<BackendDatasetResult>>>,
504    /// Dataset statistics
505    pub dataset_stats: DatasetStatsSummary,
506}
507
508/// Summary statistics about the evaluated datasets.
509#[derive(Debug, Clone, Serialize, Deserialize, Default)]
510pub struct DatasetStatsSummary {
511    /// Total examples
512    pub total_examples: usize,
513    /// Total gold entities
514    pub total_entities: usize,
515    /// Entity type distribution
516    pub entity_type_distribution: HashMap<String, usize>,
517    /// Domain distribution (for synthetic)
518    pub domain_distribution: HashMap<String, usize>,
519    /// Difficulty distribution (for synthetic)
520    pub difficulty_distribution: HashMap<String, usize>,
521}
522
523// =============================================================================
524// Evaluation Harness
525// =============================================================================
526
527/// Main evaluation harness.
528pub struct EvalHarness {
529    config: EvalConfig,
530    registry: BackendRegistry,
531    loader: Option<DatasetLoader>,
532}
533
534impl EvalHarness {
535    /// Create a new harness with given config.
536    pub fn new(config: EvalConfig) -> Result<Self> {
537        let loader = if let Some(ref dir) = config.cache_dir {
538            Some(DatasetLoader::with_cache_dir(dir)?)
539        } else {
540            DatasetLoader::new().ok()
541        };
542
543        Ok(Self {
544            config,
545            registry: BackendRegistry::new(),
546            loader,
547        })
548    }
549
550    /// Create with default config and default backends.
551    pub fn with_defaults() -> Result<Self> {
552        let mut harness = Self::new(EvalConfig::default())?;
553        harness.registry = BackendRegistry::default();
554        Ok(harness)
555    }
556
557    /// Create with custom config and default backends.
558    pub fn with_config(config: EvalConfig) -> Result<Self> {
559        let mut harness = Self::new(config)?;
560        harness.registry = BackendRegistry::default();
561        Ok(harness)
562    }
563
564    /// Register a backend.
565    pub fn register(
566        &mut self,
567        name: impl Into<String>,
568        description: impl Into<String>,
569        model: Box<dyn Model>,
570    ) {
571        self.registry.register(name, description, model);
572    }
573
574    /// Register all default backends.
575    pub fn register_defaults(&mut self) {
576        self.registry.register_defaults();
577    }
578
579    /// Get read-only reference to registry.
580    pub fn registry(&self) -> &BackendRegistry {
581        &self.registry
582    }
583
584    /// Get mutable reference to registry.
585    pub fn registry_mut(&mut self) -> &mut BackendRegistry {
586        &mut self.registry
587    }
588
589    /// Get number of registered backends.
590    pub fn backend_count(&self) -> usize {
591        self.registry.len()
592    }
593
594    /// Run evaluation on synthetic datasets.
595    pub fn run_synthetic(&self) -> Result<EvalResults> {
596        if self.registry.is_empty() {
597            return Err(Error::InvalidInput(
598                "No backends registered for evaluation".to_string(),
599            ));
600        }
601
602        let all_examples = all_datasets();
603        let test_cases: Vec<_> = all_examples
604            .iter()
605            .filter(|ex| !ex.text.is_empty())
606            .take(if self.config.max_examples_per_dataset > 0 {
607                self.config.max_examples_per_dataset
608            } else {
609                usize::MAX
610            })
611            .map(|ex| (ex.text.clone(), ex.entities.clone()))
612            .collect();
613
614        let dataset_stats = compute_dataset_stats(&all_examples);
615
616        let mut backends_results = Vec::new();
617
618        for (name, desc, model) in self.registry.iter() {
619            let result = self.evaluate_model_on_cases(model, name, "synthetic", &test_cases)?;
620            backends_results.push((name.to_string(), desc.to_string(), vec![result]));
621        }
622
623        // Aggregate results
624        let backends = backends_results
625            .into_iter()
626            .map(|(name, desc, results)| aggregate_backend_results(&name, &desc, results))
627            .collect();
628
629        // Breakdowns
630        let by_difficulty = if self.config.breakdown_by_difficulty {
631            Some(self.compute_difficulty_breakdown()?)
632        } else {
633            None
634        };
635
636        let by_domain = if self.config.breakdown_by_domain {
637            Some(self.compute_domain_breakdown()?)
638        } else {
639            None
640        };
641
642        Ok(EvalResults {
643            timestamp: chrono::Utc::now().to_rfc3339(),
644            config: self.config.clone(),
645            backends,
646            by_difficulty,
647            by_domain,
648            dataset_stats,
649        })
650    }
651
652    /// Run evaluation on real datasets (requires `eval` feature for downloading).
653    #[cfg(feature = "eval")]
654    pub fn run_real_datasets(&self, datasets: &[DatasetId]) -> Result<EvalResults> {
655        if self.registry.is_empty() {
656            return Err(Error::InvalidInput(
657                "No backends registered for evaluation".to_string(),
658            ));
659        }
660
661        let loader = self
662            .loader
663            .as_ref()
664            .ok_or_else(|| Error::InvalidInput("Dataset loader not initialized".to_string()))?;
665
666        let mut all_test_cases: Vec<(String, Vec<GoldEntity>)> = Vec::new();
667        let mut dataset_results: HashMap<String, Vec<(String, Vec<GoldEntity>)>> = HashMap::new();
668
669        for dataset_id in datasets {
670            let loadable = match crate::eval::LoadableDatasetId::try_from(*dataset_id) {
671                Ok(id) => id,
672                Err(e) => {
673                    log::warn!("Skipping {} (not loadable): {}", dataset_id.name(), e);
674                    continue;
675                }
676            };
677
678            match loader.load_or_download(loadable) {
679                Ok(loaded) => {
680                    let cases = loaded.to_test_cases();
681                    let limited: Vec<_> = if self.config.max_examples_per_dataset > 0 {
682                        cases
683                            .into_iter()
684                            .take(self.config.max_examples_per_dataset)
685                            .collect()
686                    } else {
687                        cases
688                    };
689                    dataset_results.insert(dataset_id.name().to_string(), limited.clone());
690                    all_test_cases.extend(limited);
691                }
692                Err(e) => {
693                    log::warn!("Failed to load {}: {}", dataset_id.name(), e);
694                }
695            }
696        }
697
698        if all_test_cases.is_empty() {
699            return Err(Error::InvalidInput("No datasets loaded".to_string()));
700        }
701
702        let mut backends_results = Vec::new();
703
704        for (name, desc, model) in self.registry.iter() {
705            let mut per_dataset_results = Vec::new();
706
707            for (dataset_name, cases) in &dataset_results {
708                let result = self.evaluate_model_on_cases(model, name, dataset_name, cases)?;
709                per_dataset_results.push(result);
710            }
711
712            backends_results.push((name.to_string(), desc.to_string(), per_dataset_results));
713        }
714
715        let backends = backends_results
716            .into_iter()
717            .map(|(name, desc, results)| aggregate_backend_results(&name, &desc, results))
718            .collect();
719
720        Ok(EvalResults {
721            timestamp: chrono::Utc::now().to_rfc3339(),
722            config: self.config.clone(),
723            backends,
724            by_difficulty: None,
725            by_domain: None,
726            dataset_stats: DatasetStatsSummary::default(),
727        })
728    }
729
730    /// Load cached dataset without downloading.
731    pub fn run_cached_datasets(&self, datasets: &[DatasetId]) -> Result<EvalResults> {
732        if self.registry.is_empty() {
733            return Err(Error::InvalidInput(
734                "No backends registered for evaluation".to_string(),
735            ));
736        }
737
738        let loader = self
739            .loader
740            .as_ref()
741            .ok_or_else(|| Error::InvalidInput("Dataset loader not initialized".to_string()))?;
742
743        let mut all_test_cases: Vec<(String, Vec<GoldEntity>)> = Vec::new();
744        let mut dataset_results: HashMap<String, Vec<(String, Vec<GoldEntity>)>> = HashMap::new();
745
746        for dataset_id in datasets {
747            let loadable = match crate::eval::LoadableDatasetId::try_from(*dataset_id) {
748                Ok(id) => id,
749                Err(_) => continue,
750            };
751
752            if loader.is_cached(loadable) {
753                match loader.load(loadable) {
754                    Ok(loaded) => {
755                        let cases = loaded.to_test_cases();
756                        let limited: Vec<_> = if self.config.max_examples_per_dataset > 0 {
757                            cases
758                                .into_iter()
759                                .take(self.config.max_examples_per_dataset)
760                                .collect()
761                        } else {
762                            cases
763                        };
764                        dataset_results.insert(dataset_id.name().to_string(), limited.clone());
765                        all_test_cases.extend(limited);
766                    }
767                    Err(e) => {
768                        log::warn!("Failed to load cached {}: {}", dataset_id.name(), e);
769                    }
770                }
771            }
772        }
773
774        if all_test_cases.is_empty() {
775            return Err(Error::InvalidInput(
776                "No cached datasets available".to_string(),
777            ));
778        }
779
780        let mut backends_results = Vec::new();
781
782        for (name, desc, model) in self.registry.iter() {
783            let mut per_dataset_results = Vec::new();
784
785            for (dataset_name, cases) in &dataset_results {
786                let result = self.evaluate_model_on_cases(model, name, dataset_name, cases)?;
787                per_dataset_results.push(result);
788            }
789
790            backends_results.push((name.to_string(), desc.to_string(), per_dataset_results));
791        }
792
793        let backends = backends_results
794            .into_iter()
795            .map(|(name, desc, results)| aggregate_backend_results(&name, &desc, results))
796            .collect();
797
798        Ok(EvalResults {
799            timestamp: chrono::Utc::now().to_rfc3339(),
800            config: self.config.clone(),
801            backends,
802            by_difficulty: None,
803            by_domain: None,
804            dataset_stats: DatasetStatsSummary::default(),
805        })
806    }
807
808    /// Evaluate a single backend on test cases.
809    fn evaluate_model_on_cases(
810        &self,
811        model: &dyn Model,
812        backend_name: &str,
813        dataset_name: &str,
814        test_cases: &[(String, Vec<GoldEntity>)],
815    ) -> Result<BackendDatasetResult> {
816        // Warmup
817        if self.config.warmup && !test_cases.is_empty() {
818            for _ in 0..self.config.warmup_iterations {
819                let _ = model.extract_entities(&test_cases[0].0, None);
820            }
821        }
822
823        let start = Instant::now();
824        let results = evaluate_ner_model(model, test_cases)?;
825        let duration = start.elapsed();
826
827        let total_gold: usize = test_cases.iter().map(|(_, gold)| gold.len()).sum();
828
829        Ok(BackendDatasetResult {
830            backend_name: backend_name.to_string(),
831            dataset_name: dataset_name.to_string(),
832            num_examples: test_cases.len(),
833            num_gold_entities: total_gold,
834            precision: results.precision,
835            recall: results.recall,
836            f1: results.f1,
837            macro_f1: results.macro_f1,
838            found: results.found,
839            expected: results.expected,
840            per_type: results.per_type,
841            duration_ms: duration.as_secs_f64() * 1000.0,
842            tokens_per_second: results.tokens_per_second,
843        })
844    }
845
846    /// Compute breakdown by difficulty.
847    fn compute_difficulty_breakdown(&self) -> Result<HashMap<String, Vec<BackendDatasetResult>>> {
848        let difficulties = [
849            Difficulty::Easy,
850            Difficulty::Medium,
851            Difficulty::Hard,
852            Difficulty::Adversarial,
853        ];
854
855        let mut breakdown = HashMap::new();
856
857        for difficulty in difficulties {
858            let subset: Vec<_> = datasets_by_difficulty(difficulty)
859                .into_iter()
860                .filter(|ex| !ex.text.is_empty())
861                .map(|ex| (ex.text, ex.entities))
862                .collect();
863
864            if subset.is_empty() {
865                continue;
866            }
867
868            let difficulty_name = format!("{:?}", difficulty);
869            let mut difficulty_results = Vec::new();
870
871            for (name, _desc, model) in self.registry.iter() {
872                let result =
873                    self.evaluate_model_on_cases(model, name, &difficulty_name, &subset)?;
874                difficulty_results.push(result);
875            }
876
877            breakdown.insert(difficulty_name, difficulty_results);
878        }
879
880        Ok(breakdown)
881    }
882
883    /// Compute breakdown by domain.
884    fn compute_domain_breakdown(&self) -> Result<HashMap<String, Vec<BackendDatasetResult>>> {
885        let domains = [
886            Domain::News,
887            Domain::Financial,
888            Domain::Technical,
889            Domain::Sports,
890            Domain::Entertainment,
891            Domain::Politics,
892            Domain::Ecommerce,
893            Domain::Travel,
894            Domain::Weather,
895            Domain::Academic,
896            Domain::Historical,
897            Domain::Food,
898            Domain::RealEstate,
899            Domain::Conversational,
900            Domain::SocialMedia,
901            Domain::Biomedical,
902            Domain::Legal,
903            Domain::Scientific,
904        ];
905
906        let mut breakdown = HashMap::new();
907
908        for domain in domains {
909            let subset: Vec<_> = datasets_by_domain(domain)
910                .into_iter()
911                .filter(|ex| !ex.text.is_empty())
912                .map(|ex| (ex.text, ex.entities))
913                .collect();
914
915            if subset.is_empty() {
916                continue;
917            }
918
919            let domain_name = format!("{:?}", domain);
920            let mut domain_results = Vec::new();
921
922            for (name, _desc, model) in self.registry.iter() {
923                let result = self.evaluate_model_on_cases(model, name, &domain_name, &subset)?;
924                domain_results.push(result);
925            }
926
927            breakdown.insert(domain_name, domain_results);
928        }
929
930        Ok(breakdown)
931    }
932}
933
934// =============================================================================
935// HTML Report Generation
936// =============================================================================
937
938impl EvalResults {
939    /// Generate HTML report.
940    pub fn to_html(&self) -> String {
941        let mut html = String::new();
942
943        html.push_str(HTML_HEAD);
944        html.push_str("<body>\n");
945        html.push_str("<div class=\"container\">\n");
946
947        // Header
948        html.push_str("<h1>NER Evaluation Report</h1>\n");
949        html.push_str(&format!(
950            "<p class=\"timestamp\">Generated: {}</p>\n",
951            self.timestamp
952        ));
953
954        // Dataset stats
955        html.push_str("<h2>Dataset Summary</h2>\n");
956        html.push_str("<div class=\"stats-grid\">\n");
957        html.push_str(&format!(
958            "<div class=\"stat-box\"><span class=\"stat-value\">{}</span><span class=\"stat-label\">Examples</span></div>\n",
959            self.dataset_stats.total_examples
960        ));
961        html.push_str(&format!(
962            "<div class=\"stat-box\"><span class=\"stat-value\">{}</span><span class=\"stat-label\">Entities</span></div>\n",
963            self.dataset_stats.total_entities
964        ));
965        html.push_str(&format!(
966            "<div class=\"stat-box\"><span class=\"stat-value\">{}</span><span class=\"stat-label\">Backends</span></div>\n",
967            self.backends.len()
968        ));
969        html.push_str("</div>\n");
970
971        // Overall results table
972        html.push_str("<h2>Overall Results</h2>\n");
973        html.push_str("<table>\n");
974        html.push_str("<thead><tr><th>Backend</th><th>F1</th><th>Precision</th><th>Recall</th><th>Found/Expected</th><th>Time</th></tr></thead>\n");
975        html.push_str("<tbody>\n");
976
977        for backend in &self.backends {
978            let f1_class = if backend.f1.mean > 0.8 {
979                "good"
980            } else if backend.f1.mean > 0.5 {
981                "ok"
982            } else {
983                "poor"
984            };
985
986            html.push_str(&format!(
987                "<tr><td><strong>{}</strong><br><small>{}</small></td>\
988                 <td class=\"{}\"><strong>{:.1}%</strong><br><small>{}</small></td>\
989                 <td>{:.1}%</td>\
990                 <td>{:.1}%</td>\
991                 <td>{} / {}</td>\
992                 <td>{:.1}ms</td></tr>\n",
993                backend.backend_name,
994                backend.description,
995                f1_class,
996                backend.f1.mean * 100.0,
997                backend.f1.format_with_ci(),
998                backend.precision.mean * 100.0,
999                backend.recall.mean * 100.0,
1000                backend.total_found,
1001                backend.total_expected,
1002                backend.total_duration_ms,
1003            ));
1004        }
1005        html.push_str("</tbody></table>\n");
1006
1007        // Difficulty breakdown
1008        if let Some(ref by_diff) = self.by_difficulty {
1009            html.push_str("<h2>Results by Difficulty</h2>\n");
1010            html.push_str(&self.render_breakdown_table(by_diff));
1011        }
1012
1013        // Domain breakdown
1014        if let Some(ref by_dom) = self.by_domain {
1015            html.push_str("<h2>Results by Domain</h2>\n");
1016            html.push_str(&self.render_breakdown_table(by_dom));
1017        }
1018
1019        // Per-type metrics for best backend
1020        if let Some(best) = self.backends.iter().max_by(|a, b| {
1021            a.f1.mean
1022                .partial_cmp(&b.f1.mean)
1023                .unwrap_or(std::cmp::Ordering::Equal)
1024        }) {
1025            if !best.per_dataset.is_empty() {
1026                html.push_str(&format!(
1027                    "<h2>Per-Type Metrics ({})</h2>\n",
1028                    best.backend_name
1029                ));
1030
1031                // Aggregate per-type from all datasets
1032                let mut type_metrics: HashMap<String, Vec<&TypeMetrics>> = HashMap::new();
1033                for ds_result in &best.per_dataset {
1034                    for (type_name, metrics) in &ds_result.per_type {
1035                        type_metrics
1036                            .entry(type_name.clone())
1037                            .or_default()
1038                            .push(metrics);
1039                    }
1040                }
1041
1042                html.push_str("<table>\n");
1043                html.push_str("<thead><tr><th>Type</th><th>F1</th><th>Precision</th><th>Recall</th><th>Correct/Expected</th></tr></thead>\n");
1044                html.push_str("<tbody>\n");
1045
1046                let mut sorted_types: Vec<_> = type_metrics.iter().collect();
1047                sorted_types.sort_by(|a, b| {
1048                    let avg_f1_a = a.1.iter().map(|m| m.f1).sum::<f64>() / a.1.len() as f64;
1049                    let avg_f1_b = b.1.iter().map(|m| m.f1).sum::<f64>() / b.1.len() as f64;
1050                    avg_f1_b
1051                        .partial_cmp(&avg_f1_a)
1052                        .unwrap_or(std::cmp::Ordering::Equal)
1053                });
1054
1055                for (type_name, metrics_list) in sorted_types {
1056                    let avg_f1 =
1057                        metrics_list.iter().map(|m| m.f1).sum::<f64>() / metrics_list.len() as f64;
1058                    let avg_p = metrics_list.iter().map(|m| m.precision).sum::<f64>()
1059                        / metrics_list.len() as f64;
1060                    let avg_r = metrics_list.iter().map(|m| m.recall).sum::<f64>()
1061                        / metrics_list.len() as f64;
1062                    let total_correct: usize = metrics_list.iter().map(|m| m.correct).sum();
1063                    let total_expected: usize = metrics_list.iter().map(|m| m.expected).sum();
1064
1065                    let f1_class = if avg_f1 > 0.8 {
1066                        "good"
1067                    } else if avg_f1 > 0.5 {
1068                        "ok"
1069                    } else {
1070                        "poor"
1071                    };
1072
1073                    html.push_str(&format!(
1074                        "<tr><td>{}</td><td class=\"{}\">{:.1}%</td><td>{:.1}%</td><td>{:.1}%</td><td>{}/{}</td></tr>\n",
1075                        type_name,
1076                        f1_class,
1077                        avg_f1 * 100.0,
1078                        avg_p * 100.0,
1079                        avg_r * 100.0,
1080                        total_correct,
1081                        total_expected,
1082                    ));
1083                }
1084                html.push_str("</tbody></table>\n");
1085            }
1086        }
1087
1088        // Entity type distribution
1089        if !self.dataset_stats.entity_type_distribution.is_empty() {
1090            html.push_str("<h2>Entity Type Distribution</h2>\n");
1091            html.push_str("<table>\n");
1092            html.push_str("<thead><tr><th>Type</th><th>Count</th><th>Percent</th></tr></thead>\n");
1093            html.push_str("<tbody>\n");
1094
1095            let total: usize = self.dataset_stats.entity_type_distribution.values().sum();
1096            let mut sorted: Vec<_> = self.dataset_stats.entity_type_distribution.iter().collect();
1097            sorted.sort_by(|a, b| b.1.cmp(a.1));
1098
1099            for (type_name, count) in sorted {
1100                let pct = (*count as f64 / total as f64) * 100.0;
1101                html.push_str(&format!(
1102                    "<tr><td>{}</td><td>{}</td><td>{:.1}%</td></tr>\n",
1103                    type_name, count, pct
1104                ));
1105            }
1106            html.push_str("</tbody></table>\n");
1107        }
1108
1109        html.push_str("</div>\n</body></html>\n");
1110        html
1111    }
1112
1113    /// Render a breakdown table.
1114    fn render_breakdown_table(
1115        &self,
1116        breakdown: &HashMap<String, Vec<BackendDatasetResult>>,
1117    ) -> String {
1118        let mut html = String::new();
1119        html.push_str("<table>\n");
1120
1121        // Get backend names from first entry
1122        let backend_names: Vec<_> = breakdown
1123            .values()
1124            .next()
1125            .map(|results| results.iter().map(|r| r.backend_name.as_str()).collect())
1126            .unwrap_or_default();
1127
1128        // Header
1129        html.push_str("<thead><tr><th>Category</th>");
1130        for name in &backend_names {
1131            html.push_str(&format!("<th>{}</th>", name));
1132        }
1133        html.push_str("</tr></thead>\n");
1134
1135        // Body
1136        html.push_str("<tbody>\n");
1137        let mut sorted_keys: Vec<_> = breakdown.keys().collect();
1138        sorted_keys.sort();
1139
1140        for key in sorted_keys {
1141            let results = &breakdown[key];
1142            html.push_str(&format!("<tr><td>{}</td>", key));
1143
1144            for backend_name in &backend_names {
1145                if let Some(result) = results.iter().find(|r| r.backend_name == *backend_name) {
1146                    let f1_class = if result.f1 > 0.8 {
1147                        "good"
1148                    } else if result.f1 > 0.5 {
1149                        "ok"
1150                    } else {
1151                        "poor"
1152                    };
1153                    html.push_str(&format!(
1154                        "<td class=\"{}\">{:.1}%</td>",
1155                        f1_class,
1156                        result.f1 * 100.0
1157                    ));
1158                } else {
1159                    html.push_str("<td>-</td>");
1160                }
1161            }
1162            html.push_str("</tr>\n");
1163        }
1164        html.push_str("</tbody></table>\n");
1165
1166        html
1167    }
1168}
1169
1170// =============================================================================
1171// Helpers
1172// =============================================================================
1173
1174/// Compute dataset statistics from examples.
1175fn compute_dataset_stats(examples: &[AnnotatedExample]) -> DatasetStatsSummary {
1176    let mut entity_type_dist: HashMap<String, usize> = HashMap::new();
1177    let mut domain_dist: HashMap<String, usize> = HashMap::new();
1178    let mut difficulty_dist: HashMap<String, usize> = HashMap::new();
1179
1180    let mut total_entities = 0;
1181
1182    for ex in examples {
1183        *domain_dist.entry(format!("{:?}", ex.domain)).or_insert(0) += 1;
1184        *difficulty_dist
1185            .entry(format!("{:?}", ex.difficulty))
1186            .or_insert(0) += 1;
1187
1188        for entity in &ex.entities {
1189            let type_name = format!("{:?}", entity.entity_type);
1190            *entity_type_dist.entry(type_name).or_insert(0) += 1;
1191            total_entities += 1;
1192        }
1193    }
1194
1195    DatasetStatsSummary {
1196        total_examples: examples.len(),
1197        total_entities,
1198        entity_type_distribution: entity_type_dist,
1199        domain_distribution: domain_dist,
1200        difficulty_distribution: difficulty_dist,
1201    }
1202}
1203
1204/// Aggregate results from multiple datasets into a single backend result.
1205fn aggregate_backend_results(
1206    name: &str,
1207    desc: &str,
1208    results: Vec<BackendDatasetResult>,
1209) -> BackendAggregateResult {
1210    if results.is_empty() {
1211        return BackendAggregateResult {
1212            backend_name: name.to_string(),
1213            description: desc.to_string(),
1214            f1: MetricWithVariance::default(),
1215            precision: MetricWithVariance::default(),
1216            recall: MetricWithVariance::default(),
1217            total_examples: 0,
1218            total_found: 0,
1219            total_expected: 0,
1220            total_duration_ms: 0.0,
1221            per_dataset: vec![],
1222        };
1223    }
1224
1225    let f1s: Vec<f64> = results.iter().map(|r| r.f1).collect();
1226    let precisions: Vec<f64> = results.iter().map(|r| r.precision).collect();
1227    let recalls: Vec<f64> = results.iter().map(|r| r.recall).collect();
1228
1229    BackendAggregateResult {
1230        backend_name: name.to_string(),
1231        description: desc.to_string(),
1232        f1: MetricWithVariance::from_samples(&f1s),
1233        precision: MetricWithVariance::from_samples(&precisions),
1234        recall: MetricWithVariance::from_samples(&recalls),
1235        total_examples: results.iter().map(|r| r.num_examples).sum(),
1236        total_found: results.iter().map(|r| r.found).sum(),
1237        total_expected: results.iter().map(|r| r.expected).sum(),
1238        total_duration_ms: results.iter().map(|r| r.duration_ms).sum(),
1239        per_dataset: results,
1240    }
1241}
1242
1243// =============================================================================
1244// HTML Template
1245// =============================================================================
1246
1247const HTML_HEAD: &str = r#"<!DOCTYPE html>
1248<html lang="en">
1249<head>
1250<meta charset="UTF-8">
1251<meta name="viewport" content="width=device-width, initial-scale=1.0">
1252<title>NER Evaluation Report</title>
1253<style>
1254:root {
1255    --bg: #0d1117;
1256    --fg: #c9d1d9;
1257    --accent: #58a6ff;
1258    --good: #3fb950;
1259    --ok: #d29922;
1260    --poor: #f85149;
1261    --border: #30363d;
1262    --surface: #161b22;
1263}
1264* { box-sizing: border-box; }
1265body {
1266    font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, Oxygen, Ubuntu, sans-serif;
1267    background: var(--bg);
1268    color: var(--fg);
1269    margin: 0;
1270    padding: 0;
1271    line-height: 1.6;
1272}
1273.container { max-width: 1200px; margin: 0 auto; padding: 2rem; }
1274h1 { color: var(--accent); border-bottom: 2px solid var(--border); padding-bottom: 0.5rem; }
1275h2 { color: var(--fg); margin-top: 2rem; }
1276.timestamp { color: #8b949e; font-size: 0.9rem; }
1277.stats-grid { display: grid; grid-template-columns: repeat(auto-fit, minmax(150px, 1fr)); gap: 1rem; margin: 1rem 0; }
1278.stat-box {
1279    background: var(--surface);
1280    border: 1px solid var(--border);
1281    border-radius: 8px;
1282    padding: 1rem;
1283    text-align: center;
1284}
1285.stat-value { display: block; font-size: 2rem; font-weight: bold; color: var(--accent); }
1286.stat-label { display: block; font-size: 0.8rem; color: #8b949e; text-transform: uppercase; }
1287table {
1288    width: 100%;
1289    border-collapse: collapse;
1290    background: var(--surface);
1291    border-radius: 8px;
1292    overflow: hidden;
1293    margin: 1rem 0;
1294}
1295th, td { padding: 0.75rem 1rem; text-align: left; border-bottom: 1px solid var(--border); }
1296th { background: #21262d; color: var(--fg); font-weight: 600; }
1297tr:hover { background: #21262d; }
1298.good { color: var(--good); }
1299.ok { color: var(--ok); }
1300.poor { color: var(--poor); }
1301small { color: #8b949e; }
1302</style>
1303</head>
1304"#;
1305
1306// =============================================================================
1307// Tests
1308// =============================================================================
1309
1310#[cfg(test)]
1311mod tests {
1312    use super::*;
1313
1314    #[test]
1315    fn test_eval_config_default() {
1316        let config = EvalConfig::default();
1317        assert!(config.breakdown_by_difficulty);
1318        assert!(config.breakdown_by_domain);
1319    }
1320
1321    #[test]
1322    fn test_backend_registry() {
1323        let mut registry = BackendRegistry::new();
1324        assert!(registry.is_empty());
1325
1326        registry.register_defaults();
1327        assert!(!registry.is_empty());
1328        // Number of backends depends on feature flags (onnx, candle, etc.)
1329        assert!(
1330            !registry.is_empty(),
1331            "Expected at least 1 backend, got {}",
1332            registry.len()
1333        );
1334    }
1335
1336    #[test]
1337    fn test_synthetic_eval() {
1338        // Keep this test fast and deterministic: avoid `with_defaults()` which may
1339        // initialize optional ML backends depending on feature flags.
1340        let mut harness = EvalHarness::new(EvalConfig {
1341            max_examples_per_dataset: 5,
1342            breakdown_by_difficulty: false,
1343            breakdown_by_domain: false,
1344            breakdown_by_type: true,
1345            warmup: false,
1346            warmup_iterations: 0,
1347            min_confidence: None,
1348            cache_dir: None,
1349            normalize_types: false,
1350        })
1351        .unwrap();
1352        harness.registry.register_defaults();
1353        let results = harness.run_synthetic();
1354        assert!(results.is_ok());
1355
1356        let results = results.unwrap();
1357        assert!(!results.backends.is_empty());
1358    }
1359
1360    #[test]
1361    fn test_html_generation() {
1362        // Same rationale as test_synthetic_eval(): keep this fast and local-only.
1363        let mut harness = EvalHarness::new(EvalConfig {
1364            max_examples_per_dataset: 5,
1365            breakdown_by_difficulty: false,
1366            breakdown_by_domain: false,
1367            breakdown_by_type: true,
1368            warmup: false,
1369            warmup_iterations: 0,
1370            min_confidence: None,
1371            cache_dir: None,
1372            normalize_types: false,
1373        })
1374        .unwrap();
1375        harness.registry.register_defaults();
1376        let results = harness.run_synthetic().unwrap();
1377        let html = results.to_html();
1378
1379        assert!(html.contains("<html"));
1380        assert!(html.contains("NER Evaluation Report"));
1381        assert!(html.contains("RegexNER"));
1382    }
1383}
anno_eval/eval/harness.rs

anno_eval/eval/
harness.rs
