anno_eval/eval/dataset/
mod.rs

1//! Dataset API for NER evaluation.
2//!
3//! This module provides a unified interface for working with NER datasets,
4//! whether loaded from files or synthetically generated.
5//!
6//! # Design Philosophy
7//!
8//! Following Rust idioms (burntsushi, shepmaster patterns):
9//! - Use concrete types over trait objects where possible
10//! - Implement standard traits (`IntoIterator`, `Index`, etc.)
11//! - Keep the API surface small and predictable
12//! - Prefer composition over inheritance
13//!
14//! # Quick Start
15//!
16//! ```rust
17//! use anno_eval::eval::dataset::{NERDataset, Domain, Difficulty};
18//!
19//! // Create from synthetic data
20//! let dataset = NERDataset::synthetic();
21//! println!("Total examples: {}", dataset.len());
22//!
23//! // Filter by domain
24//! let news = dataset.filter_domain(Domain::News);
25//! println!("News examples: {}", news.len());
26//!
27//! // Filter by difficulty
28//! let hard = dataset.filter_difficulty(Difficulty::Hard);
29//!
30//! // Iterate
31//! for example in &dataset {
32//!     println!("{}: {} entities", example.text.len(), example.entity_count());
33//! }
34//! ```
35//!
36//! # Module Structure
37//!
38//! - `types`: Core types (`AnnotatedExample`, `Domain`, `Difficulty`)
39//! - `synthetic`: Synthetic dataset generation by domain
40//! - Main module: `NERDataset` struct and operations
41
42pub mod synthetic;
43pub mod types;
44
45pub use types::{AnnotatedExample, Difficulty, Domain};
46
47use crate::eval::GoldEntity;
48use anno::Result;
49use serde::{Deserialize, Serialize};
50use std::collections::HashMap;
51use std::ops::Index;
52use std::path::Path;
53
54// ============================================================================
55// NERDataset - Core Dataset Container
56// ============================================================================
57
58/// A collection of annotated NER examples with metadata and filtering.
59///
60/// `NERDataset` is the primary type for working with NER evaluation data.
61/// It wraps a `Vec<AnnotatedExample>` with convenient methods for filtering,
62/// statistics, and conversion to evaluation formats.
63///
64/// # Example
65///
66/// ```rust
67/// use anno_eval::eval::dataset::{NERDataset, Domain};
68///
69/// // Load synthetic data
70/// let mut dataset = NERDataset::synthetic();
71///
72/// // Filter to specific domain
73/// let biomedical = dataset.filter_domain(Domain::Biomedical);
74///
75/// // Get statistics
76/// let stats = biomedical.stats();
77/// println!("Examples: {}, Entities: {}", stats.total_examples, stats.total_entities);
78///
79/// // Convert to test cases for evaluation
80/// let test_cases = biomedical.to_test_cases();
81/// ```
82#[derive(Debug, Clone, Serialize, Deserialize)]
83pub struct NERDataset {
84    /// The underlying examples.
85    examples: Vec<AnnotatedExample>,
86    /// Dataset name/identifier.
87    name: String,
88    /// Optional source information (file path, URL, etc.).
89    source: Option<String>,
90}
91
92impl NERDataset {
93    // ========================================================================
94    // Constructors
95    // ========================================================================
96
97    /// Create an empty dataset with a name.
98    pub fn new(name: impl Into<String>) -> Self {
99        Self {
100            examples: Vec::new(),
101            name: name.into(),
102            source: None,
103        }
104    }
105
106    /// Create a dataset from a vector of examples.
107    pub fn from_examples(name: impl Into<String>, examples: Vec<AnnotatedExample>) -> Self {
108        Self {
109            examples,
110            name: name.into(),
111            source: None,
112        }
113    }
114
115    /// Create a dataset with source information.
116    pub fn with_source(mut self, source: impl Into<String>) -> Self {
117        self.source = Some(source.into());
118        self
119    }
120
121    /// Load the full synthetic dataset (all domains, all difficulties).
122    ///
123    /// This is the primary constructor for testing and development.
124    /// The synthetic data covers many entity types and edge cases.
125    ///
126    /// # Example
127    ///
128    /// ```rust
129    /// use anno_eval::eval::dataset::NERDataset;
130    ///
131    /// let dataset = NERDataset::synthetic();
132    /// assert!(!dataset.is_empty());
133    /// ```
134    pub fn synthetic() -> Self {
135        Self {
136            examples: synthetic::all_datasets(),
137            name: "synthetic".to_string(),
138            source: Some("anno_eval::eval::dataset::synthetic".to_string()),
139        }
140    }
141
142    /// Load synthetic data for a specific domain.
143    pub fn synthetic_domain(domain: Domain) -> Self {
144        Self {
145            examples: synthetic::by_domain(domain),
146            name: format!("synthetic_{:?}", domain).to_lowercase(),
147            source: Some("anno_eval::eval::dataset::synthetic".to_string()),
148        }
149    }
150
151    /// Load a dataset from a JSON file.
152    ///
153    /// Supports both JSON arrays and JSONL (one object per line).
154    ///
155    /// # Format
156    ///
157    /// ```json
158    /// [
159    ///   {
160    ///     "text": "John works at Google.",
161    ///     "entities": [
162    ///       {"text": "John", "label": "PER", "start": 0, "end": 4},
163    ///       {"text": "Google", "label": "ORG", "start": 14, "end": 20}
164    ///     ]
165    ///   }
166    /// ]
167    /// ```
168    pub fn from_json<P: AsRef<Path>>(path: P) -> Result<Self> {
169        let path = path.as_ref();
170        let name = path
171            .file_stem()
172            .and_then(|s| s.to_str())
173            .unwrap_or("unknown")
174            .to_string();
175
176        let test_cases = crate::eval::datasets::load_json_ner_dataset(path)?;
177        let examples = test_cases
178            .into_iter()
179            .map(|(text, entities)| AnnotatedExample::new(text, entities))
180            .collect();
181
182        Ok(Self {
183            examples,
184            name,
185            source: Some(path.display().to_string()),
186        })
187    }
188
189    /// Load a dataset from CoNLL-2003 format.
190    ///
191    /// # Format
192    ///
193    /// Each line: `word POS chunk NER-tag`
194    /// Empty lines separate sentences.
195    pub fn from_conll<P: AsRef<Path>>(path: P) -> Result<Self> {
196        let path = path.as_ref();
197        let name = path
198            .file_stem()
199            .and_then(|s| s.to_str())
200            .unwrap_or("unknown")
201            .to_string();
202
203        let test_cases = crate::eval::load_conll2003(path)?;
204        let examples = test_cases
205            .into_iter()
206            .map(|(text, entities)| AnnotatedExample::new(text, entities).with_domain(Domain::News))
207            .collect();
208
209        Ok(Self {
210            examples,
211            name,
212            source: Some(path.display().to_string()),
213        })
214    }
215
216    /// Auto-detect format and load from file.
217    ///
218    /// Tries CoNLL first, then JSON/JSONL.
219    pub fn from_file<P: AsRef<Path>>(path: P) -> Result<Self> {
220        let path = path.as_ref();
221        let name = path
222            .file_stem()
223            .and_then(|s| s.to_str())
224            .unwrap_or("unknown")
225            .to_string();
226
227        let test_cases = crate::eval::datasets::load_ner_dataset(path)?;
228        let examples = test_cases
229            .into_iter()
230            .map(|(text, entities)| AnnotatedExample::new(text, entities))
231            .collect();
232
233        Ok(Self {
234            examples,
235            name,
236            source: Some(path.display().to_string()),
237        })
238    }
239
240    // ========================================================================
241    // Accessors
242    // ========================================================================
243
244    /// Returns the dataset name.
245    pub fn name(&self) -> &str {
246        &self.name
247    }
248
249    /// Returns the source (file path, URL, etc.) if known.
250    pub fn source(&self) -> Option<&str> {
251        self.source.as_deref()
252    }
253
254    /// Returns the number of examples.
255    pub fn len(&self) -> usize {
256        self.examples.len()
257    }
258
259    /// Returns true if the dataset is empty.
260    pub fn is_empty(&self) -> bool {
261        self.examples.is_empty()
262    }
263
264    /// Get an example by index.
265    pub fn get(&self, index: usize) -> Option<&AnnotatedExample> {
266        self.examples.get(index)
267    }
268
269    /// Returns a slice of all examples.
270    pub fn as_slice(&self) -> &[AnnotatedExample] {
271        &self.examples
272    }
273
274    /// Returns a mutable slice of all examples.
275    pub fn as_mut_slice(&mut self) -> &mut [AnnotatedExample] {
276        &mut self.examples
277    }
278
279    /// Returns an iterator over the examples.
280    pub fn iter(&self) -> impl Iterator<Item = &AnnotatedExample> {
281        self.examples.iter()
282    }
283
284    /// Returns a mutable iterator over the examples.
285    pub fn iter_mut(&mut self) -> impl Iterator<Item = &mut AnnotatedExample> {
286        self.examples.iter_mut()
287    }
288
289    // ========================================================================
290    // Filtering
291    // ========================================================================
292
293    /// Filter to a specific domain, returning a new dataset.
294    pub fn filter_domain(&self, domain: Domain) -> Self {
295        let examples = self
296            .examples
297            .iter()
298            .filter(|ex| ex.domain == domain)
299            .cloned()
300            .collect();
301
302        Self {
303            examples,
304            name: format!("{}_{:?}", self.name, domain).to_lowercase(),
305            source: self.source.clone(),
306        }
307    }
308
309    /// Filter to a specific difficulty, returning a new dataset.
310    pub fn filter_difficulty(&self, difficulty: Difficulty) -> Self {
311        let examples = self
312            .examples
313            .iter()
314            .filter(|ex| ex.difficulty == difficulty)
315            .cloned()
316            .collect();
317
318        Self {
319            examples,
320            name: format!("{}_{:?}", self.name, difficulty).to_lowercase(),
321            source: self.source.clone(),
322        }
323    }
324
325    /// Filter to examples containing a specific entity type.
326    pub fn filter_entity_type(&self, entity_type: &anno::EntityType) -> Self {
327        let examples = self
328            .examples
329            .iter()
330            .filter(|ex| ex.entities.iter().any(|e| &e.entity_type == entity_type))
331            .cloned()
332            .collect();
333
334        Self {
335            examples,
336            name: format!("{}_filtered", self.name),
337            source: self.source.clone(),
338        }
339    }
340
341    /// Filter using a custom predicate.
342    pub fn filter<F>(&self, predicate: F) -> Self
343    where
344        F: Fn(&AnnotatedExample) -> bool,
345    {
346        let examples = self
347            .examples
348            .iter()
349            .filter(|ex| predicate(ex))
350            .cloned()
351            .collect();
352
353        Self {
354            examples,
355            name: format!("{}_filtered", self.name),
356            source: self.source.clone(),
357        }
358    }
359
360    /// Take the first n examples.
361    pub fn take(&self, n: usize) -> Self {
362        Self {
363            examples: self.examples.iter().take(n).cloned().collect(),
364            name: format!("{}_head_{}", self.name, n),
365            source: self.source.clone(),
366        }
367    }
368
369    /// Skip the first n examples.
370    pub fn skip(&self, n: usize) -> Self {
371        Self {
372            examples: self.examples.iter().skip(n).cloned().collect(),
373            name: format!("{}_tail", self.name),
374            source: self.source.clone(),
375        }
376    }
377
378    // ========================================================================
379    // Mutation
380    // ========================================================================
381
382    /// Add an example to the dataset.
383    pub fn push(&mut self, example: AnnotatedExample) {
384        self.examples.push(example);
385    }
386
387    /// Extend with examples from another source.
388    pub fn extend<I: IntoIterator<Item = AnnotatedExample>>(&mut self, iter: I) {
389        self.examples.extend(iter);
390    }
391
392    /// Merge another dataset into this one.
393    pub fn merge(&mut self, other: Self) {
394        self.examples.extend(other.examples);
395    }
396
397    // ========================================================================
398    // Conversion
399    // ========================================================================
400
401    /// Convert to test cases for evaluation functions.
402    ///
403    /// This is the bridge to `evaluate_ner_model()` and similar functions.
404    pub fn to_test_cases(&self) -> Vec<(String, Vec<GoldEntity>)> {
405        self.examples
406            .iter()
407            .map(|ex| (ex.text.clone(), ex.entities.clone()))
408            .collect()
409    }
410
411    /// Consume and convert to test cases.
412    pub fn into_test_cases(self) -> Vec<(String, Vec<GoldEntity>)> {
413        self.examples
414            .into_iter()
415            .map(|ex| ex.into_test_case())
416            .collect()
417    }
418
419    /// Convert to owned examples vec.
420    pub fn into_inner(self) -> Vec<AnnotatedExample> {
421        self.examples
422    }
423
424    // ========================================================================
425    // Statistics
426    // ========================================================================
427
428    /// Compute dataset statistics.
429    pub fn stats(&self) -> DatasetStats {
430        let total_examples = self.examples.len();
431        let total_entities: usize = self.examples.iter().map(|ex| ex.entities.len()).sum();
432
433        let mut domains = HashMap::new();
434        let mut difficulties = HashMap::new();
435        let mut entity_types = HashMap::new();
436
437        for ex in &self.examples {
438            *domains.entry(ex.domain).or_insert(0) += 1;
439            *difficulties.entry(ex.difficulty).or_insert(0) += 1;
440
441            for entity in &ex.entities {
442                let type_str = crate::eval::entity_type_to_string(&entity.entity_type);
443                *entity_types.entry(type_str).or_insert(0) += 1;
444            }
445        }
446
447        DatasetStats {
448            total_examples,
449            total_entities,
450            avg_entities_per_example: if total_examples > 0 {
451                total_entities as f64 / total_examples as f64
452            } else {
453                0.0
454            },
455            domains,
456            difficulties,
457            entity_types,
458        }
459    }
460}
461
462// ============================================================================
463// Standard Trait Implementations
464// ============================================================================
465
466impl Default for NERDataset {
467    fn default() -> Self {
468        Self::new("default")
469    }
470}
471
472impl Index<usize> for NERDataset {
473    type Output = AnnotatedExample;
474
475    fn index(&self, index: usize) -> &Self::Output {
476        &self.examples[index]
477    }
478}
479
480impl<'a> IntoIterator for &'a NERDataset {
481    type Item = &'a AnnotatedExample;
482    type IntoIter = std::slice::Iter<'a, AnnotatedExample>;
483
484    fn into_iter(self) -> Self::IntoIter {
485        self.examples.iter()
486    }
487}
488
489impl IntoIterator for NERDataset {
490    type Item = AnnotatedExample;
491    type IntoIter = std::vec::IntoIter<AnnotatedExample>;
492
493    fn into_iter(self) -> Self::IntoIter {
494        self.examples.into_iter()
495    }
496}
497
498impl FromIterator<AnnotatedExample> for NERDataset {
499    fn from_iter<I: IntoIterator<Item = AnnotatedExample>>(iter: I) -> Self {
500        Self {
501            examples: iter.into_iter().collect(),
502            name: "collected".to_string(),
503            source: None,
504        }
505    }
506}
507
508impl Extend<AnnotatedExample> for NERDataset {
509    fn extend<I: IntoIterator<Item = AnnotatedExample>>(&mut self, iter: I) {
510        self.examples.extend(iter);
511    }
512}
513
514// ============================================================================
515// Dataset Statistics
516// ============================================================================
517
518/// Statistics about a dataset.
519#[derive(Debug, Clone, Serialize, Deserialize)]
520pub struct DatasetStats {
521    /// Total number of examples.
522    pub total_examples: usize,
523    /// Total number of entities across all examples.
524    pub total_entities: usize,
525    /// Average entities per example.
526    pub avg_entities_per_example: f64,
527    /// Count by domain.
528    pub domains: HashMap<Domain, usize>,
529    /// Count by difficulty.
530    pub difficulties: HashMap<Difficulty, usize>,
531    /// Count by entity type (as string labels).
532    pub entity_types: HashMap<String, usize>,
533}
534
535impl DatasetStats {
536    /// Format as a human-readable summary.
537    pub fn summary(&self) -> String {
538        let mut s = String::new();
539        s.push_str(&format!("Examples: {}\n", self.total_examples));
540        s.push_str(&format!("Entities: {}\n", self.total_entities));
541        s.push_str(&format!(
542            "Avg entities/example: {:.1}\n",
543            self.avg_entities_per_example
544        ));
545
546        s.push_str("\nDomains:\n");
547        let mut domains: Vec<_> = self.domains.iter().collect();
548        domains.sort_by(|a, b| b.1.cmp(a.1));
549        for (domain, count) in domains.iter().take(5) {
550            s.push_str(&format!("  {:?}: {}\n", domain, count));
551        }
552        if domains.len() > 5 {
553            s.push_str(&format!("  ... and {} more\n", domains.len() - 5));
554        }
555
556        s.push_str("\nEntity Types:\n");
557        let mut types: Vec<_> = self.entity_types.iter().collect();
558        types.sort_by(|a, b| b.1.cmp(a.1));
559        for (etype, count) in types.iter().take(10) {
560            s.push_str(&format!("  {}: {}\n", etype, count));
561        }
562        if types.len() > 10 {
563            s.push_str(&format!("  ... and {} more\n", types.len() - 10));
564        }
565
566        s
567    }
568}
569
570// ============================================================================
571// Tests
572// ============================================================================
573
574#[cfg(test)]
575mod tests {
576    use super::*;
577
578    #[test]
579    fn test_synthetic_not_empty() {
580        let dataset = NERDataset::synthetic();
581        assert!(!dataset.is_empty());
582        assert!(
583            dataset.len() >= 50,
584            "Should have substantial synthetic data"
585        );
586    }
587
588    #[test]
589    fn test_filter_domain() {
590        let dataset = NERDataset::synthetic();
591        let news = dataset.filter_domain(Domain::News);
592
593        assert!(!news.is_empty());
594        for ex in &news {
595            assert_eq!(ex.domain, Domain::News);
596        }
597    }
598
599    #[test]
600    fn test_filter_difficulty() {
601        let dataset = NERDataset::synthetic();
602        let hard = dataset.filter_difficulty(Difficulty::Hard);
603
604        assert!(!hard.is_empty());
605        for ex in &hard {
606            assert_eq!(ex.difficulty, Difficulty::Hard);
607        }
608    }
609
610    #[test]
611    fn test_to_test_cases() {
612        let dataset = NERDataset::synthetic().take(5);
613        let test_cases = dataset.to_test_cases();
614
615        assert_eq!(test_cases.len(), 5);
616        for (text, entities) in &test_cases {
617            assert!(!text.is_empty());
618            // Some might have no entities (negative examples)
619            let _ = entities;
620        }
621    }
622
623    #[test]
624    fn test_stats() {
625        let dataset = NERDataset::synthetic();
626        let stats = dataset.stats();
627
628        assert!(stats.total_examples > 0);
629        assert!(stats.total_entities > 0);
630        assert!(!stats.domains.is_empty());
631        assert!(!stats.entity_types.is_empty());
632    }
633
634    #[test]
635    fn test_indexing() {
636        let dataset = NERDataset::synthetic();
637        let first = &dataset[0];
638        assert!(!first.text.is_empty());
639    }
640
641    #[test]
642    fn test_into_iterator() {
643        let dataset = NERDataset::synthetic().take(3);
644        let mut count = 0;
645        for _ex in &dataset {
646            count += 1;
647        }
648        assert_eq!(count, 3);
649    }
650
651    #[test]
652    fn test_from_iterator() {
653        let examples = vec![
654            AnnotatedExample::from_tuples("John works", vec![("John", "PER")]),
655            AnnotatedExample::from_tuples("At Google", vec![("Google", "ORG")]),
656        ];
657
658        let dataset: NERDataset = examples.into_iter().collect();
659        assert_eq!(dataset.len(), 2);
660    }
661}
anno_eval/eval/dataset/mod.rs

anno_eval/eval/dataset/
mod.rs
