Skip to main content

anno_eval/eval/
dataset_metadata.rs

1//! Table-driven dataset metadata.
2//!
3//! This module provides a single source of truth for all dataset metadata,
4//! replacing the 15+ separate `match` statements that previously existed.
5//!
6//! # Design
7//!
8//! Instead of:
9//! ```rust,ignore
10//! fn domain(&self) -> &'static str {
11//!     match self {
12//!         DatasetId::WikiGold => "news",
13//!         DatasetId::BC5CDR => "biomedical",
14//!         // ... 450+ more cases
15//!     }
16//! }
17//! ```
18//!
19//! We use:
20//! ```rust,ignore
21//! fn domain(&self) -> &'static str {
22//!     METADATA[self.index()].domain
23//! }
24//! ```
25//!
26//! This reduces ~18K lines to ~3K lines and makes adding datasets O(1) instead of O(n).
27
28use bitflags::bitflags;
29
30bitflags! {
31    /// Dataset capability/category flags.
32    ///
33    /// Using bitflags allows O(1) category checks and easy combination queries.
34    #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
35    pub struct DatasetFlags: u32 {
36        /// Named Entity Recognition dataset
37        const NER = 1 << 0;
38        /// Coreference resolution dataset
39        const COREFERENCE = 1 << 1;
40        /// Intra-document coreference
41        const INTRA_DOC_COREF = 1 << 2;
42        /// Cross-document / inter-document coreference
43        const INTER_DOC_COREF = 1 << 3;
44        /// Temporal NER (dates, durations, temporal expressions)
45        const TEMPORAL_NER = 1 << 4;
46        /// Biomedical/clinical domain
47        const BIOMEDICAL = 1 << 5;
48        /// Social media domain (Twitter, Reddit, etc.)
49        const SOCIAL_MEDIA = 1 << 6;
50        /// Specialized domain (not general news)
51        const SPECIALIZED_DOMAIN = 1 << 7;
52        /// Relation extraction dataset
53        const RELATION_EXTRACTION = 1 << 8;
54        /// Historical texts
55        const HISTORICAL = 1 << 9;
56        /// Bias evaluation dataset
57        const BIAS_EVALUATION = 1 << 10;
58        /// Dialogue/conversational coreference
59        const DIALOGUE_COREF = 1 << 11;
60        /// Joint NER + Relation Extraction
61        const JOINT_NER_RE = 1 << 12;
62        /// Discontinuous/nested NER
63        const DISCONTINUOUS_NER = 1 << 13;
64        /// Few-shot learning dataset
65        const FEW_SHOT = 1 << 14;
66        /// Multilingual dataset
67        const MULTILINGUAL = 1 << 15;
68        /// Constructed/artificial language
69        const CONSTRUCTED_LANGUAGE = 1 << 16;
70        /// Code-switching dataset
71        const CODE_SWITCHING = 1 << 17;
72        /// African language dataset
73        const AFRICAN_LANGUAGE = 1 << 18;
74        /// Entity linking dataset
75        const ENTITY_LINKING = 1 << 19;
76        /// Event extraction dataset
77        const EVENT_EXTRACTION = 1 << 20;
78        /// Legal domain
79        const LEGAL = 1 << 21;
80        /// Financial domain
81        const FINANCIAL = 1 << 22;
82        /// Scientific domain
83        const SCIENTIFIC = 1 << 23;
84        /// Literary/fiction domain
85        const LITERARY = 1 << 24;
86        /// News domain
87        const NEWS = 1 << 25;
88        /// Low-resource language
89        const LOW_RESOURCE = 1 << 26;
90    }
91}
92
93impl Default for DatasetFlags {
94    fn default() -> Self {
95        Self::NER
96    }
97}
98
99/// Complete metadata for a single dataset.
100///
101/// All fields are static references for zero-cost access.
102#[derive(Debug, Clone, Copy)]
103pub struct DatasetMetadata {
104    /// Human-readable name
105    pub name: &'static str,
106    /// Short description (1-2 sentences)
107    pub description: &'static str,
108    /// Download URL (HuggingFace, GitHub, etc.)
109    pub download_url: &'static str,
110    /// Primary domain (news, biomedical, social-media, etc.)
111    pub domain: &'static str,
112    /// ISO 639-1/3 language code (en, de, zh, multilingual, etc.)
113    pub language: &'static str,
114    /// Entity types in this dataset
115    pub entity_types: &'static [&'static str],
116    /// Capability flags (replaces all is_* methods)
117    pub flags: DatasetFlags,
118    /// Academic citation (if available)
119    pub citation: Option<&'static str>,
120    /// License (MIT, CC-BY, etc.)
121    pub license: Option<&'static str>,
122    /// Publication year
123    pub year: Option<u16>,
124    /// Paper URL (arXiv, ACL Anthology, etc.)
125    pub paper_url: Option<&'static str>,
126}
127
128impl DatasetMetadata {
129    /// Create metadata with required fields only.
130    #[must_use]
131    pub const fn new(
132        name: &'static str,
133        description: &'static str,
134        download_url: &'static str,
135    ) -> Self {
136        Self {
137            name,
138            description,
139            download_url,
140            domain: "general",
141            language: "en",
142            entity_types: &[],
143            flags: DatasetFlags::NER,
144            citation: None,
145            license: None,
146            year: None,
147            paper_url: None,
148        }
149    }
150
151    /// Builder: set domain.
152    #[must_use]
153    pub const fn domain(mut self, domain: &'static str) -> Self {
154        self.domain = domain;
155        self
156    }
157
158    /// Builder: set language.
159    #[must_use]
160    pub const fn language(mut self, language: &'static str) -> Self {
161        self.language = language;
162        self
163    }
164
165    /// Builder: set entity types.
166    #[must_use]
167    pub const fn entity_types(mut self, types: &'static [&'static str]) -> Self {
168        self.entity_types = types;
169        self
170    }
171
172    /// Builder: set flags.
173    #[must_use]
174    pub const fn flags(mut self, flags: DatasetFlags) -> Self {
175        self.flags = flags;
176        self
177    }
178
179    /// Builder: set citation.
180    #[must_use]
181    pub const fn citation(mut self, citation: &'static str) -> Self {
182        self.citation = Some(citation);
183        self
184    }
185
186    /// Builder: set license.
187    #[must_use]
188    pub const fn license(mut self, license: &'static str) -> Self {
189        self.license = Some(license);
190        self
191    }
192
193    /// Builder: set year.
194    #[must_use]
195    pub const fn year(mut self, year: u16) -> Self {
196        self.year = Some(year);
197        self
198    }
199
200    /// Builder: set paper URL.
201    #[must_use]
202    pub const fn paper_url(mut self, url: &'static str) -> Self {
203        self.paper_url = Some(url);
204        self
205    }
206
207    // Flag query methods (all O(1))
208
209    /// Returns `true` if this dataset supports named entity recognition.
210    ///
211    /// NER datasets provide entity span annotations with type labels.
212    #[inline]
213    pub const fn is_ner(&self) -> bool {
214        self.flags.contains(DatasetFlags::NER)
215    }
216
217    /// Returns `true` if this dataset contains coreference annotations.
218    ///
219    /// Coreference datasets provide mention clusters or entity chains
220    /// indicating which mentions refer to the same entity.
221    #[inline]
222    pub const fn is_coreference(&self) -> bool {
223        self.flags.contains(DatasetFlags::COREFERENCE)
224    }
225
226    /// Returns `true` if this dataset contains within-document coreference.
227    ///
228    /// Intra-document coreference datasets link mentions within a single
229    /// document, forming coreference chains or clusters.
230    #[inline]
231    pub const fn is_intra_doc_coref(&self) -> bool {
232        self.flags.contains(DatasetFlags::INTRA_DOC_COREF)
233    }
234
235    /// Returns `true` if this dataset contains cross-document coreference.
236    ///
237    /// Inter-document coreference datasets link entities across multiple
238    /// documents, enabling cross-document entity resolution.
239    #[inline]
240    pub const fn is_inter_doc_coref(&self) -> bool {
241        self.flags.contains(DatasetFlags::INTER_DOC_COREF)
242    }
243
244    /// Returns `true` if this dataset contains temporal entity annotations.
245    ///
246    /// Temporal NER datasets include time expressions, events with temporal
247    /// anchors, or entities with temporal attributes (e.g., historical entities).
248    #[inline]
249    pub const fn is_temporal_ner(&self) -> bool {
250        self.flags.contains(DatasetFlags::TEMPORAL_NER)
251    }
252
253    /// Returns `true` if this dataset is from the biomedical domain.
254    ///
255    /// Biomedical datasets include medical, clinical, or life science text
256    /// with specialized entity types (diseases, genes, proteins, chemicals).
257    #[inline]
258    pub const fn is_biomedical(&self) -> bool {
259        self.flags.contains(DatasetFlags::BIOMEDICAL)
260    }
261
262    /// Returns `true` if this dataset contains social media text.
263    ///
264    /// Social media datasets include Twitter, Reddit, or other informal text
265    /// with non-standard capitalization, abbreviations, and informal language.
266    #[inline]
267    pub const fn is_social_media(&self) -> bool {
268        self.flags.contains(DatasetFlags::SOCIAL_MEDIA)
269    }
270
271    /// Returns `true` if this dataset is from a specialized domain.
272    ///
273    /// Specialized domain datasets include technical, legal, financial, or
274    /// other domain-specific text requiring specialized entity recognition.
275    #[inline]
276    pub const fn is_specialized_domain(&self) -> bool {
277        self.flags.contains(DatasetFlags::SPECIALIZED_DOMAIN)
278    }
279
280    /// Returns `true` if this dataset supports relation extraction.
281    ///
282    /// Relation extraction datasets provide entity-relation-entity triples,
283    /// enabling evaluation of models that extract structured relationships.
284    #[inline]
285    pub const fn is_relation_extraction(&self) -> bool {
286        self.flags.contains(DatasetFlags::RELATION_EXTRACTION)
287    }
288
289    /// Returns `true` if this dataset contains historical text.
290    ///
291    /// Historical datasets include ancient texts, historical documents, or
292    /// diachronic corpora that test model robustness to language evolution.
293    #[inline]
294    pub const fn is_historical(&self) -> bool {
295        self.flags.contains(DatasetFlags::HISTORICAL)
296    }
297
298    /// Returns `true` if this dataset is designed for bias evaluation.
299    ///
300    /// Bias evaluation datasets test for gender, demographic, or other biases
301    /// in entity recognition and coreference resolution.
302    #[inline]
303    pub const fn is_bias_evaluation(&self) -> bool {
304        self.flags.contains(DatasetFlags::BIAS_EVALUATION)
305    }
306
307    /// Returns `true` if this dataset contains dialogue coreference annotations.
308    ///
309    /// Dialogue coreference datasets include multi-party conversations, meetings,
310    /// or interviews where coreference resolution must handle speaker turns and
311    /// dialogue-specific phenomena (e.g., prosody, gestures).
312    #[inline]
313    pub const fn is_dialogue_coref(&self) -> bool {
314        self.flags.contains(DatasetFlags::DIALOGUE_COREF)
315    }
316
317    /// Returns `true` if this dataset supports joint NER and relation extraction.
318    ///
319    /// Joint datasets provide both entity annotations and relation triples,
320    /// enabling evaluation of models that perform both tasks simultaneously.
321    #[inline]
322    pub const fn is_joint_ner_re(&self) -> bool {
323        self.flags.contains(DatasetFlags::JOINT_NER_RE)
324    }
325
326    /// Returns `true` if this dataset contains discontinuous entity annotations.
327    ///
328    /// Discontinuous entities span non-contiguous tokens (e.g., "left and right
329    /// ventricle" where "ventricle" is split). Requires specialized evaluation
330    /// metrics beyond standard span-based NER.
331    #[inline]
332    pub const fn is_discontinuous_ner(&self) -> bool {
333        self.flags.contains(DatasetFlags::DISCONTINUOUS_NER)
334    }
335
336    /// Returns `true` if this dataset is designed for few-shot learning evaluation.
337    ///
338    /// Few-shot datasets typically have small training sets or are used to test
339    /// zero-shot transfer from related domains.
340    #[inline]
341    pub const fn is_few_shot(&self) -> bool {
342        self.flags.contains(DatasetFlags::FEW_SHOT)
343    }
344
345    /// Returns `true` if this dataset covers multiple languages.
346    ///
347    /// Multilingual datasets enable cross-lingual evaluation and testing of
348    /// zero-shot transfer between languages.
349    #[inline]
350    pub const fn is_multilingual(&self) -> bool {
351        self.flags.contains(DatasetFlags::MULTILINGUAL)
352    }
353
354    /// Returns `true` if this dataset contains constructed/artificial languages.
355    ///
356    /// Constructed languages (e.g., Esperanto, Klingon) test model generalization
357    /// to languages with different structural properties than natural languages.
358    #[inline]
359    pub const fn is_constructed_language(&self) -> bool {
360        self.flags.contains(DatasetFlags::CONSTRUCTED_LANGUAGE)
361    }
362
363    /// Returns `true` if this dataset contains code-switched text.
364    ///
365    /// Code-switching datasets include text where speakers mix multiple languages
366    /// within the same utterance, requiring models to handle language boundaries.
367    #[inline]
368    pub const fn is_code_switching(&self) -> bool {
369        self.flags.contains(DatasetFlags::CODE_SWITCHING)
370    }
371
372    /// Returns `true` if this dataset contains African languages.
373    ///
374    /// African language datasets are important for evaluating model performance
375    /// on under-resourced languages and diverse linguistic structures.
376    #[inline]
377    pub const fn is_african_language(&self) -> bool {
378        self.flags.contains(DatasetFlags::AFRICAN_LANGUAGE)
379    }
380}
381
382// =============================================================================
383// Standard entity type sets (reduces duplication)
384// =============================================================================
385
386/// CoNLL-style entity types (PER, LOC, ORG, MISC)
387pub static CONLL_TYPES: &[&str] = &["PER", "LOC", "ORG", "MISC"];
388
389/// OntoNotes entity types (18 types)
390pub static ONTONOTES_TYPES: &[&str] = &[
391    "PERSON",
392    "NORP",
393    "FAC",
394    "ORG",
395    "GPE",
396    "LOC",
397    "PRODUCT",
398    "EVENT",
399    "WORK_OF_ART",
400    "LAW",
401    "LANGUAGE",
402    "DATE",
403    "TIME",
404    "PERCENT",
405    "MONEY",
406    "QUANTITY",
407    "ORDINAL",
408    "CARDINAL",
409];
410
411/// Biomedical entity types
412pub static BIO_TYPES: &[&str] = &["Chemical", "Disease", "Gene", "Species"];
413
414/// ACE entity types
415pub static ACE_TYPES: &[&str] = &["PER", "ORG", "GPE", "LOC", "FAC", "VEH", "WEA"];
416
417// =============================================================================
418// Static metadata table (proof of concept - first 20 datasets)
419// =============================================================================
420//
421// This table replaces the 15+ match statements in loader.rs.
422// Full migration will happen incrementally.
423
424/// Metadata for WikiGold dataset
425pub static WIKIGOLD: DatasetMetadata = DatasetMetadata::new(
426    "WikiGold",
427    "Wikipedia-based NER (PER, LOC, ORG, MISC)",
428    "https://huggingface.co/datasets/wikigold",
429)
430.domain("news")
431.language("en")
432.entity_types(CONLL_TYPES)
433.flags(DatasetFlags::NER.union(DatasetFlags::NEWS))
434.year(2009);
435
436/// Metadata for WNUT-17 dataset
437pub static WNUT17: DatasetMetadata = DatasetMetadata::new(
438    "WNUT-17",
439    "Social media NER (emerging entities)",
440    "https://huggingface.co/datasets/wnut_17",
441)
442.domain("social-media")
443.language("en")
444.entity_types(&[
445    "person",
446    "location",
447    "corporation",
448    "product",
449    "creative-work",
450    "group",
451])
452.flags(DatasetFlags::NER.union(DatasetFlags::SOCIAL_MEDIA))
453.year(2017);
454
455/// Metadata for BC5CDR dataset
456pub static BC5CDR: DatasetMetadata = DatasetMetadata::new(
457    "BC5CDR",
458    "Biomedical NER (chemicals, diseases)",
459    "https://huggingface.co/datasets/bc5cdr",
460)
461.domain("biomedical")
462.language("en")
463.entity_types(&["Chemical", "Disease"])
464.flags(
465    DatasetFlags::NER
466        .union(DatasetFlags::BIOMEDICAL)
467        .union(DatasetFlags::SPECIALIZED_DOMAIN),
468)
469.year(2015);
470
471/// Metadata for GAP coreference dataset
472pub static GAP: DatasetMetadata = DatasetMetadata::new(
473    "GAP",
474    "Gendered Ambiguous Pronouns",
475    "https://huggingface.co/datasets/gap",
476)
477.domain("coreference")
478.language("en")
479.entity_types(&["Pronoun", "Name"])
480.flags(
481    DatasetFlags::COREFERENCE
482        .union(DatasetFlags::INTRA_DOC_COREF)
483        .union(DatasetFlags::BIAS_EVALUATION),
484)
485.year(2018);
486
487/// Metadata for OntoNotes 5.0 coreference dataset
488pub static ONTONOTES_COREF: DatasetMetadata = DatasetMetadata::new(
489    "OntoNotes 5.0 (Coreference)",
490    "Standard coreference benchmark",
491    "https://catalog.ldc.upenn.edu/LDC2013T19",
492)
493.domain("coreference")
494.language("en")
495.entity_types(ONTONOTES_TYPES)
496.flags(
497    DatasetFlags::COREFERENCE
498        .union(DatasetFlags::INTRA_DOC_COREF)
499        .union(DatasetFlags::NER),
500)
501.year(2013);
502
503/// Metadata for ECB+ cross-document coreference dataset
504pub static ECBPLUS: DatasetMetadata = DatasetMetadata::new(
505    "ECB+",
506    "Event Coreference Bank Plus",
507    "http://www.newsreader-project.eu/results/data/the-ecb-corpus/",
508)
509.domain("coreference")
510.language("en")
511.entity_types(&["Event", "Entity"])
512.flags(
513    DatasetFlags::COREFERENCE
514        .union(DatasetFlags::INTER_DOC_COREF)
515        .union(DatasetFlags::EVENT_EXTRACTION),
516)
517.year(2014);
518
519/// Metadata for MultiNERD multilingual dataset
520pub static MULTINERD: DatasetMetadata = DatasetMetadata::new(
521    "MultiNERD",
522    "Multilingual NER (10 languages)",
523    "https://huggingface.co/datasets/Babelscape/multinerd",
524)
525.domain("multilingual")
526.language("multilingual")
527.entity_types(&[
528    "PER", "LOC", "ORG", "ANIM", "BIO", "CEL", "DIS", "EVE", "FOOD", "INST", "MEDIA", "MYTH",
529    "PLANT", "TIME", "VEHI",
530])
531.flags(DatasetFlags::NER.union(DatasetFlags::MULTILINGUAL))
532.year(2022);
533
534/// Metadata for FewNERD dataset
535pub static FEWNERD: DatasetMetadata = DatasetMetadata::new(
536    "FewNERD",
537    "Few-shot NER with fine-grained types",
538    "https://huggingface.co/datasets/DFKI-SLT/few-nerd",
539)
540.domain("general")
541.language("en")
542.entity_types(&[
543    "person",
544    "location",
545    "organization",
546    "building",
547    "art",
548    "product",
549    "event",
550    "other",
551])
552.flags(DatasetFlags::NER.union(DatasetFlags::FEW_SHOT))
553.year(2021);
554
555/// Metadata for MasakhaNER African languages dataset
556pub static MASAKHANER: DatasetMetadata = DatasetMetadata::new(
557    "MasakhaNER",
558    "NER for African languages",
559    "https://huggingface.co/datasets/masakhaner",
560)
561.domain("low-resource")
562.language("multilingual")
563.entity_types(CONLL_TYPES)
564.flags(
565    DatasetFlags::NER
566        .union(DatasetFlags::MULTILINGUAL)
567        .union(DatasetFlags::AFRICAN_LANGUAGE)
568        .union(DatasetFlags::LOW_RESOURCE),
569)
570.year(2021);
571
572/// Metadata for GENIA biomedical dataset
573pub static GENIA: DatasetMetadata = DatasetMetadata::new(
574    "GENIA",
575    "Biomedical NER (genes, proteins)",
576    "http://www.geniaproject.org/",
577)
578.domain("biomedical")
579.language("en")
580.entity_types(&["DNA", "RNA", "protein", "cell_line", "cell_type"])
581.flags(
582    DatasetFlags::NER
583        .union(DatasetFlags::BIOMEDICAL)
584        .union(DatasetFlags::SPECIALIZED_DOMAIN),
585)
586.year(2003);
587
588#[cfg(test)]
589mod tests {
590    use super::*;
591
592    #[test]
593    fn test_flags_operations() {
594        let flags = DatasetFlags::NER | DatasetFlags::BIOMEDICAL | DatasetFlags::SPECIALIZED_DOMAIN;
595        assert!(flags.contains(DatasetFlags::NER));
596        assert!(flags.contains(DatasetFlags::BIOMEDICAL));
597        assert!(!flags.contains(DatasetFlags::SOCIAL_MEDIA));
598    }
599
600    #[test]
601    fn test_metadata_builder() {
602        let meta = DatasetMetadata::new("Test", "A test dataset", "https://example.com")
603            .domain("biomedical")
604            .language("en")
605            .entity_types(&["Disease", "Drug"])
606            .flags(DatasetFlags::NER | DatasetFlags::BIOMEDICAL)
607            .year(2023);
608
609        assert_eq!(meta.name, "Test");
610        assert_eq!(meta.domain, "biomedical");
611        assert!(meta.is_biomedical());
612        assert!(!meta.is_social_media());
613        assert_eq!(meta.year, Some(2023));
614    }
615
616    #[test]
617    fn test_const_construction() {
618        // Verify const construction works (for static tables)
619        const META: DatasetMetadata = DatasetMetadata::new("Const", "Desc", "url")
620            .domain("test")
621            .language("en");
622
623        assert_eq!(META.name, "Const");
624        assert_eq!(META.domain, "test");
625    }
626
627    #[test]
628    fn test_static_wikigold() {
629        assert_eq!(WIKIGOLD.name, "WikiGold");
630        assert_eq!(WIKIGOLD.domain, "news");
631        assert!(WIKIGOLD.is_ner());
632        assert!(!WIKIGOLD.is_coreference());
633    }
634
635    #[test]
636    fn test_static_bc5cdr() {
637        assert!(BC5CDR.is_biomedical());
638        assert!(BC5CDR.is_specialized_domain());
639        assert_eq!(BC5CDR.entity_types.len(), 2);
640    }
641
642    #[test]
643    fn test_static_gap() {
644        assert!(GAP.is_coreference());
645        assert!(GAP.is_intra_doc_coref());
646        assert!(GAP.is_bias_evaluation());
647        assert!(!GAP.is_ner());
648    }
649
650    #[test]
651    fn test_static_ecbplus() {
652        assert!(ECBPLUS.is_inter_doc_coref());
653        assert!(!ECBPLUS.is_intra_doc_coref());
654    }
655
656    #[test]
657    fn test_static_masakhaner() {
658        assert!(MASAKHANER.is_african_language());
659        assert!(MASAKHANER.is_multilingual());
660    }
661}