anno_eval/eval/
task_mapping.rs

1//! Task-Dataset-Backend Mapping System
2//!
3//! This module provides a cohesive system for mapping:
4//! - Tasks (NER, NED, Coreference, etc.) → Datasets
5//! - Datasets → Backends that can evaluate them
6//! - Backends → Tasks they support (via trait inspection)
7//!
8//! # Design Philosophy
9//!
10//! - **Trait-based capabilities**: Backend capabilities are determined by trait implementations
11//! - **Many-to-many relationships**: A dataset can support multiple tasks, a backend can support multiple tasks
12//! - **Explicit capabilities**: Each backend declares what tasks it supports via traits
13//! - **Dataset metadata**: Each dataset declares what tasks it can evaluate
14//! - **Task requirements**: Each task declares what datasets are suitable
15//!
16//! # Trait-Based Capability Detection
17//!
18//! Backends are queried for capabilities using trait bounds:
19//! - `Model` → NER capability
20//! - `ZeroShotNER` → Zero-shot NER capability
21//! - `RelationExtractor` → Relation extraction capability
22//! - `DiscontinuousNER` → Discontinuous NER capability
23//! - `CoreferenceResolver` → Coreference resolution capability
24
25use crate::eval::loader::DatasetId;
26use serde::{Deserialize, Serialize};
27use std::collections::HashMap;
28
29// Re-export traits for capability detection
30pub use crate::eval::coref_resolver::CoreferenceResolver as CoreferenceResolverTrait;
31pub use anno::backends::inference::{
32    DiscontinuousNER as DiscontinuousNERTrait, RelationExtractor as RelationExtractorTrait,
33    ZeroShotNER as ZeroShotNERTrait,
34};
35
36/// Information extraction and NLP tasks supported by anno.
37#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
38#[non_exhaustive]
39pub enum Task {
40    /// Named Entity Recognition: extract entity spans with types
41    NER,
42    /// Named Entity Disambiguation: link entities to knowledge bases
43    NED,
44    /// Relation Extraction: extract entity-relation-entity triples
45    RelationExtraction,
46    /// Intra-document Coreference: resolve mentions within a document
47    IntraDocCoref,
48    /// Inter-document Coreference: resolve mentions across documents
49    InterDocCoref,
50    /// Abstract Anaphora: resolve pronouns to events/propositions
51    AbstractAnaphora,
52    /// Discontinuous NER: extract non-contiguous entity spans
53    DiscontinuousNER,
54    /// Event Extraction: extract event triggers and arguments
55    EventExtraction,
56    /// Text Classification: classify entire text or spans
57    TextClassification,
58    /// Sentiment analysis (a specialization of text classification)
59    SentimentAnalysis,
60    /// Part-of-speech tagging
61    PosTagging,
62    /// Machine translation
63    MachineTranslation,
64    /// Language modeling
65    LanguageModeling,
66    /// Temporal extraction / temporal information (catalogued from registries).
67    ///
68    /// Note: this is distinct from `HierarchicalExtraction`. Some registries use "temporal"
69    /// to mean time expressions, event ordering, or temporal relations; we keep it separate
70    /// so we don't mislabel temporal datasets as hierarchical.
71    Temporal,
72    /// Hierarchical Structure Extraction: extract nested structures
73    HierarchicalExtraction,
74    /// Discourse relations (e.g., PDTB-style)
75    DiscourseRelations,
76    /// Discourse coherence / coherence modeling
77    DiscourseCoherence,
78    /// Discourse segmentation
79    DiscourseSegmentation,
80    /// Speech act classification (dialogue acts)
81    SpeechActClassification,
82}
83
84impl Task {
85    /// All supported tasks.
86    pub fn all() -> &'static [Task] {
87        &[
88            Task::NER,
89            Task::NED,
90            Task::RelationExtraction,
91            Task::IntraDocCoref,
92            Task::InterDocCoref,
93            Task::AbstractAnaphora,
94            Task::DiscontinuousNER,
95            Task::EventExtraction,
96            Task::TextClassification,
97            Task::SentimentAnalysis,
98            Task::PosTagging,
99            Task::MachineTranslation,
100            Task::LanguageModeling,
101            Task::Temporal,
102            Task::HierarchicalExtraction,
103            Task::DiscourseRelations,
104            Task::DiscourseCoherence,
105            Task::DiscourseSegmentation,
106            Task::SpeechActClassification,
107        ]
108    }
109
110    /// Human-readable name for this task.
111    pub fn name(&self) -> &'static str {
112        match self {
113            Task::NER => "Named Entity Recognition",
114            Task::NED => "Named Entity Disambiguation",
115            Task::RelationExtraction => "Relation Extraction",
116            Task::IntraDocCoref => "Intra-document Coreference",
117            Task::InterDocCoref => "Inter-document Coreference",
118            Task::AbstractAnaphora => "Abstract Anaphora Resolution",
119            Task::DiscontinuousNER => "Discontinuous NER",
120            Task::EventExtraction => "Event Extraction",
121            Task::TextClassification => "Text Classification",
122            Task::SentimentAnalysis => "Sentiment Analysis",
123            Task::PosTagging => "Part-of-speech Tagging",
124            Task::MachineTranslation => "Machine Translation",
125            Task::LanguageModeling => "Language Modeling",
126            Task::Temporal => "Temporal",
127            Task::HierarchicalExtraction => "Hierarchical Structure Extraction",
128            Task::DiscourseRelations => "Discourse Relations",
129            Task::DiscourseCoherence => "Discourse Coherence",
130            Task::DiscourseSegmentation => "Discourse Segmentation",
131            Task::SpeechActClassification => "Speech Act Classification",
132        }
133    }
134
135    /// Short code for this task (for CLI/config).
136    pub fn code(&self) -> &'static str {
137        match self {
138            Task::NER => "ner",
139            Task::NED => "ned",
140            Task::RelationExtraction => "re",
141            Task::IntraDocCoref => "intra-coref",
142            Task::InterDocCoref => "inter-coref",
143            Task::AbstractAnaphora => "abstract-anaphora",
144            Task::DiscontinuousNER => "discontinuous-ner",
145            Task::EventExtraction => "events",
146            Task::TextClassification => "classification",
147            Task::SentimentAnalysis => "sentiment",
148            Task::PosTagging => "pos",
149            Task::MachineTranslation => "mt",
150            Task::LanguageModeling => "lm",
151            Task::Temporal => "temporal",
152            Task::HierarchicalExtraction => "hierarchical",
153            Task::DiscourseRelations => "discourse-relations",
154            Task::DiscourseCoherence => "discourse-coherence",
155            Task::DiscourseSegmentation => "discourse-segmentation",
156            Task::SpeechActClassification => "speech-act-classification",
157        }
158    }
159
160    /// Parse task from short code string.
161    ///
162    /// Supports many common aliases used in dataset registry:
163    /// - NER: "ner", "sequence_labeling", "nested-ner", "mner", "pii_detection", "slot_filling"
164    /// - NED: "ned", "el", "entity_linking", "entity-linking"
165    /// - RelationExtraction: "re", "relation_extraction", "relation-extraction"
166    /// - IntraDocCoref: "coref", "intra-coref", "intra_coref"
167    /// - InterDocCoref: "inter-coref", "inter_coref", "cdcr", "event_coref"
168    /// - AbstractAnaphora: "abstract-anaphora", "abstract_anaphora", "bridging", "discourse_deixis"
169    /// - DiscontinuousNER: "discontinuous-ner", "discontinuous_ner", "dner"
170    /// - EventExtraction: "events", "event_extraction", "event-extraction"
171    /// - TextClassification: "classification", "text-classification", "bias_evaluation", "qa", "harmonic_analysis"
172    /// - Temporal: "temporal", "timex"
173    /// - HierarchicalExtraction: "hierarchical"
174    pub fn from_code(code: &str) -> Option<Self> {
175        match code.to_lowercase().as_str() {
176            // NER family
177            "ner" | "sequence_labeling" | "nested-ner" | "mner" | "pii_detection"
178            | "slot_filling" => Some(Task::NER),
179
180            // Entity Linking / Disambiguation
181            "ned" | "el" | "entity_linking" | "entity-linking" => Some(Task::NED),
182
183            // Relation Extraction
184            "re" | "relation" | "relation_extraction" | "relation-extraction" => {
185                Some(Task::RelationExtraction)
186            }
187
188            // Intra-document coreference
189            "coref" | "intra-coref" | "intra_coref" | "intracoref" => Some(Task::IntraDocCoref),
190
191            // Inter-document coreference (CDCR)
192            "inter-coref" | "inter_coref" | "intercoref" | "cdcr" | "event_coref" => {
193                Some(Task::InterDocCoref)
194            }
195
196            // Abstract Anaphora (includes bridging and discourse deixis)
197            "abstract-anaphora" | "abstract_anaphora" | "bridging" | "discourse_deixis" => {
198                Some(Task::AbstractAnaphora)
199            }
200
201            // Discontinuous NER
202            "discontinuous-ner" | "discontinuous_ner" | "disc-ner" | "dner" => {
203                Some(Task::DiscontinuousNER)
204            }
205
206            // Event Extraction
207            "events" | "event" | "event_extraction" | "event-extraction" | "ee" => {
208                Some(Task::EventExtraction)
209            }
210
211            // Text Classification (includes bias evaluation, QA, harmonic analysis)
212            "classification"
213            | "text_classification"
214            | "text-classification"
215            | "tc"
216            | "bias_evaluation"
217            | "qa"
218            | "harmonic_analysis" => Some(Task::TextClassification),
219
220            // Sentiment analysis (often listed separately in registries)
221            "sentiment" | "sentiment_analysis" | "sentiment-analysis" => {
222                Some(Task::SentimentAnalysis)
223            }
224
225            // POS tagging
226            "pos" | "pos_tagging" | "pos-tagging" => Some(Task::PosTagging),
227
228            // Machine translation
229            "mt" | "machine_translation" | "machine-translation" | "translation" => {
230                Some(Task::MachineTranslation)
231            }
232
233            // Language modeling
234            "lm" | "language_modeling" | "language-modeling" => Some(Task::LanguageModeling),
235
236            // Temporal
237            "temporal" | "timex" | "time_expressions" | "time-expressions" => Some(Task::Temporal),
238
239            // Hierarchical Extraction
240            "hierarchical" | "hierarchical-extraction" | "he" => Some(Task::HierarchicalExtraction),
241
242            // Discourse tasks (catalogued in the registry; evaluation may be stubbed for now)
243            "discourse_relations" | "discourse-relations" => Some(Task::DiscourseRelations),
244            "discourse_coherence" | "discourse-coherence" => Some(Task::DiscourseCoherence),
245            "discourse_segmentation" | "discourse-segmentation" => {
246                Some(Task::DiscourseSegmentation)
247            }
248            "speech_act_classification" | "speech-act-classification" => {
249                Some(Task::SpeechActClassification)
250            }
251
252            _ => None,
253        }
254    }
255
256    /// Check if this task is in the NER family (NER, discontinuous NER, NED).
257    pub fn is_ner_family(&self) -> bool {
258        matches!(self, Task::NER | Task::DiscontinuousNER | Task::NED)
259    }
260
261    /// Check if this task is in the coreference family.
262    pub fn is_coref_family(&self) -> bool {
263        matches!(
264            self,
265            Task::IntraDocCoref | Task::InterDocCoref | Task::AbstractAnaphora
266        )
267    }
268}
269
270/// Tasks that a dataset supports *for evaluation*.
271///
272/// Derived from the dataset registry's `DatasetId::tasks()` strings and mapped through
273/// `Task::from_code` (which supports many aliases used in the registry).
274///
275/// We intentionally **do not** default unknown datasets to NER: that hides missing metadata and
276/// creates misleading task×dataset×backend matrices.
277pub fn dataset_tasks(dataset: DatasetId) -> Vec<Task> {
278    // Keep the registry as the single source of truth for the string task codes,
279    // but use its typed view to avoid duplicating parsing/alias logic here.
280    dataset.tasks_typed()
281}
282
283/// Mapping from tasks to suitable datasets.
284pub fn task_datasets(task: Task) -> &'static [DatasetId] {
285    match task {
286        Task::NER => &[
287            DatasetId::WikiGold,
288            DatasetId::Wnut17,
289            DatasetId::MitMovie,
290            DatasetId::MitRestaurant,
291            DatasetId::CoNLL2003Sample,
292            DatasetId::OntoNotesSample,
293            DatasetId::MultiNERD,
294            DatasetId::BC5CDR,
295            DatasetId::NCBIDisease,
296            DatasetId::GENIA,
297            DatasetId::AnatEM,
298            DatasetId::BC2GM,
299            DatasetId::BC4CHEMD,
300            DatasetId::TweetNER7,
301            DatasetId::BroadTwitterCorpus,
302            DatasetId::FabNER,
303            DatasetId::FewNERD,
304            DatasetId::CrossNER,
305            DatasetId::UniversalNERBench,
306            DatasetId::WikiANN,
307            DatasetId::MultiCoNER,
308            DatasetId::MultiCoNERv2,
309            DatasetId::WikiNeural,
310            DatasetId::PolyglotNER,
311            DatasetId::UniversalNER,
312            DatasetId::UNER,
313            DatasetId::MSNER,
314            DatasetId::BioMNER,
315            DatasetId::LegNER,
316            DatasetId::OntoNotes50,
317            DatasetId::GermEval2014,
318            DatasetId::HAREM,
319            DatasetId::SemEval2013Task91,
320            DatasetId::MUC6,
321            DatasetId::MUC7,
322            DatasetId::JNLPBA,
323            DatasetId::BC2GMFull,
324            DatasetId::CRAFT,
325            DatasetId::FinNER,
326            DatasetId::LegalNER,
327            DatasetId::SciERCNER,
328            // Constructed languages (CoNLL-U)
329            DatasetId::TaggedPBCEsperanto,
330            DatasetId::TaggedPBCKlingon,
331            // Note: These variants were referenced but not added to enum
332            // Using existing variants: CoNLL2003Sample, Wnut17, BC5CDR, NCBIDisease
333        ],
334        Task::DiscontinuousNER => &[DatasetId::CADEC, DatasetId::ShARe13, DatasetId::ShARe14],
335        Task::RelationExtraction => &[
336            DatasetId::DocRED,
337            DatasetId::ReTACRED,
338            DatasetId::NYTFB,
339            DatasetId::WEBNLG,
340            DatasetId::GoogleRE,
341            DatasetId::BioRED,
342            DatasetId::SciER,
343            DatasetId::MixRED,
344            DatasetId::CovEReD,
345        ],
346        Task::IntraDocCoref => &[
347            DatasetId::GAP,
348            DatasetId::PreCo,
349            DatasetId::LitBank,
350            DatasetId::HumanVoiceAgentInteraction,
351        ],
352        Task::InterDocCoref => &[DatasetId::ECBPlus, DatasetId::WikiCoref],
353        Task::AbstractAnaphora => &[
354            DatasetId::GAP,
355            DatasetId::PreCo,
356            DatasetId::LitBank,
357            DatasetId::HumanVoiceAgentInteraction,
358        ],
359        Task::EventExtraction => &[
360            DatasetId::ACE2005,
361            DatasetId::MAVEN,
362            DatasetId::MAVENArg,
363            DatasetId::CASIE,
364            DatasetId::RAMS,
365        ],
366        Task::NED => &[DatasetId::AIDA, DatasetId::TACKBP],
367        Task::TextClassification => &[
368            DatasetId::MasakhaNEWS,
369            DatasetId::AGNews,
370            DatasetId::DBPedia14,
371            DatasetId::YahooAnswers,
372            DatasetId::TREC,
373            DatasetId::TweetTopic,
374        ],
375        // We catalog these tasks/datasets, but we don't currently run end-to-end evaluation
376        // pipelines for them.
377        Task::SentimentAnalysis => &[],
378        Task::PosTagging => &[],
379        Task::MachineTranslation => &[],
380        Task::LanguageModeling => &[],
381        Task::Temporal => &[DatasetId::TimexRecognitionSentenceOriginal],
382        Task::HierarchicalExtraction => {
383            // GLiNER multi-task can do hierarchical extraction, but we don't have dedicated datasets yet
384            &[]
385        }
386        Task::DiscourseRelations => &[
387            DatasetId::DisrptEngDepScidtbRels,
388            DatasetId::DisrptDeuRstPccRels,
389        ],
390        Task::SpeechActClassification => &[DatasetId::ViraDialogActsLive],
391        Task::DiscourseSegmentation => &[
392            DatasetId::DisrptEngDepScidtbConlluSeg,
393            DatasetId::DisrptDeuRstPccConlluSeg,
394        ],
395        Task::DiscourseCoherence => &[],
396    }
397}
398
399/// Detect backend capabilities via trait inspection.
400///
401/// This function attempts to determine what tasks a backend supports
402/// by checking if it implements relevant traits. For runtime detection,
403/// use `detect_backend_capabilities` instead.
404pub fn backend_tasks(backend_name: &str) -> &'static [Task] {
405    match backend_name {
406        // Regex-based backends
407        "pattern" | "RegexNER" => &[Task::NER], // Only structured entities
408        "heuristic" | "HeuristicNER" => &[Task::NER],
409        "stacked" | "StackedNER" => &[Task::NER],
410        "crf" | "CrfNER" => &[Task::NER],
411        "hmm" | "HmmNER" => &[Task::NER],
412        "ensemble" | "EnsembleNER" => &[Task::NER],
413        "heuristic_crf" | "heuristic-crf" | "HeuristicCrfNER" => &[Task::NER],
414
415        // ML-based NER backends (all implement Model)
416        "bert_onnx" | "BertNEROnnx" => &[Task::NER],
417        "candle_ner" | "CandleNER" => &[Task::NER],
418        "nuner" | "NuNER" | "nuner_4k" | "nunerzero4k" => &[Task::NER], // Also implements ZeroShotNER
419        "b2ner" | "B2NER" => &[Task::NER],
420        "deberta_v3" | "DeBERTaV3NER" => &[Task::NER],
421        "albert" | "ALBERTNER" => &[Task::NER],
422
423        // Zero-shot NER backends (implement Model + ZeroShotNER)
424        "gliner_onnx" | "GLiNEROnnx" => &[Task::NER],
425        "gliner_candle" | "GLiNERCandle" => &[Task::NER],
426        "gliner_poly" | "GLiNERPoly" => &[Task::NER],
427        "gliner_pii" | "pii_ml" => &[Task::NER], // PII entity types
428        "gliner_relex" | "relex" => &[Task::NER, Task::RelationExtraction],
429        "universal_ner" | "UniversalNER" => &[Task::NER],
430
431        // Multi-task backends (GLiNER multi-task implements multiple traits)
432        "gliner_multitask"
433        | "GLiNERMultitask"
434        | "GLiNERMultitaskOnnx"
435        | "GLiNERMultitaskCandle" => &[
436            Task::NER,
437            Task::TextClassification,
438            Task::SpeechActClassification,
439            Task::Temporal,
440            Task::HierarchicalExtraction,
441            Task::DiscourseRelations,
442            Task::DiscourseSegmentation,
443            // Treat event trigger extraction as “entities” with event-type labels.
444            // This makes event datasets runnable even before a dedicated event pipeline exists.
445            Task::EventExtraction,
446            Task::RelationExtraction, // Via RelationExtractor trait
447        ],
448
449        // Discontinuous NER backends (implement DiscontinuousNER trait)
450        "w2ner" | "W2NER" => &[Task::NER, Task::DiscontinuousNER],
451
452        // Joint entity-relation backends
453        // TPLinker is a relation extraction model. It is not a general NER backend.
454        "tplinker" | "TPLinker" => &[Task::RelationExtraction],
455
456        // Neural coreference backends (implement CorefBackend -- text-based)
457        "fcoref" | "f-coref" => &[Task::IntraDocCoref],
458
459        // Coreference backends (implement CoreferenceResolver trait)
460        //
461        // Note: the same resolver interface is used for both intra-doc and inter-doc eval in
462        // `TaskEvaluator` (it is invoked per document, and spans are offset to avoid collisions).
463        "coref_resolver" | "CorefResolver" | "SimpleCorefResolver" | "DiscourseAwareResolver" => &[
464            Task::IntraDocCoref,
465            Task::InterDocCoref,
466            Task::AbstractAnaphora,
467        ],
468        "mention_ranking" | "MentionRankingCoref" => &[Task::IntraDocCoref, Task::InterDocCoref],
469
470        _ => &[],
471    }
472}
473
474/// Runtime capability detection for a backend instance.
475///
476/// Uses backend name to determine capabilities (fallback when type_id isn't available).
477/// For more accurate detection, use `detect_backend_capabilities_by_name` with the backend's name.
478///
479/// Note: Full runtime detection with trait objects is limited by Rust's type system.
480/// A capability registry pattern would be more reliable but requires backend registration.
481pub fn detect_backend_capabilities(backend: &dyn crate::Model) -> Vec<Task> {
482    // Use the backend's name to determine capabilities
483    // This is a pragmatic approach given Rust's trait object limitations
484    let backend_name = backend.name();
485    detect_backend_capabilities_by_name(backend_name)
486}
487
488/// Capability detection using backend name (fallback when type_id isn't available).
489///
490/// This is less accurate than `detect_backend_capabilities` but works with trait objects.
491pub fn detect_backend_capabilities_by_name(backend_name: &str) -> Vec<Task> {
492    backend_tasks(backend_name).to_vec()
493}
494
495/// Get all tasks that a dataset supports.
496pub fn get_dataset_tasks(dataset: DatasetId) -> Vec<Task> {
497    dataset_tasks(dataset)
498}
499
500/// Get all datasets suitable for a task.
501pub fn get_task_datasets(task: Task) -> Vec<DatasetId> {
502    // Defensive: keep the curated task→datasets list, but only return datasets that
503    // *actually* declare support for the task in the registry metadata.
504    //
505    // This prevents the benchmark/matrix from turning red due to mapping drift
506    // (e.g., dataset listed for a task but missing `tasks: [...]` in the registry).
507    task_datasets(task)
508        .iter()
509        .copied()
510        .filter(|d| dataset_tasks(*d).contains(&task))
511        .collect()
512}
513
514/// Get all backends that support a task.
515///
516/// For benchmarking, only returns "stacked" (which combines pattern+heuristic)
517/// and ML backends, since individual pattern/heuristic backends are incomplete.
518pub fn get_task_backends(task: Task) -> Vec<&'static str> {
519    let mut backends = Vec::new();
520    for backend in [
521        // Builtins
522        "stacked",
523        "crf",
524        "hmm",
525        "heuristic",
526        "ensemble",
527        "heuristic_crf",
528        // Pattern is intentionally excluded from NER eval by default (structured-only),
529        // but it still advertises NER capability via `backend_tasks` and can be enabled by callers.
530        // "pattern",
531        // ML backends
532        "bert_onnx",
533        "candle_ner",
534        "nuner",
535        "nuner_4k",
536        "b2ner",
537        "gliner_onnx",
538        "gliner_candle",
539        "gliner_multitask",
540        "gliner_pii",
541        "gliner_relex",
542        "w2ner",
543        // New backends
544        "tplinker",
545        "gliner_poly",
546        "deberta_v3",
547        "albert",
548        "universal_ner",
549        // Special backends
550        "coref_resolver",
551        "mention_ranking",
552        "fcoref",
553    ] {
554        if backend_tasks(backend).contains(&task) {
555            backends.push(backend);
556        }
557    }
558    backends
559}
560
561/// Comprehensive task-dataset-backend mapping.
562#[derive(Debug, Clone, Serialize, Deserialize)]
563pub struct TaskMapping {
564    /// Tasks → Datasets
565    pub task_to_datasets: HashMap<String, Vec<String>>,
566    /// Datasets → Tasks
567    pub dataset_to_tasks: HashMap<String, Vec<String>>,
568    /// Backends → Tasks
569    pub backend_to_tasks: HashMap<String, Vec<String>>,
570    /// Tasks → Backends
571    pub task_to_backends: HashMap<String, Vec<String>>,
572}
573
574impl TaskMapping {
575    /// Build a complete mapping from all available data.
576    pub fn build() -> Self {
577        let mut task_to_datasets = HashMap::new();
578        let mut dataset_to_tasks = HashMap::new();
579        let mut backend_to_tasks = HashMap::new();
580        let mut task_to_backends = HashMap::new();
581
582        // Build task → datasets
583        for task in Task::all() {
584            let datasets = get_task_datasets(*task)
585                .iter()
586                .map(|d| format!("{:?}", d))
587                .collect();
588            task_to_datasets.insert(task.code().to_string(), datasets);
589        }
590
591        // Build dataset → tasks
592        for dataset in DatasetId::all() {
593            let tasks = get_dataset_tasks(*dataset)
594                .iter()
595                .map(|t| t.code().to_string())
596                .collect();
597            dataset_to_tasks.insert(format!("{:?}", dataset), tasks);
598        }
599
600        // Build backend → tasks
601        for backend in [
602            "pattern",
603            "heuristic",
604            "stacked",
605            "crf",
606            "hmm",
607            "ensemble",
608            "heuristic_crf",
609            "bert_onnx",
610            "candle_ner",
611            "nuner",
612            "deberta_v3",
613            "albert",
614            "gliner_onnx",
615            "gliner_candle",
616            "gliner_poly",
617            "gliner_multitask",
618            "universal_ner",
619            "w2ner",
620            "tplinker",
621            "coref_resolver",
622            "mention_ranking",
623        ] {
624            let tasks = backend_tasks(backend)
625                .iter()
626                .map(|t| t.code().to_string())
627                .collect();
628            backend_to_tasks.insert(backend.to_string(), tasks);
629        }
630
631        // Build task → backends
632        for task in Task::all() {
633            let backends: Vec<String> = get_task_backends(*task)
634                .iter()
635                .map(|s| s.to_string())
636                .collect();
637            task_to_backends.insert(task.code().to_string(), backends);
638        }
639
640        Self {
641            task_to_datasets,
642            dataset_to_tasks,
643            backend_to_tasks,
644            task_to_backends,
645        }
646    }
647
648    /// Get datasets for a task.
649    pub fn datasets_for_task(&self, task: &str) -> Option<&Vec<String>> {
650        self.task_to_datasets.get(task)
651    }
652
653    /// Get tasks for a dataset.
654    pub fn tasks_for_dataset(&self, dataset: &str) -> Option<&Vec<String>> {
655        self.dataset_to_tasks.get(dataset)
656    }
657
658    /// Get tasks for a backend.
659    pub fn tasks_for_backend(&self, backend: &str) -> Option<&Vec<String>> {
660        self.backend_to_tasks.get(backend)
661    }
662
663    /// Get backends for a task.
664    pub fn backends_for_task(&self, task: &str) -> Option<&Vec<String>> {
665        self.task_to_backends.get(task)
666    }
667}
668
669#[cfg(test)]
670mod tests {
671    use super::*;
672
673    #[test]
674    fn test_task_mapping() {
675        let mapping = TaskMapping::build();
676        assert!(mapping.datasets_for_task("ner").is_some());
677        assert!(mapping.tasks_for_dataset("WikiGold").is_some());
678        assert!(mapping.tasks_for_backend("gliner_multitask").is_some());
679        assert!(mapping.backends_for_task("ner").is_some());
680    }
681
682    #[test]
683    fn test_gliner_multitask_capabilities() {
684        let tasks = backend_tasks("gliner_multitask");
685        assert!(tasks.contains(&Task::NER));
686        assert!(tasks.contains(&Task::TextClassification));
687        assert!(tasks.contains(&Task::HierarchicalExtraction));
688        assert!(tasks.contains(&Task::RelationExtraction));
689    }
690
691    #[test]
692    fn test_dataset_tasks_are_deduced_from_registry() {
693        // Basic sanity: we should be able to derive at least one eval task from
694        // the registry for a known dataset with explicit tasks.
695        let tasks = dataset_tasks(DatasetId::WikiGold);
696        assert!(tasks.contains(&Task::NER));
697    }
698
699    #[test]
700    fn test_registry_task_codes_are_parseable() {
701        // Ensures the task strings in the dataset registry don't silently drift away
702        // from our enum mapping (which would make task selection/eval inconsistent).
703        //
704        // We only enforce this for *in-scope* task codes. The registry currently catalogs
705        // some out-of-scope/non-text tasks (e.g., audio captioning) that we do not intend
706        // to model in the eval task enum.
707        fn is_in_scope_task_code(task_str: &str) -> bool {
708            matches!(
709                task_str.to_lowercase().as_str(),
710                // NER family
711                "ner"
712                    | "sequence_labeling"
713                    | "nested-ner"
714                    | "mner"
715                    | "pii_detection"
716                    | "slot_filling"
717                    // Coref family
718                    | "coref"
719                    | "intra-coref"
720                    | "intra_coref"
721                    | "intracoref"
722                    | "inter-coref"
723                    | "inter_coref"
724                    | "intercoref"
725                    | "cdcr"
726                    | "event_coref"
727                    // RE / events
728                    | "re"
729                    | "relation"
730                    | "relation_extraction"
731                    | "relation-extraction"
732                    | "events"
733                    | "event"
734                    | "event_extraction"
735                    | "event-extraction"
736                    | "ee"
737                    // Discourse / anaphora
738                    | "abstract-anaphora"
739                    | "abstract_anaphora"
740                    | "bridging"
741                    | "discourse_deixis"
742                    | "discourse_relations"
743                    | "discourse-relations"
744                    | "discourse_coherence"
745                    | "discourse-coherence"
746                    | "discourse_segmentation"
747                    | "discourse-segmentation"
748                    | "speech_act_classification"
749                    | "speech-act-classification"
750                    // Classification-ish (text-only; not core scope but supported in Task enum)
751                    | "classification"
752                    | "text_classification"
753                    | "text-classification"
754                    | "tc"
755                    | "bias_evaluation"
756                    | "qa"
757                    | "harmonic_analysis"
758                    | "sentiment"
759                    | "sentiment_analysis"
760                    | "sentiment-analysis"
761                    | "pos"
762                    | "pos_tagging"
763                    | "pos-tagging"
764                    // MT/LM (text-only but not evaluated end-to-end yet)
765                    | "mt"
766                    | "machine_translation"
767                    | "machine-translation"
768                    | "translation"
769                    | "lm"
770                    | "language_modeling"
771                    | "language-modeling"
772            )
773        }
774
775        for ds in DatasetId::all() {
776            for &task_str in ds.tasks() {
777                if !is_in_scope_task_code(task_str) {
778                    continue;
779                }
780                assert!(
781                    Task::from_code(task_str).is_some(),
782                    "Unrecognized task code '{}' on dataset {:?}",
783                    task_str,
784                    ds
785                );
786            }
787        }
788    }
789}
anno_eval/eval/task_mapping.rs

anno_eval/eval/
task_mapping.rs
