anno_eval/eval/
datasets.rs

1//! Dataset loading for NER evaluation.
2#![allow(missing_docs)] // Internal evaluation types
3//!
4//! Supports multiple dataset formats:
5//! - CoNLL-2003 (classic BIO tagging format)
6//! - JSON/JSONL (format used by OpenNER, MultiNERD, Wikiann)
7//! - HuggingFace Datasets format
8//!
9//! Recent datasets (2024-2026):
10//! - OpenNER 1.0: 52 languages, standardized JSON
11//! - MultiNERD: Multi-ontology, JSON format
12//! - Wikiann: 282 languages, JSONL format
13
14use anno::EntityType;
15use anno::{Error, Result};
16use serde::{Deserialize, Serialize};
17use std::path::Path;
18
19// =============================================================================
20// Gold Entity Types
21// =============================================================================
22
23/// Gold standard entity annotation.
24///
25/// This is the canonical entity type for all NER evaluation.
26/// Use this instead of creating local entity structs.
27///
28/// # Example
29/// ```rust
30/// use anno_eval::eval::GoldEntity;
31/// use anno::EntityType;
32///
33/// let entity = GoldEntity::new("John Doe", EntityType::Person, 0);
34/// assert_eq!(entity.end, 8);
35/// ```
36#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
37pub struct GoldEntity {
38    /// The entity text (surface form)
39    pub text: String,
40    /// Normalized entity type
41    pub entity_type: EntityType,
42    /// Original label from dataset (e.g., "B-PER", "ACTOR")
43    pub original_label: String,
44    /// Character offset (start)
45    pub start: usize,
46    /// Character offset (end, exclusive)
47    pub end: usize,
48}
49
50impl GoldEntity {
51    /// Create a new gold entity with computed end offset.
52    ///
53    /// Note: Uses character count (not byte count) for Unicode correctness.
54    #[must_use]
55    pub fn new(text: impl Into<String>, entity_type: EntityType, start: usize) -> Self {
56        let text = text.into();
57        // Use char count for Unicode correctness
58        let end = start + text.chars().count();
59        Self {
60            text,
61            entity_type,
62            original_label: String::new(),
63            start,
64            end,
65        }
66    }
67
68    /// Create with explicit span (start, end).
69    pub fn with_span(
70        text: impl Into<String>,
71        entity_type: EntityType,
72        start: usize,
73        end: usize,
74    ) -> Self {
75        Self {
76            text: text.into(),
77            entity_type,
78            original_label: String::new(),
79            start,
80            end,
81        }
82    }
83
84    /// Create with explicit original label.
85    ///
86    /// Note: Uses character count (not byte count) for Unicode correctness.
87    pub fn with_label(
88        text: impl Into<String>,
89        entity_type: EntityType,
90        original_label: impl Into<String>,
91        start: usize,
92    ) -> Self {
93        let text = text.into();
94        let end = start + text.chars().count();
95        Self {
96            text,
97            entity_type,
98            original_label: original_label.into(),
99            start,
100            end,
101        }
102    }
103
104    /// Create with all fields explicit.
105    pub fn full(
106        text: impl Into<String>,
107        entity_type: EntityType,
108        original_label: impl Into<String>,
109        start: usize,
110        end: usize,
111    ) -> Self {
112        Self {
113            text: text.into(),
114            entity_type,
115            original_label: original_label.into(),
116            start,
117            end,
118        }
119    }
120
121    /// Check if this entity overlaps with another.
122    pub fn overlaps(&self, other: &Self) -> bool {
123        self.start < other.end && other.start < self.end
124    }
125
126    /// Safely extract text from source using character offsets.
127    ///
128    /// GoldEntity stores character offsets, not byte offsets. This method
129    /// correctly extracts text by iterating over characters.
130    ///
131    /// # Arguments
132    /// * `source_text` - The original text from which this entity was extracted
133    ///
134    /// # Returns
135    /// The extracted text, or empty string if offsets are invalid
136    #[must_use]
137    pub fn extract_text(&self, source_text: &str) -> String {
138        let char_count = source_text.chars().count();
139        if self.start >= char_count || self.end > char_count || self.start >= self.end {
140            return String::new();
141        }
142        source_text
143            .chars()
144            .skip(self.start)
145            .take(self.end - self.start)
146            .collect()
147    }
148
149    /// Check if spans match exactly.
150    pub fn span_matches(&self, other: &Self) -> bool {
151        self.start == other.start && self.end == other.end
152    }
153
154    /// Check if this is an exact match (span + type).
155    pub fn exact_matches(&self, other: &Self) -> bool {
156        self.span_matches(other) && self.entity_type == other.entity_type
157    }
158}
159
160/// JSON format for NER datasets (OpenNER, MultiNERD style).
161///
162/// Format:
163/// ```json
164/// {
165///   "text": "John Smith works at Acme Corp.",
166///   "entities": [
167///     {"text": "John Smith", "label": "PER", "start": 0, "end": 10},
168///     {"text": "Acme Corp", "label": "ORG", "start": 20, "end": 29}
169///   ]
170/// }
171/// ```
172#[derive(Debug, Clone, Serialize, Deserialize)]
173pub struct JSONNERExample {
174    pub text: String,
175    pub entities: Vec<JSONEntity>,
176}
177
178#[derive(Debug, Clone, Serialize, Deserialize)]
179pub struct JSONEntity {
180    pub text: String,
181    pub label: String,
182    pub start: usize,
183    pub end: usize,
184    #[serde(default)]
185    pub confidence: Option<f64>,
186}
187
188/// JSONL format (one JSON object per line).
189///
190/// Used by Wikiann and many recent datasets.
191pub type JSONLNERExample = JSONNERExample;
192
193/// Load JSON format NER dataset.
194///
195/// Supports:
196/// - Single JSON file with array of examples
197/// - JSONL file (one JSON object per line)
198pub fn load_json_ner_dataset<P: AsRef<Path>>(path: P) -> Result<Vec<(String, Vec<GoldEntity>)>> {
199    let content = std::fs::read_to_string(path.as_ref()).map_err(Error::Io)?;
200
201    let mut test_cases = Vec::new();
202
203    // Try JSONL first (one object per line)
204    let is_jsonl = content.lines().count() > 1
205        && content
206            .lines()
207            .all(|line| line.trim().starts_with('{') && line.trim().ends_with('}'));
208
209    if is_jsonl {
210        // JSONL format
211        for (line_num, line) in content.lines().enumerate() {
212            let line = line.trim();
213            if line.is_empty() {
214                continue;
215            }
216
217            let example: JSONLNERExample = serde_json::from_str(line).map_err(|e| {
218                Error::Parse(format!(
219                    "Failed to parse JSONL line {}: {}",
220                    line_num + 1,
221                    e
222                ))
223            })?;
224
225            let entities: Vec<GoldEntity> = example
226                .entities
227                .into_iter()
228                .map(|e| {
229                    let entity_type = map_label_to_entity_type(&e.label);
230                    GoldEntity::full(e.text, entity_type, &e.label, e.start, e.end)
231                })
232                .collect();
233
234            // Validate entities against text
235            let validation = crate::eval::validation::validate_ground_truth_entities(
236                &example.text,
237                &entities,
238                false, // Warnings for overlaps, not errors
239            );
240            if !validation.is_valid {
241                return Err(Error::InvalidInput(format!(
242                    "Invalid entities in dataset: {}",
243                    validation.errors.join("; ")
244                )));
245            }
246
247            test_cases.push((example.text, entities));
248        }
249    } else {
250        // Single JSON file (array of examples)
251        let examples: Vec<JSONNERExample> = serde_json::from_str(&content)
252            .map_err(|e| Error::Parse(format!("Failed to parse JSON: {}", e)))?;
253
254        for example in examples {
255            let entities: Vec<GoldEntity> = example
256                .entities
257                .into_iter()
258                .map(|e| {
259                    let entity_type = map_label_to_entity_type(&e.label);
260                    GoldEntity::full(e.text, entity_type, &e.label, e.start, e.end)
261                })
262                .collect();
263
264            // Validate entities against text
265            let validation = crate::eval::validation::validate_ground_truth_entities(
266                &example.text,
267                &entities,
268                false, // Warnings for overlaps, not errors
269            );
270            if !validation.is_valid {
271                return Err(Error::InvalidInput(format!(
272                    "Invalid entities in dataset: {}",
273                    validation.errors.join("; ")
274                )));
275            }
276
277            test_cases.push((example.text, entities));
278        }
279    }
280
281    Ok(test_cases)
282}
283
284/// Load HuggingFace Datasets format.
285///
286/// HuggingFace datasets are typically stored as JSON/JSONL with specific structure.
287/// This function handles common HuggingFace NER dataset formats.
288pub fn load_hf_ner_dataset<P: AsRef<Path>>(path: P) -> Result<Vec<(String, Vec<GoldEntity>)>> {
289    // HuggingFace datasets are often JSONL or JSON arrays
290    // Try JSONL first, then fall back to JSON array
291    load_json_ner_dataset(path)
292}
293
294/// Map label string to EntityType.
295///
296/// **Prefer `crate::schema::map_to_canonical()` for new code** - it handles
297/// more types correctly and preserves semantic distinctions.
298///
299/// Handles various label formats:
300/// - CoNLL: "PER", "ORG", "LOC", "MISC"
301/// - OpenNER: Standardized labels
302/// - MultiNERD: Extended labels (PER, ORG, LOC, ANIM, BIO, CEL, etc.)
303/// - Wikiann: "PER", "ORG", "LOC", "MISC"
304fn map_label_to_entity_type(label: &str) -> EntityType {
305    // Use the new canonical mapper for consistent semantics
306    anno::schema::map_to_canonical(label, None)
307}
308
309/// Auto-detect dataset format and load.
310///
311/// Tries to detect format based on file extension and content:
312/// - `.conll`, `.conll2003` → CoNLL-2003 format
313/// - `.json`, `.jsonl` → JSON/JSONL format
314/// - `.txt` → Try CoNLL first, then JSON
315pub fn load_ner_dataset<P: AsRef<Path>>(path: P) -> Result<Vec<(String, Vec<GoldEntity>)>> {
316    let path = path.as_ref();
317    let extension = path
318        .extension()
319        .and_then(|ext| ext.to_str())
320        .unwrap_or("")
321        .to_lowercase();
322
323    match extension.as_str() {
324        "conll" | "conll2003" | "txt" => {
325            // Try CoNLL format first
326            load_conll_2003_dataset_internal(path).or_else(|_| {
327                // If CoNLL fails, try JSON
328                load_json_ner_dataset(path)
329            })
330        }
331        "json" | "jsonl" => load_json_ner_dataset(path),
332        _ => {
333            // Try CoNLL first (most common), then JSON
334            load_conll_2003_dataset_internal(path).or_else(|_| load_json_ner_dataset(path))
335        }
336    }
337}
338
339/// Internal CoNLL-2003 loader (used by datasets module).
340fn load_conll_2003_dataset_internal<P: AsRef<Path>>(
341    path: P,
342) -> Result<Vec<(String, Vec<GoldEntity>)>> {
343    // Re-export the public function
344    crate::eval::load_conll2003(path)
345}
346
347/// Dataset metadata for tracking dataset information.
348#[derive(Debug, Clone, Serialize, Deserialize)]
349pub struct DatasetMetadata {
350    pub name: String,
351    pub format: String,
352    pub language: Option<String>,
353    pub entity_types: Vec<String>,
354    pub num_examples: usize,
355    pub source: Option<String>,
356    pub year: Option<u32>,
357}
358
359/// Extract dataset metadata from loaded examples.
360pub fn extract_dataset_metadata(
361    examples: &[(String, Vec<GoldEntity>)],
362    name: &str,
363) -> DatasetMetadata {
364    let mut entity_types = std::collections::HashSet::new();
365    for (_, entities) in examples {
366        for entity in entities {
367            let type_str = crate::eval::entity_type_to_string(&entity.entity_type);
368            entity_types.insert(type_str);
369        }
370    }
371
372    DatasetMetadata {
373        name: name.to_string(),
374        format: "auto-detected".to_string(),
375        language: None,
376        entity_types: entity_types.into_iter().collect(),
377        num_examples: examples.len(),
378        source: None,
379        year: None,
380    }
381}
382
383#[cfg(test)]
384mod tests {
385    use super::*;
386    use std::fs::File;
387    use std::io::Write;
388
389    #[test]
390    fn test_load_json_ner_dataset() {
391        let json_content = r#"[
392            {
393                "text": "John Smith works at Acme Corp.",
394                "entities": [
395                    {"text": "John Smith", "label": "PER", "start": 0, "end": 10},
396                    {"text": "Acme Corp", "label": "ORG", "start": 20, "end": 29}
397                ]
398            }
399        ]"#;
400
401        let temp_dir = std::env::temp_dir();
402        let file_path = temp_dir.join("test_ner.json");
403        let mut file = File::create(&file_path).expect("should create test file");
404        file.write_all(json_content.as_bytes())
405            .expect("should write test file");
406        file.flush().expect("should flush test file");
407
408        let result = load_json_ner_dataset(&file_path).expect("should load test dataset");
409        assert_eq!(result.len(), 1);
410        assert_eq!(result[0].0, "John Smith works at Acme Corp.");
411        assert_eq!(result[0].1.len(), 2);
412
413        std::fs::remove_file(&file_path).ok();
414    }
415
416    #[test]
417    fn test_load_jsonl_ner_dataset() {
418        let jsonl_content = r#"{"text": "John Smith works.", "entities": [{"text": "John Smith", "label": "PER", "start": 0, "end": 10}]}
419{"text": "Acme Corp is hiring.", "entities": [{"text": "Acme Corp", "label": "ORG", "start": 0, "end": 9}]}
420"#;
421
422        let temp_dir = std::env::temp_dir();
423        let file_path = temp_dir.join("test_ner.jsonl");
424        let mut file = File::create(&file_path).unwrap();
425        file.write_all(jsonl_content.as_bytes()).unwrap();
426        file.flush().unwrap();
427
428        let result = load_json_ner_dataset(&file_path).expect("should load test dataset");
429        assert_eq!(result.len(), 2);
430
431        std::fs::remove_file(&file_path).ok();
432    }
433
434    #[test]
435    fn test_map_label_to_entity_type() {
436        // Core types
437        assert!(matches!(
438            map_label_to_entity_type("PER"),
439            EntityType::Person
440        ));
441        assert!(matches!(
442            map_label_to_entity_type("ORG"),
443            EntityType::Organization
444        ));
445        assert!(matches!(
446            map_label_to_entity_type("LOC"),
447            EntityType::Location
448        ));
449
450        // MISC -> Custom or Other
451        assert!(matches!(
452            map_label_to_entity_type("MISC"),
453            EntityType::Custom { .. }
454        ));
455
456        // ANIM now preserves semantics as Custom type
457        assert!(matches!(
458            map_label_to_entity_type("ANIM"),
459            EntityType::Custom { .. }
460        ));
461    }
462
463    #[test]
464    fn test_load_ner_dataset_auto_detect() {
465        // Test JSON detection
466        let json_content = r#"[{"text": "Test", "entities": []}]"#;
467        let temp_dir = std::env::temp_dir();
468        let file_path = temp_dir.join("test_auto.json");
469        let mut file = File::create(&file_path).expect("should create test file");
470        file.write_all(json_content.as_bytes())
471            .expect("should write test file");
472        file.flush().expect("should flush test file");
473
474        let result = load_ner_dataset(&file_path);
475        assert!(result.is_ok());
476
477        std::fs::remove_file(&file_path).ok();
478    }
479}
anno_eval/eval/datasets.rs

anno_eval/eval/
datasets.rs
