Skip to main content

anno_eval/eval/
loader.rs

1//! Dataset downloading and caching for NER evaluation.
2//!
3//! Downloads, caches, and parses real NER datasets from public sources.
4//! Follows burntsushi's philosophy: real-world data, not toy examples.
5//!
6//! # Quick Start
7//!
8//! ```rust,ignore
9//! use anno_eval::eval::{DatasetLoader, LoadableDatasetId};
10//!
11//! let loader = DatasetLoader::new()?;
12//! let dataset = loader.load(LoadableDatasetId::WikiGold)?;
13//! println!("Loaded {} sentences", dataset.len());
14//! ```
15//!
16//! # Dataset IDs: catalog vs loadable
17//!
18//! The full dataset catalog lives in [`dataset_registry::DatasetId`]. Not every dataset
19//! in the catalog can be downloaded/parsed by this crate (some require licenses or have
20//! unimplemented formats).
21//!
22//! This module provides [`LoadableDatasetId`], a wrapper that guarantees a dataset has
23//! a loading implementation. Use it for `DatasetLoader::{load, load_or_download}`.
24//!
25//! # Supported Datasets
26//!
27//! | Dataset | Source | License | Entities |
28//! |---------|--------|---------|----------|
29//! | WikiGold | Wikipedia | CC-BY | PER, LOC, ORG, MISC |
30//! | WNUT-17 | Social Media | Open | person, location, corporation, etc. |
31//! | MIT Movie | MIT | Research | actor, director, genre, title, etc. |
32//! | MIT Restaurant | MIT | Research | amenity, cuisine, dish, etc. |
33//! | CoNLL-2003 Sample | Public | Research | PER, LOC, ORG, MISC |
34//! | OntoNotes Sample | Public | Research | 18 entity types |
35//! | BC5CDR | PubMed | Research | Disease, Chemical |
36//! | NCBI Disease | PubMed | Research | Disease |
37//!
38//! # Design Philosophy
39//!
40//! - **Lazy downloading**: Only fetch what's needed
41//! - **Persistent caching**: Never re-download unchanged data
42//! - **Integrity verification**: SHA256 checksums for all downloads
43//! - **Graceful degradation**: Work offline with cached data
44//! - **Clear errors**: Explain exactly what went wrong
45//!
46//! # Extended Example
47//!
48//! ```rust,ignore
49//! use anno_eval::eval::{DatasetLoader, LoadableDatasetId};
50//!
51//! let loader = DatasetLoader::new()?;
52//!
53//! // Check cache status before loading
54//! if loader.is_cached(LoadableDatasetId::WikiGold) {
55//!     println!("WikiGold is cached, will load from disk");
56//! }
57//!
58//! // Load dataset (downloads if not cached, verifies checksum)
59//! let dataset = loader.load(LoadableDatasetId::WikiGold)?;
60//! println!("Loaded {} sentences with {} entities",
61//!     dataset.len(), dataset.entity_count());
62//!
63//! // Iterate over examples
64//! for example in dataset.iter() {
65//!     println!("Text: {}", example.text);
66//!     for entity in &example.entities {
67//!         println!("  {} [{}]", entity.text, entity.entity_type);
68//!     }
69//! }
70//! ```
71//!
72//! [`dataset_registry::DatasetId`]: super::dataset_registry::DatasetId
73
74#[cfg(test)]
75use anno::EntityType;
76use anno::{Error, Result};
77use serde::{Deserialize, Serialize};
78use std::collections::HashMap;
79use std::fs;
80use std::path::PathBuf;
81
82// =============================================================================
83// Dataset Identification
84// =============================================================================
85
86/// Dataset identifier (full catalog).
87///
88/// This is the single source of truth for dataset metadata. Not all datasets are
89/// loadable by `DatasetLoader`. Use [`LoadableDatasetId`] when you want that guarantee.
90///
91/// # Usage
92///
93/// ```rust,ignore
94/// use anno_eval::eval::{DatasetLoader, LoadableDatasetId};
95///
96/// let loader = DatasetLoader::new()?;
97/// let dataset = loader.load(LoadableDatasetId::WikiGold)?;
98/// ```
99///
100/// # Available Datasets
101///
102/// ## NER Datasets
103///
104/// | Dataset | Size | Domain | Entity Types |
105/// |---------|------|--------|--------------|
106/// | `WikiGold` | ~3.5k entities | Wikipedia | PER, LOC, ORG, MISC |
107/// | `Wnut17` | ~2k entities | Social media | person, location, etc. |
108/// | `MitMovie` | ~10k entities | Movies | actor, director, genre |
109/// | `MitRestaurant` | ~8k entities | Restaurants | cuisine, dish, etc. |
110/// | `CoNLL2003Sample` | ~20k entities | News | PER, LOC, ORG, MISC |
111/// | `OntoNotesSample` | ~18k entities | Mixed | 18 types |
112/// | `BC5CDR` | ~28k entities | Biomedical | Disease, Chemical |
113/// | `NCBIDisease` | ~6k entities | Biomedical | Disease |
114///
115/// ## Coreference Datasets
116///
117/// | Dataset | Size | Domain | Features |
118/// |---------|------|--------|----------|
119/// | `GAP` | 8,908 pairs | Wikipedia | Gender-balanced pronouns |
120/// | `PreCo` | 38k docs | Reading | Includes singletons |
121/// | `LitBank` | 100 works | Literature | Literary coreference |
122///
123/// # Extending
124///
125/// To add a new loadable dataset:
126/// 1. Add the variant here
127/// 2. Implement loading in `DatasetLoader::load()`
128/// 3. Add to `DatasetId::all()` iterator
129/// 4. Ensure it exists in `dataset_registry::DatasetId` (metadata catalog)
130///
131pub use super::dataset_registry::DatasetId;
132
133/// Dataset identifier guaranteed to be loadable by [`DatasetLoader`].
134///
135/// This is a thin wrapper over [`DatasetId`] that enforces loader support via
136/// `TryFrom<DatasetId>` / `FromStr`.
137#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
138pub struct LoadableDatasetId(DatasetId);
139
140#[derive(Debug, Copy, Clone, PartialEq, Eq)]
141enum DatasetParsePlan {
142    Conll,
143    JsonlNer,
144    WikiannJson,
145    TweetNer7,
146    DocredJson,
147    GoogleReCorpus,
148    ChisiecJson,
149    CadecHybrid,
150    Bc5cdr,
151    NcbiDisease,
152    GapTsv,
153    PrecoJsonl,
154    Litbank,
155    EcbPlus,
156    AfriSenti,
157    AfriQa,
158    MasakhaNews,
159    Conllu,
160    AgNews,
161    Dbpedia14,
162    YahooAnswers,
163    Trec,
164    TweetTopic,
165    Maven,
166    MavenArg,
167    Casie,
168    Rams,
169    HfApiResponse,
170    /// TSV-based NER format (HIPE-2022 style)
171    TsvNer,
172    /// CSV-based NER format (E-NER/EDGAR-NER style: Token,Tag)
173    CsvNer,
174}
175
176impl LoadableDatasetId {
177    /// All datasets that have loading implementations.
178    #[must_use]
179    pub fn all() -> Vec<Self> {
180        DatasetId::all()
181            .iter()
182            .copied()
183            .filter(|id| Self::is_loadable_dataset(*id))
184            .map(Self)
185            .collect()
186    }
187
188    /// Access the underlying catalog ID.
189    #[must_use]
190    pub fn into_inner(self) -> DatasetId {
191        self.0
192    }
193
194    #[must_use]
195    fn is_loadable_dataset(id: DatasetId) -> bool {
196        // Loader truth: some catalog entries are intentionally non-loadable (blocked, deprecated,
197        // or require manual access) even if they have a parse plan.
198        //
199        // Keep this list small and motivated. Prefer expressing “non-loadable” in the registry
200        // metadata first, then enforcing it here so CLI counts and `LoadableDatasetId` remain
201        // honest.
202        if matches!(
203            id,
204            DatasetId::CoNLL2002 | DatasetId::CoNLL2002Spanish | DatasetId::CoNLL2002Dutch
205        ) {
206            return false;
207        }
208
209        Self::parse_plan(id).is_some()
210    }
211
212    #[must_use]
213    fn parse_plan(id: DatasetId) -> Option<DatasetParsePlan> {
214        if let Some(hint) = Self::registry_hint_plan(id) {
215            return Some(hint);
216        }
217
218        Some(match id {
219            // ===================================================================
220            // CoNLL/BIO format datasets
221            // ===================================================================
222            // These datasets use standard CoNLL-2003 or BIO tagging format.
223            // The CoNLL parser handles: space/tab-separated columns, IOB2/BIO schemes,
224            // and standard entity type mappings (PER, ORG, LOC, etc.).
225            //
226            // Note: Some datasets listed here are also handled by registry_hint_plan()
227            // but are kept for explicit control or historical reasons.
228            //
229            // To add more: If a dataset has format="CoNLL" or format="BIO" in the registry
230            // and is NER-only, it should be auto-detected. Add here only if auto-detection
231            // fails (e.g., multi-task datasets, ambiguous metadata).
232            DatasetId::WikiGold
233            | DatasetId::Wnut17
234            | DatasetId::MitMovie
235            | DatasetId::MitRestaurant
236            | DatasetId::CoNLL2003Sample
237            | DatasetId::OntoNotesSample
238            | DatasetId::UniversalNERBench
239            | DatasetId::LegNER
240            | DatasetId::TwiConv
241            | DatasetId::AIDA
242            | DatasetId::TACKBP
243            | DatasetId::BroadTwitterCorpus
244            | DatasetId::BioMNER
245            | DatasetId::MasakhaNER
246            | DatasetId::MasakhaNER2
247            | DatasetId::OntoNotes50
248            | DatasetId::GermEval2014
249            | DatasetId::HAREM
250            | DatasetId::SemEval2013Task91
251            | DatasetId::MUC6
252            | DatasetId::MUC7
253            | DatasetId::JNLPBA
254            | DatasetId::BC2GMFull
255            | DatasetId::CRAFT
256            | DatasetId::FinNER
257            | DatasetId::LegalNER
258            | DatasetId::AGRONER
259            | DatasetId::ATISFlightBooking
260            | DatasetId::AerospaceNERDataset
261            | DatasetId::AstrologyNER
262            | DatasetId::AutomotiveNER
263            | DatasetId::AviationProductsNER
264            | DatasetId::BeekeepingNER
265            | DatasetId::BrewingNER
266            | DatasetId::ChineseEngineeringGeologyNER
267            | DatasetId::CocktailNER
268            | DatasetId::ConstructionNER
269            | DatasetId::CropDiseaseNER
270            | DatasetId::CryptoNER
271            | DatasetId::CyNERAptner
272            | DatasetId::DnDNERBenchmark
273            | DatasetId::EFGC
274            | DatasetId::EnergyNER
275            | DatasetId::EquestrianNER
276            | DatasetId::EsportsNER
277            | DatasetId::FINERFood
278            | DatasetId::FitnessNER
279            | DatasetId::FourRegionsGeologyNER
280            | DatasetId::GNERGeoscience
281            | DatasetId::GardeningNER
282            | DatasetId::HealthcareAdminNER
283            | DatasetId::HindiEnglishSocialMediaNER
284            | DatasetId::JobPostingNER
285            | DatasetId::LogisticsNER
286            | DatasetId::ArabicEventCoref
287            | DatasetId::FrenchFullLengthFictionCoref
288            | DatasetId::LongtoNotes
289            | DatasetId::ManufacturingNER
290            | DatasetId::MaritimeNER
291            | DatasetId::MovieCoref
292            | DatasetId::MusicNER
293            | DatasetId::NDNER
294            | DatasetId::OrigamiNER
295            | DatasetId::PaleontologyNER
296            | DatasetId::PharmaNER
297            | DatasetId::PhotographyNER
298            | DatasetId::RealEstateNER
299            | DatasetId::RitterTwitterNER
300            | DatasetId::SECFilingsNER
301            | DatasetId::SanskritNERBhagavadGita
302            | DatasetId::ScubaNER
303            | DatasetId::SportsNERGeneral
304            | DatasetId::TVShowMultilingualCoref
305            | DatasetId::TarotNER
306            | DatasetId::TelecomNER
307            | DatasetId::TelenovelaNER
308            | DatasetId::TourismNER
309            | DatasetId::VeterinaryNER
310            | DatasetId::WaterResourceNER
311            | DatasetId::WeatherNER
312            | DatasetId::WineNER
313            | DatasetId::WoodworkingNER
314            // Added 2025-12: More CoNLL-format datasets with direct download URLs
315            | DatasetId::QxoRef
316            | DatasetId::GICoref
317            | DatasetId::WNUT16
318            | DatasetId::NoiseBench
319            | DatasetId::CrossWeigh
320            | DatasetId::ZELDA
321            | DatasetId::GENIANested => DatasetParsePlan::Conll,
322
323            // ===================================================================
324            // JSONL format (HuggingFace style)
325            // ===================================================================
326            // These datasets use JSON Lines format with standard fields:
327            // - `tokens`: Array of token strings
328            // - `ner_tags`: Array of BIO/IOB2 tags
329            // - Optional: `text` (original sentence), `entities` (structured)
330            //
331            // The JsonlNer parser is flexible and handles variations in field names.
332            DatasetId::MultiNERD
333            | DatasetId::ACORD
334            | DatasetId::ARFFiction
335            | DatasetId::AgMNER
336            | DatasetId::AgriNER
337            | DatasetId::AnimeMangaNER
338            | DatasetId::AntiquesNER
339            | DatasetId::AstronomicalTelegramKEE
340            | DatasetId::BirdwatchingNER
341            | DatasetId::BoardGameNER
342            | DatasetId::BookCorefBamman
343            | DatasetId::CUAD
344            | DatasetId::ChessNER
345            | DatasetId::CoMTA
346            | DatasetId::DeepFashion2
347            | DatasetId::FIREBALL
348            | DatasetId::FashionIQ
349            | DatasetId::FragranceNER
350            | DatasetId::IMDbSemiStructuredRE
351            | DatasetId::InsuranceNER
352            | DatasetId::KnittingNER
353            | DatasetId::LLMRocMinNER
354            | DatasetId::MOFDataset
355            | DatasetId::MathDial
356            | DatasetId::NHKRecipeDataset
357            | DatasetId::NaturalProductsRE
358            | DatasetId::NumismaticsNER
359            | DatasetId::PhilatelyNER
360            | DatasetId::PolyIE
361            | DatasetId::RecipeDBAnnotated
362            | DatasetId::ResumeNER
363            | DatasetId::RetailInventoryNER
364            | DatasetId::SPoRC
365            | DatasetId::Saraga
366            | DatasetId::SolidStateDoping
367            | DatasetId::SpotifyPodcastsDataset
368            | DatasetId::TattooNER
369            | DatasetId::ThemeParkNER
370            | DatasetId::VREN
371            | DatasetId::VisDialCoref
372            // Added 2025-12: More JSONL-format datasets with direct download URLs
373            | DatasetId::REBEL
374            | DatasetId::BBQ
375            | DatasetId::RealToxicityPrompts
376            | DatasetId::BookCoref
377            | DatasetId::BookCorefSplit
378            | DatasetId::WIESP2022NER
379            | DatasetId::FewRel
380            | DatasetId::PIIMasking200k
381            | DatasetId::B2NERD
382            | DatasetId::OpenNER
383            | DatasetId::FictionNER750M => DatasetParsePlan::JsonlNer,
384
385            // WikiANN-ish JSON format (from download_hf_datasets.py)
386            DatasetId::UNER | DatasetId::MSNER => DatasetParsePlan::WikiannJson,
387
388            // TweetNER7 uses JSON format with integer tags
389            DatasetId::TweetNER7 => DatasetParsePlan::TweetNer7,
390
391            // ===================================================================
392            // JSON format (Relation extraction datasets)
393            // ===================================================================
394            // Relation extraction datasets use JSON format with entity pairs and relations.
395            // Common structure: {entities: [...], relations: [...], text: "..."}
396            //
397            // The DocredJson parser handles DocRED, ReTACRED, and similar formats.
398            // Some datasets (SciERC, PolyIE, EnzChemRED) are auto-detected via
399            // registry_hint_plan() but listed here for explicit control.
400            //
401            // To add more: RE datasets with format="JSON" or format="JSONL" and is_re=true
402            // should be auto-detected. Add here only for special cases.
403            DatasetId::DocRED
404            | DatasetId::ReTACRED
405            | DatasetId::NYTFB
406            | DatasetId::WEBNLG
407            | DatasetId::GoogleRE
408            | DatasetId::BioRED
409            | DatasetId::SciER
410            | DatasetId::MixRED
411            | DatasetId::CovEReD
412            | DatasetId::ACE2005
413            | DatasetId::MuDoCo
414            | DatasetId::SciERCNER
415            // Added 2025-01-27: RE datasets with public URLs but missing hints
416            // Note: SciERC, PolyIE, EnzChemRED are now handled by registry_hint_plan() auto-detection
417            => DatasetParsePlan::DocredJson,
418
419            // CHisIEC: Ancient Chinese historical NER+RE (custom JSON format)
420            DatasetId::CHisIEC => DatasetParsePlan::ChisiecJson,
421
422            // Discontinuous NER (HF datasets-server API or JSONL format)
423            DatasetId::CADEC | DatasetId::ShARe13 | DatasetId::ShARe14 => {
424                DatasetParsePlan::CadecHybrid
425            }
426
427            // Biomedical formats
428            DatasetId::BC5CDR => DatasetParsePlan::Bc5cdr,
429            DatasetId::NCBIDisease => DatasetParsePlan::NcbiDisease,
430
431            // Coreference formats
432            // Note: GUM uses CoNLL format and is handled by registry_hint_plan()
433            DatasetId::GAP | DatasetId::WikiCoref => DatasetParsePlan::GapTsv,
434            DatasetId::WinoBias => DatasetParsePlan::GapTsv,
435            DatasetId::PreCo | DatasetId::SciCo => DatasetParsePlan::PrecoJsonl,
436            DatasetId::LitBank => DatasetParsePlan::Litbank,
437            DatasetId::ECBPlus => DatasetParsePlan::EcbPlus,
438
439            // African language datasets
440            DatasetId::AfriSenti => DatasetParsePlan::AfriSenti,
441            DatasetId::AfriQA => DatasetParsePlan::AfriQa,
442            DatasetId::MasakhaNEWS => DatasetParsePlan::MasakhaNews,
443            DatasetId::MasakhaPOS => DatasetParsePlan::Conllu,
444
445            // Text classification datasets
446            DatasetId::AGNews => DatasetParsePlan::AgNews,
447            DatasetId::DBPedia14 => DatasetParsePlan::Dbpedia14,
448            DatasetId::YahooAnswers => DatasetParsePlan::YahooAnswers,
449            DatasetId::TREC => DatasetParsePlan::Trec,
450            DatasetId::TweetTopic => DatasetParsePlan::TweetTopic,
451
452            // Event extraction datasets
453            DatasetId::MAVEN => DatasetParsePlan::Maven,
454            DatasetId::MAVENArg => DatasetParsePlan::MavenArg,
455            DatasetId::CASIE => DatasetParsePlan::Casie,
456            DatasetId::RAMS => DatasetParsePlan::Rams,
457
458            // HF datasets-server API responses
459            DatasetId::GENIA
460            | DatasetId::AnatEM
461            | DatasetId::BC2GM
462            | DatasetId::BC4CHEMD
463            | DatasetId::FewNERD
464            | DatasetId::CrossNER
465            | DatasetId::FabNER
466            | DatasetId::WikiNeural
467            | DatasetId::WikiANN
468            | DatasetId::MultiCoNER
469            | DatasetId::MultiCoNERv2
470            | DatasetId::PolyglotNER
471            | DatasetId::UniversalNER => DatasetParsePlan::HfApiResponse,
472
473            _ => return None,
474        })
475    }
476
477    /// Best-effort "hint" derived from registry metadata.
478    ///
479    /// This is intentionally conservative: it returns `None` if the registry
480    /// doesn't provide enough signal to choose the same plan as `parse_plan`.
481    ///
482    /// The long-term goal is to have this cover most datasets so `parse_plan`
483    /// shrinks to a small exception table.
484    ///
485    /// # Architecture Notes
486    ///
487    /// This function enables automatic format detection for datasets with:
488    /// - Clear format metadata (`format:` field in registry)
489    /// - Single primary task (NER, RE, coref, etc.)
490    /// - Standard annotation schemes (BIO, IOB2, CoNLL, JSONL)
491    ///
492    /// Datasets with multiple tasks (e.g., NER+RE) or ambiguous formats
493    /// should be added to explicit matches in `parse_plan()` instead.
494    ///
495    /// # Adding New Datasets
496    ///
497    /// 1. **Common formats (CoNLL, JSONL, TSV)**: Add format metadata to registry,
498    ///    ensure single primary task, and this function will auto-detect.
499    /// 2. **Special formats**: Add explicit match in `parse_plan()` match statement.
500    /// 3. **Multi-task datasets**: Add explicit match (auto-detection is conservative).
501    #[must_use]
502    fn registry_hint_plan(id: DatasetId) -> Option<DatasetParsePlan> {
503        // Hybrids / special cases first.
504        if id == DatasetId::CHisIEC {
505            // CHisIEC is multi-task and uses a custom JSON array format (not DocRED/CrossRE).
506            // It MUST not be auto-detected as DocredJson.
507            return Some(DatasetParsePlan::ChisiecJson);
508        }
509        if matches!(
510            id,
511            DatasetId::CADEC | DatasetId::ShARe13 | DatasetId::ShARe14
512        ) {
513            return Some(DatasetParsePlan::CadecHybrid);
514        }
515        if id == DatasetId::GoogleRE {
516            // The Google relation-extraction-corpus files are not DocRED/CrossRE.
517            // They are JSONL records with fields like {pred, sub, obj, evidences[], judgments[]}.
518            return Some(DatasetParsePlan::GoogleReCorpus);
519        }
520        if id == DatasetId::TweetNER7 {
521            return Some(DatasetParsePlan::TweetNer7);
522        }
523
524        // HuggingFace-hosted datasets with bespoke (non-generic) parse plans.
525        //
526        // These are marked `access_status: HuggingFace` in the registry, so they must be
527        // hintable; but their parse plans are not derivable from `format` alone (or use
528        // specialized parsing despite being hosted on the Hub).
529        if id == DatasetId::BC5CDR {
530            return Some(DatasetParsePlan::Bc5cdr);
531        }
532        if id == DatasetId::NCBIDisease {
533            return Some(DatasetParsePlan::NcbiDisease);
534        }
535        if id == DatasetId::AfriSenti {
536            return Some(DatasetParsePlan::AfriSenti);
537        }
538        if id == DatasetId::AfriQA {
539            return Some(DatasetParsePlan::AfriQa);
540        }
541        if id == DatasetId::MasakhaNEWS {
542            return Some(DatasetParsePlan::MasakhaNews);
543        }
544        if id == DatasetId::AGNews {
545            return Some(DatasetParsePlan::AgNews);
546        }
547        if id == DatasetId::DBPedia14 {
548            return Some(DatasetParsePlan::Dbpedia14);
549        }
550        if id == DatasetId::YahooAnswers {
551            return Some(DatasetParsePlan::YahooAnswers);
552        }
553        if id == DatasetId::TREC {
554            return Some(DatasetParsePlan::Trec);
555        }
556        if id == DatasetId::TweetTopic {
557            return Some(DatasetParsePlan::TweetTopic);
558        }
559
560        // Coref datasets in CoNLL format (not GAP TSV): parse as regular CoNLL.
561        // These have NER annotations alongside coref, so the CoNLL parser works.
562        if matches!(
563            id,
564            DatasetId::QxoRef
565                | DatasetId::GICoref
566                | DatasetId::WNUT16
567                | DatasetId::NoiseBench
568                | DatasetId::CrossWeigh
569                | DatasetId::ZELDA
570                | DatasetId::GENIANested
571                // Added 2025-12: More CoNLL NER datasets
572                | DatasetId::HistNERo
573                | DatasetId::DutchArchaeology
574                | DatasetId::FINER
575                | DatasetId::CALCS2018
576                | DatasetId::MedievalCharterNER
577                | DatasetId::RockNER
578                | DatasetId::AIDACoNLL
579                | DatasetId::NNE
580                | DatasetId::GermEvalDiscontinuous
581                | DatasetId::PubMedDiscontinuous
582                | DatasetId::IndicNER
583                | DatasetId::NorNE
584                | DatasetId::CHEMDNER
585                | DatasetId::GeoWebNews
586                | DatasetId::TASTEset
587                | DatasetId::RecipeNER
588                | DatasetId::AstroNER
589                | DatasetId::FinanceNER
590                | DatasetId::TechNER
591                | DatasetId::CALCS
592                | DatasetId::LinCE
593                | DatasetId::FinTechPatent
594                | DatasetId::WaterAgriNER
595                | DatasetId::NERsocialFood
596                | DatasetId::RussianCulturalNER
597                | DatasetId::EighteenthCenturyNER
598                | DatasetId::GuaraniNER
599                | DatasetId::ShipiboKoniboNER
600                | DatasetId::BASHI
601                | DatasetId::ENER
602                // Added 2025-01-27: CoNLL NER datasets with public URLs but missing hints
603                | DatasetId::OntoNotes50
604                | DatasetId::GUM
605                // CoNLL04RE is CoNLL format (works with CoNLL parser even though it's for RE)
606                | DatasetId::CoNLL04RE
607        ) {
608            return Some(DatasetParsePlan::Conll);
609        }
610
611        // JSONL datasets that should use JsonlNer despite having coref/RE tasks.
612        // The JsonlNer parser handles tokens + ner_tags fields generically.
613        if matches!(
614            id,
615            DatasetId::REBEL
616                | DatasetId::BBQ
617                | DatasetId::RealToxicityPrompts
618                | DatasetId::BookCoref
619                | DatasetId::BookCorefSplit
620                | DatasetId::WIESP2022NER
621                | DatasetId::FewRel
622                | DatasetId::PIIMasking200k
623                | DatasetId::B2NERD
624                | DatasetId::OpenNER
625                | DatasetId::FictionNER750M
626                // Added 2025-12: More JSONL NER datasets
627                | DatasetId::MultiWOZNER
628                | DatasetId::HinglishNER
629                | DatasetId::ChineseNestedNER
630                | DatasetId::AgCNER
631                | DatasetId::LongDocNER
632                | DatasetId::MultiBioNERLong
633                | DatasetId::ReasoningNER
634                | DatasetId::BioNERLLaMA
635                | DatasetId::LexGLUENER
636                | DatasetId::FinBenNER
637                | DatasetId::FiNER139
638                | DatasetId::SciNER
639                | DatasetId::CharacterCodex
640                | DatasetId::AIONER
641                | DatasetId::WIESPAstro
642                | DatasetId::CEREC
643                | DatasetId::DELICATE
644                | DatasetId::CSN
645                // Added 2025-01-27: JSONL NER datasets with public URLs but missing hints
646                | DatasetId::SCINERNested
647                | DatasetId::AgriNER
648                | DatasetId::MOFDataset
649                | DatasetId::SolidStateDoping
650        ) {
651            return Some(DatasetParsePlan::JsonlNer);
652        }
653
654        // CoNLLU datasets (Universal Dependencies treebanks with NER annotations)
655        if matches!(
656            id,
657            DatasetId::AncientGreekUD
658                | DatasetId::LatinUD
659                | DatasetId::SanskritUD
660                | DatasetId::OldEnglishUD
661                | DatasetId::OldNorseUD
662                | DatasetId::UDEsperantoCairo
663                // Added 2025-12: More UD treebanks (ancient/historical/constructed)
664                | DatasetId::CopticScriptorium
665                | DatasetId::TaggedPBCEsperanto
666                | DatasetId::TaggedPBCKlingon
667                | DatasetId::AkkadianUD
668                | DatasetId::AncientHebrewUD
669                | DatasetId::ClassicalChineseUD
670                | DatasetId::CopticUD
671                | DatasetId::GothicUD
672                | DatasetId::HittiteUD
673                | DatasetId::OldChurchSlavonicUD
674                | DatasetId::LatinITTB
675                | DatasetId::LatinPROIEL
676                | DatasetId::EsperantoUD
677                | DatasetId::NavajoMorph
678        ) {
679            return Some(DatasetParsePlan::Conllu);
680        }
681
682        // TSV NER datasets (HIPE-2022 style)
683        if matches!(id, DatasetId::HIPE2022) {
684            return Some(DatasetParsePlan::TsvNer);
685        }
686
687        // CSV NER datasets
688        if matches!(id, DatasetId::ENer) {
689            return Some(DatasetParsePlan::CsvNer);
690        }
691
692        // Some datasets are fetched via HF datasets-server (JSON response), regardless of their
693        // “canonical” source format in the registry.
694        if matches!(
695            id,
696            DatasetId::GENIA
697                | DatasetId::AnatEM
698                | DatasetId::BC2GM
699                | DatasetId::BC4CHEMD
700                | DatasetId::FewNERD
701                | DatasetId::CrossNER
702                | DatasetId::FabNER
703                | DatasetId::WikiNeural
704                | DatasetId::WikiANN
705                | DatasetId::MultiCoNER
706                | DatasetId::MultiCoNERv2
707                | DatasetId::PolyglotNER
708                | DatasetId::UniversalNER
709        ) {
710            return Some(DatasetParsePlan::HfApiResponse);
711        }
712
713        // Use tasks_or_inferred() to support datasets that have task-like categories
714        // (e.g., `categories: [ner, ancient]`) but no explicit `tasks` field yet.
715        // This enables automatic loading for datasets with correct format + category.
716        let inferred_tasks = id.tasks_or_inferred();
717        let is_ner = inferred_tasks.contains(&"ner");
718        let is_coref = inferred_tasks.contains(&"coref");
719        let is_re = inferred_tasks.contains(&"re");
720        let is_event = inferred_tasks.contains(&"event_extraction");
721
722        let annotation_scheme = id.annotation_scheme().unwrap_or("");
723
724        // Format string is about the *source* dataset, not necessarily our cached
725        // representation; still useful for many direct-download datasets.
726        if let Some(format) = id.format() {
727            let plan = match format {
728                "CoNLL" | "BIO" | "IOB2" if is_ner && !is_coref && !is_re && !is_event => {
729                    Some(DatasetParsePlan::Conll)
730                }
731                "CoNLL-U" | "CoNLLU" if is_ner && !is_coref && !is_re && !is_event => {
732                    Some(DatasetParsePlan::Conllu)
733                }
734                "JSONL" if is_ner && !is_coref && !is_re && !is_event => {
735                    Some(DatasetParsePlan::JsonlNer)
736                }
737                // Coreference: we only treat these as unambiguous when the dataset is *only* coref.
738                // Check both format field and annotation_scheme for CoNLLCoref datasets.
739                "TSV" | "ConllCoref" | "CoNLLCoref"
740                    if is_coref && !is_ner && !is_re && !is_event =>
741                {
742                    // GAP-style TSV and similar "coref-only" exports.
743                    Some(DatasetParsePlan::GapTsv)
744                }
745                // CoNLL format with CoNLLCoref annotation scheme (e.g., GICoref, qxoRef)
746                "CoNLL"
747                    if annotation_scheme == "CoNLLCoref"
748                        && is_coref
749                        && !is_ner
750                        && !is_re
751                        && !is_event =>
752                {
753                    Some(DatasetParsePlan::GapTsv)
754                }
755                "JSONL" if is_coref && !is_ner && !is_re && !is_event => {
756                    // PreCo/SciCo style JSONL exports.
757                    Some(DatasetParsePlan::PrecoJsonl)
758                }
759
760                // Relation extraction: unambiguous only when the dataset is *only* RE.
761                // Also handle RE datasets that might have NER as secondary task (e.g., SciERCNER).
762                "JSON" | "JSONL" if is_re && !is_coref && !is_event => {
763                    Some(DatasetParsePlan::DocredJson)
764                }
765
766                // Event extraction: treat JSONL as MAVEN-ish only when the dataset is event-only.
767                "JSONL" if is_event && !is_ner && !is_coref && !is_re => {
768                    Some(DatasetParsePlan::Maven)
769                }
770
771                // TSV NER format (HIPE-2022 style)
772                "TSV" if is_ner && !is_coref && !is_re && !is_event => {
773                    Some(DatasetParsePlan::TsvNer)
774                }
775
776                // CSV NER format (E-NER/EDGAR-NER style: Token,Tag)
777                "CSV" if is_ner && !is_coref && !is_re && !is_event => {
778                    Some(DatasetParsePlan::CsvNer)
779                }
780
781                // These are too ambiguous across sources; treat as unknown for now.
782                _ => None,
783            };
784
785            if plan.is_some() {
786                return plan;
787            }
788        }
789
790        // HuggingFace datasets: if the registry provides an HF id/config and marks the
791        // dataset as HuggingFace-accessible, we can treat it as an HfApiResponse plan.
792        //
793        // Keep this as a *fallback*: registry format/task signals (and explicit hint lists
794        // above) take precedence, so we don't override e.g. CoNLL/JSONL datasets that are
795        // hosted on HuggingFace but should be parsed as raw files.
796        if id.hf_id().is_some()
797            && id.access_status()
798                == crate::eval::dataset_registry::DatasetAccessibility::HuggingFace
799        {
800            return Some(DatasetParsePlan::HfApiResponse);
801        }
802
803        None
804    }
805}
806
807impl std::ops::Deref for LoadableDatasetId {
808    type Target = DatasetId;
809
810    fn deref(&self) -> &Self::Target {
811        &self.0
812    }
813}
814
815impl From<LoadableDatasetId> for DatasetId {
816    fn from(value: LoadableDatasetId) -> Self {
817        value.0
818    }
819}
820
821impl std::str::FromStr for LoadableDatasetId {
822    type Err = Error;
823
824    fn from_str(s: &str) -> std::result::Result<Self, Self::Err> {
825        let id: DatasetId = s
826            .parse()
827            .map_err(|e| Error::InvalidInput(format!("Invalid dataset id '{}': {}", s, e)))?;
828        Self::try_from(id)
829    }
830}
831
832impl TryFrom<DatasetId> for LoadableDatasetId {
833    type Error = Error;
834
835    fn try_from(value: DatasetId) -> std::result::Result<Self, Self::Error> {
836        if Self::is_loadable_dataset(value) {
837            Ok(Self(value))
838        } else {
839            Err(Error::InvalidInput(format!(
840                "Dataset {:?} is in the registry but has no loading implementation",
841                value
842            )))
843        }
844    }
845}
846
847// =============================================================================
848// Loader-specific extensions for `DatasetId`
849// =============================================================================
850//
851// `DatasetId` is defined in `dataset_registry` (the full catalog of datasets).
852// This module is intentionally split:
853// - `DatasetId`: metadata/catalog (can include datasets that are not loadable here)
854// - `LoadableDatasetId`: the subset that this loader+parsers support
855
856// Extension methods for DatasetId that delegate to registry
857impl DatasetId {
858    /// Get default metadata for this dataset.
859    ///
860    /// Delegates to registry methods.
861    #[must_use]
862    pub fn default_metadata(&self) -> DatasetMetadata {
863        DatasetMetadata {
864            description: Some(self.description().to_string()),
865            license: self.license().map(|s| s.to_string()),
866            citation: self.citation().map(|s| s.to_string()),
867            split: Some("test".to_string()), // Default to test split
868            language: Some(self.language().to_string()),
869            domain: Some(self.domain().to_string()),
870            original_source: Some(self.download_url().to_string()),
871            version: None,
872            annotators: None,
873            guidelines_url: None,
874        }
875    }
876}
877
878// =============================================================================
879// Cache Manifest
880// =============================================================================
881
882/// Entry in the cache manifest tracking a downloaded dataset file.
883#[derive(Debug, Clone, Serialize, Deserialize)]
884pub struct CacheManifestEntry {
885    /// Dataset identifier.
886    pub dataset_id: String,
887    /// Source URL from which the data was downloaded.
888    ///
889    /// This is the *logical* registry URL (what the dataset entry says).
890    pub source_url: String,
891    /// The URL that was actually fetched (after any rewrites / mirrors).
892    ///
893    /// This is helpful for reproducibility (and debugging) when the registry URL is a
894    /// homepage (e.g., HuggingFace dataset page) but we fetched via datasets-server
895    /// or `resolve/main/...`.
896    #[serde(default)]
897    pub resolved_url: Option<String>,
898    /// SHA-256 checksum of the downloaded file.
899    pub sha256: String,
900    /// File size in bytes.
901    pub file_size: u64,
902    /// ISO 8601 timestamp of when the file was downloaded.
903    pub downloaded_at: String,
904    /// Number of sentences parsed from the dataset.
905    pub sentence_count: usize,
906    /// Number of entities parsed from the dataset.
907    pub entity_count: usize,
908    /// Anno version that downloaded this file.
909    pub anno_version: String,
910}
911
912/// Cache manifest tracking all downloaded datasets.
913///
914/// Stored as `manifest.json` in the cache directory.
915/// Enables:
916/// - Reproducibility: know exactly what was downloaded
917/// - Validation: verify checksums match
918/// - Debugging: trace when data was fetched
919#[derive(Debug, Clone, Default, Serialize, Deserialize)]
920pub struct CacheManifest {
921    /// Version of the manifest format.
922    pub version: u32,
923    /// Entries for each cached dataset.
924    pub entries: HashMap<String, CacheManifestEntry>,
925}
926
927impl CacheManifest {
928    /// Current manifest format version.
929    pub const CURRENT_VERSION: u32 = 1;
930
931    /// Create a new empty manifest.
932    #[must_use]
933    pub fn new() -> Self {
934        Self {
935            version: Self::CURRENT_VERSION,
936            entries: HashMap::new(),
937        }
938    }
939
940    /// Load manifest from a cache directory.
941    ///
942    /// Returns empty manifest if file doesn't exist.
943    pub fn load(cache_dir: &std::path::Path) -> Result<Self> {
944        let manifest_path = cache_dir.join("manifest.json");
945        if !manifest_path.exists() {
946            return Ok(Self::new());
947        }
948
949        let content = fs::read_to_string(&manifest_path)
950            .map_err(|e| Error::InvalidInput(format!("Failed to read manifest: {}", e)))?;
951
952        serde_json::from_str(&content)
953            .map_err(|e| Error::InvalidInput(format!("Failed to parse manifest: {}", e)))
954    }
955
956    /// Save manifest to a cache directory.
957    pub fn save(&self, cache_dir: &std::path::Path) -> Result<()> {
958        let manifest_path = cache_dir.join("manifest.json");
959        let content = serde_json::to_string_pretty(self)
960            .map_err(|e| Error::InvalidInput(format!("Failed to serialize manifest: {}", e)))?;
961
962        fs::write(&manifest_path, content)
963            .map_err(|e| Error::InvalidInput(format!("Failed to write manifest: {}", e)))
964    }
965
966    /// Add or update an entry in the manifest.
967    pub fn update_entry(&mut self, entry: CacheManifestEntry) {
968        self.entries.insert(entry.dataset_id.clone(), entry);
969    }
970
971    /// Get entry for a dataset.
972    #[must_use]
973    pub fn get(&self, dataset_id: &str) -> Option<&CacheManifestEntry> {
974        self.entries.get(dataset_id)
975    }
976
977    /// Check if a dataset's cached file matches the manifest.
978    ///
979    /// Returns `true` if the file exists and checksum matches.
980    #[must_use]
981    pub fn verify_entry(&self, dataset_id: &str, cache_dir: &std::path::Path) -> bool {
982        let Some(entry) = self.entries.get(dataset_id) else {
983            return false;
984        };
985
986        // Check file exists with expected size
987        let file_path = cache_dir.join(&entry.dataset_id);
988        let Ok(metadata) = fs::metadata(&file_path) else {
989            return false;
990        };
991
992        metadata.len() == entry.file_size
993        // Note: Full SHA-256 verification is expensive; size check is usually sufficient
994    }
995}
996
997// =============================================================================
998// Temporal Metadata
999// =============================================================================
1000
1001/// Temporal metadata for datasets (optional).
1002///
1003/// Used for temporal stratification of evaluation metrics.
1004/// Most datasets don't have temporal metadata, so this is optional.
1005#[derive(Debug, Clone, Default, Serialize, Deserialize)]
1006pub struct TemporalMetadata {
1007    /// KB version used for entity linking (if applicable)
1008    pub kb_version: Option<String>,
1009    /// Temporal cutoff date (entities before this date are "old", after are "new")
1010    pub temporal_cutoff: Option<String>, // ISO 8601 date string
1011    /// Entity creation dates (if available)
1012    pub entity_creation_dates: Option<HashMap<String, String>>, // entity_id -> ISO 8601 date
1013}
1014
1015// =============================================================================
1016// Annotated Structures
1017// =============================================================================
1018
1019/// Annotated token with NER tag.
1020#[derive(Debug, Clone, Serialize, Deserialize)]
1021pub struct AnnotatedToken {
1022    /// Token text
1023    pub text: String,
1024    /// NER tag (BIO format)
1025    pub ner_tag: String,
1026}
1027
1028/// Annotated sentence with tokens.
1029#[derive(Debug, Clone, Serialize, Deserialize)]
1030pub struct AnnotatedSentence {
1031    /// Tokens in the sentence
1032    pub tokens: Vec<AnnotatedToken>,
1033    /// Source dataset
1034    pub source_dataset: DatasetId,
1035}
1036
1037impl AnnotatedSentence {
1038    /// Get text of the sentence.
1039    #[must_use]
1040    pub fn text(&self) -> String {
1041        self.tokens
1042            .iter()
1043            .map(|t| t.text.as_str())
1044            .collect::<Vec<_>>()
1045            .join(" ")
1046    }
1047
1048    /// Get entities from this sentence.
1049    #[must_use]
1050    pub fn entities(&self) -> Vec<super::datasets::GoldEntity> {
1051        // Precompute character offsets for each token in `self.text()` (tokens joined by a
1052        // single ASCII space). Offsets are in *character* units, not bytes.
1053        let mut token_starts: Vec<usize> = Vec::with_capacity(self.tokens.len());
1054        let mut pos: usize = 0;
1055        for (i, tok) in self.tokens.iter().enumerate() {
1056            token_starts.push(pos);
1057            pos += tok.text.chars().count();
1058            if i + 1 < self.tokens.len() {
1059                pos += 1; // the space inserted by `text()`
1060            }
1061        }
1062
1063        let mut entities = Vec::new();
1064        let mut i = 0;
1065        while i < self.tokens.len() {
1066            let tag = &self.tokens[i].ner_tag;
1067
1068            // Check for BIO format (B-TYPE, I-TYPE)
1069            if tag.starts_with("B-") {
1070                let entity_type = tag.trim_start_matches("B-");
1071                let mut end = i + 1;
1072                // Only continue with I-tags that match the same entity type
1073                while end < self.tokens.len()
1074                    && self.tokens[end].ner_tag.starts_with("I-")
1075                    && self.tokens[end].ner_tag.trim_start_matches("I-") == entity_type
1076                {
1077                    end += 1;
1078                }
1079                let entity_text: String = self.tokens[i..end]
1080                    .iter()
1081                    .map(|t| t.text.as_str())
1082                    .collect::<Vec<_>>()
1083                    .join(" ");
1084
1085                let start = token_starts.get(i).copied().unwrap_or(0);
1086                let end_char = if end <= i {
1087                    start
1088                } else {
1089                    let last = end - 1;
1090                    token_starts.get(last).copied().unwrap_or(start)
1091                        + self.tokens[last].text.chars().count()
1092                };
1093
1094                entities.push(super::datasets::GoldEntity {
1095                    text: entity_text,
1096                    original_label: entity_type.to_string(),
1097                    entity_type: anno::EntityType::from_label(entity_type),
1098                    start,
1099                    end: end_char,
1100                });
1101                i = end;
1102            }
1103            // Fallback: non-BIO format (simple type names like "person", "organization")
1104            // Used by datasets like FewNERD that use ClassLabels
1105            else if tag != "O" && !tag.starts_with("I-") && !tag.starts_with("TAG_") {
1106                let entity_type = tag.as_str();
1107                let mut end = i + 1;
1108                // Continue while same entity type (treating consecutive same labels as one entity)
1109                while end < self.tokens.len() && self.tokens[end].ner_tag == *tag {
1110                    end += 1;
1111                }
1112                let entity_text: String = self.tokens[i..end]
1113                    .iter()
1114                    .map(|t| t.text.as_str())
1115                    .collect::<Vec<_>>()
1116                    .join(" ");
1117
1118                let start = token_starts.get(i).copied().unwrap_or(0);
1119                let end_char = if end <= i {
1120                    start
1121                } else {
1122                    let last = end - 1;
1123                    token_starts.get(last).copied().unwrap_or(start)
1124                        + self.tokens[last].text.chars().count()
1125                };
1126
1127                entities.push(super::datasets::GoldEntity {
1128                    text: entity_text,
1129                    original_label: entity_type.to_string(),
1130                    entity_type: anno::EntityType::from_label(entity_type),
1131                    start,
1132                    end: end_char,
1133                });
1134                i = end;
1135            } else {
1136                i += 1;
1137            }
1138        }
1139        entities
1140    }
1141}
1142
1143/// Dataset metadata for provenance and reproducibility.
1144///
1145/// Captures all information needed to understand where data came from
1146/// and how to attribute/license it correctly.
1147#[derive(Debug, Clone, Default, Serialize, Deserialize)]
1148pub struct DatasetMetadata {
1149    /// Human-readable description.
1150    pub description: Option<String>,
1151    /// License (e.g., "CC-BY-4.0", "Apache-2.0", "Research-only").
1152    pub license: Option<String>,
1153    /// Citation reference (BibTeX key or paper title).
1154    pub citation: Option<String>,
1155    /// Split type (e.g., "train", "dev", "test").
1156    pub split: Option<String>,
1157    /// Language code (e.g., "en", "de", "zh").
1158    pub language: Option<String>,
1159    /// Domain (e.g., "news", "biomedical", "social-media").
1160    pub domain: Option<String>,
1161    /// Original source repository or homepage.
1162    pub original_source: Option<String>,
1163    /// Version string if the dataset has versions.
1164    pub version: Option<String>,
1165    /// Number of annotators (for IAA reporting).
1166    pub annotators: Option<u32>,
1167    /// Annotation guidelines URL if available.
1168    pub guidelines_url: Option<String>,
1169}
1170
1171/// Where the dataset was loaded from.
1172#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
1173pub enum DataSource {
1174    /// Loaded from S3 cache (fast, reliable).
1175    S3Cache,
1176    /// Loaded from local file cache.
1177    LocalCache,
1178    /// Downloaded from original URL (may be slow or unavailable).
1179    OriginalUrl,
1180    /// Dataset was skipped (URL broken or unavailable).
1181    #[default]
1182    Skipped,
1183    /// Embedded in binary (test fixtures).
1184    Embedded,
1185}
1186
1187impl DataSource {
1188    /// Human-readable description of the source.
1189    #[must_use]
1190    pub fn description(&self) -> &'static str {
1191        match self {
1192            Self::S3Cache => "S3 cache",
1193            Self::LocalCache => "local cache",
1194            Self::OriginalUrl => "original URL",
1195            Self::Skipped => "skipped (unavailable)",
1196            Self::Embedded => "embedded",
1197        }
1198    }
1199
1200    /// Whether this source is cached (fast).
1201    #[must_use]
1202    pub fn is_cached(&self) -> bool {
1203        matches!(self, Self::S3Cache | Self::LocalCache | Self::Embedded)
1204    }
1205
1206    /// Whether this source indicates the data is unavailable.
1207    #[must_use]
1208    pub fn is_unavailable(&self) -> bool {
1209        matches!(self, Self::Skipped)
1210    }
1211}
1212
1213impl std::fmt::Display for DataSource {
1214    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1215        write!(f, "{}", self.description())
1216    }
1217}
1218
1219/// A loaded dataset with metadata.
1220#[derive(Debug, Clone, Serialize, Deserialize)]
1221pub struct LoadedDataset {
1222    /// Dataset identifier.
1223    pub id: DatasetId,
1224    /// Annotated sentences.
1225    pub sentences: Vec<AnnotatedSentence>,
1226    /// When the dataset was loaded.
1227    pub loaded_at: String, // ISO 8601
1228    /// Source URL.
1229    pub source_url: String,
1230    /// Where the data was actually loaded from.
1231    #[serde(default)]
1232    pub data_source: DataSource,
1233    /// Optional temporal metadata for temporal stratification
1234    #[serde(default)]
1235    pub temporal_metadata: Option<TemporalMetadata>,
1236    /// Dataset provenance and attribution metadata.
1237    #[serde(default)]
1238    pub metadata: DatasetMetadata,
1239}
1240
1241impl LoadedDataset {
1242    /// Total number of sentences.
1243    #[must_use]
1244    pub fn len(&self) -> usize {
1245        self.sentences.len()
1246    }
1247
1248    /// Check if empty.
1249    #[must_use]
1250    pub fn is_empty(&self) -> bool {
1251        self.sentences.is_empty()
1252    }
1253
1254    /// Total number of entities.
1255    #[must_use]
1256    pub fn entity_count(&self) -> usize {
1257        self.sentences.iter().map(|s| s.entities().len()).sum()
1258    }
1259
1260    /// Count entities by type.
1261    #[must_use]
1262    pub fn entity_counts_by_type(&self) -> HashMap<String, usize> {
1263        let mut counts = HashMap::new();
1264        for sentence in &self.sentences {
1265            for entity in sentence.entities() {
1266                *counts.entry(entity.original_label.clone()).or_insert(0) += 1;
1267            }
1268        }
1269        counts
1270    }
1271
1272    /// Get statistics summary.
1273    #[must_use]
1274    pub fn stats(&self) -> DatasetStats {
1275        DatasetStats {
1276            name: self.id.name().to_string(),
1277            sentences: self.len(),
1278            tokens: self.sentences.iter().map(|s| s.tokens.len()).sum(),
1279            entities: self.entity_count(),
1280            entities_by_type: self.entity_counts_by_type(),
1281        }
1282    }
1283
1284    /// Convert to test cases format for evaluation.
1285    #[must_use]
1286    pub fn to_test_cases(&self) -> Vec<(String, Vec<super::datasets::GoldEntity>)> {
1287        self.sentences
1288            .iter()
1289            .map(|s| (s.text(), s.entities()))
1290            .collect()
1291    }
1292}
1293
1294// =============================================================================
1295// Dataset Statistics
1296// =============================================================================
1297
1298/// Dataset statistics.
1299#[derive(Debug, Clone, Serialize, Deserialize)]
1300pub struct DatasetStats {
1301    /// Dataset name.
1302    pub name: String,
1303    /// Number of sentences.
1304    pub sentences: usize,
1305    /// Total token count.
1306    pub tokens: usize,
1307    /// Total entity count.
1308    pub entities: usize,
1309    /// Entities by type.
1310    pub entities_by_type: HashMap<String, usize>,
1311}
1312
1313/// A document with text and gold relations for relation extraction evaluation.
1314#[derive(Debug, Clone)]
1315pub struct RelationDocument {
1316    /// Document text
1317    pub text: String,
1318    /// Gold standard relations
1319    pub relations: Vec<super::relation::RelationGold>,
1320}
1321
1322impl std::fmt::Display for DatasetStats {
1323    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1324        writeln!(f, "Dataset: {}", self.name)?;
1325        writeln!(f, "  Sentences: {}", self.sentences)?;
1326        writeln!(f, "  Tokens: {}", self.tokens)?;
1327        writeln!(f, "  Entities: {}", self.entities)?;
1328        writeln!(f, "  Entity types:")?;
1329        let mut types: Vec<_> = self.entities_by_type.iter().collect();
1330        types.sort_by(|a, b| b.1.cmp(a.1));
1331        for (etype, count) in types {
1332            writeln!(f, "    {}: {}", etype, count)?;
1333        }
1334        Ok(())
1335    }
1336}
1337
1338// =============================================================================
1339// Dataset Loader
1340// =============================================================================
1341
1342/// Loads and caches NER datasets.
1343///
1344/// # Caching Strategy (Tiered)
1345///
1346/// 1. **Local cache** (`~/.cache/anno/datasets/`): Checked first, fastest
1347/// 2. **S3 cache** (`s3://anno-data/`): Checked if `ANNO_S3_CACHE=1` and AWS credentials available
1348/// 3. **URL download**: Original source URL, last resort
1349///
1350/// This tiered approach:
1351/// - Maximizes speed (local cache hit is instant)
1352/// - Provides redundancy (S3 backup if original URLs go down)
1353/// - Allows sharing datasets across machines via S3
1354///
1355/// # Environment Variables
1356///
1357/// - `ANNO_CACHE_DIR`: Override default cache location
1358/// - `ANNO_S3_CACHE`: Set to "1" to enable S3 fallback (requires AWS credentials)
1359/// - `ANNO_S3_BUCKET`: Override S3 bucket name (default: `arc-anno-data`)
1360///
1361/// # Example
1362///
1363/// ```bash
1364/// # Enable S3 caching
1365/// export ANNO_S3_CACHE=1
1366/// export AWS_PROFILE=default  # or set AWS_ACCESS_KEY_ID, etc.
1367///
1368/// # Now load_or_download() will try S3 before URL
1369/// ```
1370pub struct DatasetLoader {
1371    cache_dir: PathBuf,
1372    /// S3 bucket name for fallback caching (if enabled)
1373    s3_bucket: Option<String>,
1374    /// Cache manifest for tracking downloads.
1375    ///
1376    /// This is shared across evaluation threads, so it must be thread-safe.
1377    manifest: std::sync::RwLock<CacheManifest>,
1378}
1379
1380impl DatasetLoader {
1381    // Default download cap when `ANNO_MAX_DOWNLOAD_BYTES` is unset.
1382    //
1383    // Rationale: unset should be *usable* but *safe* by default. This cap is meant to prevent
1384    // accidental multi-GB downloads while still allowing many evaluation datasets.
1385    #[cfg(feature = "eval")]
1386    const DEFAULT_MAX_DOWNLOAD_BYTES: u64 = 50 * 1024 * 1024; // 50 MiB
1387
1388    #[cfg(feature = "eval")]
1389    fn max_download_bytes() -> Option<u64> {
1390        match std::env::var("ANNO_MAX_DOWNLOAD_BYTES").ok() {
1391            Some(s) => {
1392                let s = s.trim();
1393                if s.is_empty() {
1394                    return Some(Self::DEFAULT_MAX_DOWNLOAD_BYTES);
1395                }
1396                let Ok(v) = s.parse::<u64>() else {
1397                    return Some(Self::DEFAULT_MAX_DOWNLOAD_BYTES);
1398                };
1399                if v == 0 {
1400                    None // explicit opt-out
1401                } else {
1402                    Some(v)
1403                }
1404            }
1405            None => Some(Self::DEFAULT_MAX_DOWNLOAD_BYTES),
1406        }
1407    }
1408
1409    #[cfg(feature = "eval")]
1410    fn enforce_max_download_bytes(content_len: usize, source: &str) -> Result<()> {
1411        let Some(limit) = Self::max_download_bytes() else {
1412            return Ok(());
1413        };
1414        let len = content_len as u64;
1415        if len > limit {
1416            return Err(Error::InvalidInput(format!(
1417                "Download rejected ({} bytes > ANNO_MAX_DOWNLOAD_BYTES={} bytes) from {}",
1418                len, limit, source
1419            )));
1420        }
1421        Ok(())
1422    }
1423
1424    #[cfg(feature = "eval")]
1425    fn hf_dataset_from_rows_url(url: &str) -> Option<String> {
1426        // Example:
1427        // https://datasets-server.huggingface.co/rows?dataset=masakhane%2Fmasakhaner2&config=bam&split=test...
1428        let (_, query) = url.split_once('?')?;
1429        for part in query.split('&') {
1430            let (k, v) = part.split_once('=')?;
1431            if k == "dataset" {
1432                // Minimal percent-decoding sufficient for HF dataset ids.
1433                let decoded = v.replace("%2F", "/").replace("%2f", "/");
1434                return Some(decoded);
1435            }
1436        }
1437        None
1438    }
1439
1440    /// Create a new loader with default cache directory.
1441    ///
1442    /// Default location: `~/.cache/anno/datasets` (platform cache via `dirs` crate)
1443    /// Falls back to `.anno/datasets` in current directory if `dirs` crate unavailable.
1444    ///
1445    /// S3 fallback is enabled if `ANNO_S3_CACHE=1` environment variable is set.
1446    pub fn new() -> Result<Self> {
1447        // Load workspace `.env` (idempotent, does not override existing env vars).
1448        // This keeps eval tooling usable without manual exporting of env vars.
1449        anno::env::load_dotenv();
1450
1451        // Check for custom cache directory
1452        let cache_dir = if let Ok(custom_dir) = std::env::var("ANNO_CACHE_DIR") {
1453            PathBuf::from(custom_dir).join("datasets")
1454        } else {
1455            let base_dir = dirs::cache_dir().unwrap_or_else(|| PathBuf::from("."));
1456            base_dir.join("anno").join("datasets")
1457        };
1458
1459        // Check for S3 caching
1460        let s3_bucket = if std::env::var("ANNO_S3_CACHE").unwrap_or_default() == "1" {
1461            Some(std::env::var("ANNO_S3_BUCKET").unwrap_or_else(|_| "arc-anno-data".to_string()))
1462        } else {
1463            None
1464        };
1465
1466        fs::create_dir_all(&cache_dir).map_err(|e| {
1467            Error::InvalidInput(format!("Failed to create cache dir {:?}: {}", cache_dir, e))
1468        })?;
1469
1470        let manifest = CacheManifest::load(&cache_dir)?;
1471
1472        Ok(Self {
1473            cache_dir,
1474            s3_bucket,
1475            manifest: std::sync::RwLock::new(manifest),
1476        })
1477    }
1478
1479    /// Update the cache manifest with a new entry and save it.
1480    #[cfg(feature = "eval")]
1481    fn update_manifest(&self, entry: CacheManifestEntry) -> Result<()> {
1482        let mut manifest = self
1483            .manifest
1484            .write()
1485            .map_err(|_| Error::InvalidInput("cache manifest lock poisoned".to_string()))?;
1486        manifest.update_entry(entry);
1487        manifest.save(&self.cache_dir)?;
1488        Ok(())
1489    }
1490
1491    /// Create a loader with a custom cache directory.
1492    pub fn with_cache_dir(cache_dir: impl Into<PathBuf>) -> Result<Self> {
1493        let cache_dir = cache_dir.into();
1494        let s3_bucket = if std::env::var("ANNO_S3_CACHE").unwrap_or_default() == "1" {
1495            Some(std::env::var("ANNO_S3_BUCKET").unwrap_or_else(|_| "arc-anno-data".to_string()))
1496        } else {
1497            None
1498        };
1499
1500        fs::create_dir_all(&cache_dir).map_err(|e| {
1501            Error::InvalidInput(format!("Failed to create cache dir {:?}: {}", cache_dir, e))
1502        })?;
1503
1504        let manifest = CacheManifest::load(&cache_dir)?;
1505
1506        Ok(Self {
1507            cache_dir,
1508            s3_bucket,
1509            manifest: std::sync::RwLock::new(manifest),
1510        })
1511    }
1512
1513    /// Base directory for cached datasets.
1514    #[must_use]
1515    pub fn cache_dir(&self) -> &std::path::Path {
1516        &self.cache_dir
1517    }
1518
1519    /// Whether S3 fallback caching is enabled.
1520    #[must_use]
1521    pub fn s3_enabled(&self) -> bool {
1522        self.s3_bucket.is_some()
1523    }
1524
1525    /// The configured S3 bucket name (if enabled).
1526    #[must_use]
1527    pub fn s3_bucket(&self) -> Option<&str> {
1528        self.s3_bucket.as_deref()
1529    }
1530
1531    /// Snapshot of cached dataset manifest entries (cloned).
1532    ///
1533    /// This is intended for tooling (CLI/scripts) that wants to iterate cached datasets
1534    /// without holding the manifest lock for a long time.
1535    #[must_use]
1536    pub fn cached_manifest_entries(&self) -> Vec<CacheManifestEntry> {
1537        let Ok(manifest) = self.manifest.read() else {
1538            return Vec::new();
1539        };
1540        let mut out: Vec<CacheManifestEntry> = manifest.entries.values().cloned().collect();
1541        out.sort_by(|a, b| a.dataset_id.cmp(&b.dataset_id));
1542        out
1543    }
1544
1545    /// Upload a locally cached dataset (by `DatasetId`) to S3, using its manifest entry.
1546    ///
1547    /// This uses the same object layout as `load_or_download()`'s best-effort upload:
1548    /// - `datasets/<cache_filename>` (legacy/mutable key)
1549    /// - `datasets/by-sha256/<sha>/<cache_filename>` (immutable snapshot)
1550    /// - `datasets/<cache_filename>.latest.json` (pointer)
1551    /// - `datasets/<cache_filename>.manifest.json` (sidecar metadata)
1552    #[cfg(feature = "eval")]
1553    pub fn upload_cached_dataset_to_s3(&self, bucket: &str, id: DatasetId) -> Result<()> {
1554        let key = id.cache_filename();
1555        let entry = {
1556            let guard = self
1557                .manifest
1558                .read()
1559                .map_err(|_| Error::InvalidInput("cache manifest lock poisoned".to_string()))?;
1560            guard
1561                .get(key)
1562                .cloned()
1563                .ok_or_else(|| Error::InvalidInput(format!("No manifest entry for {}", key)))?
1564        };
1565        let path = self.cache_path_for(id);
1566        let content = std::fs::read_to_string(&path).map_err(|e| {
1567            Error::InvalidInput(format!(
1568                "Failed to read cached dataset {}: {}",
1569                path.display(),
1570                e
1571            ))
1572        })?;
1573        Self::enforce_max_download_bytes(content.len(), "local cache (sync-s3)")?;
1574        self.upload_to_s3(bucket, id, &content, &entry)
1575    }
1576
1577    #[must_use]
1578    fn cache_path_for(&self, id: DatasetId) -> PathBuf {
1579        self.cache_dir.join(id.cache_filename())
1580    }
1581
1582    #[must_use]
1583    fn is_cached_for(&self, id: DatasetId) -> bool {
1584        if !self.cache_path_for(id).exists() {
1585            return false;
1586        }
1587
1588        // Best-effort cache invalidation: if the registry URL changes, the on-disk cached
1589        // file may no longer correspond to the dataset id. Prefer re-download over silently
1590        // using stale/incorrect data.
1591        //
1592        // This is intentionally lightweight (no checksum re-hash here).
1593        if let Ok(manifest) = self.manifest.read() {
1594            if let Some(entry) = manifest.get(id.cache_filename()) {
1595                if entry.source_url != id.download_url() {
1596                    return false;
1597                }
1598                // Treat “cached but empty” as invalid: this is almost always a bad download
1599                // (HTML, auth wall, or format mismatch) and should be re-fetched.
1600                if entry.sentence_count == 0 {
1601                    return false;
1602                }
1603            }
1604        }
1605
1606        true
1607    }
1608
1609    /// Get the cache path for a dataset.
1610    #[must_use]
1611    pub fn cache_path(&self, id: LoadableDatasetId) -> PathBuf {
1612        self.cache_path_for(id.0)
1613    }
1614
1615    /// Check if a dataset is cached locally.
1616    #[must_use]
1617    pub fn is_cached(&self, id: LoadableDatasetId) -> bool {
1618        self.is_cached_for(id.0)
1619    }
1620
1621    /// Load a dataset from cache.
1622    ///
1623    /// Returns an error if the dataset is not cached.
1624    pub fn load(&self, id: LoadableDatasetId) -> Result<LoadedDataset> {
1625        let dataset_id = id.0;
1626        let cache_path = self.cache_path(id);
1627        if !cache_path.exists() {
1628            return Err(Error::InvalidInput(format!(
1629                "Dataset {:?} not cached at {:?}",
1630                dataset_id, cache_path
1631            )));
1632        }
1633
1634        let content = fs::read_to_string(&cache_path).map_err(|e| {
1635            Error::InvalidInput(format!("Failed to read cache {:?}: {}", cache_path, e))
1636        })?;
1637
1638        let mut dataset = self.parse_content_impl(&content, dataset_id)?;
1639        if dataset.sentences.is_empty() {
1640            return Err(Error::InvalidInput(format!(
1641                "Cached dataset '{}' parsed to 0 sentences (cache_path={:?})",
1642                dataset_id.name(),
1643                cache_path
1644            )));
1645        }
1646        dataset.data_source = DataSource::LocalCache;
1647        Ok(dataset)
1648    }
1649
1650    /// Load or download a dataset.
1651    ///
1652    /// Tries cache first, then S3 (if enabled), then downloads from URL.
1653    #[cfg(feature = "eval")]
1654    pub fn load_or_download(&self, id: LoadableDatasetId) -> Result<LoadedDataset> {
1655        let dataset_id = id.0;
1656        // 1. Check local cache first
1657        if self.is_cached(id) {
1658            return self.load(id);
1659        }
1660
1661        // 2. Try S3 cache if enabled
1662        if let Some(ref bucket) = self.s3_bucket {
1663            if let Ok((content, manifest_entry)) = self.download_from_s3(bucket, dataset_id) {
1664                Self::enforce_max_download_bytes(content.len(), "S3")?;
1665                // Cache locally for future use
1666                let cache_path = self.cache_path(id);
1667                fs::write(&cache_path, &content).map_err(|e| {
1668                    Error::InvalidInput(format!("Failed to write cache {:?}: {}", cache_path, e))
1669                })?;
1670
1671                let mut dataset = self.parse_content_impl(&content, dataset_id)?;
1672                dataset.data_source = DataSource::S3Cache;
1673
1674                // Best-effort: if S3 provides a manifest entry, record it locally.
1675                if let Some(entry) = manifest_entry {
1676                    let _ = self.update_manifest(entry);
1677                }
1678                return Ok(dataset);
1679            }
1680        }
1681
1682        // 3. Download from original URL
1683        let (content, resolved_url) = self.download_with_resolved_url(dataset_id)?;
1684        Self::enforce_max_download_bytes(content.len(), &resolved_url)?;
1685        let file_size = content.len() as u64;
1686        let sha256 = self.compute_sha256(&content);
1687
1688        // 4. Cache the downloaded content locally
1689        let cache_path = self.cache_path(id);
1690        fs::write(&cache_path, &content).map_err(|e| {
1691            Error::InvalidInput(format!("Failed to write cache {:?}: {}", cache_path, e))
1692        })?;
1693
1694        // 5. Parse the content
1695        let mut dataset = self.parse_content_impl(&content, dataset_id)?;
1696        dataset.data_source = DataSource::OriginalUrl;
1697
1698        // 6. Update manifest with download metadata
1699        let entry = CacheManifestEntry {
1700            dataset_id: dataset_id.cache_filename().to_string(),
1701            source_url: dataset_id.download_url().to_string(),
1702            resolved_url: Some(resolved_url.clone()),
1703            sha256,
1704            file_size,
1705            downloaded_at: chrono::Utc::now().to_rfc3339(),
1706            sentence_count: dataset.sentences.len(),
1707            entity_count: dataset.entity_count(),
1708            anno_version: env!("CARGO_PKG_VERSION").to_string(),
1709        };
1710        let _ = self.update_manifest(entry); // Best effort
1711
1712        // 7. Optionally upload to S3 for future use (best effort)
1713        if let Some(ref bucket) = self.s3_bucket {
1714            let entry = CacheManifestEntry {
1715                dataset_id: dataset_id.cache_filename().to_string(),
1716                source_url: dataset_id.download_url().to_string(),
1717                resolved_url: Some(resolved_url),
1718                sha256: self.compute_sha256(&content),
1719                file_size: content.len() as u64,
1720                downloaded_at: chrono::Utc::now().to_rfc3339(),
1721                sentence_count: dataset.sentences.len(),
1722                entity_count: dataset.entity_count(),
1723                anno_version: env!("CARGO_PKG_VERSION").to_string(),
1724            };
1725            let _ = self.upload_to_s3(bucket, dataset_id, &content, &entry);
1726        }
1727
1728        Ok(dataset)
1729    }
1730
1731    /// Load-or-download shim when `eval` is disabled.
1732    ///
1733    /// This method will load from cache if present, otherwise return a helpful error.
1734    #[cfg(not(feature = "eval"))]
1735    pub fn load_or_download(&self, id: LoadableDatasetId) -> Result<LoadedDataset> {
1736        if self.is_cached(id) {
1737            return self.load(id);
1738        }
1739
1740        Err(Error::InvalidInput(
1741            "Dataset is not cached. Rebuild with feature `eval` to enable downloading.".to_string(),
1742        ))
1743    }
1744
1745    /// Download dataset from S3 cache bucket.
1746    ///
1747    /// Uses AWS CLI under the hood (requires `aws` command in PATH and valid credentials).
1748    #[cfg(feature = "eval")]
1749    fn download_from_s3(
1750        &self,
1751        bucket: &str,
1752        id: DatasetId,
1753    ) -> Result<(String, Option<CacheManifestEntry>)> {
1754        use std::process::Command;
1755
1756        // Prefer a content-addressed snapshot if the pointer exists; fall back to the legacy key.
1757        let pointer_key = format!("datasets/{}.latest.json", id.cache_filename());
1758        let pointer_uri = format!("s3://{}/{}", bucket, pointer_key);
1759
1760        let mut content: Option<String> = None;
1761
1762        // Attempt: pointer -> by-sha256 snapshot
1763        let pointer = Command::new("aws")
1764            .args(["s3", "cp", &pointer_uri, "-"])
1765            .output()
1766            .ok()
1767            .and_then(|o| {
1768                if !o.status.success() {
1769                    // Pointer missing or unreadable; fall back to legacy key.
1770                    return None;
1771                }
1772                String::from_utf8(o.stdout).ok()
1773            })
1774            .and_then(|json| serde_json::from_str::<serde_json::Value>(&json).ok());
1775
1776        if let Some(pointer) = pointer {
1777            if let Some(sha) = pointer.get("sha256").and_then(|v| v.as_str()) {
1778                let by_sha_key = format!("datasets/by-sha256/{}/{}", sha, id.cache_filename());
1779                let by_sha_uri = format!("s3://{}/{}", bucket, by_sha_key);
1780                let output = Command::new("aws")
1781                    .args(["s3", "cp", &by_sha_uri, "-"])
1782                    .output()
1783                    .ok();
1784
1785                if let Some(output) = output {
1786                    if output.status.success() {
1787                        content = String::from_utf8(output.stdout).ok();
1788                    }
1789                }
1790            }
1791        }
1792
1793        // Fall back: legacy key (mutable)
1794        if content.is_none() {
1795            let s3_key = format!("datasets/{}", id.cache_filename());
1796            let s3_uri = format!("s3://{}/{}", bucket, s3_key);
1797
1798            // Use aws s3 cp to download to stdout
1799            let output = Command::new("aws")
1800                .args(["s3", "cp", &s3_uri, "-"])
1801                .output()
1802                .map_err(|e| Error::InvalidInput(format!("Failed to run aws s3 cp: {}", e)))?;
1803
1804            if output.status.success() {
1805                content = Some(String::from_utf8(output.stdout).map_err(|e| {
1806                    Error::InvalidInput(format!("S3 content not valid UTF-8: {}", e))
1807                })?);
1808            } else {
1809                return Err(Error::InvalidInput(format!(
1810                    "S3 download failed: {}",
1811                    String::from_utf8_lossy(&output.stderr)
1812                )));
1813            }
1814        }
1815
1816        let Some(content) = content else {
1817            return Err(Error::InvalidInput("S3 download failed".to_string()));
1818        };
1819
1820        // Best-effort: attempt to fetch per-dataset manifest entry.
1821        let manifest = self.download_manifest_entry_from_s3(bucket, id).ok();
1822
1823        Ok((content, manifest))
1824    }
1825
1826    /// Upload dataset to S3 cache bucket.
1827    ///
1828    /// Best effort - failures are logged but don't stop execution.
1829    #[cfg(feature = "eval")]
1830    fn upload_to_s3(
1831        &self,
1832        bucket: &str,
1833        id: DatasetId,
1834        content: &str,
1835        manifest_entry: &CacheManifestEntry,
1836    ) -> Result<()> {
1837        use std::io::Write;
1838        use std::process::{Command, Stdio};
1839
1840        let s3_key = format!("datasets/{}", id.cache_filename());
1841        let s3_uri = format!("s3://{}/{}", bucket, s3_key);
1842
1843        // Use aws s3 cp to upload from stdin
1844        let mut child = Command::new("aws")
1845            .args(["s3", "cp", "-", &s3_uri])
1846            .stdin(Stdio::piped())
1847            .stdout(Stdio::null())
1848            .stderr(Stdio::null())
1849            .spawn()
1850            .map_err(|e| Error::InvalidInput(format!("Failed to spawn aws s3 cp: {}", e)))?;
1851
1852        if let Some(ref mut stdin) = child.stdin {
1853            let _ = stdin.write_all(content.as_bytes());
1854        }
1855
1856        let status = child
1857            .wait()
1858            .map_err(|e| Error::InvalidInput(format!("Failed to wait for aws s3 cp: {}", e)))?;
1859
1860        if !status.success() {
1861            return Err(Error::InvalidInput("S3 upload failed".to_string()));
1862        }
1863
1864        // Also upload a content-addressed snapshot. This supports the "durable mirror" use case:
1865        // once uploaded, snapshots are stable even if upstream changes.
1866        let by_sha_key = format!(
1867            "datasets/by-sha256/{}/{}",
1868            manifest_entry.sha256,
1869            id.cache_filename()
1870        );
1871        let by_sha_uri = format!("s3://{}/{}", bucket, by_sha_key);
1872
1873        let mut child = Command::new("aws")
1874            .args(["s3", "cp", "-", &by_sha_uri])
1875            .stdin(Stdio::piped())
1876            .stdout(Stdio::null())
1877            .stderr(Stdio::null())
1878            .spawn()
1879            .map_err(|e| Error::InvalidInput(format!("Failed to spawn aws s3 cp: {}", e)))?;
1880
1881        if let Some(ref mut stdin) = child.stdin {
1882            let _ = stdin.write_all(content.as_bytes());
1883        }
1884
1885        let _ = child.wait();
1886
1887        // Update a pointer file for the latest snapshot.
1888        let pointer_key = format!("datasets/{}.latest.json", id.cache_filename());
1889        let pointer_uri = format!("s3://{}/{}", bucket, pointer_key);
1890        let pointer_json = serde_json::json!({
1891            "dataset": id.cache_filename(),
1892            "sha256": manifest_entry.sha256,
1893            "updated_at": chrono::Utc::now().to_rfc3339(),
1894            "source_url": manifest_entry.source_url,
1895            "resolved_url": manifest_entry.resolved_url,
1896            "file_size": manifest_entry.file_size,
1897        })
1898        .to_string();
1899
1900        let mut child = Command::new("aws")
1901            .args(["s3", "cp", "-", &pointer_uri])
1902            .stdin(Stdio::piped())
1903            .stdout(Stdio::null())
1904            .stderr(Stdio::null())
1905            .spawn()
1906            .map_err(|e| Error::InvalidInput(format!("Failed to spawn aws s3 cp: {}", e)))?;
1907
1908        if let Some(ref mut stdin) = child.stdin {
1909            let _ = stdin.write_all(pointer_json.as_bytes());
1910        }
1911
1912        let _ = child.wait();
1913
1914        // Upload a sidecar manifest entry (JSON). This makes the S3 cache self-describing.
1915        let manifest_key = format!("datasets/{}.manifest.json", id.cache_filename());
1916        let manifest_uri = format!("s3://{}/{}", bucket, manifest_key);
1917        let json = serde_json::to_string_pretty(manifest_entry)
1918            .map_err(|e| Error::InvalidInput(format!("Failed to serialize S3 manifest: {}", e)))?;
1919
1920        let mut child = Command::new("aws")
1921            .args(["s3", "cp", "-", &manifest_uri])
1922            .stdin(Stdio::piped())
1923            .stdout(Stdio::null())
1924            .stderr(Stdio::null())
1925            .spawn()
1926            .map_err(|e| Error::InvalidInput(format!("Failed to spawn aws s3 cp: {}", e)))?;
1927
1928        if let Some(ref mut stdin) = child.stdin {
1929            let _ = stdin.write_all(json.as_bytes());
1930        }
1931
1932        let status = child
1933            .wait()
1934            .map_err(|e| Error::InvalidInput(format!("Failed to wait for aws s3 cp: {}", e)))?;
1935
1936        if status.success() {
1937            Ok(())
1938        } else {
1939            Err(Error::InvalidInput("S3 manifest upload failed".to_string()))
1940        }
1941    }
1942
1943    /// Best-effort: download the sidecar manifest entry for a dataset from S3.
1944    #[cfg(feature = "eval")]
1945    fn download_manifest_entry_from_s3(
1946        &self,
1947        bucket: &str,
1948        id: DatasetId,
1949    ) -> Result<CacheManifestEntry> {
1950        use std::process::Command;
1951
1952        let manifest_key = format!("datasets/{}.manifest.json", id.cache_filename());
1953        let manifest_uri = format!("s3://{}/{}", bucket, manifest_key);
1954
1955        let output = Command::new("aws")
1956            .args(["s3", "cp", &manifest_uri, "-"])
1957            .output()
1958            .map_err(|e| Error::InvalidInput(format!("Failed to run aws s3 cp: {}", e)))?;
1959
1960        if !output.status.success() {
1961            return Err(Error::InvalidInput(
1962                "S3 manifest download failed".to_string(),
1963            ));
1964        }
1965
1966        let json = String::from_utf8(output.stdout)
1967            .map_err(|e| Error::InvalidInput(format!("S3 manifest not valid UTF-8: {}", e)))?;
1968
1969        serde_json::from_str(&json)
1970            .map_err(|e| Error::InvalidInput(format!("Failed to parse S3 manifest JSON: {}", e)))
1971    }
1972
1973    /// Download dataset from source with retry logic and pagination support.
1974    ///
1975    /// Implements exponential backoff retry strategy:
1976    /// - 3 retries maximum
1977    /// - Initial delay: 1 second
1978    /// - Exponential backoff: 2^attempt seconds
1979    /// - Max delay: 10 seconds
1980    ///
1981    /// For HuggingFace datasets-server API, automatically paginates to download full dataset.
1982    #[cfg(feature = "eval")]
1983    fn download_with_resolved_url(&self, id: DatasetId) -> Result<(String, String)> {
1984        let url = id.download_url().to_string();
1985
1986        fn env_single_csv(keys: &[&str]) -> Option<String> {
1987            for &k in keys {
1988                let Ok(raw) = std::env::var(k) else {
1989                    continue;
1990                };
1991                let mut parts = raw
1992                    .split(',')
1993                    .map(|s| s.trim().to_ascii_lowercase())
1994                    .filter(|s| !s.is_empty())
1995                    .collect::<Vec<_>>();
1996                parts.sort();
1997                parts.dedup();
1998                if parts.len() == 1 {
1999                    return Some(parts[0].clone());
2000                }
2001            }
2002            None
2003        }
2004
2005        // Check if URL is empty (dataset not available for download)
2006        if url.is_empty() {
2007            let access_status = id.access_status();
2008            let notes = id.notes();
2009            let mirror_url = id.mirror_url();
2010
2011            let mut error_msg = format!(
2012                "Dataset '{}' ({:?}) has no download URL available.",
2013                id.name(),
2014                id
2015            );
2016
2017            // Add access status information
2018            match access_status {
2019                super::dataset_registry::DatasetAccessibility::Public => {
2020                    error_msg.push_str("\n\nStatus: Public (but URL not configured)");
2021                }
2022                super::dataset_registry::DatasetAccessibility::HuggingFace => {
2023                    if let Some(hf_id) = id.hf_id() {
2024                        error_msg.push_str(&format!(
2025                            "\n\nStatus: Available on HuggingFace\nDataset ID: {}",
2026                            hf_id
2027                        ));
2028                    } else {
2029                        error_msg
2030                            .push_str("\n\nStatus: Available on HuggingFace (ID not configured)");
2031                    }
2032                }
2033                super::dataset_registry::DatasetAccessibility::Local => {
2034                    error_msg.push_str("\n\nStatus: Available locally (check testdata/ or cache)");
2035                }
2036                super::dataset_registry::DatasetAccessibility::Registration => {
2037                    error_msg
2038                        .push_str("\n\nStatus: Requires registration (e.g., LDC for academics)");
2039                }
2040                super::dataset_registry::DatasetAccessibility::ContactAuthors => {
2041                    error_msg.push_str("\n\nStatus: Contact dataset authors for access");
2042                }
2043                super::dataset_registry::DatasetAccessibility::NotYetReleased => {
2044                    error_msg.push_str("\n\nStatus: Not yet publicly released");
2045                }
2046                super::dataset_registry::DatasetAccessibility::DependsOnOther => {
2047                    if let Some(dep) = id.depends_on() {
2048                        error_msg
2049                            .push_str(&format!("\n\nStatus: Depends on another dataset: {}", dep));
2050                    } else {
2051                        error_msg.push_str("\n\nStatus: Depends on another dataset");
2052                    }
2053                }
2054                super::dataset_registry::DatasetAccessibility::Deprecated => {
2055                    error_msg.push_str("\n\nStatus: Deprecated or no longer available");
2056                }
2057            }
2058
2059            if let Some(note) = notes {
2060                error_msg.push_str(&format!("\n\nNotes: {}", note));
2061            }
2062
2063            if let Some(mirror) = mirror_url {
2064                error_msg.push_str(&format!("\n\nMirror URL: {}", mirror));
2065            }
2066
2067            return Err(Error::InvalidInput(error_msg));
2068        }
2069
2070        // For multilingual datasets hosted on HuggingFace with per-language configs (e.g. WikiANN),
2071        // let callers bias the HF "config" selection via `ANNO_MUXER_PIN_LANG`.
2072        //
2073        // This keeps muxer facet pins honest: a `.lang=de` run should ideally download a German
2074        // config, not a random/first config.
2075        let preferred_hf_config: Option<String> = if id.language().eq_ignore_ascii_case("mul") {
2076            // Explicit override wins if set.
2077            env_single_csv(&["ANNO_HF_DATASET_CONFIG"])
2078                .or_else(|| env_single_csv(&["ANNO_MUXER_PIN_LANG", "ANNO_MUXER_FILTER_LANG"]))
2079        } else {
2080            env_single_csv(&["ANNO_HF_DATASET_CONFIG"])
2081        };
2082
2083        // If the registry marks this dataset as HuggingFace-accessible and provides an HF id,
2084        // prefer fetching via datasets-server (even if the registry `url` is a paper/GitHub homepage).
2085        if id.access_status() == super::dataset_registry::DatasetAccessibility::HuggingFace {
2086            if let Some(hf_ds) = id.hf_id().map(|s| s.to_string()) {
2087                if let Ok((config, split)) =
2088                    self.resolve_hf_config_split_prefer(&hf_ds, preferred_hf_config.as_deref())
2089                {
2090                    let base = Self::hf_rows_url(&hf_ds, &config, &split);
2091                    if std::env::var("ANNO_MUXER_VERBOSE")
2092                        .ok()
2093                        .is_some_and(|v| v == "1" || v.eq_ignore_ascii_case("true"))
2094                    {
2095                        eprintln!(
2096                            "dataset-loader: hf dataset={} config={} split={} (preferred={:?})",
2097                            hf_ds, config, split, preferred_hf_config
2098                        );
2099                    }
2100                    match self.download_hf_dataset_paginated(id, &base) {
2101                        Ok(content) => return Ok((content, base)),
2102                        Err(_e) => {
2103                            // Fall back to downloading a raw file from the dataset repo.
2104                            if let Ok((content, resolved)) =
2105                                self.download_hf_dataset_file_from_hub(&hf_ds)
2106                            {
2107                                return Ok((content, resolved));
2108                            }
2109                        }
2110                    }
2111                }
2112            }
2113        }
2114
2115        // HuggingFace dataset pages are common in the registry. Convert them into a
2116        // datasets-server API download so we can actually fetch data (and avoid caching HTML).
2117        //
2118        // Example:
2119        // - page:  https://huggingface.co/datasets/KevinSpaghetti/cadec
2120        // - API:   https://datasets-server.huggingface.co/rows?dataset=KevinSpaghetti%2Fcadec&...
2121        if url.contains("huggingface.co/datasets/") {
2122            // Prefer the registry's `hf_id` when available (more reliable than parsing the URL).
2123            let hf_ds = id
2124                .hf_id()
2125                .map(|s| s.to_string())
2126                .or_else(|| Self::extract_hf_dataset_name(&url));
2127
2128            if let Some(hf_ds) = hf_ds {
2129                match self.resolve_hf_config_split_prefer(&hf_ds, preferred_hf_config.as_deref()) {
2130                    Ok((config, split)) => {
2131                        let base = Self::hf_rows_url(&hf_ds, &config, &split);
2132                        if std::env::var("ANNO_MUXER_VERBOSE")
2133                            .ok()
2134                            .is_some_and(|v| v == "1" || v.eq_ignore_ascii_case("true"))
2135                        {
2136                            eprintln!(
2137                                "dataset-loader: hf dataset={} config={} split={} (preferred={:?})",
2138                                hf_ds, config, split, preferred_hf_config
2139                            );
2140                        }
2141                        match self.download_hf_dataset_paginated(id, &base) {
2142                            Ok(content) => return Ok((content, base)),
2143                            Err(e) => {
2144                                // Some HF datasets don't support datasets-server row export.
2145                                // Fall back to downloading a raw file from the dataset repo.
2146                                if let Ok((content, resolved)) =
2147                                    self.download_hf_dataset_file_from_hub(&hf_ds)
2148                                {
2149                                    return Ok((content, resolved));
2150                                }
2151                                return Err(e);
2152                            }
2153                        }
2154                    }
2155                    Err(_e) => {
2156                        // If the datasets-server endpoint is unavailable (rate limits, transient issues),
2157                        // prefer falling back to the hub-file path rather than downloading HTML.
2158                        if let Ok((content, resolved)) =
2159                            self.download_hf_dataset_file_from_hub(&hf_ds)
2160                        {
2161                            return Ok((content, resolved));
2162                        }
2163                    }
2164                }
2165            }
2166        }
2167
2168        // Check if this is a HuggingFace datasets-server API URL
2169        if url.contains("datasets-server.huggingface.co/rows") {
2170            match self.download_hf_dataset_paginated(id, &url) {
2171                Ok(content) => return Ok((content, url)),
2172                Err(e) => {
2173                    // Some datasets fail row export (422). If we can infer the HF dataset id,
2174                    // fall back to the hub file download path.
2175                    if let Some(hf_ds) = Self::hf_dataset_from_rows_url(&url) {
2176                        if let Ok((content, resolved)) =
2177                            self.download_hf_dataset_file_from_hub(&hf_ds)
2178                        {
2179                            return Ok((content, resolved));
2180                        }
2181                    }
2182                    return Err(e);
2183                }
2184            }
2185        }
2186
2187        // Regular download with retry logic
2188        const MAX_RETRIES: u32 = 3;
2189        const INITIAL_DELAY_SECS: u64 = 1;
2190
2191        let mut last_error = None;
2192
2193        for attempt in 0..=MAX_RETRIES {
2194            match self.download_attempt(&url) {
2195                Ok(content) => {
2196                    // Checksum validation removed - registry doesn't provide expected_checksum
2197                    // Future: Add checksum validation if registry provides it
2198
2199                    return Ok((content, url.clone()));
2200                }
2201                Err(e) => {
2202                    last_error = Some(e);
2203                    if attempt < MAX_RETRIES {
2204                        let delay_secs = (INITIAL_DELAY_SECS * (1 << attempt)).min(10);
2205                        log::warn!(
2206                            "Download attempt {} failed for {}, retrying in {}s...",
2207                            attempt + 1,
2208                            &url,
2209                            delay_secs
2210                        );
2211                        std::thread::sleep(std::time::Duration::from_secs(delay_secs));
2212                    }
2213                }
2214            }
2215        }
2216
2217        Err(last_error.unwrap_or_else(|| {
2218            Error::InvalidInput(format!(
2219                "Failed to download {} after {} retries. \
2220                 Check network connection and try again. \
2221                 URL: {}",
2222                id.name(),
2223                MAX_RETRIES + 1,
2224                &url
2225            ))
2226        }))
2227    }
2228
2229    /// Extract `<org>/<dataset>` from `https://huggingface.co/datasets/<org>/<dataset>(/...)`.
2230    #[cfg(feature = "eval")]
2231    fn extract_hf_dataset_name(url: &str) -> Option<String> {
2232        let marker = "huggingface.co/datasets/";
2233        let idx = url.find(marker)? + marker.len();
2234        let rest = &url[idx..];
2235        let rest = rest.split('?').next().unwrap_or(rest);
2236        let rest = rest.trim_matches('/');
2237        let mut parts = rest.split('/');
2238        let org = parts.next()?;
2239        let name = parts.next()?;
2240        Some(format!("{}/{}", org, name))
2241    }
2242
2243    #[cfg(feature = "eval")]
2244    fn url_encode_component(s: &str) -> String {
2245        // Minimal encoding for HF datasets-server query parameters.
2246        // In practice we mainly need `/` -> `%2F`.
2247        s.replace('/', "%2F").replace(' ', "%20")
2248    }
2249
2250    #[cfg(feature = "eval")]
2251    fn hf_rows_url(dataset: &str, config: &str, split: &str) -> String {
2252        format!(
2253            "https://datasets-server.huggingface.co/rows?dataset={}&config={}&split={}&offset=0&length=100",
2254            Self::url_encode_component(dataset),
2255            Self::url_encode_component(config),
2256            Self::url_encode_component(split),
2257        )
2258    }
2259
2260    /// Resolve a usable (config, split) pair for a HF dataset via datasets-server, with an optional
2261    /// preferred config.
2262    #[cfg(feature = "eval")]
2263    fn resolve_hf_config_split_prefer(
2264        &self,
2265        dataset: &str,
2266        preferred_config: Option<&str>,
2267    ) -> Result<(String, String)> {
2268        let url = format!(
2269            "https://datasets-server.huggingface.co/splits?dataset={}",
2270            Self::url_encode_component(dataset)
2271        );
2272        let response = ureq::get(&url)
2273            .timeout(std::time::Duration::from_secs(30))
2274            .call()
2275            .map_err(|e| Error::InvalidInput(format!("Failed to query HF splits: {}", e)))?;
2276
2277        if response.status() != 200 {
2278            return Err(Error::InvalidInput(format!(
2279                "HF splits query returned HTTP {} for dataset {}",
2280                response.status(),
2281                dataset
2282            )));
2283        }
2284
2285        let body = response.into_string().map_err(|e| {
2286            Error::InvalidInput(format!("Failed to read HF splits response: {}", e))
2287        })?;
2288
2289        let json: serde_json::Value = serde_json::from_str(&body).map_err(|e| {
2290            Error::InvalidInput(format!("Invalid JSON from HF splits endpoint: {}", e))
2291        })?;
2292
2293        let splits = json
2294            .get("splits")
2295            .and_then(|v| v.as_array())
2296            .ok_or_else(|| {
2297                Error::InvalidInput("HF splits response missing `splits`".to_string())
2298            })?;
2299
2300        // Prefer smaller/eval splits if present, otherwise pick the first split.
2301        // If `preferred_config` is provided, try to pick a split within that config first.
2302        let mut chosen = None;
2303        let prefer = preferred_config.map(|s| s.trim()).filter(|s| !s.is_empty());
2304
2305        for pass in 0..2 {
2306            for s in splits {
2307                let config = s.get("config").and_then(|v| v.as_str());
2308                let split = s.get("split").and_then(|v| v.as_str());
2309                if let (Some(config), Some(split)) = (config, split) {
2310                    // Pass 0: only consider preferred config. Pass 1: consider any config.
2311                    if pass == 0 {
2312                        if let Some(p) = prefer {
2313                            if config != p {
2314                                continue;
2315                            }
2316                        } else {
2317                            // No preference: skip pass 0.
2318                            break;
2319                        }
2320                    }
2321
2322                    if split == "test" {
2323                        chosen = Some((config.to_string(), split.to_string()));
2324                        break;
2325                    }
2326                    if chosen.is_none() {
2327                        chosen = Some((config.to_string(), split.to_string()));
2328                    }
2329                }
2330            }
2331            if chosen.is_some() {
2332                break;
2333            }
2334        }
2335
2336        chosen.ok_or_else(|| {
2337            Error::InvalidInput(format!(
2338                "HF splits endpoint returned no usable (config, split) for dataset {}",
2339                dataset
2340            ))
2341        })
2342    }
2343
2344    /// Best-effort fallback: download a raw dataset file from the HuggingFace Hub.
2345    ///
2346    /// This is a fallback for datasets that do not support the datasets-server row export.
2347    #[cfg(feature = "eval")]
2348    fn download_hf_dataset_file_from_hub(&self, dataset: &str) -> Result<(String, String)> {
2349        // Example API: https://huggingface.co/api/datasets/coref-data/preco_raw
2350        let api_url = format!("https://huggingface.co/api/datasets/{}", dataset);
2351        let response = ureq::get(&api_url)
2352            .timeout(std::time::Duration::from_secs(30))
2353            .call()
2354            .map_err(|e| {
2355                Error::InvalidInput(format!(
2356                    "Failed to query HuggingFace dataset metadata for {}: {}",
2357                    dataset, e
2358                ))
2359            })?;
2360
2361        if response.status() != 200 {
2362            return Err(Error::InvalidInput(format!(
2363                "HuggingFace dataset metadata request returned HTTP {} for {}",
2364                response.status(),
2365                dataset
2366            )));
2367        }
2368
2369        let body = response.into_string().map_err(|e| {
2370            Error::InvalidInput(format!(
2371                "Failed to read HuggingFace dataset metadata: {}",
2372                e
2373            ))
2374        })?;
2375        let json: serde_json::Value = serde_json::from_str(&body).map_err(|e| {
2376            Error::InvalidInput(format!(
2377                "Invalid JSON from HuggingFace dataset metadata: {}",
2378                e
2379            ))
2380        })?;
2381
2382        let siblings = json
2383            .get("siblings")
2384            .and_then(|v| v.as_array())
2385            .ok_or_else(|| {
2386                Error::InvalidInput("HF dataset metadata missing `siblings`".to_string())
2387            })?;
2388
2389        // Prefer smaller / eval-friendly splits, but fall back to whatever exists.
2390        //
2391        // Many HF dataset repos store raw source files (CoNLL/TSV/etc.) instead of JSONL,
2392        // often under subdirectories. We handle both.
2393        let preferred = [
2394            // JSONL (common)
2395            "dev.jsonl",
2396            "validation.jsonl",
2397            "test.jsonl",
2398            "train.jsonl",
2399        ];
2400        let mut chosen: Option<String> = None;
2401
2402        for want in preferred {
2403            if siblings.iter().any(|s| {
2404                s.get("rfilename")
2405                    .and_then(|v| v.as_str())
2406                    .is_some_and(|f| f == want)
2407            }) {
2408                chosen = Some(want.to_string());
2409                break;
2410            }
2411        }
2412
2413        if chosen.is_none() {
2414            // Otherwise, pick a file we can plausibly parse (allow subdirectories).
2415            //
2416            // Priority:
2417            // - contains split hint: test > dev/valid > train
2418            // - extension: jsonl > conll/bio/iob/tsv/txt
2419            //
2420            // This is still best-effort; the parser selection is handled elsewhere via the dataset id.
2421            fn split_rank(name: &str) -> u8 {
2422                let n = name.to_lowercase();
2423                if n.contains("test") {
2424                    0
2425                } else if n.contains("dev") || n.contains("valid") || n.contains("validation") {
2426                    1
2427                } else if n.contains("train") {
2428                    2
2429                } else {
2430                    3
2431                }
2432            }
2433            fn ext_rank(name: &str) -> u8 {
2434                let n = name.to_lowercase();
2435                if n.ends_with(".jsonl") {
2436                    0
2437                } else if n.ends_with(".conll")
2438                    || n.ends_with(".bio")
2439                    || n.ends_with(".iob")
2440                    || n.ends_with(".iob2")
2441                {
2442                    1
2443                } else if n.ends_with(".tsv") {
2444                    2
2445                } else if n.ends_with(".txt") {
2446                    3
2447                } else {
2448                    9
2449                }
2450            }
2451
2452            let mut candidates: Vec<String> = siblings
2453                .iter()
2454                .filter_map(|s| {
2455                    s.get("rfilename")
2456                        .and_then(|v| v.as_str())
2457                        .map(|f| f.to_string())
2458                })
2459                .filter(|f| ext_rank(f) < 9)
2460                .collect();
2461
2462            candidates.sort_by(|a, b| {
2463                (split_rank(a), ext_rank(a), a.len()).cmp(&(split_rank(b), ext_rank(b), b.len()))
2464            });
2465
2466            chosen = candidates.first().cloned();
2467        }
2468
2469        let Some(filename) = chosen else {
2470            return Err(Error::InvalidInput(format!(
2471                "No downloadable .jsonl file discovered for HF dataset {}",
2472                dataset
2473            )));
2474        };
2475
2476        // Download file via resolve (use main; can be made content-addressed later via `sha`).
2477        let file_url = format!(
2478            "https://huggingface.co/datasets/{}/resolve/main/{}",
2479            dataset, filename
2480        );
2481
2482        // Best-effort: avoid downloading huge files if a max byte limit is configured.
2483        if let Some(limit) = Self::max_download_bytes() {
2484            if let Ok(resp) = ureq::head(&file_url)
2485                .timeout(std::time::Duration::from_secs(30))
2486                .call()
2487            {
2488                if resp.status() == 200 {
2489                    if let Some(len) = resp
2490                        .header("Content-Length")
2491                        .and_then(|s| s.parse::<u64>().ok())
2492                    {
2493                        if len > limit {
2494                            return Err(Error::InvalidInput(format!(
2495                                "Download rejected ({} bytes > ANNO_MAX_DOWNLOAD_BYTES={} bytes) from {}",
2496                                len, limit, file_url
2497                            )));
2498                        }
2499                    }
2500                }
2501            }
2502        }
2503
2504        let content = self.download_attempt(&file_url)?;
2505        Ok((content, file_url))
2506    }
2507
2508    /// Placeholder when `eval` is disabled.
2509    #[cfg(not(feature = "eval"))]
2510    #[allow(dead_code)]
2511    fn download_hf_dataset_file_from_hub(&self, _dataset: &str) -> Result<(String, String)> {
2512        Err(Error::InvalidInput(
2513            "HuggingFace file fallback requires feature `eval`".to_string(),
2514        ))
2515    }
2516
2517    /// Download HuggingFace dataset with pagination support.
2518    ///
2519    /// HF datasets-server API limits responses to 100 rows by default.
2520    /// This function automatically paginates to download the full dataset.
2521    ///
2522    /// For datasets available on HuggingFace Hub, can also use hf-hub crate
2523    /// for direct file downloads (faster, no pagination needed).
2524    #[cfg(feature = "eval")]
2525    fn download_hf_dataset_paginated(&self, id: DatasetId, base_url: &str) -> Result<String> {
2526        // Try hf-hub direct download first (if available and dataset is on HF)
2527        #[cfg(feature = "onnx")] // hf-hub is available with onnx feature
2528        {
2529            if let Ok(content) = self.try_hf_hub_download(id) {
2530                return Ok(content);
2531            }
2532        }
2533
2534        // Fall back to paginated API download
2535        // HuggingFace datasets-server API limits to 100 rows per request
2536        const PAGE_SIZE: usize = 100;
2537        let mut all_rows = Vec::new();
2538        let mut features = None;
2539        let mut offset: usize = 0;
2540        let mut total_rows = None;
2541
2542        log::info!(
2543            "Downloading {} with pagination (page size: {})",
2544            id.name(),
2545            PAGE_SIZE
2546        );
2547
2548        loop {
2549            // Build paginated URL
2550            let url = if base_url.contains("offset=") {
2551                // Replace existing offset parameter
2552                let prev_offset = offset.saturating_sub(PAGE_SIZE);
2553                base_url
2554                    .replace(
2555                        &format!("offset={}", prev_offset),
2556                        &format!("offset={}", offset),
2557                    )
2558                    .replace("length=100", &format!("length={}", PAGE_SIZE))
2559            } else {
2560                // Add pagination parameters
2561                let separator = if base_url.contains('?') { "&" } else { "?" };
2562                format!(
2563                    "{}{}offset={}&length={}",
2564                    base_url, separator, offset, PAGE_SIZE
2565                )
2566            };
2567
2568            match self.download_attempt(&url) {
2569                Ok(content) => {
2570                    let parsed: serde_json::Value =
2571                        serde_json::from_str(&content).map_err(|e| {
2572                            Error::InvalidInput(format!("Invalid JSON response: {}", e))
2573                        })?;
2574
2575                    // Extract features (only from first page)
2576                    if features.is_none() {
2577                        features = parsed.get("features").cloned();
2578                    }
2579
2580                    // Extract total number of rows (if available)
2581                    if total_rows.is_none() {
2582                        total_rows = parsed
2583                            .get("num_rows_total")
2584                            .and_then(|v| v.as_u64())
2585                            .map(|n| n as usize);
2586                    }
2587
2588                    // Extract rows from this page
2589                    if let Some(rows) = parsed.get("rows").and_then(|v| v.as_array()) {
2590                        if rows.is_empty() {
2591                            break; // No more rows
2592                        }
2593                        all_rows.extend_from_slice(rows);
2594                        log::debug!(
2595                            "Downloaded {} rows (total so far: {})",
2596                            rows.len(),
2597                            all_rows.len()
2598                        );
2599
2600                        // Check if we've got all rows
2601                        if let Some(total) = total_rows {
2602                            if all_rows.len() >= total {
2603                                break;
2604                            }
2605                        } else if rows.len() < PAGE_SIZE {
2606                            // No total available, but got fewer rows than requested = last page
2607                            break;
2608                        }
2609
2610                        offset += PAGE_SIZE;
2611                    } else {
2612                        // No rows in response, might be error or empty dataset
2613                        break;
2614                    }
2615                }
2616                Err(e) => {
2617                    // If we got some rows, return partial dataset with warning
2618                    if !all_rows.is_empty() {
2619                        log::warn!(
2620                            "Failed to download full {} dataset (got {} rows before error: {}). \
2621                             Returning partial dataset.",
2622                            id.name(),
2623                            all_rows.len(),
2624                            e
2625                        );
2626                        break;
2627                    } else {
2628                        return Err(e);
2629                    }
2630                }
2631            }
2632
2633            // Safety limit: prevent infinite loops
2634            if offset > 1_000_000 {
2635                log::warn!(
2636                    "Reached safety limit (1M rows) for {}. Returning partial dataset ({} rows).",
2637                    id.name(),
2638                    all_rows.len()
2639                );
2640                break;
2641            }
2642        }
2643
2644        // Reconstruct full API response format
2645        let mut response: serde_json::Value = serde_json::json!({
2646            "rows": all_rows,
2647        });
2648
2649        if let Some(features_val) = features {
2650            response["features"] = features_val;
2651        }
2652
2653        if let Some(total) = total_rows {
2654            response["num_rows_total"] = serde_json::json!(total);
2655        }
2656
2657        serde_json::to_string(&response).map_err(|e| {
2658            Error::InvalidInput(format!("Failed to serialize paginated response: {}", e))
2659        })
2660    }
2661
2662    /// Try downloading dataset directly from HuggingFace Hub using hf-hub crate.
2663    ///
2664    /// This is faster than paginated API calls for datasets available on HF Hub.
2665    /// Returns Ok(content) if successful, Err if not available or hf-hub not enabled.
2666    ///
2667    /// Uses `HF_TOKEN` environment variable if set for accessing gated datasets.
2668    #[cfg(all(feature = "eval", feature = "onnx"))]
2669    fn try_hf_hub_download(&self, id: DatasetId) -> Result<String> {
2670        use hf_hub::api::sync::{Api, ApiBuilder};
2671
2672        // Map dataset IDs to HuggingFace dataset names and file paths
2673        let (dataset_name, file_path) = match id {
2674            DatasetId::MultiNERD => ("Babelscape/multinerd", "test/test_en.jsonl"),
2675            DatasetId::TweetNER7 => ("tner/tweetner7", "dataset/2020.dev.json"),
2676            DatasetId::BroadTwitterCorpus => ("GateNLP/broad_twitter_corpus", "test/a.conll"),
2677            DatasetId::CADEC => ("KevinSpaghetti/cadec", "data/test.jsonl"),
2678            DatasetId::PreCo => ("coref-data/preco", "data/test.jsonl"),
2679            // Gated datasets that require HF_TOKEN
2680            DatasetId::MultiCoNER => ("MultiCoNER/multiconer_v1", "en/test.conll"),
2681            DatasetId::MultiCoNERv2 => ("MultiCoNER/multiconer_v2", "en/test.conll"),
2682            _ => {
2683                return Err(Error::InvalidInput(
2684                    "Dataset not available via hf-hub".to_string(),
2685                ))
2686            }
2687        };
2688
2689        // Load .env for HF_TOKEN if not already set
2690        anno::env::load_dotenv();
2691
2692        // Use HF_TOKEN from environment if available (for gated datasets)
2693        let api = if let Some(token) = anno::env::hf_token() {
2694            ApiBuilder::new()
2695                .with_token(Some(token))
2696                .build()
2697                .map_err(|e| {
2698                    Error::InvalidInput(format!(
2699                        "Failed to initialize HuggingFace API with token: {}",
2700                        e
2701                    ))
2702                })?
2703        } else {
2704            Api::new().map_err(|e| {
2705                Error::InvalidInput(format!("Failed to initialize HuggingFace API: {}", e))
2706            })?
2707        };
2708
2709        let repo = api.dataset(dataset_name.to_string());
2710        let file_path_buf = repo.get(file_path).map_err(|e| {
2711            Error::InvalidInput(format!(
2712                "Failed to download {} from HuggingFace Hub: {}. \
2713                 Falling back to HTTP download.",
2714                file_path, e
2715            ))
2716        })?;
2717
2718        std::fs::read_to_string(&file_path_buf)
2719            .map_err(|e| Error::InvalidInput(format!("Failed to read downloaded file: {}", e)))
2720    }
2721
2722    /// Placeholder for when hf-hub is not available.
2723    #[cfg(not(all(feature = "eval", feature = "onnx")))]
2724    #[allow(dead_code)] // Part of trait interface, may be unused in some feature combinations
2725    fn try_hf_hub_download(&self, _id: DatasetId) -> Result<String> {
2726        Err(Error::InvalidInput("hf-hub not available".to_string()))
2727    }
2728
2729    /// Single download attempt with retry logic.
2730    ///
2731    /// Retries up to 3 times with exponential backoff for transient failures (timeouts, 5xx errors).
2732    #[cfg(feature = "eval")]
2733    fn download_attempt(&self, url: &str) -> Result<String> {
2734        const MAX_RETRIES: usize = 3;
2735        const INITIAL_TIMEOUT_SECS: u64 = 30;
2736        const MAX_TIMEOUT_SECS: u64 = 120;
2737
2738        let mut last_error = None;
2739
2740        for attempt in 0..=MAX_RETRIES {
2741            let timeout_secs = (INITIAL_TIMEOUT_SECS * (1 << attempt.min(2))).min(MAX_TIMEOUT_SECS);
2742
2743            match ureq::get(url)
2744                .timeout(std::time::Duration::from_secs(timeout_secs))
2745                .call()
2746            {
2747                Ok(response) => {
2748                    if response.status() == 200 {
2749                        // Success - read content
2750                        let content = response.into_string().map_err(|e| {
2751                            Error::InvalidInput(format!(
2752                                "Failed to read response from {}: {}. \
2753                                 Response may be too large or corrupted.",
2754                                url, e
2755                            ))
2756                        })?;
2757
2758                        // Heuristic guardrail: many registry URLs are dataset homepages. If we fetch HTML,
2759                        // treat it as a download failure so we don't silently cache garbage and then parse
2760                        // "0 sentences".
2761                        let lower = content
2762                            .chars()
2763                            .take(2048)
2764                            .collect::<String>()
2765                            .to_lowercase();
2766                        if lower.contains("<html") || lower.contains("<!doctype html") {
2767                            return Err(Error::InvalidInput(format!(
2768                                "Downloaded HTML from {}. This URL looks like a webpage, not a raw dataset file.",
2769                                url
2770                            )));
2771                        }
2772
2773                        return Ok(content);
2774                    }
2775
2776                    let status = response.status();
2777                    let body = response.into_string().unwrap_or_default();
2778                    let body = body.trim();
2779                    let body_preview = if body.len() > 800 {
2780                        format!("{}…", &body[..800])
2781                    } else {
2782                        body.to_string()
2783                    };
2784
2785                    // Retry on 5xx server errors (transient)
2786                    if (500..600).contains(&status) && attempt < MAX_RETRIES {
2787                        let wait_ms = 1000 * (1 << attempt); // Exponential backoff: 1s, 2s, 4s
2788                        log::debug!(
2789                            "Server error {} downloading {} (attempt {}/{}), retrying in {}ms...",
2790                            status,
2791                            url,
2792                            attempt + 1,
2793                            MAX_RETRIES + 1,
2794                            wait_ms
2795                        );
2796                        std::thread::sleep(std::time::Duration::from_millis(wait_ms));
2797                        last_error = Some(format!(
2798                            "HTTP {} downloading {}. Server returned error status. {}{}",
2799                            status,
2800                            url,
2801                            if body_preview.is_empty() {
2802                                ""
2803                            } else {
2804                                "Response body: "
2805                            },
2806                            body_preview
2807                        ));
2808                        continue;
2809                    }
2810
2811                    // Non-retryable error (4xx client errors)
2812                    return Err(Error::InvalidInput(format!(
2813                        "HTTP {} downloading {}. \
2814                         Server returned error status. \
2815                         Dataset may be temporarily unavailable or URL changed. {}{}",
2816                        status,
2817                        url,
2818                        if body_preview.is_empty() {
2819                            ""
2820                        } else {
2821                            "Response body: "
2822                        },
2823                        body_preview
2824                    )));
2825                }
2826                Err(ureq::Error::Transport(e)) => {
2827                    // Network errors (timeouts, connection failures) - retry
2828                    let error_msg = format!("{}", e);
2829                    let is_timeout =
2830                        error_msg.contains("timeout") || error_msg.contains("timed out");
2831
2832                    if is_timeout && attempt < MAX_RETRIES {
2833                        let wait_ms = 1000 * (1 << attempt); // Exponential backoff
2834                        log::debug!(
2835                            "Timeout downloading {} (attempt {}/{}), retrying in {}ms...",
2836                            url,
2837                            attempt + 1,
2838                            MAX_RETRIES + 1,
2839                            wait_ms
2840                        );
2841                        std::thread::sleep(std::time::Duration::from_millis(wait_ms));
2842                        last_error = Some(error_msg);
2843                        continue;
2844                    }
2845
2846                    // Final attempt failed or non-timeout error
2847                    return Err(Error::InvalidInput(format!(
2848                        "Network error downloading {}: {}. \
2849                         Check your internet connection and try again. \
2850                         {}",
2851                        url,
2852                        error_msg,
2853                        if attempt > 0 {
2854                            format!("(Failed after {} retries)", attempt)
2855                        } else {
2856                            String::new()
2857                        }
2858                    )));
2859                }
2860                Err(e) => {
2861                    // Other errors (non-retryable)
2862                    let error_msg = format!("{}", e);
2863                    return Err(Error::InvalidInput(format!(
2864                        "Error downloading {}: {}. \
2865                         Check your internet connection and try again.",
2866                        url, error_msg
2867                    )));
2868                }
2869            }
2870        }
2871
2872        // All retries exhausted
2873        Err(Error::InvalidInput(format!(
2874            "Failed to download {} after {} attempts. Last error: {}",
2875            url,
2876            MAX_RETRIES + 1,
2877            last_error.unwrap_or_else(|| "unknown error".to_string())
2878        )))
2879    }
2880
2881    /// Single download attempt returning raw bytes (for binary formats like ZIP).
2882    ///
2883    /// Same retry/backoff logic as `download_attempt` but reads response as bytes.
2884    #[cfg(feature = "eval")]
2885    fn download_attempt_bytes(&self, url: &str) -> Result<Vec<u8>> {
2886        const MAX_RETRIES: usize = 3;
2887        const INITIAL_TIMEOUT_SECS: u64 = 30;
2888        const MAX_TIMEOUT_SECS: u64 = 120;
2889        const MAX_BYTES: usize = 50 * 1024 * 1024; // 50 MB limit
2890
2891        let mut last_error = None;
2892
2893        for attempt in 0..=MAX_RETRIES {
2894            let timeout_secs = (INITIAL_TIMEOUT_SECS * (1 << attempt.min(2))).min(MAX_TIMEOUT_SECS);
2895
2896            match ureq::get(url)
2897                .timeout(std::time::Duration::from_secs(timeout_secs))
2898                .call()
2899            {
2900                Ok(response) => {
2901                    if response.status() == 200 {
2902                        use std::io::Read as _;
2903                        let mut bytes = Vec::new();
2904                        response
2905                            .into_reader()
2906                            .take(MAX_BYTES as u64)
2907                            .read_to_end(&mut bytes)
2908                            .map_err(|e| {
2909                                Error::InvalidInput(format!(
2910                                    "Failed to read bytes from {}: {}",
2911                                    url, e
2912                                ))
2913                            })?;
2914                        return Ok(bytes);
2915                    }
2916
2917                    let status = response.status();
2918                    if (500..600).contains(&status) && attempt < MAX_RETRIES {
2919                        let wait_ms = 1000 * (1 << attempt);
2920                        std::thread::sleep(std::time::Duration::from_millis(wait_ms));
2921                        last_error = Some(format!("HTTP {} from {}", status, url));
2922                        continue;
2923                    }
2924
2925                    return Err(Error::InvalidInput(format!(
2926                        "HTTP {} downloading binary from {}",
2927                        status, url
2928                    )));
2929                }
2930                Err(ureq::Error::Transport(e)) => {
2931                    let msg = format!("{}", e);
2932                    if (msg.contains("timeout") || msg.contains("timed out"))
2933                        && attempt < MAX_RETRIES
2934                    {
2935                        let wait_ms = 1000 * (1 << attempt);
2936                        std::thread::sleep(std::time::Duration::from_millis(wait_ms));
2937                        last_error = Some(msg);
2938                        continue;
2939                    }
2940                    return Err(Error::InvalidInput(format!(
2941                        "Network error downloading {}: {}",
2942                        url, msg
2943                    )));
2944                }
2945                Err(e) => {
2946                    return Err(Error::InvalidInput(format!(
2947                        "Error downloading {}: {}",
2948                        url, e
2949                    )));
2950                }
2951            }
2952        }
2953
2954        Err(Error::InvalidInput(format!(
2955            "Failed to download {} after {} attempts. Last error: {}",
2956            url,
2957            MAX_RETRIES + 1,
2958            last_error.unwrap_or_else(|| "unknown".to_string())
2959        )))
2960    }
2961
2962    /// Compute SHA256 checksum of content.
2963    #[cfg(feature = "eval")]
2964    fn compute_sha256(&self, content: &str) -> String {
2965        #[cfg(feature = "eval")]
2966        {
2967            use sha2::{Digest, Sha256};
2968            let mut hasher = Sha256::new();
2969            hasher.update(content.as_bytes());
2970            format!("{:x}", hasher.finalize())
2971        }
2972        #[cfg(not(feature = "eval"))]
2973        {
2974            // Fallback if sha2 not available
2975            use std::collections::hash_map::DefaultHasher;
2976            use std::hash::{Hash, Hasher};
2977            let mut hasher = DefaultHasher::new();
2978            content.hash(&mut hasher);
2979            format!("{:x}", hasher.finish())
2980        }
2981    }
2982
2983    /// Get temporal metadata for a dataset if available.
2984    fn get_temporal_metadata(id: DatasetId) -> Option<TemporalMetadata> {
2985        match id {
2986            DatasetId::TweetNER7 => {
2987                // TweetNER7 has temporal entity recognition - use dataset creation date as cutoff
2988                Some(TemporalMetadata {
2989                    kb_version: None,                                // No KB linking in TweetNER7
2990                    temporal_cutoff: Some("2017-01-01".to_string()), // Approximate dataset creation
2991                    entity_creation_dates: None,                     // Would need entity linking
2992                })
2993            }
2994            DatasetId::BroadTwitterCorpus => {
2995                // BroadTwitterCorpus is stratified across times - use approximate cutoff
2996                Some(TemporalMetadata {
2997                    kb_version: None,
2998                    temporal_cutoff: Some("2018-01-01".to_string()), // Approximate
2999                    entity_creation_dates: None,
3000                })
3001            }
3002            DatasetId::BC5CDR
3003            | DatasetId::NCBIDisease
3004            | DatasetId::GENIA
3005            | DatasetId::AnatEM
3006            | DatasetId::BC2GM
3007            | DatasetId::BC4CHEMD => {
3008                // Biomedical datasets might have KB versions (UMLS, etc.)
3009                Some(TemporalMetadata {
3010                    kb_version: None,
3011                    temporal_cutoff: None,
3012                    entity_creation_dates: None,
3013                })
3014            }
3015            _ => None, // Most datasets don't have temporal metadata
3016        }
3017    }
3018
3019    /// Parse content based on dataset format.
3020    ///
3021    /// This is intentionally `pub` (behind the `eval` module) to enable offline evaluation
3022    /// workflows and integration tests that want to exercise parsing without network I/O.
3023    ///
3024    /// It does **not** download or read from cache. For typical usage, prefer
3025    /// [`DatasetLoader::load_or_download`] / [`DatasetLoader::load`].
3026    pub fn parse_content_str(&self, content: &str, id: DatasetId) -> Result<LoadedDataset> {
3027        self.parse_content_impl(content, id)
3028    }
3029
3030    /// Internal implementation for parsing a dataset payload.
3031    fn parse_content_impl(&self, content: &str, id: DatasetId) -> Result<LoadedDataset> {
3032        if content.trim().is_empty() {
3033            return Err(Error::InvalidInput(format!(
3034                "Dataset {:?} file is empty",
3035                id
3036            )));
3037        }
3038
3039        let plan = LoadableDatasetId::parse_plan(id).ok_or_else(|| {
3040            Error::InvalidInput(format!("No parser configured for dataset {:?}", id))
3041        })?;
3042
3043        let result = match plan {
3044            DatasetParsePlan::Conll => self.parse_conll(content, id),
3045            DatasetParsePlan::JsonlNer => self.parse_jsonl_ner(content, id),
3046            DatasetParsePlan::WikiannJson => self.parse_wikiann_json(content, id),
3047            DatasetParsePlan::TweetNer7 => self.parse_tweetner7(content, id),
3048            DatasetParsePlan::DocredJson => self.parse_docred(content, id),
3049            DatasetParsePlan::GoogleReCorpus => self.parse_google_re_corpus(content, id),
3050            DatasetParsePlan::ChisiecJson => self.parse_chisiec(content, id),
3051            DatasetParsePlan::CadecHybrid => {
3052                if self.is_hf_api_response(content) {
3053                    self.parse_cadec_hf_api(content, id)
3054                } else {
3055                    self.parse_cadec_jsonl(content, id)
3056                }
3057            }
3058            DatasetParsePlan::Bc5cdr => self.parse_bc5cdr(content, id),
3059            DatasetParsePlan::NcbiDisease => self.parse_ncbi_disease(content, id),
3060            DatasetParsePlan::GapTsv => self.parse_gap(content, id),
3061            DatasetParsePlan::PrecoJsonl => self.parse_preco_jsonl(content, id),
3062            DatasetParsePlan::Litbank => self.parse_litbank(content, id),
3063            DatasetParsePlan::EcbPlus => self.parse_ecb_plus(content, id),
3064            DatasetParsePlan::AfriSenti => self.parse_afrisenti(content, id),
3065            DatasetParsePlan::AfriQa => self.parse_afriqa(content, id),
3066            DatasetParsePlan::MasakhaNews => self.parse_masakhanews(content, id),
3067            DatasetParsePlan::Conllu => self.parse_conllu(content, id),
3068            DatasetParsePlan::AgNews => self.parse_agnews(content, id),
3069            DatasetParsePlan::Dbpedia14 => self.parse_dbpedia14(content, id),
3070            DatasetParsePlan::YahooAnswers => self.parse_yahoo_answers(content, id),
3071            DatasetParsePlan::Trec => self.parse_trec(content, id),
3072            DatasetParsePlan::TweetTopic => self.parse_tweettopic(content, id),
3073            DatasetParsePlan::Maven => self.parse_maven(content, id),
3074            DatasetParsePlan::MavenArg => self.parse_maven_arg(content, id),
3075            DatasetParsePlan::Casie => self.parse_casie(content, id),
3076            DatasetParsePlan::Rams => self.parse_rams(content, id),
3077            DatasetParsePlan::HfApiResponse => self.parse_hf_api_response(content, id),
3078            DatasetParsePlan::TsvNer => self.parse_tsv_ner(content, id),
3079            DatasetParsePlan::CsvNer => self.parse_csv_ner(content, id),
3080        }?;
3081
3082        // Validate parsed dataset is not empty
3083        if result.sentences.is_empty() {
3084            return Err(Error::InvalidInput(format!(
3085                "Dataset {:?} parsed successfully but contains no sentences. \
3086                 This may indicate a parsing issue or empty dataset file.",
3087                id
3088            )));
3089        }
3090
3091        Ok(result)
3092    }
3093
3094    /// Check if content is HuggingFace datasets-server API response.
3095    fn is_hf_api_response(&self, content: &str) -> bool {
3096        // Check for HF datasets-server API response structure
3097        let trimmed = content.trim_start();
3098        trimmed.starts_with("{\"rows\":")
3099            || trimmed.starts_with("{\"features\":")
3100            || (trimmed.starts_with("{")
3101                && trimmed.contains("\"rows\":[")
3102                && trimmed.contains("\"features\":["))
3103    }
3104
3105    /// Parse CoNLL/BIO format content.
3106    fn parse_conll(&self, content: &str, id: DatasetId) -> Result<LoadedDataset> {
3107        let mut sentences = Vec::new();
3108        let mut current_tokens = Vec::new();
3109
3110        // Detect format: MIT datasets use TAB separator with TAG first
3111        let is_mit_format = matches!(id, DatasetId::MitMovie | DatasetId::MitRestaurant);
3112
3113        for line in content.lines() {
3114            let line = line.trim();
3115
3116            // Empty line = sentence boundary
3117            if line.is_empty() {
3118                if !current_tokens.is_empty() {
3119                    sentences.push(AnnotatedSentence {
3120                        tokens: std::mem::take(&mut current_tokens),
3121                        source_dataset: id,
3122                    });
3123                }
3124                continue;
3125            }
3126
3127            // Skip document markers
3128            if line.starts_with("-DOCSTART-") {
3129                continue;
3130            }
3131
3132            // Parse based on format
3133            let (text, ner_tag) = if is_mit_format {
3134                // MIT format: TAG\tword (tab-separated)
3135                let parts: Vec<&str> = line.split('\t').collect();
3136                if parts.len() >= 2 {
3137                    (parts[1].to_string(), parts[0].to_string())
3138                } else {
3139                    continue;
3140                }
3141            } else {
3142                // Standard CoNLL/BIO format (space-separated)
3143                let parts: Vec<&str> = line.split_whitespace().collect();
3144                if parts.is_empty() {
3145                    continue;
3146                }
3147
3148                if parts.len() >= 4 {
3149                    // CoNLL-2003 format: word POS chunk NER
3150                    (parts[0].to_string(), parts[3].to_string())
3151                } else if parts.len() >= 2 {
3152                    // BIO format: word NER
3153                    (parts[0].to_string(), parts[parts.len() - 1].to_string())
3154                } else {
3155                    // Single column - assume O tag
3156                    (parts[0].to_string(), "O".to_string())
3157                }
3158            };
3159
3160            // Normalize BIO tags: some corpora (e.g., WikiGold) use `I-XXX` even at the
3161            // beginning of an entity. Treat such `I-` tags as `B-` when they don't
3162            // continue an entity of the same type.
3163            let ner_tag = if let Some(label) = ner_tag.strip_prefix("I-") {
3164                let continues_same = current_tokens
3165                    .last()
3166                    .map(|t| t.ner_tag.as_str())
3167                    .is_some_and(|prev| {
3168                        (prev.starts_with("B-") || prev.starts_with("I-"))
3169                            && prev.get(2..).is_some_and(|prev_label| prev_label == label)
3170                    });
3171
3172                if continues_same {
3173                    ner_tag
3174                } else {
3175                    format!("B-{}", label)
3176                }
3177            } else {
3178                ner_tag
3179            };
3180
3181            current_tokens.push(AnnotatedToken { text, ner_tag });
3182        }
3183
3184        // Don't forget last sentence
3185        if !current_tokens.is_empty() {
3186            sentences.push(AnnotatedSentence {
3187                tokens: current_tokens,
3188                source_dataset: id,
3189            });
3190        }
3191
3192        if sentences.is_empty() {
3193            return Err(Error::InvalidInput(format!(
3194                "CoNLL file for {:?} contains no valid sentences",
3195                id
3196            )));
3197        }
3198
3199        let now = chrono::Utc::now().to_rfc3339();
3200
3201        Ok(LoadedDataset {
3202            id,
3203            sentences,
3204            loaded_at: now,
3205            source_url: id.download_url().to_string(),
3206            data_source: DataSource::LocalCache,
3207            temporal_metadata: Self::get_temporal_metadata(id),
3208            metadata: id.default_metadata(),
3209        })
3210    }
3211
3212    /// Parse JSONL NER format (HuggingFace style, e.g., MultiNERD).
3213    ///
3214    /// Expected format: `{"tokens": ["word1", "word2"], "ner_tags": [0, 1, 0]}`
3215    fn parse_jsonl_ner(&self, content: &str, id: DatasetId) -> Result<LoadedDataset> {
3216        let mut sentences = Vec::new();
3217
3218        // MultiNERD tag mapping (index -> label)
3219        let tag_labels = [
3220            "O", "B-PER", "I-PER", "B-ORG", "I-ORG", "B-LOC", "I-LOC", "B-ANIM", "I-ANIM", "B-BIO",
3221            "I-BIO", "B-CEL", "I-CEL", "B-DIS", "I-DIS", "B-EVE", "I-EVE", "B-FOOD", "I-FOOD",
3222            "B-INST", "I-INST", "B-MEDIA", "I-MEDIA", "B-MYTH", "I-MYTH", "B-PLANT", "I-PLANT",
3223            "B-TIME", "I-TIME", "B-VEHI", "I-VEHI",
3224        ];
3225
3226        for line in content.lines() {
3227            let line = line.trim();
3228            if line.is_empty() {
3229                continue;
3230            }
3231
3232            // Parse JSON line
3233            let parsed: serde_json::Value = match serde_json::from_str(line) {
3234                Ok(v) => v,
3235                Err(_) => continue, // Skip malformed lines
3236            };
3237
3238            let tokens = match parsed.get("tokens").and_then(|v| v.as_array()) {
3239                Some(t) => t,
3240                None => continue,
3241            };
3242
3243            let ner_tags = match parsed.get("ner_tags").and_then(|v| v.as_array()) {
3244                Some(t) => t,
3245                None => continue,
3246            };
3247
3248            if tokens.len() != ner_tags.len() {
3249                continue; // Skip malformed entries
3250            }
3251
3252            let mut annotated_tokens = Vec::new();
3253            for (token, tag) in tokens.iter().zip(ner_tags.iter()) {
3254                let text = token.as_str().unwrap_or("").to_string();
3255                let tag_idx = tag.as_u64().unwrap_or(0) as usize;
3256                let ner_tag = tag_labels.get(tag_idx).unwrap_or(&"O").to_string();
3257                annotated_tokens.push(AnnotatedToken { text, ner_tag });
3258            }
3259
3260            if !annotated_tokens.is_empty() {
3261                sentences.push(AnnotatedSentence {
3262                    tokens: annotated_tokens,
3263                    source_dataset: id,
3264                });
3265            }
3266        }
3267
3268        if sentences.is_empty() {
3269            return Err(Error::InvalidInput(format!(
3270                "JSONL NER file for {:?} contains no valid sentences",
3271                id
3272            )));
3273        }
3274
3275        let now = chrono::Utc::now().to_rfc3339();
3276        Ok(LoadedDataset {
3277            id,
3278            sentences,
3279            loaded_at: now,
3280            source_url: id.download_url().to_string(),
3281            data_source: DataSource::LocalCache,
3282            temporal_metadata: Self::get_temporal_metadata(id),
3283            metadata: id.default_metadata(),
3284        })
3285    }
3286
3287    /// Parse HuggingFace datasets-server API response.
3288    ///
3289    /// Expected format:
3290    /// ```json
3291    /// {
3292    ///   "features": [{"name": "tokens", ...}, {"name": "ner_tags", ...}],
3293    ///   "rows": [{"row_idx": 0, "row": {"tokens": [...], "ner_tags": [...]}}, ...]
3294    /// }
3295    /// ```
3296    fn parse_hf_api_response(&self, content: &str, id: DatasetId) -> Result<LoadedDataset> {
3297        let parsed: serde_json::Value = serde_json::from_str(content)
3298            .map_err(|e| Error::InvalidInput(format!("Failed to parse HF API response: {}", e)))?;
3299
3300        let mut sentences = Vec::new();
3301
3302        // Extract tag names from features if available (for integer tag mapping)
3303        let tag_names = self.extract_tag_names_from_features(&parsed);
3304        let class_names = self.extract_class_names_from_features(&parsed);
3305
3306        let rows = parsed
3307            .get("rows")
3308            .and_then(|v| v.as_array())
3309            .ok_or_else(|| Error::InvalidInput("No 'rows' array in HF API response".to_string()))?;
3310
3311        for row_obj in rows {
3312            let row = match row_obj.get("row") {
3313                Some(r) => r,
3314                None => continue,
3315            };
3316
3317            // Primary path: token-level NER rows (`tokens` + `ner_tags`)
3318            if let (Some(tokens), Some(ner_tags)) = (
3319                row.get("tokens").and_then(|v| v.as_array()),
3320                row.get("ner_tags").and_then(|v| v.as_array()),
3321            ) {
3322                if tokens.len() != ner_tags.len() {
3323                    continue;
3324                }
3325
3326                let mut annotated_tokens = Vec::new();
3327                for (token, tag) in tokens.iter().zip(ner_tags.iter()) {
3328                    let text = token.as_str().unwrap_or("").to_string();
3329
3330                    // Handle both integer and string tags
3331                    let ner_tag = if let Some(tag_idx) = tag.as_u64() {
3332                        // Integer tag - map using feature names or default
3333                        tag_names
3334                            .get(tag_idx as usize)
3335                            .cloned()
3336                            .unwrap_or_else(|| format!("TAG_{}", tag_idx))
3337                    } else if let Some(tag_str) = tag.as_str() {
3338                        // String tag - use directly
3339                        tag_str.to_string()
3340                    } else {
3341                        "O".to_string()
3342                    };
3343
3344                    annotated_tokens.push(AnnotatedToken { text, ner_tag });
3345                }
3346
3347                if !annotated_tokens.is_empty() {
3348                    sentences.push(AnnotatedSentence {
3349                        tokens: annotated_tokens,
3350                        source_dataset: id,
3351                    });
3352                }
3353                continue;
3354            }
3355
3356            // Fallback: temporal standoff rows (`text` + `*_expressions` with char offsets).
3357            //
3358            // Some HF datasets expose TIMEX/EVENT spans as character ranges in lists like:
3359            // - `time_expressions`: [{start_char, end_char, ...}, ...]
3360            // - `event_expressions`: ...
3361            // - `signal_expressions`: ...
3362            //
3363            // We tokenize `text` by whitespace and assign BIO tags based on overlap with any span.
3364            if let Some(text) = row.get("text").and_then(|v| v.as_str()).map(str::trim) {
3365                let has_temporal_spans = row
3366                    .get("time_expressions")
3367                    .and_then(|v| v.as_array())
3368                    .is_some()
3369                    || row
3370                        .get("event_expressions")
3371                        .and_then(|v| v.as_array())
3372                        .is_some()
3373                    || row
3374                        .get("signal_expressions")
3375                        .and_then(|v| v.as_array())
3376                        .is_some();
3377
3378                if has_temporal_spans && !text.is_empty() {
3379                    fn spans_from_array(
3380                        arr: Option<&Vec<serde_json::Value>>,
3381                    ) -> Vec<(usize, usize)> {
3382                        let mut out = Vec::new();
3383                        let Some(arr) = arr else { return out };
3384                        for item in arr {
3385                            let Some(s) = item.get("start_char").and_then(|v| v.as_u64()) else {
3386                                continue;
3387                            };
3388                            let Some(e) = item.get("end_char").and_then(|v| v.as_u64()) else {
3389                                continue;
3390                            };
3391                            let s = s as usize;
3392                            let e = e as usize;
3393                            if e > s {
3394                                out.push((s, e));
3395                            }
3396                        }
3397                        out
3398                    }
3399
3400                    fn overlaps(token_s: usize, token_e: usize, spans: &[(usize, usize)]) -> bool {
3401                        spans.iter().any(|(s, e)| token_s < *e && token_e > *s)
3402                    }
3403
3404                    // Tokenize by whitespace, tracking char offsets (Rust `char` count).
3405                    let mut tokens: Vec<(String, usize, usize)> = Vec::new();
3406                    let mut cur = String::new();
3407                    let mut cur_start: Option<usize> = None;
3408
3409                    for (i, ch) in text.chars().enumerate() {
3410                        if ch.is_whitespace() {
3411                            if let Some(s) = cur_start.take() {
3412                                let e = i;
3413                                if !cur.is_empty() {
3414                                    tokens.push((std::mem::take(&mut cur), s, e));
3415                                } else {
3416                                    cur.clear();
3417                                }
3418                            }
3419                        } else {
3420                            if cur_start.is_none() {
3421                                cur_start = Some(i);
3422                            }
3423                            cur.push(ch);
3424                        }
3425                    }
3426                    if let Some(s) = cur_start.take() {
3427                        let e = text.chars().count();
3428                        if !cur.is_empty() {
3429                            tokens.push((cur, s, e));
3430                        }
3431                    }
3432
3433                    if tokens.is_empty() {
3434                        continue;
3435                    }
3436
3437                    let timex_spans =
3438                        spans_from_array(row.get("time_expressions").and_then(|v| v.as_array()));
3439                    let event_spans =
3440                        spans_from_array(row.get("event_expressions").and_then(|v| v.as_array()));
3441                    let signal_spans =
3442                        spans_from_array(row.get("signal_expressions").and_then(|v| v.as_array()));
3443
3444                    let mut annotated_tokens = Vec::with_capacity(tokens.len());
3445                    let mut prev_label: Option<&'static str> = None;
3446                    for (tok, s, e) in tokens {
3447                        let label = if overlaps(s, e, &timex_spans) {
3448                            Some("TIMEX")
3449                        } else if overlaps(s, e, &event_spans) {
3450                            Some("EVENT")
3451                        } else if overlaps(s, e, &signal_spans) {
3452                            Some("SIGNAL")
3453                        } else {
3454                            None
3455                        };
3456
3457                        let ner_tag = match (label, prev_label) {
3458                            (None, _) => {
3459                                prev_label = None;
3460                                "O".to_string()
3461                            }
3462                            (Some(l), Some(p)) if l == p => format!("I-{}", l),
3463                            (Some(l), _) => {
3464                                prev_label = Some(l);
3465                                format!("B-{}", l)
3466                            }
3467                        };
3468
3469                        annotated_tokens.push(AnnotatedToken { text: tok, ner_tag });
3470                    }
3471
3472                    sentences.push(AnnotatedSentence {
3473                        tokens: annotated_tokens,
3474                        source_dataset: id,
3475                    });
3476                    continue;
3477                }
3478            }
3479
3480            // Fallback: DISRPT-style CoNLL-U rows for discourse segmentation.
3481            //
3482            // DISRPT `*.conllu` configs are exported via HF datasets-server as arrays:
3483            // - `form`: token surface strings
3484            // - `misc`: per-token features, including `Seg=B-seg` (segment boundary) or `Seg=O`
3485            //
3486            // We convert boundaries into BIO tags of a single entity type (`SEG`), treating each
3487            // segment (EDU) as an entity span.
3488            if let (Some(forms), Some(misc)) = (
3489                row.get("form").and_then(|v| v.as_array()),
3490                row.get("misc").and_then(|v| v.as_array()),
3491            ) {
3492                if !forms.is_empty() && forms.len() == misc.len() {
3493                    let mut annotated_tokens = Vec::with_capacity(forms.len());
3494                    let mut in_seg = false;
3495                    for (i, (f, m)) in forms.iter().zip(misc.iter()).enumerate() {
3496                        let tok = f.as_str().unwrap_or("").to_string();
3497                        let misc_s = m.as_str().unwrap_or("");
3498                        let start = misc_s.contains("Seg=B-seg");
3499                        let ner_tag = if i == 0 || start {
3500                            in_seg = true;
3501                            "B-SEG".to_string()
3502                        } else if in_seg {
3503                            "I-SEG".to_string()
3504                        } else {
3505                            "O".to_string()
3506                        };
3507                        annotated_tokens.push(AnnotatedToken { text: tok, ner_tag });
3508                    }
3509
3510                    sentences.push(AnnotatedSentence {
3511                        tokens: annotated_tokens,
3512                        source_dataset: id,
3513                    });
3514                    continue;
3515                }
3516            }
3517
3518            // Fallback: classification-ish rows (`<text>` + `label`).
3519            //
3520            // We encode the gold label as `B-<LABEL>` on a single token, matching the loader's
3521            // convention for other classification datasets.
3522            let text = if let Some(s) = row.get("text").and_then(|v| v.as_str()) {
3523                s.trim().to_string()
3524            } else if let (Some(a), Some(b)) = (
3525                row.get("unit1_txt").and_then(|v| v.as_str()),
3526                row.get("unit2_txt").and_then(|v| v.as_str()),
3527            ) {
3528                format!("{} [SEP] {}", a.trim(), b.trim())
3529            } else if let (Some(a), Some(b)) = (
3530                row.get("sentence1").and_then(|v| v.as_str()),
3531                row.get("sentence2").and_then(|v| v.as_str()),
3532            ) {
3533                format!("{} [SEP] {}", a.trim(), b.trim())
3534            } else if let (Some(a), Some(b)) = (
3535                row.get("premise").and_then(|v| v.as_str()),
3536                row.get("hypothesis").and_then(|v| v.as_str()),
3537            ) {
3538                format!("{} [SEP] {}", a.trim(), b.trim())
3539            } else {
3540                continue;
3541            };
3542            if text.trim().is_empty() {
3543                continue;
3544            }
3545
3546            let label_value = row.get("label").or_else(|| row.get("labels"));
3547            let label = match label_value {
3548                Some(v) if v.is_string() => v.as_str().unwrap_or("").to_string(),
3549                Some(v) if v.is_number() => {
3550                    let idx = v.as_u64().unwrap_or(0) as usize;
3551                    class_names
3552                        .get(idx)
3553                        .cloned()
3554                        .unwrap_or_else(|| format!("LABEL_{}", idx))
3555                }
3556                _ => "label".to_string(),
3557            };
3558            if label.trim().is_empty() {
3559                continue;
3560            }
3561
3562            sentences.push(AnnotatedSentence {
3563                tokens: vec![AnnotatedToken {
3564                    text,
3565                    ner_tag: format!("B-{}", label),
3566                }],
3567                source_dataset: id,
3568            });
3569        }
3570
3571        if sentences.is_empty() {
3572            return Err(Error::InvalidInput(format!(
3573                "HF API response for {:?} contains no valid sentences",
3574                id
3575            )));
3576        }
3577
3578        let now = chrono::Utc::now().to_rfc3339();
3579        Ok(LoadedDataset {
3580            id,
3581            sentences,
3582            loaded_at: now,
3583            source_url: id.download_url().to_string(),
3584            data_source: DataSource::LocalCache,
3585            temporal_metadata: Self::get_temporal_metadata(id),
3586            metadata: id.default_metadata(),
3587        })
3588    }
3589
3590    /// Extract tag names from HF API features metadata.
3591    fn extract_tag_names_from_features(&self, parsed: &serde_json::Value) -> Vec<String> {
3592        let mut tag_names = Vec::new();
3593
3594        if let Some(features) = parsed.get("features").and_then(|v| v.as_array()) {
3595            for feature in features {
3596                let name = feature.get("name").and_then(|v| v.as_str());
3597                if name == Some("ner_tags") {
3598                    // Look for ClassLabel names
3599                    if let Some(names) = feature
3600                        .get("type")
3601                        .and_then(|t| t.get("feature"))
3602                        .and_then(|f| f.get("names"))
3603                        .and_then(|n| n.as_array())
3604                    {
3605                        for name in names {
3606                            if let Some(s) = name.as_str() {
3607                                tag_names.push(s.to_string());
3608                            }
3609                        }
3610                    }
3611                    break;
3612                }
3613            }
3614        }
3615
3616        tag_names
3617    }
3618
3619    /// Extract class label names for `label` fields in HF API features metadata.
3620    fn extract_class_names_from_features(&self, parsed: &serde_json::Value) -> Vec<String> {
3621        let mut names = Vec::new();
3622        if let Some(features) = parsed.get("features").and_then(|v| v.as_array()) {
3623            for feature in features {
3624                let name = feature.get("name").and_then(|v| v.as_str());
3625                if name == Some("label") {
3626                    if let Some(label_names) = feature
3627                        .get("type")
3628                        .and_then(|t| t.get("names"))
3629                        .and_then(|n| n.as_array())
3630                    {
3631                        for n in label_names {
3632                            if let Some(s) = n.as_str() {
3633                                names.push(s.to_string());
3634                            }
3635                        }
3636                    }
3637                    break;
3638                }
3639            }
3640        }
3641        names
3642    }
3643
3644    /// Parse TweetNER7 JSON format.
3645    ///
3646    /// TweetNER7 is JSONL with each line: {"tokens": [...], "tags": [...]}
3647    /// Tag mapping from label.json (tag -> id format, we need id -> tag):
3648    fn parse_tweetner7(&self, content: &str, id: DatasetId) -> Result<LoadedDataset> {
3649        // TweetNER7 tag mapping from label.json (index order!)
3650        // {"B-corporation": 0, "B-creative_work": 1, "B-event": 2, "B-group": 3,
3651        //  "B-location": 4, "B-person": 5, "B-product": 6, "I-corporation": 7,
3652        //  "I-creative_work": 8, "I-event": 9, "I-group": 10, "I-location": 11,
3653        //  "I-person": 12, "I-product": 13, "O": 14}
3654        let tag_labels = [
3655            "B-corporation",   // 0
3656            "B-creative_work", // 1
3657            "B-event",         // 2
3658            "B-group",         // 3
3659            "B-location",      // 4
3660            "B-person",        // 5
3661            "B-product",       // 6
3662            "I-corporation",   // 7
3663            "I-creative_work", // 8
3664            "I-event",         // 9
3665            "I-group",         // 10
3666            "I-location",      // 11
3667            "I-person",        // 12
3668            "I-product",       // 13
3669            "O",               // 14
3670        ];
3671
3672        let mut sentences = Vec::new();
3673
3674        // Parse as JSONL (one JSON object per line)
3675        for line in content.lines() {
3676            let line = line.trim();
3677            if line.is_empty() {
3678                continue;
3679            }
3680
3681            let parsed: serde_json::Value = match serde_json::from_str(line) {
3682                Ok(v) => v,
3683                Err(_) => continue, // Skip malformed lines
3684            };
3685
3686            let tokens = match parsed.get("tokens").and_then(|v| v.as_array()) {
3687                Some(t) => t,
3688                None => continue,
3689            };
3690
3691            let tags = match parsed.get("tags").and_then(|v| v.as_array()) {
3692                Some(t) => t,
3693                None => continue,
3694            };
3695
3696            if tokens.len() != tags.len() {
3697                continue;
3698            }
3699
3700            let mut annotated_tokens = Vec::new();
3701            for (token, tag) in tokens.iter().zip(tags.iter()) {
3702                let text = token.as_str().unwrap_or("").to_string();
3703                let tag_idx = tag.as_u64().unwrap_or(0) as usize;
3704                let ner_tag = tag_labels.get(tag_idx).unwrap_or(&"O").to_string();
3705                annotated_tokens.push(AnnotatedToken { text, ner_tag });
3706            }
3707
3708            if !annotated_tokens.is_empty() {
3709                sentences.push(AnnotatedSentence {
3710                    tokens: annotated_tokens,
3711                    source_dataset: id,
3712                });
3713            }
3714        }
3715
3716        let now = chrono::Utc::now().to_rfc3339();
3717        Ok(LoadedDataset {
3718            id,
3719            sentences,
3720            loaded_at: now,
3721            source_url: id.download_url().to_string(),
3722            data_source: DataSource::LocalCache,
3723            temporal_metadata: Self::get_temporal_metadata(id),
3724            metadata: id.default_metadata(),
3725        })
3726    }
3727
3728    /// Parse HIPE-2022 style TSV NER format.
3729    ///
3730    /// Expected format:
3731    /// ```text
3732    /// TOKEN    NE-COARSE-LIT   NE-COARSE-METO   NE-FINE-LIT   ...
3733    /// # hipe2022:document_id = doc123
3734    /// word1    B-PER           _                B-pers.author ...
3735    /// word2    I-PER           _                I-pers.author ...
3736    /// word3    O               _                O             ...
3737    /// ```
3738    ///
3739    /// - First line is header (starts with TOKEN)
3740    /// - Lines starting with `#` are metadata comments
3741    /// - Data lines are tab-separated with token in first column
3742    /// - NE-COARSE-LIT (column 2) contains BIO-tagged NER labels
3743    /// - `_` means no annotation
3744    fn parse_tsv_ner(&self, content: &str, id: DatasetId) -> Result<LoadedDataset> {
3745        let mut sentences = Vec::new();
3746        let mut current_tokens = Vec::new();
3747
3748        for line in content.lines() {
3749            let line = line.trim();
3750
3751            // Skip empty lines - they indicate sentence boundaries
3752            if line.is_empty() {
3753                if !current_tokens.is_empty() {
3754                    sentences.push(AnnotatedSentence {
3755                        tokens: std::mem::take(&mut current_tokens),
3756                        source_dataset: id,
3757                    });
3758                }
3759                continue;
3760            }
3761
3762            // Skip header line
3763            if line.starts_with("TOKEN\t") || line.starts_with("TOKEN ") {
3764                continue;
3765            }
3766
3767            // Skip metadata comments
3768            if line.starts_with('#') {
3769                // Document boundary comment can also serve as sentence boundary
3770                if line.contains("document_id") && !current_tokens.is_empty() {
3771                    sentences.push(AnnotatedSentence {
3772                        tokens: std::mem::take(&mut current_tokens),
3773                        source_dataset: id,
3774                    });
3775                }
3776                continue;
3777            }
3778
3779            // Parse data line
3780            let parts: Vec<&str> = line.split('\t').collect();
3781            if parts.len() < 2 {
3782                continue; // Malformed line
3783            }
3784
3785            let token_text = parts[0].to_string();
3786            let ner_label = parts.get(1).unwrap_or(&"O");
3787
3788            // Convert underscore to O (no annotation)
3789            let ner_tag = if *ner_label == "_" || ner_label.is_empty() {
3790                "O".to_string()
3791            } else {
3792                ner_label.to_string()
3793            };
3794
3795            current_tokens.push(AnnotatedToken {
3796                text: token_text,
3797                ner_tag,
3798            });
3799        }
3800
3801        // Don't forget the last sentence
3802        if !current_tokens.is_empty() {
3803            sentences.push(AnnotatedSentence {
3804                tokens: current_tokens,
3805                source_dataset: id,
3806            });
3807        }
3808
3809        if sentences.is_empty() {
3810            return Err(Error::InvalidInput(format!(
3811                "TSV NER file for {:?} contains no valid sentences",
3812                id
3813            )));
3814        }
3815
3816        let now = chrono::Utc::now().to_rfc3339();
3817        Ok(LoadedDataset {
3818            id,
3819            sentences,
3820            loaded_at: now,
3821            source_url: id.download_url().to_string(),
3822            data_source: DataSource::LocalCache,
3823            temporal_metadata: Self::get_temporal_metadata(id),
3824            metadata: id.default_metadata(),
3825        })
3826    }
3827
3828    /// Parse CSV NER format (E-NER/EDGAR-NER style).
3829    ///
3830    /// Expected format: `Token,Tag` (comma-separated)
3831    /// Uses `-DOCSTART-` for document boundaries and empty lines for sentence boundaries.
3832    /// Tags use BIO scheme (e.g., O, B-PERSON, I-BUSINESS).
3833    ///
3834    /// # Errors
3835    ///
3836    /// Returns error if CSV is malformed or has no valid sentences.
3837    fn parse_csv_ner(&self, content: &str, id: DatasetId) -> Result<LoadedDataset> {
3838        let mut sentences = Vec::new();
3839        let mut current_tokens = Vec::new();
3840
3841        for line in content.lines() {
3842            let line = line.trim();
3843
3844            // Empty lines indicate sentence boundaries
3845            if line.is_empty() {
3846                if !current_tokens.is_empty() {
3847                    sentences.push(AnnotatedSentence {
3848                        tokens: std::mem::take(&mut current_tokens),
3849                        source_dataset: id,
3850                    });
3851                }
3852                continue;
3853            }
3854
3855            // Handle -DOCSTART- markers (document boundary, also a sentence boundary)
3856            if line.starts_with("-DOCSTART-") {
3857                if !current_tokens.is_empty() {
3858                    sentences.push(AnnotatedSentence {
3859                        tokens: std::mem::take(&mut current_tokens),
3860                        source_dataset: id,
3861                    });
3862                }
3863                continue;
3864            }
3865
3866            // Skip header line if present
3867            if line.eq_ignore_ascii_case("token,tag")
3868                || line.eq_ignore_ascii_case("text,label")
3869                || line.eq_ignore_ascii_case("word,ner")
3870            {
3871                continue;
3872            }
3873
3874            // Parse comma-separated line
3875            // Handle cases where the token itself might be a comma: ",O" means token is comma
3876            let (token_text, ner_tag) = if let Some(rest) = line.strip_prefix(',') {
3877                // Token is empty or the comma itself
3878                if let Some(idx) = rest.find(',') {
3879                    // Format: ,tag (empty token) or ,,tag (comma token)
3880                    if idx == 0 {
3881                        // ",,tag" -> token is comma, tag follows
3882                        (",".to_string(), rest[1..].to_string())
3883                    } else {
3884                        // ",tag" -> empty token
3885                        (String::new(), rest[..idx].to_string())
3886                    }
3887                } else {
3888                    // ",tag" with no more commas
3889                    (String::new(), rest.to_string())
3890                }
3891            } else if let Some(idx) = line.rfind(',') {
3892                // Normal case: Token,Tag
3893                let token = line[..idx].to_string();
3894                let tag = line[idx + 1..].to_string();
3895                (token, tag)
3896            } else {
3897                // Malformed line, skip
3898                continue;
3899            };
3900
3901            // Skip if we got empty results
3902            if token_text.is_empty() && ner_tag.is_empty() {
3903                continue;
3904            }
3905
3906            // Convert empty tag to O
3907            let ner_tag = if ner_tag.is_empty() || ner_tag == "_" {
3908                "O".to_string()
3909            } else {
3910                ner_tag
3911            };
3912
3913            current_tokens.push(AnnotatedToken {
3914                text: token_text,
3915                ner_tag,
3916            });
3917        }
3918
3919        // Don't forget the last sentence
3920        if !current_tokens.is_empty() {
3921            sentences.push(AnnotatedSentence {
3922                tokens: current_tokens,
3923                source_dataset: id,
3924            });
3925        }
3926
3927        if sentences.is_empty() {
3928            return Err(Error::InvalidInput(format!(
3929                "CSV NER file for {:?} contains no valid sentences",
3930                id
3931            )));
3932        }
3933
3934        let now = chrono::Utc::now().to_rfc3339();
3935        Ok(LoadedDataset {
3936            id,
3937            sentences,
3938            loaded_at: now,
3939            source_url: id.download_url().to_string(),
3940            data_source: DataSource::LocalCache,
3941            temporal_metadata: Self::get_temporal_metadata(id),
3942            metadata: id.default_metadata(),
3943        })
3944    }
3945
3946    /// Parse WikiANN-style JSON array format.
3947    ///
3948    /// Expected format: `[{"text": "...", "tokens": ["word1", "word2"], "ner_tags": ["O", "B-LOC"]}, ...]`
3949    /// Used by UNER and MSNER datasets converted via download_hf_datasets.py.
3950    fn parse_wikiann_json(&self, content: &str, id: DatasetId) -> Result<LoadedDataset> {
3951        let parsed: serde_json::Value = serde_json::from_str(content)
3952            .map_err(|e| Error::InvalidInput(format!("Failed to parse JSON: {}", e)))?;
3953
3954        let mut sentences = Vec::new();
3955
3956        // Handle JSON array format
3957        if let Some(items) = parsed.as_array() {
3958            for item in items {
3959                let tokens = match item.get("tokens").and_then(|v| v.as_array()) {
3960                    Some(t) => t,
3961                    None => continue,
3962                };
3963
3964                let ner_tags = match item.get("ner_tags").and_then(|v| v.as_array()) {
3965                    Some(t) => t,
3966                    None => continue,
3967                };
3968
3969                if tokens.len() != ner_tags.len() {
3970                    continue;
3971                }
3972
3973                let mut annotated_tokens = Vec::new();
3974                for (token, tag) in tokens.iter().zip(ner_tags.iter()) {
3975                    let text = token.as_str().unwrap_or("").to_string();
3976                    let ner_tag = tag.as_str().unwrap_or("O").to_string();
3977                    annotated_tokens.push(AnnotatedToken { text, ner_tag });
3978                }
3979
3980                if !annotated_tokens.is_empty() {
3981                    sentences.push(AnnotatedSentence {
3982                        tokens: annotated_tokens,
3983                        source_dataset: id,
3984                    });
3985                }
3986            }
3987        }
3988
3989        let now = chrono::Utc::now().to_rfc3339();
3990        Ok(LoadedDataset {
3991            id,
3992            sentences,
3993            loaded_at: now,
3994            source_url: id.download_url().to_string(),
3995            data_source: DataSource::LocalCache,
3996            temporal_metadata: Self::get_temporal_metadata(id),
3997            metadata: id.default_metadata(),
3998        })
3999    }
4000
4001    /// Parse DocRED document-level relation extraction JSON format.
4002    ///
4003    /// DocRED is primarily for relation extraction but includes NER annotations.
4004    /// Parse DocRED/CrossRE JSON format for relation extraction.
4005    ///
4006    /// CrossRE format: {"doc_key": "...", "sentence": [...], "ner": [[start, end, type], ...], "relations": [...]}
4007    fn parse_docred(&self, content: &str, id: DatasetId) -> Result<LoadedDataset> {
4008        let mut sentences = Vec::new();
4009
4010        // Parse as JSONL (one JSON object per line)
4011        for line in content.lines() {
4012            let line = line.trim();
4013            if line.is_empty() {
4014                continue;
4015            }
4016
4017            let doc: serde_json::Value = match serde_json::from_str(line) {
4018                Ok(v) => v,
4019                Err(_) => continue,
4020            };
4021
4022            // CrossRE format: sentence array + ner array
4023            let tokens_arr = match doc.get("sentence").and_then(|v| v.as_array()) {
4024                Some(t) => t,
4025                None => continue,
4026            };
4027
4028            let ner_spans = doc.get("ner").and_then(|v| v.as_array());
4029
4030            // Build token list with entity annotations
4031            let mut tokens: Vec<AnnotatedToken> = tokens_arr
4032                .iter()
4033                .filter_map(|t| t.as_str())
4034                .map(|word| AnnotatedToken {
4035                    text: word.to_string(),
4036                    ner_tag: "O".to_string(),
4037                })
4038                .collect();
4039
4040            // Apply NER annotations: [start, end, type]
4041            if let Some(ner) = ner_spans {
4042                for span in ner {
4043                    if let Some(arr) = span.as_array() {
4044                        if arr.len() >= 3 {
4045                            let start = arr[0].as_u64().unwrap_or(0) as usize;
4046                            let end = arr[1].as_u64().unwrap_or(0) as usize;
4047                            let ent_type = arr[2].as_str().unwrap_or("ENTITY");
4048
4049                            // Apply BIO tags
4050                            for idx in start..=end {
4051                                if idx < tokens.len() {
4052                                    tokens[idx].ner_tag = if idx == start {
4053                                        format!("B-{}", ent_type.to_uppercase())
4054                                    } else {
4055                                        format!("I-{}", ent_type.to_uppercase())
4056                                    };
4057                                }
4058                            }
4059                        }
4060                    }
4061                }
4062            }
4063
4064            if !tokens.is_empty() {
4065                sentences.push(AnnotatedSentence {
4066                    tokens,
4067                    source_dataset: id,
4068                });
4069            }
4070        }
4071
4072        if sentences.is_empty() {
4073            return Err(Error::InvalidInput(format!(
4074                "DocRED/CrossRE JSON for {:?} contains no valid sentences",
4075                id
4076            )));
4077        }
4078
4079        let now = chrono::Utc::now().to_rfc3339();
4080        Ok(LoadedDataset {
4081            id,
4082            sentences,
4083            loaded_at: now,
4084            source_url: id.download_url().to_string(),
4085            data_source: DataSource::LocalCache,
4086            temporal_metadata: Self::get_temporal_metadata(id),
4087            metadata: id.default_metadata(),
4088        })
4089    }
4090
4091    /// Parse Google's relation-extraction-corpus JSONL format.
4092    ///
4093    /// Example record (one per line):
4094    /// ```json
4095    /// {"pred":"/people/person/place_of_birth","sub":"/m/...","obj":"/m/...","evidences":[{"url":"...","snippet":"..."}],"judgments":[{"rater":"...","judgment":"yes"}]}
4096    /// ```
4097    ///
4098    /// This dataset does **not** include token-level entity spans. For now we treat each
4099    /// evidence snippet as a plain sentence with all tokens tagged as `O`, which is
4100    /// sufficient for sanity evaluation plumbing (and avoids failing dataset loads).
4101    fn parse_google_re_corpus(&self, content: &str, id: DatasetId) -> Result<LoadedDataset> {
4102        let mut sentences: Vec<AnnotatedSentence> = Vec::new();
4103
4104        for line in content.lines() {
4105            let line = line.trim();
4106            if line.is_empty() {
4107                continue;
4108            }
4109
4110            let rec: serde_json::Value = match serde_json::from_str(line) {
4111                Ok(v) => v,
4112                Err(_) => continue,
4113            };
4114
4115            let snippet = rec
4116                .get("evidences")
4117                .and_then(|v| v.as_array())
4118                .and_then(|arr| arr.first())
4119                .and_then(|ev| ev.get("snippet"))
4120                .and_then(|s| s.as_str());
4121            let Some(snippet) = snippet else {
4122                continue;
4123            };
4124
4125            let tokens: Vec<AnnotatedToken> = snippet
4126                .split_whitespace()
4127                .filter(|t| !t.is_empty())
4128                .map(|t| AnnotatedToken {
4129                    text: t.to_string(),
4130                    ner_tag: "O".to_string(),
4131                })
4132                .collect();
4133
4134            if tokens.is_empty() {
4135                continue;
4136            }
4137
4138            sentences.push(AnnotatedSentence {
4139                tokens,
4140                source_dataset: id,
4141            });
4142        }
4143
4144        if sentences.is_empty() {
4145            return Err(Error::InvalidInput(format!(
4146                "Google relation-extraction-corpus file for {:?} contains no usable evidence snippets",
4147                id
4148            )));
4149        }
4150
4151        let now = chrono::Utc::now().to_rfc3339();
4152        Ok(LoadedDataset {
4153            id,
4154            sentences,
4155            loaded_at: now,
4156            source_url: id.download_url().to_string(),
4157            data_source: DataSource::LocalCache,
4158            temporal_metadata: Self::get_temporal_metadata(id),
4159            metadata: id.default_metadata(),
4160        })
4161    }
4162
4163    /// Parse CADEC from HuggingFace datasets-server API.
4164    ///
4165    /// CADEC HF API format: {"text": "...", "ade": "...", "term_PT": "..."}
4166    /// Each row is a text-ADE pair (one sentence per ADE mention).
4167    /// The `ade` field contains the adverse drug event mention within `text`.
4168    fn parse_cadec_hf_api(&self, content: &str, id: DatasetId) -> Result<LoadedDataset> {
4169        let parsed: serde_json::Value = serde_json::from_str(content).map_err(|e| {
4170            Error::InvalidInput(format!("Failed to parse CADEC HF API response: {}", e))
4171        })?;
4172
4173        let mut sentences = Vec::new();
4174
4175        let rows = parsed
4176            .get("rows")
4177            .and_then(|v| v.as_array())
4178            .ok_or_else(|| {
4179                Error::InvalidInput("No 'rows' array in CADEC HF API response".to_string())
4180            })?;
4181
4182        for row_obj in rows {
4183            let row = match row_obj.get("row") {
4184                Some(r) => r,
4185                None => continue,
4186            };
4187
4188            let text = match row.get("text").and_then(|v| v.as_str()) {
4189                Some(t) => t,
4190                None => continue,
4191            };
4192
4193            let ade_text = match row.get("ade").and_then(|v| v.as_str()) {
4194                Some(a) => a,
4195                None => continue,
4196            };
4197
4198            // Find ADE span in text.
4199            //
4200            // IMPORTANT: anno uses **character offsets** globally, but this parser assigns BIO tags
4201            // token-by-token and uses byte spans internally. We must preserve indices and avoid
4202            // Unicode casefolding (which can change string length).
4203            //
4204            // Strategy:
4205            // 1) Try exact match (fast, preserves byte indices).
4206            // 2) Fall back to ASCII case-insensitive match (ADE strings are ASCII-centric).
4207            let ade_start_byte = text.find(ade_text).or_else(|| {
4208                let needle_len = ade_text.len();
4209                if needle_len == 0 {
4210                    return None;
4211                }
4212                for (b, _) in text.char_indices() {
4213                    if let Some(hay) = text.get(b..b + needle_len) {
4214                        if hay.eq_ignore_ascii_case(ade_text) {
4215                            return Some(b);
4216                        }
4217                    }
4218                }
4219                None
4220            });
4221            let Some(ade_start_byte) = ade_start_byte else {
4222                continue; // ADE not found in text
4223            };
4224            let ade_end_byte = ade_start_byte + ade_text.len();
4225
4226            // Tokenize text preserving byte offsets.
4227            let mut tokens: Vec<AnnotatedToken> = Vec::new();
4228            let mut byte_idx = 0;
4229            let words: Vec<&str> = text.split_whitespace().collect();
4230
4231            for word in words {
4232                let word_start =
4233                    text.get(byte_idx..).and_then(|s| s.find(word)).unwrap_or(0) + byte_idx;
4234                let word_end = word_start + word.len();
4235
4236                // Check if this word overlaps with ADE span (byte spans)
4237                let ner_tag = if word_start >= ade_start_byte && word_end <= ade_end_byte {
4238                    // Check if this is the first word of the ADE
4239                    if word_start == ade_start_byte
4240                        || tokens.is_empty()
4241                        || !tokens
4242                            .last()
4243                            .expect("tokens.is_empty() checked above")
4244                            .ner_tag
4245                            .starts_with("I-")
4246                    {
4247                        "B-adverse_drug_event".to_string()
4248                    } else {
4249                        "I-adverse_drug_event".to_string()
4250                    }
4251                } else {
4252                    "O".to_string()
4253                };
4254
4255                tokens.push(AnnotatedToken {
4256                    text: word.to_string(),
4257                    ner_tag,
4258                });
4259
4260                // Update byte position to after this word (including trailing space)
4261                byte_idx = word_end;
4262                if byte_idx < text.len() && text.as_bytes().get(byte_idx) == Some(&b' ') {
4263                    byte_idx += 1;
4264                }
4265            }
4266
4267            if !tokens.is_empty() {
4268                sentences.push(AnnotatedSentence {
4269                    tokens,
4270                    source_dataset: id,
4271                });
4272            }
4273        }
4274
4275        let now = chrono::Utc::now().to_rfc3339();
4276        Ok(LoadedDataset {
4277            id,
4278            sentences,
4279            loaded_at: now,
4280            source_url: id.download_url().to_string(),
4281            data_source: DataSource::LocalCache,
4282            temporal_metadata: Self::get_temporal_metadata(id),
4283            metadata: id.default_metadata(),
4284        })
4285    }
4286
4287    /// Parse CADEC JSONL format with support for discontinuous entities.
4288    ///
4289    /// CADEC format can include:
4290    /// - Standard BIO tags: `{"tokens": [...], "ner_tags": [...]}`
4291    /// - Entity spans: `{"tokens": [...], "entities": [{"text": "...", "label": "...", "start": 0, "end": 10}]}`
4292    /// - Discontinuous entities: `{"entities": [{"text": "...", "label": "...", "spans": [[0, 5], [10, 15]]}]}`
4293    ///
4294    /// For discontinuous entities, we convert them to BIO tags by marking all tokens
4295    /// within any span as part of the entity.
4296    fn parse_cadec_jsonl(&self, content: &str, id: DatasetId) -> Result<LoadedDataset> {
4297        let mut sentences = Vec::new();
4298
4299        for line in content.lines() {
4300            let line = line.trim();
4301            if line.is_empty() {
4302                continue;
4303            }
4304
4305            // Parse JSON line
4306            let parsed: serde_json::Value = match serde_json::from_str(line) {
4307                Ok(v) => v,
4308                Err(_) => continue, // Skip malformed lines
4309            };
4310
4311            // Try to get tokens
4312            let tokens = match parsed.get("tokens").and_then(|v| v.as_array()) {
4313                Some(t) => t,
4314                None => continue,
4315            };
4316
4317            let mut annotated_tokens = Vec::new();
4318            let mut char_offset = 0;
4319
4320            // Build token list with character offsets
4321            let mut token_offsets = Vec::new();
4322            for token in tokens {
4323                let text = token.as_str().unwrap_or("").to_string();
4324                let start = char_offset;
4325                char_offset += text.chars().count() + 1; // +1 for space
4326                let end = char_offset - 1;
4327                token_offsets.push((text, start, end));
4328            }
4329
4330            // Initialize all tokens as "O"
4331            for (text, _, _) in &token_offsets {
4332                annotated_tokens.push(AnnotatedToken {
4333                    text: text.clone(),
4334                    ner_tag: "O".to_string(),
4335                });
4336            }
4337
4338            // Try to parse entities (for discontinuous support)
4339            if let Some(entities) = parsed.get("entities").and_then(|v| v.as_array()) {
4340                for entity in entities {
4341                    let label = entity
4342                        .get("label")
4343                        .or_else(|| entity.get("entity_type"))
4344                        .and_then(|v| v.as_str())
4345                        .unwrap_or("UNKNOWN")
4346                        .to_string();
4347
4348                    // Check for discontinuous spans
4349                    if let Some(spans) = entity.get("spans").and_then(|v| v.as_array()) {
4350                        // Discontinuous entity with multiple spans
4351                        for span in spans {
4352                            if let Some(span_array) = span.as_array() {
4353                                if span_array.len() >= 2 {
4354                                    let start = span_array[0].as_u64().unwrap_or(0) as usize;
4355                                    let end = span_array[1].as_u64().unwrap_or(0) as usize;
4356
4357                                    // Mark tokens within this span
4358                                    for (idx, (_, token_start, token_end)) in
4359                                        token_offsets.iter().enumerate()
4360                                    {
4361                                        if *token_start >= start && *token_end <= end {
4362                                            if idx > 0
4363                                                && annotated_tokens[idx - 1]
4364                                                    .ner_tag
4365                                                    .starts_with(&format!("I-{}", label))
4366                                                || annotated_tokens[idx - 1]
4367                                                    .ner_tag
4368                                                    .starts_with(&format!("B-{}", label))
4369                                            {
4370                                                annotated_tokens[idx].ner_tag =
4371                                                    format!("I-{}", label);
4372                                            } else {
4373                                                annotated_tokens[idx].ner_tag =
4374                                                    format!("B-{}", label);
4375                                            }
4376                                        }
4377                                    }
4378                                }
4379                            }
4380                        }
4381                    } else if let (Some(start_val), Some(end_val)) = (
4382                        entity.get("start").and_then(|v| v.as_u64()),
4383                        entity.get("end").and_then(|v| v.as_u64()),
4384                    ) {
4385                        // Contiguous entity
4386                        let start = start_val as usize;
4387                        let end = end_val as usize;
4388
4389                        // Mark tokens within this span
4390                        for (idx, (_, token_start, token_end)) in token_offsets.iter().enumerate() {
4391                            if *token_start >= start && *token_end <= end {
4392                                if idx > 0
4393                                    && (annotated_tokens[idx - 1]
4394                                        .ner_tag
4395                                        .starts_with(&format!("I-{}", label))
4396                                        || annotated_tokens[idx - 1]
4397                                            .ner_tag
4398                                            .starts_with(&format!("B-{}", label)))
4399                                {
4400                                    annotated_tokens[idx].ner_tag = format!("I-{}", label);
4401                                } else {
4402                                    annotated_tokens[idx].ner_tag = format!("B-{}", label);
4403                                }
4404                            }
4405                        }
4406                    }
4407                }
4408            } else if let Some(ner_tags) = parsed.get("ner_tags").and_then(|v| v.as_array()) {
4409                // Fallback to standard BIO tags
4410                let tag_labels = [
4411                    "O",
4412                    "B-PER",
4413                    "I-PER",
4414                    "B-ORG",
4415                    "I-ORG",
4416                    "B-LOC",
4417                    "I-LOC",
4418                    "B-MISC",
4419                    "I-MISC",
4420                    "B-DRUG",
4421                    "I-DRUG",
4422                    "B-ADR",
4423                    "I-ADR",
4424                    "B-DISEASE",
4425                    "I-DISEASE",
4426                ];
4427
4428                for (idx, (text, _, _)) in token_offsets.iter().enumerate() {
4429                    if let Some(tag_val) = ner_tags.get(idx) {
4430                        let tag_idx = tag_val.as_u64().unwrap_or(0) as usize;
4431                        let ner_tag = tag_labels.get(tag_idx).unwrap_or(&"O").to_string();
4432                        annotated_tokens[idx] = AnnotatedToken {
4433                            text: text.clone(),
4434                            ner_tag,
4435                        };
4436                    }
4437                }
4438            }
4439
4440            if !annotated_tokens.is_empty() {
4441                sentences.push(AnnotatedSentence {
4442                    tokens: annotated_tokens,
4443                    source_dataset: id,
4444                });
4445            }
4446        }
4447
4448        if sentences.is_empty() {
4449            return Err(Error::InvalidInput(format!(
4450                "CADEC JSONL file for {:?} contains no valid sentences",
4451                id
4452            )));
4453        }
4454
4455        let now = chrono::Utc::now().to_rfc3339();
4456        Ok(LoadedDataset {
4457            id,
4458            sentences,
4459            loaded_at: now,
4460            source_url: id.download_url().to_string(),
4461            data_source: DataSource::LocalCache,
4462            temporal_metadata: Self::get_temporal_metadata(id),
4463            metadata: id.default_metadata(),
4464        })
4465    }
4466
4467    /// Parse BC5CDR BioC XML format.
4468    ///
4469    /// Note: This is a simplified parser that extracts text passages.
4470    /// Full annotation extraction would require proper XML parsing.
4471    /// Parse BC5CDR dataset in CoNLL format from BioFLAIR.
4472    ///
4473    /// Format: WORD\tPOS\tCHUNK\tNER_TAG
4474    fn parse_bc5cdr(&self, content: &str, id: DatasetId) -> Result<LoadedDataset> {
4475        let mut sentences = Vec::new();
4476        let mut current_tokens = Vec::new();
4477
4478        for line in content.lines() {
4479            let line = line.trim();
4480
4481            // Skip DOCSTART lines
4482            if line.starts_with("-DOCSTART-") {
4483                continue;
4484            }
4485
4486            if line.is_empty() {
4487                // End of sentence
4488                if !current_tokens.is_empty() {
4489                    sentences.push(AnnotatedSentence {
4490                        tokens: std::mem::take(&mut current_tokens),
4491                        source_dataset: id,
4492                    });
4493                }
4494                continue;
4495            }
4496
4497            // Parse CoNLL line: WORD\tPOS\tCHUNK\tNER_TAG
4498            let parts: Vec<&str> = line.split('\t').collect();
4499            if parts.len() >= 4 {
4500                let word = parts[0].to_string();
4501                let ner_tag = parts[3].to_string();
4502
4503                // Map Entity tags to BIO format
4504                let normalized_tag = if ner_tag.contains("Entity")
4505                    || ner_tag.contains("CHEMICAL")
4506                    || ner_tag.contains("DISEASE")
4507                {
4508                    // Convert I-Entity to B-CHEMICAL or I-CHEMICAL based on context
4509                    if ner_tag.starts_with("B-") {
4510                        "B-CHEMICAL".to_string()
4511                    } else if ner_tag.starts_with("I-") {
4512                        "I-CHEMICAL".to_string()
4513                    } else {
4514                        "O".to_string()
4515                    }
4516                } else {
4517                    ner_tag
4518                };
4519
4520                current_tokens.push(AnnotatedToken {
4521                    text: word,
4522                    ner_tag: normalized_tag,
4523                });
4524            }
4525        }
4526
4527        // Don't forget the last sentence
4528        if !current_tokens.is_empty() {
4529            sentences.push(AnnotatedSentence {
4530                tokens: current_tokens,
4531                source_dataset: id,
4532            });
4533        }
4534
4535        let now = chrono::Utc::now().to_rfc3339();
4536        Ok(LoadedDataset {
4537            id,
4538            sentences,
4539            loaded_at: now,
4540            source_url: id.download_url().to_string(),
4541            data_source: DataSource::LocalCache,
4542            temporal_metadata: Self::get_temporal_metadata(id),
4543            metadata: id.default_metadata(),
4544        })
4545    }
4546
4547    /// Parse NCBI Disease corpus format.
4548    ///
4549    /// Format: PMID|t|Title or PMID|a|Abstract, followed by annotation lines.
4550    /// Parse NCBI Disease dataset in CoNLL format from BioFLAIR.
4551    ///
4552    /// Format: WORD\tPOS\tCHUNK\tNER_TAG
4553    fn parse_ncbi_disease(&self, content: &str, id: DatasetId) -> Result<LoadedDataset> {
4554        let mut sentences = Vec::new();
4555        let mut current_tokens = Vec::new();
4556
4557        for line in content.lines() {
4558            let line = line.trim();
4559
4560            if line.is_empty() {
4561                // End of sentence
4562                if !current_tokens.is_empty() {
4563                    sentences.push(AnnotatedSentence {
4564                        tokens: std::mem::take(&mut current_tokens),
4565                        source_dataset: id,
4566                    });
4567                }
4568                continue;
4569            }
4570
4571            // Parse CoNLL line: WORD\tPOS\tCHUNK\tNER_TAG
4572            let parts: Vec<&str> = line.split('\t').collect();
4573            if parts.len() >= 4 {
4574                let word = parts[0].to_string();
4575                let ner_tag = parts[3].to_string();
4576
4577                current_tokens.push(AnnotatedToken {
4578                    text: word,
4579                    ner_tag,
4580                });
4581            }
4582        }
4583
4584        // Don't forget the last sentence
4585        if !current_tokens.is_empty() {
4586            sentences.push(AnnotatedSentence {
4587                tokens: current_tokens,
4588                source_dataset: id,
4589            });
4590        }
4591
4592        if sentences.is_empty() {
4593            return Err(Error::InvalidInput(format!(
4594                "NCBI Disease file for {:?} contains no valid sentences",
4595                id
4596            )));
4597        }
4598
4599        let now = chrono::Utc::now().to_rfc3339();
4600        Ok(LoadedDataset {
4601            id,
4602            sentences,
4603            loaded_at: now,
4604            source_url: id.download_url().to_string(),
4605            data_source: DataSource::LocalCache,
4606            temporal_metadata: Self::get_temporal_metadata(id),
4607            metadata: id.default_metadata(),
4608        })
4609    }
4610
4611    /// Parse GAP coreference TSV format.
4612    ///
4613    /// GAP format columns: ID, Text, Pronoun, Pronoun-offset, A, A-offset, A-coref, B, B-offset, B-coref, URL
4614    fn parse_gap(&self, content: &str, id: DatasetId) -> Result<LoadedDataset> {
4615        let mut sentences = Vec::new();
4616        let mut first_line = true;
4617
4618        for line in content.lines() {
4619            // Skip header
4620            if first_line {
4621                first_line = false;
4622                continue;
4623            }
4624
4625            let parts: Vec<&str> = line.split('\t').collect();
4626            if parts.len() < 10 {
4627                continue;
4628            }
4629
4630            let text = parts[1];
4631
4632            // Create tokens (whitespace tokenization for simplicity)
4633            let tokens: Vec<AnnotatedToken> = text
4634                .split_whitespace()
4635                .map(|w| AnnotatedToken {
4636                    text: w.to_string(),
4637                    ner_tag: "O".to_string(),
4638                })
4639                .collect();
4640
4641            if !tokens.is_empty() {
4642                sentences.push(AnnotatedSentence {
4643                    tokens,
4644                    source_dataset: id,
4645                });
4646            }
4647        }
4648
4649        if sentences.is_empty() {
4650            return Err(Error::InvalidInput(format!(
4651                "GAP TSV file for {:?} contains no valid sentences",
4652                id
4653            )));
4654        }
4655
4656        let now = chrono::Utc::now().to_rfc3339();
4657        Ok(LoadedDataset {
4658            id,
4659            sentences,
4660            loaded_at: now,
4661            source_url: id.download_url().to_string(),
4662            data_source: DataSource::LocalCache,
4663            temporal_metadata: Self::get_temporal_metadata(id),
4664            metadata: id.default_metadata(),
4665        })
4666    }
4667
4668    /// Parse PreCo JSONL format from HuggingFace.
4669    ///
4670    /// PreCo JSONL format: One JSON object per line with "sentences" array.
4671    /// Note: PreCo is a coreference dataset, not NER. This parser extracts sentences
4672    /// for NER evaluation (which will have 0 entities). Use `load_coref()` for coreference.
4673    fn parse_preco_jsonl(&self, content: &str, id: DatasetId) -> Result<LoadedDataset> {
4674        let mut sentences = Vec::new();
4675        let mut line_count = 0usize;
4676        let mut parsed_count = 0usize;
4677
4678        for line in content.lines() {
4679            let line = line.trim();
4680            if line.is_empty() {
4681                continue;
4682            }
4683            line_count += 1;
4684
4685            let parsed: serde_json::Value = match serde_json::from_str(line) {
4686                Ok(v) => v,
4687                Err(e) => {
4688                    // Log first few parse errors for debugging
4689                    if parsed_count < 3 {
4690                        log::warn!("PreCo JSONL parse error on line {}: {}", line_count, e);
4691                    }
4692                    continue; // Skip malformed lines
4693                }
4694            };
4695
4696            // PreCo format: {"sentences": [[token1, token2, ...], ...]}
4697            if let Some(sents) = parsed.get("sentences").and_then(|v| v.as_array()) {
4698                parsed_count += 1;
4699                for sent_tokens in sents {
4700                    if let Some(token_array) = sent_tokens.as_array() {
4701                        let tokens: Vec<AnnotatedToken> = token_array
4702                            .iter()
4703                            .filter_map(|t| t.as_str())
4704                            .map(|t| AnnotatedToken {
4705                                text: t.to_string(),
4706                                ner_tag: "O".to_string(), // PreCo has no NER annotations
4707                            })
4708                            .collect();
4709
4710                        if !tokens.is_empty() {
4711                            sentences.push(AnnotatedSentence {
4712                                tokens,
4713                                source_dataset: id,
4714                            });
4715                        }
4716                    }
4717                }
4718            }
4719        }
4720
4721        if sentences.is_empty() {
4722            return Err(Error::InvalidInput(format!(
4723                "PreCo JSONL file contains no valid sentences (parsed {} of {} lines)",
4724                parsed_count, line_count
4725            )));
4726        }
4727
4728        let now = chrono::Utc::now().to_rfc3339();
4729        Ok(LoadedDataset {
4730            id,
4731            sentences,
4732            loaded_at: now,
4733            source_url: id.download_url().to_string(),
4734            data_source: DataSource::LocalCache,
4735            temporal_metadata: Self::get_temporal_metadata(id),
4736            metadata: id.default_metadata(),
4737        })
4738    }
4739
4740    /// Parse LitBank annotation format for NER.
4741    ///
4742    /// LitBank .ann format: T<id>\t<Type> <start> <end>\t<text>
4743    /// Note: LitBank is primarily a coreference dataset. This parser extracts
4744    /// entity mentions as NER annotations, but LitBank should be used with
4745    /// `load_coref()` for proper coreference evaluation.
4746    fn parse_litbank(&self, content: &str, id: DatasetId) -> Result<LoadedDataset> {
4747        // LitBank .ann format: each line is T<id>\t<Type> <start> <end>\t<text>
4748        // Extract entity mentions as NER annotations
4749        let now = chrono::Utc::now().to_rfc3339();
4750        let mut sentences = Vec::new();
4751        let mut entities: Vec<(usize, usize, String, String)> = Vec::new(); // (start, end, text, label)
4752
4753        for line in content.lines() {
4754            let line = line.trim();
4755            if line.is_empty() {
4756                continue;
4757            }
4758
4759            if line.starts_with('T') {
4760                // Entity annotation: T1\tPER 0 5\tAlice
4761                // LitBank uses ACE-style labels like PROP_PER, NOM_LOC, PRON_FAC etc.
4762                // We normalize to just the entity type (PER, LOC, ORG, GPE, FAC, VEH)
4763                // Note: LitBank also has coreference chain annotations like "character_name-ID"
4764                // which should be skipped for NER evaluation
4765                let parts: Vec<&str> = line.split('\t').collect();
4766                if parts.len() >= 3 {
4767                    let type_span: Vec<&str> = parts[1].split_whitespace().collect();
4768                    if type_span.len() >= 3 {
4769                        let raw_label = type_span[0];
4770
4771                        // Valid entity types (with or without prefix)
4772                        const VALID_ENTITY_TYPES: &[&str] = &[
4773                            "PER",
4774                            "LOC",
4775                            "ORG",
4776                            "GPE",
4777                            "FAC",
4778                            "VEH",
4779                            "PERSON",
4780                            "LOCATION",
4781                            "ORGANIZATION",
4782                        ];
4783
4784                        // Entity type annotations start with PROP_, NOM_, PRON_ OR are plain types
4785                        let is_prefixed_entity = raw_label.starts_with("PROP_")
4786                            || raw_label.starts_with("NOM_")
4787                            || raw_label.starts_with("PRON_");
4788                        let is_plain_entity = VALID_ENTITY_TYPES.contains(&raw_label);
4789
4790                        if !is_prefixed_entity && !is_plain_entity {
4791                            // Skip coreference chain annotations (e.g., "jarndyce_2-73")
4792                            continue;
4793                        }
4794
4795                        // Normalize: PROP_PER -> PER, NOM_LOC -> LOC, PRON_FAC -> FAC
4796                        let label = if is_prefixed_entity {
4797                            raw_label.split('_').next_back().unwrap_or(raw_label)
4798                        } else {
4799                            raw_label
4800                        };
4801                        let start: usize = type_span[1].parse().unwrap_or(0);
4802                        let end: usize = type_span[2].parse().unwrap_or(0);
4803                        let text = parts[2];
4804
4805                        entities.push((start, end, text.to_string(), label.to_string()));
4806                    }
4807                }
4808            }
4809        }
4810
4811        if entities.is_empty() {
4812            return Err(Error::InvalidInput(
4813                "LitBank .ann file contains no entity annotations (T lines)".to_string(),
4814            ));
4815        }
4816
4817        // Sort entities by start position
4818        entities.sort_by_key(|(start, _, _, _)| *start);
4819
4820        // Reconstruct text and create tokens with NER tags
4821        let max_end = entities
4822            .iter()
4823            .map(|(_, end, _, _)| *end)
4824            .max()
4825            .unwrap_or(0);
4826        let mut text_chars: Vec<char> = vec![' '; max_end.max(1)];
4827        let mut token_starts: Vec<usize> = Vec::new();
4828
4829        // Fill in entity text
4830        for (start, _end, text, _) in &entities {
4831            let text_chars_vec: Vec<char> = text.chars().collect();
4832            let _actual_end = (*start + text_chars_vec.len()).min(text_chars.len());
4833            if *start < text_chars.len() {
4834                for (i, ch) in text_chars_vec.iter().enumerate() {
4835                    let pos = *start + i;
4836                    if pos < text_chars.len() {
4837                        text_chars[pos] = *ch;
4838                    }
4839                }
4840            }
4841            token_starts.push(*start);
4842        }
4843
4844        // Note: text reconstruction is not used directly since we tokenize from entity text
4845        // but we keep it for potential future use (e.g., validation)
4846        let _text: String = text_chars
4847            .into_iter()
4848            .collect::<String>()
4849            .trim()
4850            .to_string();
4851
4852        // Create tokens with NER tags
4853        // Improved approach: tokenize entity text into words and apply BIO tags
4854        let mut tokens: Vec<AnnotatedToken> = Vec::new();
4855
4856        // Sort entities by start position for processing
4857        entities.sort_by_key(|(start, _, _, _)| *start);
4858
4859        let mut last_end = 0usize;
4860        for (start, end, entity_text, label) in &entities {
4861            // Add "O" tokens for gaps between entities
4862            if *start > last_end {
4863                // For gaps, we can't reconstruct the actual text without the .txt file
4864                // But we can estimate based on character distance
4865                let gap_size = *start - last_end;
4866                if gap_size > 0 {
4867                    // Estimate number of words in gap (roughly 5 chars per word + space)
4868                    let estimated_words = (gap_size / 6).max(1);
4869                    for _ in 0..estimated_words.min(10) {
4870                        // Limit gap tokens to avoid excessive placeholders
4871                        tokens.push(AnnotatedToken {
4872                            text: "[...]".to_string(),
4873                            ner_tag: "O".to_string(),
4874                        });
4875                    }
4876                }
4877            }
4878
4879            // Tokenize entity text into words and apply BIO tags
4880            let entity_words: Vec<&str> = entity_text.split_whitespace().collect();
4881            for (i, word) in entity_words.iter().enumerate() {
4882                let ner_tag = if i == 0 {
4883                    format!("B-{}", label)
4884                } else {
4885                    format!("I-{}", label)
4886                };
4887                tokens.push(AnnotatedToken {
4888                    text: word.to_string(),
4889                    ner_tag,
4890                });
4891            }
4892
4893            last_end = *end;
4894        }
4895
4896        if !tokens.is_empty() {
4897            sentences.push(AnnotatedSentence {
4898                tokens,
4899                source_dataset: id,
4900            });
4901        }
4902
4903        if sentences.is_empty() {
4904            return Err(Error::InvalidInput(
4905                "LitBank file produced no sentences after parsing".to_string(),
4906            ));
4907        }
4908
4909        Ok(LoadedDataset {
4910            id,
4911            sentences,
4912            loaded_at: now,
4913            source_url: id.download_url().to_string(),
4914            data_source: DataSource::LocalCache,
4915            temporal_metadata: Self::get_temporal_metadata(id),
4916            metadata: id.default_metadata(),
4917        })
4918    }
4919
4920    /// Parse ECB+ CSV format for event coreference.
4921    ///
4922    /// ECB+ uses CSV format with columns for event mentions and coreference links.
4923    /// For now, extracts entities as NER annotations (event triggers).
4924    fn parse_ecb_plus(&self, content: &str, id: DatasetId) -> Result<LoadedDataset> {
4925        let mut sentences = Vec::new();
4926        let mut first_line = true;
4927
4928        for line in content.lines() {
4929            // Skip header
4930            if first_line {
4931                first_line = false;
4932                continue;
4933            }
4934
4935            let parts: Vec<&str> = line.split(',').collect();
4936            if parts.len() < 3 {
4937                continue;
4938            }
4939
4940            // ECB+ CSV format: sentence_id, text, event_mention, ...
4941            // Extract text and create tokens
4942            let text = parts.get(1).unwrap_or(&"");
4943            let tokens: Vec<AnnotatedToken> = text
4944                .split_whitespace()
4945                .map(|w| AnnotatedToken {
4946                    text: w.to_string(),
4947                    ner_tag: "O".to_string(),
4948                })
4949                .collect();
4950
4951            if !tokens.is_empty() {
4952                sentences.push(AnnotatedSentence {
4953                    tokens,
4954                    source_dataset: id,
4955                });
4956            }
4957        }
4958
4959        if sentences.is_empty() {
4960            return Err(Error::InvalidInput(format!(
4961                "ECB+ CSV file for {:?} contains no valid sentences",
4962                id
4963            )));
4964        }
4965
4966        let now = chrono::Utc::now().to_rfc3339();
4967        Ok(LoadedDataset {
4968            id,
4969            sentences,
4970            loaded_at: now,
4971            source_url: id.download_url().to_string(),
4972            data_source: DataSource::LocalCache,
4973            temporal_metadata: Self::get_temporal_metadata(id),
4974            metadata: id.default_metadata(),
4975        })
4976    }
4977
4978    // =========================================================================
4979    // Coreference Loading
4980    // =========================================================================
4981
4982    /// Load coreference dataset, returning documents with chains.
4983    ///
4984    /// Use this for GAP, PreCo, and LitBank datasets.
4985    pub fn load_coref(&self, id: DatasetId) -> Result<Vec<super::coref::CorefDocument>> {
4986        if !id.is_coreference() {
4987            return Err(Error::InvalidInput(format!(
4988                "{:?} is not a coreference dataset",
4989                id
4990            )));
4991        }
4992
4993        let cache_path = self.cache_path_for(id);
4994        if !cache_path.exists() {
4995            return Err(Error::InvalidInput(format!(
4996                "Dataset {:?} not cached at {:?}. Download from {}",
4997                id,
4998                cache_path,
4999                id.download_url()
5000            )));
5001        }
5002
5003        let content = std::fs::read_to_string(&cache_path)
5004            .map_err(|e| Error::InvalidInput(format!("Failed to read {:?}: {}", cache_path, e)))?;
5005
5006        match id {
5007            DatasetId::CorefUD => super::coref_loader::parse_corefud_conllu(&content),
5008            DatasetId::GAP => {
5009                let examples = super::coref_loader::parse_gap_tsv(&content)?;
5010                Ok(examples
5011                    .into_iter()
5012                    .map(|ex| ex.to_coref_document())
5013                    .collect())
5014            }
5015            DatasetId::PreCo => {
5016                // PreCo can be JSONL (one JSON object per line) or JSON array
5017                // Try JSONL first (more common), then fall back to JSON array
5018                if content.trim().starts_with('[') {
5019                    // JSON array format
5020                    let docs = super::coref_loader::parse_preco_json(&content)?;
5021                    Ok(docs.into_iter().map(|d| d.to_coref_document()).collect())
5022                } else {
5023                    // JSONL format - parse each line and convert to JSON array format
5024                    let mut json_objects = Vec::new();
5025                    for line in content.lines() {
5026                        let line = line.trim();
5027                        if line.is_empty() {
5028                            continue;
5029                        }
5030                        // Validate it's valid JSON
5031                        if serde_json::from_str::<serde_json::Value>(line).is_ok() {
5032                            json_objects.push(line);
5033                        }
5034                    }
5035                    // Convert JSONL to JSON array format
5036                    let json_array = format!("[{}]", json_objects.join(","));
5037                    let docs = super::coref_loader::parse_preco_json(&json_array)?;
5038                    Ok(docs.into_iter().map(|d| d.to_coref_document()).collect())
5039                }
5040            }
5041            DatasetId::LitBank => {
5042                // LitBank coreference - parse .ann format for chains
5043                self.parse_litbank_coref(&content)
5044            }
5045            DatasetId::ECBPlus => {
5046                // ECB+ may be cached as either:
5047                // - ZIP binary (new: real XML annotations)
5048                // - CSV text (legacy: sentence index)
5049                let raw_bytes = std::fs::read(&cache_path).map_err(|e| {
5050                    Error::InvalidInput(format!("Failed to read {:?}: {}", cache_path, e))
5051                })?;
5052                if raw_bytes.starts_with(b"PK\x03\x04") {
5053                    return super::coref_loader::parse_ecb_plus_zip(&raw_bytes);
5054                }
5055                // Fall back to CSV parser
5056                super::coref_loader::parse_ecb_plus_coref(&content)
5057            }
5058            DatasetId::WikiCoref => {
5059                // WikiCoref uses a GAP-compatible TSV format.
5060                let examples = super::coref_loader::parse_gap_tsv(&content)?;
5061                Ok(examples
5062                    .into_iter()
5063                    .map(|ex| ex.to_coref_document())
5064                    .collect())
5065            }
5066            DatasetId::GUM
5067            | DatasetId::WinoBias
5068            | DatasetId::TwiConv
5069            | DatasetId::MuDoCo
5070            | DatasetId::SciCo => Err(Error::InvalidInput(format!(
5071                "{:?} coreference format is not yet supported (requires a dedicated parser)",
5072                id
5073            ))),
5074            DatasetId::BookCoref | DatasetId::BookCorefSplit => {
5075                // BOOKCOREF: Book-scale coreference (Martinelli et al. 2025)
5076                // Format: OntoNotes-style with character metadata
5077                // Note: Actual data requires HuggingFace datasets library to download
5078                // from Project Gutenberg. We support pre-downloaded JSONL.
5079                super::coref_loader::parse_bookcoref_json(&content)
5080            }
5081            _ => Err(Error::InvalidInput(format!(
5082                "No coreference parser for {:?}",
5083                id
5084            ))),
5085        }
5086    }
5087
5088    /// Load coreference dataset, downloading if needed.
5089    #[cfg(feature = "eval")]
5090    pub fn load_or_download_coref(
5091        &self,
5092        id: DatasetId,
5093    ) -> Result<Vec<super::coref::CorefDocument>> {
5094        if !self.is_cached_for(id) {
5095            if matches!(id, DatasetId::CorefUD) {
5096                let cache_path = self.cache_path_for(id);
5097                return Err(Error::InvalidInput(format!(
5098                    "CorefUD is not downloadable via anno yet. Please provide a local CorefUD .conllu file.\n\
5099                     - Option A: copy it to the cache path {:?}\n\
5100                     - Option B: use CorefLoader::load_corefud_from_path(<path>)",
5101                    cache_path
5102                )));
5103            }
5104
5105            let cache_path = self.cache_path_for(id);
5106            if matches!(id, DatasetId::ECBPlus) {
5107                // ECB+ is a ZIP file -- download as raw bytes
5108                let url = id.download_url();
5109                let bytes = self.download_attempt_bytes(url)?;
5110                std::fs::write(&cache_path, &bytes).map_err(|e| {
5111                    Error::InvalidInput(format!("Failed to cache {:?}: {}", cache_path, e))
5112                })?;
5113            } else {
5114                let (content, _) = self.download_with_resolved_url(id)?;
5115                std::fs::write(&cache_path, &content).map_err(|e| {
5116                    Error::InvalidInput(format!("Failed to cache {:?}: {}", cache_path, e))
5117                })?;
5118            }
5119        }
5120        self.load_coref(id)
5121    }
5122
5123    /// Parse LitBank for coreference chains.
5124    ///
5125    /// LitBank format: .ann files with T lines (mentions) and R lines (coreference relations).
5126    /// T lines: T<id>\t<Type> <start> <end>\t<text>
5127    /// R lines: R<id>\tCoref Arg1:T<id1> Arg2:T<id2>
5128    ///
5129    /// Note: LitBank .ann files reference character offsets in corresponding .txt files.
5130    /// This parser reconstructs text from mentions, but ideally should read .txt file.
5131    fn parse_litbank_coref(&self, content: &str) -> Result<Vec<super::coref::CorefDocument>> {
5132        use super::coref::{CorefChain, CorefDocument, Mention};
5133        use std::collections::HashMap;
5134
5135        // LitBank .ann format includes coreference with R lines
5136        // R1\tCoref Arg1:T1 Arg2:T2
5137        let mut mentions: HashMap<String, Mention> = HashMap::new();
5138        let mut coref_links: Vec<(String, String)> = Vec::new();
5139        let mut max_end = 0usize;
5140
5141        for line in content.lines() {
5142            let line = line.trim();
5143            if line.is_empty() {
5144                continue;
5145            }
5146
5147            if line.starts_with('T') {
5148                let parts: Vec<&str> = line.split('\t').collect();
5149                if parts.len() >= 3 {
5150                    let id = parts[0];
5151                    let type_span: Vec<&str> = parts[1].split_whitespace().collect();
5152                    if type_span.len() >= 3 {
5153                        let start: usize = type_span[1].parse().unwrap_or(0);
5154                        let end: usize = type_span[2].parse().unwrap_or(0);
5155                        let text = parts[2];
5156                        max_end = max_end.max(end);
5157                        mentions.insert(id.to_string(), Mention::new(text, start, end));
5158                    }
5159                }
5160            } else if line.starts_with('R') && line.contains("Coref") {
5161                // R1\tCoref Arg1:T1 Arg2:T2
5162                let parts: Vec<&str> = line.split_whitespace().collect();
5163                if parts.len() >= 3 {
5164                    let arg1 = parts[1].trim_start_matches("Arg1:");
5165                    let arg2 = parts[2].trim_start_matches("Arg2:");
5166                    coref_links.push((arg1.to_string(), arg2.to_string()));
5167                }
5168            }
5169        }
5170
5171        if mentions.is_empty() {
5172            return Err(Error::InvalidInput(
5173                "LitBank file contains no mentions (T lines)".to_string(),
5174            ));
5175        }
5176
5177        // Reconstruct text from mentions by sorting by start offset
5178        let mut sorted_mentions: Vec<(usize, &Mention)> =
5179            mentions.values().map(|m| (m.start, m)).collect();
5180        sorted_mentions.sort_by_key(|(start, _)| *start);
5181
5182        // Build text by inserting mentions at their offsets
5183        let mut text_chars: Vec<char> = vec![' '; max_end.max(1)];
5184        for (start, mention) in &sorted_mentions {
5185            let mention_text: Vec<char> = mention.text.chars().collect();
5186            let _end = (*start + mention_text.len()).min(text_chars.len());
5187            if *start < text_chars.len() {
5188                for (i, ch) in mention_text.iter().enumerate() {
5189                    let pos = *start + i;
5190                    if pos < text_chars.len() {
5191                        text_chars[pos] = *ch;
5192                    }
5193                }
5194            }
5195        }
5196        let text: String = text_chars
5197            .into_iter()
5198            .collect::<String>()
5199            .trim()
5200            .to_string();
5201
5202        // Build chains from links using union-find
5203        let mut chains: Vec<Vec<Mention>> = Vec::new();
5204        let mut mention_to_chain: HashMap<String, usize> = HashMap::new();
5205
5206        // First, add all mentions as singletons if they're not in any chain
5207        for (id, mention) in &mentions {
5208            if !mention_to_chain.contains_key(id) {
5209                let idx = chains.len();
5210                chains.push(vec![mention.clone()]);
5211                mention_to_chain.insert(id.clone(), idx);
5212            }
5213        }
5214
5215        // Then merge chains based on coref links
5216        for (id1, id2) in coref_links {
5217            let chain_idx = match (mention_to_chain.get(&id1), mention_to_chain.get(&id2)) {
5218                (Some(&idx1), Some(&idx2)) if idx1 != idx2 => {
5219                    // Merge chains
5220                    let to_merge = std::mem::take(&mut chains[idx2]);
5221                    chains[idx1].extend(to_merge);
5222                    // Update all mentions in merged chain to point to idx1
5223                    for m in &chains[idx1] {
5224                        // Find mention ID by matching text and position
5225                        for (mid, mref) in &mentions {
5226                            if mref.text == m.text && mref.start == m.start {
5227                                mention_to_chain.insert(mid.clone(), idx1);
5228                            }
5229                        }
5230                    }
5231                    idx1
5232                }
5233                (Some(&idx), None) => {
5234                    if let Some(m) = mentions.get(&id2) {
5235                        chains[idx].push(m.clone());
5236                        mention_to_chain.insert(id2, idx);
5237                    }
5238                    idx
5239                }
5240                (None, Some(&idx)) => {
5241                    if let Some(m) = mentions.get(&id1) {
5242                        chains[idx].push(m.clone());
5243                        mention_to_chain.insert(id1, idx);
5244                    }
5245                    idx
5246                }
5247                (None, None) => {
5248                    let idx = chains.len();
5249                    let mut chain = Vec::new();
5250                    if let Some(m) = mentions.get(&id1) {
5251                        chain.push(m.clone());
5252                        mention_to_chain.insert(id1.clone(), idx);
5253                    }
5254                    if let Some(m) = mentions.get(&id2) {
5255                        chain.push(m.clone());
5256                        mention_to_chain.insert(id2, idx);
5257                    }
5258                    if !chain.is_empty() {
5259                        chains.push(chain);
5260                    }
5261                    idx
5262                }
5263                (Some(&idx), Some(_)) => idx,
5264            };
5265            let _ = chain_idx; // Used above
5266        }
5267
5268        // Filter empty chains and convert
5269        let coref_chains: Vec<CorefChain> = chains
5270            .into_iter()
5271            .filter(|c| !c.is_empty())
5272            .enumerate()
5273            .map(|(i, mentions)| CorefChain::with_id(mentions, i as u64))
5274            .collect();
5275
5276        if coref_chains.is_empty() {
5277            return Err(Error::InvalidInput(
5278                "LitBank file contains no coreference chains".to_string(),
5279            ));
5280        }
5281
5282        // Create document with reconstructed text
5283        let doc = CorefDocument::new(&text, coref_chains);
5284        Ok(vec![doc])
5285    }
5286
5287    // =========================================================================
5288    // Relation Extraction Loading
5289    // =========================================================================
5290
5291    /// Load relation extraction dataset, returning documents with relations.
5292    ///
5293    /// Use this for DocRED and ReTACRED datasets.
5294    pub fn load_relation(&self, id: DatasetId) -> Result<Vec<RelationDocument>> {
5295        if !id.is_relation_extraction() {
5296            return Err(Error::InvalidInput(format!(
5297                "{:?} is not a relation extraction dataset",
5298                id
5299            )));
5300        }
5301
5302        let cache_path = self.cache_path_for(id);
5303        if !cache_path.exists() {
5304            return Err(Error::InvalidInput(format!(
5305                "Dataset {:?} not cached at {:?}. Download from {}",
5306                id,
5307                cache_path,
5308                id.download_url()
5309            )));
5310        }
5311
5312        let content = std::fs::read_to_string(&cache_path)
5313            .map_err(|e| Error::InvalidInput(format!("Failed to read {:?}: {}", cache_path, e)))?;
5314
5315        match id {
5316            DatasetId::DocRED
5317            | DatasetId::ReTACRED
5318            | DatasetId::NYTFB
5319            | DatasetId::WEBNLG
5320            | DatasetId::GoogleRE
5321            | DatasetId::BioRED
5322            | DatasetId::SciER
5323            | DatasetId::MixRED
5324            | DatasetId::CovEReD => {
5325                // All these datasets use the CrossRE format (same as DocRED)
5326                self.parse_docred_relations(&content)
5327            }
5328            DatasetId::CHisIEC => {
5329                // CHisIEC uses a different JSON format with entity indices
5330                self.parse_chisiec_relations(&content)
5331            }
5332            DatasetId::CADEC => {
5333                // CADEC is NER, not relation extraction
5334                Err(Error::InvalidInput(
5335                    "CADEC is a NER dataset, not relation extraction".to_string(),
5336                ))
5337            }
5338            _ => Err(Error::InvalidInput(format!(
5339                "No relation parser for {:?}",
5340                id
5341            ))),
5342        }
5343    }
5344
5345    /// Load relation extraction dataset, downloading if needed.
5346    #[cfg(feature = "eval")]
5347    pub fn load_or_download_relation(&self, id: DatasetId) -> Result<Vec<RelationDocument>> {
5348        if !self.is_cached_for(id) {
5349            let (content, _) = self.download_with_resolved_url(id)?;
5350            let cache_path = self.cache_path_for(id);
5351            std::fs::write(&cache_path, &content).map_err(|e| {
5352                Error::InvalidInput(format!("Failed to cache {:?}: {}", cache_path, e))
5353            })?;
5354        }
5355        self.load_relation(id)
5356    }
5357
5358    /// Parse DocRED/CrossRE format for relation extraction.
5359    ///
5360    /// Format: JSONL with {"sentence": [...], "ner": [[start, end, type], ...], "relations": [[id1-start, id1-end, id2-start, id2-end, rel-type, ...], ...]}
5361    fn parse_docred_relations(&self, content: &str) -> Result<Vec<RelationDocument>> {
5362        use super::relation::RelationGold;
5363
5364        let mut documents = Vec::new();
5365
5366        // Parse as JSONL (one JSON object per line)
5367        for line in content.lines() {
5368            let line = line.trim();
5369            if line.is_empty() {
5370                continue;
5371            }
5372
5373            let doc: serde_json::Value = match serde_json::from_str(line) {
5374                Ok(v) => v,
5375                Err(_) => continue,
5376            };
5377
5378            // Get sentence tokens
5379            let tokens_arr = match doc.get("sentence").and_then(|v| v.as_array()) {
5380                Some(t) => t,
5381                None => continue,
5382            };
5383
5384            // Build text from tokens (with proper spacing)
5385            let text: String = tokens_arr
5386                .iter()
5387                .filter_map(|t| t.as_str())
5388                .collect::<Vec<_>>()
5389                .join(" ");
5390
5391            // Build token-to-character offset mapping
5392            // This maps each token index to its character start position in the text
5393            let mut token_to_char: Vec<usize> = Vec::new();
5394            let mut char_pos = 0;
5395            for (i, token) in tokens_arr.iter().enumerate() {
5396                if let Some(tok_str) = token.as_str() {
5397                    token_to_char.push(char_pos);
5398                    // Add token length + 1 for space (except last token)
5399                    char_pos += tok_str.len();
5400                    if i < tokens_arr.len() - 1 {
5401                        char_pos += 1; // Space between tokens
5402                    }
5403                } else {
5404                    token_to_char.push(char_pos);
5405                }
5406            }
5407
5408            // Get NER spans: [start, end, type]
5409            let ner_spans = doc.get("ner").and_then(|v| v.as_array());
5410
5411            // Build entity map: (token_start, token_end) -> (type, text, char_start, char_end)
5412            let mut entity_map: std::collections::HashMap<
5413                (usize, usize),
5414                (String, String, usize, usize),
5415            > = std::collections::HashMap::new();
5416            if let Some(ner) = ner_spans {
5417                for span in ner {
5418                    if let Some(arr) = span.as_array() {
5419                        if arr.len() >= 3 {
5420                            let token_start = arr[0].as_u64().unwrap_or(0) as usize;
5421                            let token_end = arr[1].as_u64().unwrap_or(0) as usize;
5422                            let ent_type = arr[2].as_str().unwrap_or("ENTITY").to_string();
5423
5424                            // Extract entity text from tokens
5425                            let entity_text: String = tokens_arr
5426                                .iter()
5427                                .skip(token_start)
5428                                .take(token_end - token_start + 1)
5429                                .filter_map(|t| t.as_str())
5430                                .collect::<Vec<_>>()
5431                                .join(" ");
5432
5433                            // Calculate actual character offsets
5434                            let char_start = token_to_char.get(token_start).copied().unwrap_or(0);
5435                            let char_end = if token_end < token_to_char.len() {
5436                                // Get end position of last token
5437                                let last_token_char_start = token_to_char[token_end];
5438                                if let Some(last_token) =
5439                                    tokens_arr.get(token_end).and_then(|t| t.as_str())
5440                                {
5441                                    last_token_char_start + last_token.len()
5442                                } else {
5443                                    char_start + entity_text.len()
5444                                }
5445                            } else {
5446                                char_start + entity_text.len()
5447                            };
5448
5449                            entity_map.insert(
5450                                (token_start, token_end),
5451                                (ent_type, entity_text, char_start, char_end),
5452                            );
5453                        }
5454                    }
5455                }
5456            }
5457
5458            // Parse relations: [id1-start, id1-end, id2-start, id2-end, rel-type, ...]
5459            let relations_arr = doc.get("relations").and_then(|v| v.as_array());
5460            let mut relations = Vec::new();
5461
5462            if let Some(rels) = relations_arr {
5463                for rel in rels {
5464                    if let Some(arr) = rel.as_array() {
5465                        if arr.len() >= 5 {
5466                            let head_token_start = arr[0].as_u64().unwrap_or(0) as usize;
5467                            let head_token_end = arr[1].as_u64().unwrap_or(0) as usize;
5468                            let tail_token_start = arr[2].as_u64().unwrap_or(0) as usize;
5469                            let tail_token_end = arr[3].as_u64().unwrap_or(0) as usize;
5470                            let rel_type = arr[4].as_str().unwrap_or("RELATION").to_string();
5471
5472                            // Get entity info from map (including character offsets)
5473                            let (head_type, head_text, head_char_start, head_char_end) = entity_map
5474                                .get(&(head_token_start, head_token_end))
5475                                .cloned()
5476                                .unwrap_or_else(|| {
5477                                    // Fallback: compute from token positions
5478                                    let char_start =
5479                                        token_to_char.get(head_token_start).copied().unwrap_or(0);
5480                                    let char_end = if head_token_end < token_to_char.len() {
5481                                        let last_start = token_to_char[head_token_end];
5482                                        if let Some(last_tok) =
5483                                            tokens_arr.get(head_token_end).and_then(|t| t.as_str())
5484                                        {
5485                                            last_start + last_tok.len()
5486                                        } else {
5487                                            char_start
5488                                        }
5489                                    } else {
5490                                        char_start
5491                                    };
5492                                    ("ENTITY".to_string(), String::new(), char_start, char_end)
5493                                });
5494
5495                            let (tail_type, tail_text, tail_char_start, tail_char_end) = entity_map
5496                                .get(&(tail_token_start, tail_token_end))
5497                                .cloned()
5498                                .unwrap_or_else(|| {
5499                                    // Fallback: compute from token positions
5500                                    let char_start =
5501                                        token_to_char.get(tail_token_start).copied().unwrap_or(0);
5502                                    let char_end = if tail_token_end < token_to_char.len() {
5503                                        let last_start = token_to_char[tail_token_end];
5504                                        if let Some(last_tok) =
5505                                            tokens_arr.get(tail_token_end).and_then(|t| t.as_str())
5506                                        {
5507                                            last_start + last_tok.len()
5508                                        } else {
5509                                            char_start
5510                                        }
5511                                    } else {
5512                                        char_start
5513                                    };
5514                                    ("ENTITY".to_string(), String::new(), char_start, char_end)
5515                                });
5516
5517                            relations.push(RelationGold::new(
5518                                (head_char_start, head_char_end),
5519                                head_type,
5520                                head_text,
5521                                (tail_char_start, tail_char_end),
5522                                tail_type,
5523                                tail_text,
5524                                rel_type,
5525                            ));
5526                        }
5527                    }
5528                }
5529            }
5530
5531            if !text.is_empty() {
5532                documents.push(RelationDocument { text, relations });
5533            }
5534        }
5535
5536        Ok(documents)
5537    }
5538
5539    /// Parse CHisIEC (Chinese Historical Information Extraction Corpus) NER format.
5540    ///
5541    /// # Research Background
5542    ///
5543    /// CHisIEC (Tang et al., LREC-COLING 2024) addresses the unique challenges of
5544    /// information extraction from ancient Chinese historical texts (文言文):
5545    ///
5546    /// - **No word boundaries**: Classical Chinese has no spaces; we tokenize by character
5547    /// - **Archaic vocabulary**: Official titles (官职) differ from modern equivalents
5548    /// - **Long time span**: Covers texts from multiple dynasties across 2000+ years
5549    /// - **Domain-specific entities**: Government positions, classical texts, historical figures
5550    ///
5551    /// Paper: <https://aclanthology.org/2024.lrec-main.283/>
5552    /// GitHub: <https://github.com/tangxuemei1995/CHisIEC>
5553    ///
5554    /// # Entity Types
5555    ///
5556    /// | Type | Chinese | Meaning | Example |
5557    /// |------|---------|---------|---------|
5558    /// | PER | 人物 | Historical figure | 司馬遷, 曹操 |
5559    /// | LOC | 地点 | Location/Place | 長安, 洛陽 |
5560    /// | OFI | 官职 | Official position | 丞相, 太守 |
5561    /// | BOOK | 书籍 | Classical text | 史記, 論語 |
5562    ///
5563    /// # Format
5564    ///
5565    /// JSON array where each object contains:
5566    /// - `tokens`: String of continuous text (no spaces between characters)
5567    /// - `entities`: Array of entity objects with `type`, `start`, `end`, `span`
5568    ///
5569    /// # Distinction from HistoricalChineseNER
5570    ///
5571    /// **CHisIEC** (this dataset):
5572    /// - Ancient Chinese (文言文, pre-modern)
5573    /// - Source: 24 dynastic histories (二十四史)
5574    /// - Tasks: NER + Relation Extraction
5575    ///
5576    /// **HistoricalChineseNER** (different dataset):
5577    /// - Modern Chinese transition (1872-1949)
5578    /// - Source: ShenBao newspapers
5579    /// - Tasks: NER + Entity Linking + Coreference
5580    ///
5581    /// # Implementation Notes
5582    ///
5583    /// Chinese text is tokenized by character for NER tagging. Character offsets
5584    /// are used directly (critical for Unicode correctness).
5585    fn parse_chisiec(&self, content: &str, id: DatasetId) -> Result<LoadedDataset> {
5586        let mut sentences = Vec::new();
5587
5588        // CHisIEC RE data is a JSON array
5589        let docs: Vec<serde_json::Value> = serde_json::from_str(content)
5590            .map_err(|e| Error::InvalidInput(format!("Failed to parse CHisIEC JSON: {}", e)))?;
5591
5592        for doc in docs {
5593            // Get tokens string (characters concatenated, no spaces in Chinese)
5594            let text = match doc.get("tokens").and_then(|v| v.as_str()) {
5595                Some(t) => t.to_string(),
5596                None => continue,
5597            };
5598
5599            if text.is_empty() {
5600                continue;
5601            }
5602
5603            // Chinese text: each character is a token
5604            let chars: Vec<char> = text.chars().collect();
5605            let mut tokens: Vec<AnnotatedToken> = chars
5606                .iter()
5607                .map(|c| AnnotatedToken {
5608                    text: c.to_string(),
5609                    ner_tag: "O".to_string(),
5610                })
5611                .collect();
5612
5613            // Get entities array and build BIO tags
5614            let entities_arr = doc.get("entities").and_then(|v| v.as_array());
5615
5616            if let Some(entities) = entities_arr {
5617                for entity in entities {
5618                    let ent_type = entity
5619                        .get("type")
5620                        .and_then(|v| v.as_str())
5621                        .unwrap_or("ENTITY");
5622                    let start = entity.get("start").and_then(|v| v.as_u64()).unwrap_or(0) as usize;
5623                    let end = entity.get("end").and_then(|v| v.as_u64()).unwrap_or(0) as usize;
5624
5625                    // Apply BIO tags (CHisIEC uses character indices)
5626                    for idx in start..end {
5627                        if idx < tokens.len() {
5628                            tokens[idx].ner_tag = if idx == start {
5629                                format!("B-{}", ent_type)
5630                            } else {
5631                                format!("I-{}", ent_type)
5632                            };
5633                        }
5634                    }
5635                }
5636            }
5637
5638            if !tokens.is_empty() {
5639                sentences.push(AnnotatedSentence {
5640                    tokens,
5641                    source_dataset: id,
5642                });
5643            }
5644        }
5645
5646        if sentences.is_empty() {
5647            return Err(Error::InvalidInput(format!(
5648                "CHisIEC file for {:?} contains no valid sentences",
5649                id
5650            )));
5651        }
5652
5653        let now = chrono::Utc::now().to_rfc3339();
5654        Ok(LoadedDataset {
5655            id,
5656            sentences,
5657            loaded_at: now,
5658            source_url: id.download_url().to_string(),
5659            data_source: DataSource::LocalCache,
5660            temporal_metadata: Self::get_temporal_metadata(id),
5661            metadata: id.default_metadata(),
5662        })
5663    }
5664
5665    /// Parse CHisIEC relation extraction format.
5666    ///
5667    /// # Research Motivation
5668    ///
5669    /// Ancient Chinese historical texts encode complex socio-political relationships
5670    /// that require domain-specific relation types. The 12 relation types in CHisIEC
5671    /// reflect structures unique to dynastic China:
5672    ///
5673    /// - **Hierarchical**: 上下級 (superior-subordinate), 任職 (office-holding)
5674    /// - **Kinship**: 父母 (parents), 兄弟 (siblings)
5675    /// - **Political**: 政治奧援 (political support), 同僚 (colleagues)
5676    /// - **Military**: 敵對攻伐 (hostility/attack), 駐守 (defend/garrison)
5677    /// - **Spatial**: 到達 (arrival), 出生於某地 (birthplace), 管理 (governance)
5678    /// - **Identity**: 別名 (alias/alternative name)
5679    ///
5680    /// # Format Difference from CrossRE/DocRED
5681    ///
5682    /// **CHisIEC** uses entity indices:
5683    /// ```json
5684    /// "relations": [{"type": "上下級", "head": 0, "tail": 1, ...}]
5685    /// ```
5686    ///
5687    /// **CrossRE/DocRED** uses token spans:
5688    /// ```json
5689    /// "relations": [[0, 2, 3, 5, "relation_type"]]
5690    /// ```
5691    ///
5692    /// This difference requires a separate parser.
5693    ///
5694    /// # JSON Format
5695    ///
5696    /// ```json
5697    /// {
5698    ///   "tokens": "明日,召問其故...",
5699    ///   "entities": [
5700    ///     {"type": "PER", "start": 11, "end": 13, "span": "玉汝"},
5701    ///     {"type": "PER", "start": 14, "end": 16, "span": "嚴公"}
5702    ///   ],
5703    ///   "relations": [
5704    ///     {"type": "上下級", "head": 1, "tail": 0, "head_span": "嚴公", "tail_span": "玉汝"}
5705    ///   ],
5706    ///   "orig_id": 1260
5707    /// }
5708    /// ```
5709    ///
5710    /// # Relation Types (12 total)
5711    ///
5712    /// | Chinese | Pinyin | English | Category |
5713    /// |---------|--------|---------|----------|
5714    /// | 敵對攻伐 | dí duì gōng fá | Attack/Hostility | Military |
5715    /// | 任職 | rèn zhí | Office Holding | Political |
5716    /// | 上下級 | shàng xià jí | Superior-subordinate | Hierarchical |
5717    /// | 政治奧援 | zhèng zhì ào yuán | Political Support | Political |
5718    /// | 同僚 | tóng liáo | Colleague | Social |
5719    /// | 到達 | dào dá | Arrive | Spatial |
5720    /// | 管理 | guǎn lǐ | Manage/Govern | Administrative |
5721    /// | 出生於某地 | chū shēng yú mǒu dì | Birthplace | Spatial |
5722    /// | 駐守 | zhù shǒu | Defend/Garrison | Military |
5723    /// | 別名 | bié míng | Alias | Identity |
5724    /// | 父母 | fù mǔ | Parents | Kinship |
5725    /// | 兄弟 | xiōng dì | Siblings | Kinship |
5726    fn parse_chisiec_relations(&self, content: &str) -> Result<Vec<RelationDocument>> {
5727        use super::relation::RelationGold;
5728
5729        let mut documents = Vec::new();
5730
5731        // Parse as JSON array
5732        let docs: Vec<serde_json::Value> = serde_json::from_str(content)
5733            .map_err(|e| Error::InvalidInput(format!("Failed to parse CHisIEC JSON: {}", e)))?;
5734
5735        for doc in docs {
5736            // Get tokens string
5737            let text = match doc.get("tokens").and_then(|v| v.as_str()) {
5738                Some(t) => t.to_string(),
5739                None => continue,
5740            };
5741
5742            if text.is_empty() {
5743                continue;
5744            }
5745
5746            // Build entity list: [(type, start, end, text), ...]
5747            let entities_arr = doc.get("entities").and_then(|v| v.as_array());
5748            let mut entity_list: Vec<(String, usize, usize, String)> = Vec::new();
5749
5750            if let Some(entities) = entities_arr {
5751                for entity in entities {
5752                    let ent_type = entity
5753                        .get("type")
5754                        .and_then(|v| v.as_str())
5755                        .unwrap_or("ENTITY")
5756                        .to_string();
5757                    let start = entity.get("start").and_then(|v| v.as_u64()).unwrap_or(0) as usize;
5758                    let end = entity.get("end").and_then(|v| v.as_u64()).unwrap_or(0) as usize;
5759                    let span = entity
5760                        .get("span")
5761                        .and_then(|v| v.as_str())
5762                        .unwrap_or("")
5763                        .to_string();
5764
5765                    // Fallback to extracting from text if span is empty
5766                    let span_text = if !span.is_empty() {
5767                        span
5768                    } else {
5769                        text.chars().skip(start).take(end - start).collect()
5770                    };
5771
5772                    entity_list.push((ent_type, start, end, span_text));
5773                }
5774            }
5775
5776            // Parse relations
5777            let relations_arr = doc.get("relations").and_then(|v| v.as_array());
5778            let mut relations = Vec::new();
5779
5780            if let Some(rels) = relations_arr {
5781                for rel in rels {
5782                    let rel_type = rel
5783                        .get("type")
5784                        .and_then(|v| v.as_str())
5785                        .unwrap_or("RELATION")
5786                        .to_string();
5787                    // CHisIEC uses entity indices for head/tail
5788                    let head_idx = rel.get("head").and_then(|v| v.as_u64()).unwrap_or(0) as usize;
5789                    let tail_idx = rel.get("tail").and_then(|v| v.as_u64()).unwrap_or(0) as usize;
5790
5791                    // Look up entities by index
5792                    if head_idx < entity_list.len() && tail_idx < entity_list.len() {
5793                        let (head_type, head_start, head_end, head_text) = &entity_list[head_idx];
5794                        let (tail_type, tail_start, tail_end, tail_text) = &entity_list[tail_idx];
5795
5796                        relations.push(RelationGold::new(
5797                            (*head_start, *head_end),
5798                            head_type.clone(),
5799                            head_text.clone(),
5800                            (*tail_start, *tail_end),
5801                            tail_type.clone(),
5802                            tail_text.clone(),
5803                            rel_type,
5804                        ));
5805                    }
5806                }
5807            }
5808
5809            if !text.is_empty() {
5810                documents.push(RelationDocument { text, relations });
5811            }
5812        }
5813
5814        Ok(documents)
5815    }
5816
5817    // =========================================================================
5818    // African Language Dataset Parsers (Masakhane Community)
5819    // =========================================================================
5820
5821    /// Parse AfriSenti sentiment analysis format.
5822    ///
5823    /// AfriSenti uses TSV format: text\tlabel
5824    /// Labels: positive, neutral, negative
5825    ///
5826    /// Source: <https://github.com/afrisenti-semeval/afrisent-semeval-2023>
5827    fn parse_afrisenti(&self, content: &str, id: DatasetId) -> Result<LoadedDataset> {
5828        let mut sentences = Vec::new();
5829        let now = chrono::Utc::now().to_rfc3339();
5830
5831        for line in content.lines() {
5832            let line = line.trim();
5833            if line.is_empty() || line.starts_with('#') {
5834                continue;
5835            }
5836
5837            let parts: Vec<&str> = line.split('\t').collect();
5838            if parts.len() >= 2 {
5839                let text = parts[0].to_string();
5840                let label = parts[1].to_string();
5841
5842                // For sentiment, we create a single token annotation with the sentiment label
5843                let tokens = vec![AnnotatedToken {
5844                    text: text.clone(),
5845                    ner_tag: format!("B-{}", label),
5846                }];
5847
5848                sentences.push(AnnotatedSentence {
5849                    tokens,
5850                    source_dataset: id,
5851                });
5852            }
5853        }
5854
5855        if sentences.is_empty() {
5856            return Err(Error::InvalidInput(format!(
5857                "AfriSenti file for {:?} contains no valid sentences",
5858                id
5859            )));
5860        }
5861
5862        Ok(LoadedDataset {
5863            id,
5864            sentences,
5865            loaded_at: now,
5866            source_url: id.download_url().to_string(),
5867            data_source: DataSource::LocalCache,
5868            temporal_metadata: Self::get_temporal_metadata(id),
5869            metadata: id.default_metadata(),
5870        })
5871    }
5872
5873    /// Parse AfriQA question answering format.
5874    ///
5875    /// AfriQA uses JSON format with fields: context, question, answers
5876    /// Each answer has text and answer_start
5877    ///
5878    /// Source: <https://github.com/masakhane-io/afriqa>
5879    fn parse_afriqa(&self, content: &str, id: DatasetId) -> Result<LoadedDataset> {
5880        let mut sentences = Vec::new();
5881        let now = chrono::Utc::now().to_rfc3339();
5882
5883        // AfriQA data can be JSON array or JSONL
5884        let docs: Vec<serde_json::Value> = if content.trim().starts_with('[') {
5885            serde_json::from_str(content)
5886                .map_err(|e| Error::InvalidInput(format!("Failed to parse AfriQA JSON: {}", e)))?
5887        } else {
5888            // JSONL format
5889            content
5890                .lines()
5891                .filter(|l| !l.trim().is_empty())
5892                .filter_map(|l| serde_json::from_str(l).ok())
5893                .collect()
5894        };
5895
5896        for doc in docs {
5897            // Get context text
5898            let context = doc.get("context").and_then(|v| v.as_str()).unwrap_or("");
5899
5900            // Get answers
5901            if let Some(answers) = doc.get("answers").and_then(|v| v.as_object()) {
5902                let texts = answers.get("text").and_then(|v| v.as_array());
5903                let starts = answers.get("answer_start").and_then(|v| v.as_array());
5904
5905                if let (Some(texts), Some(starts)) = (texts, starts) {
5906                    // Create tokens for context (word-level tokenization)
5907                    let words: Vec<&str> = context.split_whitespace().collect();
5908                    let mut tokens: Vec<AnnotatedToken> = words
5909                        .iter()
5910                        .map(|w| AnnotatedToken {
5911                            text: w.to_string(),
5912                            ner_tag: "O".to_string(),
5913                        })
5914                        .collect();
5915
5916                    // Mark answer spans
5917                    for (text_val, start_val) in texts.iter().zip(starts.iter()) {
5918                        if let (Some(answer_text), Some(start)) =
5919                            (text_val.as_str(), start_val.as_u64())
5920                        {
5921                            let start = start as usize;
5922                            let answer_words: Vec<&str> = answer_text.split_whitespace().collect();
5923
5924                            // Find word index by counting words before start position
5925                            let prefix: String = context.chars().take(start).collect();
5926                            let word_idx = prefix.split_whitespace().count();
5927
5928                            // Apply BIO tags
5929                            for (i, _) in answer_words.iter().enumerate() {
5930                                let idx = word_idx + i;
5931                                if idx < tokens.len() {
5932                                    tokens[idx].ner_tag = if i == 0 {
5933                                        "B-ANSWER".to_string()
5934                                    } else {
5935                                        "I-ANSWER".to_string()
5936                                    };
5937                                }
5938                            }
5939                        }
5940                    }
5941
5942                    if !tokens.is_empty() {
5943                        sentences.push(AnnotatedSentence {
5944                            tokens,
5945                            source_dataset: id,
5946                        });
5947                    }
5948                }
5949            }
5950        }
5951
5952        if sentences.is_empty() {
5953            return Err(Error::InvalidInput(format!(
5954                "AfriQA file for {:?} contains no valid sentences",
5955                id
5956            )));
5957        }
5958
5959        Ok(LoadedDataset {
5960            id,
5961            sentences,
5962            loaded_at: now,
5963            source_url: id.download_url().to_string(),
5964            data_source: DataSource::LocalCache,
5965            temporal_metadata: Self::get_temporal_metadata(id),
5966            metadata: id.default_metadata(),
5967        })
5968    }
5969
5970    /// Parse MasakhaNEWS topic classification format.
5971    ///
5972    /// MasakhaNEWS uses TSV format: headline\tbody\tcategory
5973    /// Categories: business, entertainment, health, politics, religion, sports, technology
5974    ///
5975    /// Source: <https://github.com/masakhane-io/masakhane-news>
5976    fn parse_masakhanews(&self, content: &str, id: DatasetId) -> Result<LoadedDataset> {
5977        let mut sentences = Vec::new();
5978        let now = chrono::Utc::now().to_rfc3339();
5979
5980        for line in content.lines() {
5981            let line = line.trim();
5982            if line.is_empty() || line.starts_with('#') || line.starts_with("headline\t") {
5983                continue;
5984            }
5985
5986            let parts: Vec<&str> = line.split('\t').collect();
5987            if parts.len() >= 2 {
5988                // Use headline as text, category as label
5989                let text = parts[0].to_string();
5990                let category = if parts.len() >= 3 {
5991                    parts[2].to_string()
5992                } else {
5993                    parts[1].to_string()
5994                };
5995
5996                // For classification, create single token with category
5997                let tokens = vec![AnnotatedToken {
5998                    text,
5999                    ner_tag: format!("B-{}", category),
6000                }];
6001
6002                sentences.push(AnnotatedSentence {
6003                    tokens,
6004                    source_dataset: id,
6005                });
6006            }
6007        }
6008
6009        if sentences.is_empty() {
6010            return Err(Error::InvalidInput(format!(
6011                "MasakhaNEWS file for {:?} contains no valid sentences",
6012                id
6013            )));
6014        }
6015
6016        Ok(LoadedDataset {
6017            id,
6018            sentences,
6019            loaded_at: now,
6020            source_url: id.download_url().to_string(),
6021            data_source: DataSource::LocalCache,
6022            temporal_metadata: Self::get_temporal_metadata(id),
6023            metadata: id.default_metadata(),
6024        })
6025    }
6026
6027    /// Parse TREC question classification format.
6028    ///
6029    /// TREC format: `COARSE:fine question text`
6030    /// E.g.: `NUM:dist How far is it from Denver to Aspen ?`
6031    ///
6032    /// 6 coarse classes: ABBR, DESC, ENTY, HUM, LOC, NUM
6033    /// 50 fine-grained classes.
6034    ///
6035    /// Source: <https://cogcomp.seas.upenn.edu/Data/QA/QC/>
6036    fn parse_trec(&self, content: &str, id: DatasetId) -> Result<LoadedDataset> {
6037        let mut sentences = Vec::new();
6038        let now = chrono::Utc::now().to_rfc3339();
6039
6040        for line in content.lines() {
6041            let line = line.trim();
6042            if line.is_empty() {
6043                continue;
6044            }
6045
6046            // Format: COARSE:fine question text
6047            // Find the first space to separate label from question
6048            if let Some(space_idx) = line.find(' ') {
6049                let label = &line[..space_idx];
6050                let question = line[space_idx + 1..].trim();
6051
6052                // Extract coarse label (before colon)
6053                let coarse_label = label.split(':').next().unwrap_or(label);
6054
6055                // Create single token with the question and classification label
6056                let tokens = vec![AnnotatedToken {
6057                    text: question.to_string(),
6058                    ner_tag: format!("B-{}", coarse_label),
6059                }];
6060
6061                sentences.push(AnnotatedSentence {
6062                    tokens,
6063                    source_dataset: id,
6064                });
6065            }
6066        }
6067
6068        if sentences.is_empty() {
6069            return Err(Error::InvalidInput(format!(
6070                "TREC file for {:?} contains no valid sentences",
6071                id
6072            )));
6073        }
6074
6075        Ok(LoadedDataset {
6076            id,
6077            sentences,
6078            loaded_at: now,
6079            source_url: id.download_url().to_string(),
6080            data_source: DataSource::LocalCache,
6081            temporal_metadata: Self::get_temporal_metadata(id),
6082            metadata: id.default_metadata(),
6083        })
6084    }
6085
6086    /// Parse AG News classification format (parquet or CSV).
6087    ///
6088    /// AG News has 4 classes: World (0), Sports (1), Business (2), Sci/Tech (3)
6089    ///
6090    /// Source: <https://huggingface.co/datasets/ag_news>
6091    fn parse_agnews(&self, content: &str, id: DatasetId) -> Result<LoadedDataset> {
6092        let mut sentences = Vec::new();
6093        let now = chrono::Utc::now().to_rfc3339();
6094
6095        let label_map = ["World", "Sports", "Business", "Sci/Tech"];
6096
6097        // Try parsing as JSON (converted from parquet)
6098        for line in content.lines() {
6099            let line = line.trim();
6100            if line.is_empty() || line.starts_with('#') {
6101                continue;
6102            }
6103
6104            // Try JSON format
6105            if let Ok(obj) = serde_json::from_str::<serde_json::Value>(line) {
6106                let text = obj.get("text").and_then(|v| v.as_str()).unwrap_or_default();
6107                let label_idx = obj.get("label").and_then(|v| v.as_i64()).unwrap_or(0) as usize;
6108
6109                let label = label_map.get(label_idx).unwrap_or(&"Unknown");
6110
6111                let tokens = vec![AnnotatedToken {
6112                    text: text.to_string(),
6113                    ner_tag: format!("B-{}", label),
6114                }];
6115
6116                sentences.push(AnnotatedSentence {
6117                    tokens,
6118                    source_dataset: id,
6119                });
6120            }
6121        }
6122
6123        if sentences.is_empty() {
6124            return Err(Error::InvalidInput(format!(
6125                "AG News file for {:?} contains no valid sentences",
6126                id
6127            )));
6128        }
6129
6130        Ok(LoadedDataset {
6131            id,
6132            sentences,
6133            loaded_at: now,
6134            source_url: id.download_url().to_string(),
6135            data_source: DataSource::LocalCache,
6136            temporal_metadata: Self::get_temporal_metadata(id),
6137            metadata: id.default_metadata(),
6138        })
6139    }
6140
6141    /// Parse DBPedia-14 classification format.
6142    ///
6143    /// 14 classes from DBpedia ontology.
6144    ///
6145    /// Source: <https://huggingface.co/datasets/dbpedia_14>
6146    fn parse_dbpedia14(&self, content: &str, id: DatasetId) -> Result<LoadedDataset> {
6147        let mut sentences = Vec::new();
6148        let now = chrono::Utc::now().to_rfc3339();
6149
6150        let label_map = [
6151            "Company",
6152            "EducationalInstitution",
6153            "Artist",
6154            "Athlete",
6155            "OfficeHolder",
6156            "MeanOfTransportation",
6157            "Building",
6158            "NaturalPlace",
6159            "Village",
6160            "Animal",
6161            "Plant",
6162            "Album",
6163            "Film",
6164            "WrittenWork",
6165        ];
6166
6167        for line in content.lines() {
6168            let line = line.trim();
6169            if line.is_empty() {
6170                continue;
6171            }
6172
6173            if let Ok(obj) = serde_json::from_str::<serde_json::Value>(line) {
6174                let content_text = obj
6175                    .get("content")
6176                    .and_then(|v| v.as_str())
6177                    .unwrap_or_default();
6178                let label_idx = obj.get("label").and_then(|v| v.as_i64()).unwrap_or(0) as usize;
6179
6180                let label = label_map.get(label_idx).unwrap_or(&"Unknown");
6181
6182                let tokens = vec![AnnotatedToken {
6183                    text: content_text.to_string(),
6184                    ner_tag: format!("B-{}", label),
6185                }];
6186
6187                sentences.push(AnnotatedSentence {
6188                    tokens,
6189                    source_dataset: id,
6190                });
6191            }
6192        }
6193
6194        if sentences.is_empty() {
6195            return Err(Error::InvalidInput(format!(
6196                "DBPedia-14 file for {:?} contains no valid sentences",
6197                id
6198            )));
6199        }
6200
6201        Ok(LoadedDataset {
6202            id,
6203            sentences,
6204            loaded_at: now,
6205            source_url: id.download_url().to_string(),
6206            data_source: DataSource::LocalCache,
6207            temporal_metadata: Self::get_temporal_metadata(id),
6208            metadata: id.default_metadata(),
6209        })
6210    }
6211
6212    /// Parse Yahoo Answers topic classification format.
6213    ///
6214    /// 10 topics: Society, Science, Health, Education, Computers,
6215    /// Sports, Business, Entertainment, Family, Politics
6216    ///
6217    /// Source: <https://huggingface.co/datasets/yahoo_answers_topics>
6218    fn parse_yahoo_answers(&self, content: &str, id: DatasetId) -> Result<LoadedDataset> {
6219        let mut sentences = Vec::new();
6220        let now = chrono::Utc::now().to_rfc3339();
6221
6222        let label_map = [
6223            "Society",
6224            "Science",
6225            "Health",
6226            "Education",
6227            "Computers",
6228            "Sports",
6229            "Business",
6230            "Entertainment",
6231            "Family",
6232            "Politics",
6233        ];
6234
6235        for line in content.lines() {
6236            let line = line.trim();
6237            if line.is_empty() {
6238                continue;
6239            }
6240
6241            if let Ok(obj) = serde_json::from_str::<serde_json::Value>(line) {
6242                // Yahoo Answers has question_title, question_content, best_answer
6243                let question = obj
6244                    .get("question_title")
6245                    .and_then(|v| v.as_str())
6246                    .unwrap_or_default();
6247                let label_idx = obj.get("topic").and_then(|v| v.as_i64()).unwrap_or(0) as usize;
6248
6249                let label = label_map.get(label_idx).unwrap_or(&"Unknown");
6250
6251                let tokens = vec![AnnotatedToken {
6252                    text: question.to_string(),
6253                    ner_tag: format!("B-{}", label),
6254                }];
6255
6256                sentences.push(AnnotatedSentence {
6257                    tokens,
6258                    source_dataset: id,
6259                });
6260            }
6261        }
6262
6263        if sentences.is_empty() {
6264            return Err(Error::InvalidInput(format!(
6265                "Yahoo Answers file for {:?} contains no valid sentences",
6266                id
6267            )));
6268        }
6269
6270        Ok(LoadedDataset {
6271            id,
6272            sentences,
6273            loaded_at: now,
6274            source_url: id.download_url().to_string(),
6275            data_source: DataSource::LocalCache,
6276            temporal_metadata: Self::get_temporal_metadata(id),
6277            metadata: id.default_metadata(),
6278        })
6279    }
6280
6281    /// Parse MAVEN event detection format.
6282    ///
6283    /// MAVEN provides event triggers with 168 event types.
6284    /// Supports both:
6285    /// - Full MAVEN JSONL format (train.jsonl/valid.jsonl with events array)
6286    /// - Fallback: docid2topic.json mapping file
6287    ///
6288    /// Full format structure:
6289    /// ```json
6290    /// {
6291    ///   "id": "doc_id",
6292    ///   "content": [{"sentence": "...", "tokens": [...]}],
6293    ///   "events": [{
6294    ///     "type": "EventType",
6295    ///     "mention": [{"trigger_word": "...", "sent_id": 0, "offset": [3, 4]}]
6296    ///   }]
6297    /// }
6298    /// ```
6299    ///
6300    /// Source: <https://github.com/THU-KEG/MAVEN-dataset>
6301    fn parse_maven(&self, content: &str, id: DatasetId) -> Result<LoadedDataset> {
6302        let mut sentences = Vec::new();
6303        let now = chrono::Utc::now().to_rfc3339();
6304
6305        // Try parsing as JSONL (full MAVEN format)
6306        let mut is_jsonl = false;
6307        for line in content.lines() {
6308            let line = line.trim();
6309            if line.is_empty() {
6310                continue;
6311            }
6312
6313            if let Ok(doc) = serde_json::from_str::<serde_json::Value>(line) {
6314                // Check if this is full MAVEN format with events
6315                if let Some(events) = doc.get("events").and_then(|e| e.as_array()) {
6316                    is_jsonl = true;
6317
6318                    // Get document content for context
6319                    let doc_content: Vec<String> = doc
6320                        .get("content")
6321                        .and_then(|c| c.as_array())
6322                        .map(|sents| {
6323                            sents
6324                                .iter()
6325                                .filter_map(|s| s.get("sentence").and_then(|v| v.as_str()))
6326                                .map(|s| s.to_string())
6327                                .collect()
6328                        })
6329                        .unwrap_or_default();
6330
6331                    // Process each event
6332                    for event in events {
6333                        let event_type = event
6334                            .get("type")
6335                            .and_then(|t| t.as_str())
6336                            .unwrap_or("EVENT");
6337
6338                        // Process each mention of this event
6339                        if let Some(mentions) = event.get("mention").and_then(|m| m.as_array()) {
6340                            for mention in mentions {
6341                                let trigger_word = mention
6342                                    .get("trigger_word")
6343                                    .and_then(|t| t.as_str())
6344                                    .unwrap_or("");
6345
6346                                let sent_id =
6347                                    mention.get("sent_id").and_then(|s| s.as_u64()).unwrap_or(0)
6348                                        as usize;
6349
6350                                // Get sentence context if available
6351                                let context = doc_content
6352                                    .get(sent_id)
6353                                    .cloned()
6354                                    .unwrap_or_else(|| trigger_word.to_string());
6355
6356                                let tokens = vec![AnnotatedToken {
6357                                    text: context,
6358                                    ner_tag: format!("B-{}", event_type),
6359                                }];
6360
6361                                sentences.push(AnnotatedSentence {
6362                                    tokens,
6363                                    source_dataset: id,
6364                                });
6365                            }
6366                        }
6367                    }
6368                }
6369            }
6370        }
6371
6372        // Fallback: parse as docid2topic.json mapping
6373        if !is_jsonl {
6374            if let Ok(obj) = serde_json::from_str::<serde_json::Value>(content) {
6375                if let Some(map) = obj.as_object() {
6376                    for (doc_id, event_type) in map {
6377                        let event_type_str = event_type.as_str().unwrap_or("event");
6378
6379                        let tokens = vec![AnnotatedToken {
6380                            text: doc_id.clone(),
6381                            ner_tag: format!(
6382                                "B-EVENT_{}",
6383                                event_type_str.to_uppercase().replace(' ', "_")
6384                            ),
6385                        }];
6386
6387                        sentences.push(AnnotatedSentence {
6388                            tokens,
6389                            source_dataset: id,
6390                        });
6391                    }
6392                }
6393            }
6394        }
6395
6396        Ok(LoadedDataset {
6397            id,
6398            sentences,
6399            loaded_at: now,
6400            source_url: id.download_url().to_string(),
6401            data_source: DataSource::LocalCache,
6402            temporal_metadata: Self::get_temporal_metadata(id),
6403            metadata: id.default_metadata(),
6404        })
6405    }
6406
6407    /// Parse CASIE cybersecurity event extraction format.
6408    ///
6409    /// CASIE has 5 event types: Attack-Pattern, Vulnerability, Data-Breach, Malware, Patch
6410    ///
6411    /// Format structure:
6412    /// ```json
6413    /// {
6414    ///   "content": "document text...",
6415    ///   "cyberevent": {
6416    ///     "hopper": [{
6417    ///       "events": [{
6418    ///         "subtype": "Databreach",
6419    ///         "nugget": {"text": "trigger", "startOffset": 0, "endOffset": 10},
6420    ///         "argument": [{"text": "arg", "role": {"type": "RoleType"}}]
6421    ///       }]
6422    ///     }]
6423    ///   }
6424    /// }
6425    /// ```
6426    ///
6427    /// Source: <https://github.com/Ebiquity/CASIE>
6428    fn parse_casie(&self, content: &str, id: DatasetId) -> Result<LoadedDataset> {
6429        let mut sentences = Vec::new();
6430        let now = chrono::Utc::now().to_rfc3339();
6431
6432        // Handle JSONL format (multiple documents)
6433        for line in content.lines() {
6434            let line = line.trim();
6435            if line.is_empty() {
6436                continue;
6437            }
6438
6439            if let Ok(doc) = serde_json::from_str::<serde_json::Value>(line) {
6440                let content_text = doc
6441                    .get("content")
6442                    .and_then(|v| v.as_str())
6443                    .unwrap_or_default();
6444
6445                if content_text.is_empty() {
6446                    continue;
6447                }
6448
6449                // Extract events from cyberevent.hopper[].events[]
6450                let mut found_events = false;
6451                if let Some(hopper) = doc
6452                    .get("cyberevent")
6453                    .and_then(|ce| ce.get("hopper"))
6454                    .and_then(|h| h.as_array())
6455                {
6456                    for cluster in hopper {
6457                        if let Some(events) = cluster.get("events").and_then(|e| e.as_array()) {
6458                            for event in events {
6459                                found_events = true;
6460
6461                                // Get event subtype
6462                                let subtype = event
6463                                    .get("subtype")
6464                                    .and_then(|s| s.as_str())
6465                                    .unwrap_or("Event");
6466
6467                                // Get trigger (nugget)
6468                                let trigger_text = event
6469                                    .get("nugget")
6470                                    .and_then(|n| n.get("text"))
6471                                    .and_then(|t| t.as_str())
6472                                    .unwrap_or("");
6473
6474                                // Create entry with trigger and event type
6475                                let tokens = vec![AnnotatedToken {
6476                                    text: trigger_text.to_string(),
6477                                    ner_tag: format!("B-{}", subtype),
6478                                }];
6479
6480                                sentences.push(AnnotatedSentence {
6481                                    tokens,
6482                                    source_dataset: id,
6483                                });
6484
6485                                // Also extract arguments if present
6486                                if let Some(args) = event.get("argument").and_then(|a| a.as_array())
6487                                {
6488                                    for arg in args {
6489                                        let arg_text =
6490                                            arg.get("text").and_then(|t| t.as_str()).unwrap_or("");
6491                                        let role = arg
6492                                            .get("role")
6493                                            .and_then(|r| r.get("type"))
6494                                            .and_then(|t| t.as_str())
6495                                            .unwrap_or("Argument");
6496
6497                                        if !arg_text.is_empty() {
6498                                            let tokens = vec![AnnotatedToken {
6499                                                text: arg_text.to_string(),
6500                                                ner_tag: format!("B-ARG_{}", role),
6501                                            }];
6502
6503                                            sentences.push(AnnotatedSentence {
6504                                                tokens,
6505                                                source_dataset: id,
6506                                            });
6507                                        }
6508                                    }
6509                                }
6510                            }
6511                        }
6512                    }
6513                }
6514
6515                // If no events found, add document with O tag
6516                if !found_events {
6517                    let tokens = vec![AnnotatedToken {
6518                        text: content_text.chars().take(200).collect(),
6519                        ner_tag: "O".to_string(),
6520                    }];
6521
6522                    sentences.push(AnnotatedSentence {
6523                        tokens,
6524                        source_dataset: id,
6525                    });
6526                }
6527            }
6528        }
6529
6530        if sentences.is_empty() {
6531            return Err(Error::InvalidInput(format!(
6532                "MAVEN file for {:?} contains no valid sentences",
6533                id
6534            )));
6535        }
6536
6537        Ok(LoadedDataset {
6538            id,
6539            sentences,
6540            loaded_at: now,
6541            source_url: id.download_url().to_string(),
6542            data_source: DataSource::LocalCache,
6543            temporal_metadata: Self::get_temporal_metadata(id),
6544            metadata: id.default_metadata(),
6545        })
6546    }
6547
6548    /// Parse MAVEN-ARG event argument extraction format.
6549    ///
6550    /// MAVEN-ARG extends MAVEN with 612 argument roles across 162 event types.
6551    ///
6552    /// Format structure:
6553    /// ```json
6554    /// {
6555    ///   "id": "doc_id",
6556    ///   "document": "full document text...",
6557    ///   "events": [{
6558    ///     "type": "EventType",
6559    ///     "mention": [{"trigger_word": "...", "offset": [start, end]}],
6560    ///     "argument": {"Role": [{"content": "arg text", "offset": [s, e]}]}
6561    ///   }]
6562    /// }
6563    /// ```
6564    ///
6565    /// Source: <https://github.com/THU-KEG/MAVEN-Argument>
6566    fn parse_maven_arg(&self, content: &str, id: DatasetId) -> Result<LoadedDataset> {
6567        let mut sentences = Vec::new();
6568        let now = chrono::Utc::now().to_rfc3339();
6569
6570        for line in content.lines() {
6571            let line = line.trim();
6572            if line.is_empty() {
6573                continue;
6574            }
6575
6576            if let Ok(doc) = serde_json::from_str::<serde_json::Value>(line) {
6577                // Get document text
6578                let _doc_text = doc.get("document").and_then(|d| d.as_str()).unwrap_or("");
6579
6580                // Process events
6581                if let Some(events) = doc.get("events").and_then(|e| e.as_array()) {
6582                    for event in events {
6583                        let event_type = event
6584                            .get("type")
6585                            .and_then(|t| t.as_str())
6586                            .unwrap_or("EVENT");
6587
6588                        // Process event mentions (triggers)
6589                        if let Some(mentions) = event.get("mention").and_then(|m| m.as_array()) {
6590                            for mention in mentions {
6591                                let trigger = mention
6592                                    .get("trigger_word")
6593                                    .and_then(|t| t.as_str())
6594                                    .unwrap_or("");
6595
6596                                if !trigger.is_empty() {
6597                                    let tokens = vec![AnnotatedToken {
6598                                        text: trigger.to_string(),
6599                                        ner_tag: format!("B-{}", event_type),
6600                                    }];
6601
6602                                    sentences.push(AnnotatedSentence {
6603                                        tokens,
6604                                        source_dataset: id,
6605                                    });
6606                                }
6607                            }
6608                        }
6609
6610                        // Process arguments
6611                        if let Some(args) = event.get("argument").and_then(|a| a.as_object()) {
6612                            for (role, arg_list) in args {
6613                                if let Some(arg_arr) = arg_list.as_array() {
6614                                    for arg in arg_arr {
6615                                        // Arguments can be text or entity references
6616                                        if let Some(content) =
6617                                            arg.get("content").and_then(|c| c.as_str())
6618                                        {
6619                                            if !content.is_empty() {
6620                                                let tokens = vec![AnnotatedToken {
6621                                                    text: content.to_string(),
6622                                                    ner_tag: format!("B-ARG_{}", role),
6623                                                }];
6624
6625                                                sentences.push(AnnotatedSentence {
6626                                                    tokens,
6627                                                    source_dataset: id,
6628                                                });
6629                                            }
6630                                        }
6631                                    }
6632                                }
6633                            }
6634                        }
6635                    }
6636                }
6637            }
6638        }
6639
6640        if sentences.is_empty() {
6641            return Err(Error::InvalidInput(format!(
6642                "MAVEN-ARG file for {:?} contains no valid sentences",
6643                id
6644            )));
6645        }
6646
6647        Ok(LoadedDataset {
6648            id,
6649            sentences,
6650            loaded_at: now,
6651            source_url: id.download_url().to_string(),
6652            data_source: DataSource::LocalCache,
6653            temporal_metadata: Self::get_temporal_metadata(id),
6654            metadata: id.default_metadata(),
6655        })
6656    }
6657
6658    /// Parse RAMS event argument extraction format.
6659    ///
6660    /// RAMS (Roles Across Multiple Sentences) covers 139 event types
6661    /// with arguments that can span multiple sentences.
6662    ///
6663    /// Format structure:
6664    /// ```json
6665    /// {
6666    ///   "doc_key": "...",
6667    ///   "sentences": [["token1", "token2", ...]],
6668    ///   "evt_triggers": [[start, end, [["event.type", score]]]],
6669    ///   "gold_evt_links": [[[evt_idx], [arg_start, arg_end], "role"]]
6670    /// }
6671    /// ```
6672    ///
6673    /// Source: <https://nlp.jhu.edu/rams/>
6674    fn parse_rams(&self, content: &str, id: DatasetId) -> Result<LoadedDataset> {
6675        let mut sentences = Vec::new();
6676        let now = chrono::Utc::now().to_rfc3339();
6677
6678        for line in content.lines() {
6679            let line = line.trim();
6680            if line.is_empty() {
6681                continue;
6682            }
6683
6684            if let Ok(doc) = serde_json::from_str::<serde_json::Value>(line) {
6685                // Get all tokens (flattened sentences)
6686                let all_tokens: Vec<String> = doc
6687                    .get("sentences")
6688                    .and_then(|s| s.as_array())
6689                    .map(|sents| {
6690                        sents
6691                            .iter()
6692                            .filter_map(|s| s.as_array())
6693                            .flat_map(|toks| {
6694                                toks.iter().filter_map(|t| t.as_str().map(String::from))
6695                            })
6696                            .collect()
6697                    })
6698                    .unwrap_or_default();
6699
6700                // Process event triggers
6701                if let Some(triggers) = doc.get("evt_triggers").and_then(|t| t.as_array()) {
6702                    for trigger in triggers {
6703                        if let Some(trigger_arr) = trigger.as_array() {
6704                            if trigger_arr.len() >= 3 {
6705                                let start = trigger_arr[0].as_u64().unwrap_or(0) as usize;
6706                                let end = trigger_arr[1].as_u64().unwrap_or(0) as usize;
6707
6708                                // Get event type from nested array
6709                                let event_type = trigger_arr[2]
6710                                    .as_array()
6711                                    .and_then(|types| types.first())
6712                                    .and_then(|t| t.as_array())
6713                                    .and_then(|t| t.first())
6714                                    .and_then(|t| t.as_str())
6715                                    .unwrap_or("event");
6716
6717                                // Extract trigger text
6718                                if end <= all_tokens.len() {
6719                                    let trigger_text = all_tokens[start..=end.min(start)].join(" ");
6720                                    let tokens = vec![AnnotatedToken {
6721                                        text: trigger_text,
6722                                        ner_tag: format!("B-{}", event_type),
6723                                    }];
6724
6725                                    sentences.push(AnnotatedSentence {
6726                                        tokens,
6727                                        source_dataset: id,
6728                                    });
6729                                }
6730                            }
6731                        }
6732                    }
6733                }
6734
6735                // Process argument links
6736                if let Some(links) = doc.get("gold_evt_links").and_then(|l| l.as_array()) {
6737                    for link in links {
6738                        if let Some(link_arr) = link.as_array() {
6739                            if link_arr.len() >= 3 {
6740                                // Argument span
6741                                if let Some(span) = link_arr[1].as_array() {
6742                                    if span.len() >= 2 {
6743                                        let start = span[0].as_u64().unwrap_or(0) as usize;
6744                                        let end = span[1].as_u64().unwrap_or(0) as usize;
6745                                        let role = link_arr[2].as_str().unwrap_or("argument");
6746
6747                                        if end < all_tokens.len() {
6748                                            let arg_text =
6749                                                all_tokens[start..=end.min(start)].join(" ");
6750                                            let tokens = vec![AnnotatedToken {
6751                                                text: arg_text,
6752                                                ner_tag: format!("B-ARG_{}", role),
6753                                            }];
6754
6755                                            sentences.push(AnnotatedSentence {
6756                                                tokens,
6757                                                source_dataset: id,
6758                                            });
6759                                        }
6760                                    }
6761                                }
6762                            }
6763                        }
6764                    }
6765                }
6766            }
6767        }
6768
6769        if sentences.is_empty() {
6770            return Err(Error::InvalidInput(format!(
6771                "RAMS file for {:?} contains no valid sentences",
6772                id
6773            )));
6774        }
6775
6776        Ok(LoadedDataset {
6777            id,
6778            sentences,
6779            loaded_at: now,
6780            source_url: id.download_url().to_string(),
6781            data_source: DataSource::LocalCache,
6782            temporal_metadata: Self::get_temporal_metadata(id),
6783            metadata: id.default_metadata(),
6784        })
6785    }
6786
6787    /// Parse TweetTopic format (JSONL with text and labels).
6788    ///
6789    /// Format: {"text": "...", "label": 0, "label_name": "sports_&_gaming", ...}
6790    ///
6791    /// Source: <https://huggingface.co/datasets/cardiffnlp/tweet_topic_single>
6792    fn parse_tweettopic(&self, content: &str, id: DatasetId) -> Result<LoadedDataset> {
6793        let mut sentences = Vec::new();
6794        let now = chrono::Utc::now().to_rfc3339();
6795
6796        // Label mapping if label_name not present
6797        let label_map = [
6798            "arts_&_culture",
6799            "business_&_entrepreneurs",
6800            "pop_culture",
6801            "daily_life",
6802            "sports_&_gaming",
6803            "science_&_technology",
6804        ];
6805
6806        for line in content.lines() {
6807            let line = line.trim();
6808            if line.is_empty() {
6809                continue;
6810            }
6811
6812            if let Ok(obj) = serde_json::from_str::<serde_json::Value>(line) {
6813                let text = obj.get("text").and_then(|v| v.as_str()).unwrap_or_default();
6814
6815                // Try label_name first, then map from label index
6816                let label = obj
6817                    .get("label_name")
6818                    .and_then(|v| v.as_str())
6819                    .map(|s| s.to_string())
6820                    .or_else(|| {
6821                        obj.get("label")
6822                            .and_then(|v| v.as_i64())
6823                            .and_then(|idx| label_map.get(idx as usize))
6824                            .map(|s| s.to_string())
6825                    })
6826                    .unwrap_or_else(|| "topic".to_string());
6827
6828                let tokens = vec![AnnotatedToken {
6829                    text: text.to_string(),
6830                    ner_tag: format!("B-{}", label),
6831                }];
6832
6833                sentences.push(AnnotatedSentence {
6834                    tokens,
6835                    source_dataset: id,
6836                });
6837            }
6838        }
6839
6840        if sentences.is_empty() {
6841            return Err(Error::InvalidInput(format!(
6842                "TweetTopic file for {:?} contains no valid sentences",
6843                id
6844            )));
6845        }
6846
6847        Ok(LoadedDataset {
6848            id,
6849            sentences,
6850            loaded_at: now,
6851            source_url: id.download_url().to_string(),
6852            data_source: DataSource::LocalCache,
6853            temporal_metadata: Self::get_temporal_metadata(id),
6854            metadata: id.default_metadata(),
6855        })
6856    }
6857
6858    /// Parse CoNLL-U format (Universal Dependencies).
6859    ///
6860    /// CoNLL-U format has 10 tab-separated columns per line:
6861    /// ID FORM LEMMA UPOS XPOS FEATS HEAD DEPREL DEPS MISC
6862    ///
6863    /// Used by MasakhaPOS for African language POS tagging.
6864    ///
6865    /// Source: <https://universaldependencies.org/format.html>
6866    fn parse_conllu(&self, content: &str, id: DatasetId) -> Result<LoadedDataset> {
6867        let mut sentences = Vec::new();
6868        let now = chrono::Utc::now().to_rfc3339();
6869        let mut current_tokens = Vec::new();
6870
6871        for line in content.lines() {
6872            let line = line.trim();
6873
6874            // Skip comments
6875            if line.starts_with('#') {
6876                continue;
6877            }
6878
6879            // Blank line marks end of sentence
6880            if line.is_empty() {
6881                if !current_tokens.is_empty() {
6882                    sentences.push(AnnotatedSentence {
6883                        tokens: std::mem::take(&mut current_tokens),
6884                        source_dataset: id,
6885                    });
6886                }
6887                continue;
6888            }
6889
6890            // Parse CoNLL-U columns
6891            let fields: Vec<&str> = line.split('\t').collect();
6892            if fields.len() >= 4 {
6893                // Skip multi-word tokens (IDs with ranges like "1-2")
6894                let id_field = fields[0];
6895                if id_field.contains('-') || id_field.contains('.') {
6896                    continue;
6897                }
6898
6899                let form = fields[1]; // Word form
6900                let upos = fields[3]; // Universal POS tag
6901
6902                current_tokens.push(AnnotatedToken {
6903                    text: form.to_string(),
6904                    ner_tag: format!("B-{}", upos), // Use POS tag as entity type
6905                });
6906            }
6907        }
6908
6909        // Don't forget last sentence
6910        if !current_tokens.is_empty() {
6911            sentences.push(AnnotatedSentence {
6912                tokens: current_tokens,
6913                source_dataset: id,
6914            });
6915        }
6916
6917        if sentences.is_empty() {
6918            return Err(Error::InvalidInput(format!(
6919                "CoNLL-U file for {:?} contains no valid sentences",
6920                id
6921            )));
6922        }
6923
6924        Ok(LoadedDataset {
6925            id,
6926            sentences,
6927            loaded_at: now,
6928            source_url: id.download_url().to_string(),
6929            data_source: DataSource::LocalCache,
6930            temporal_metadata: Self::get_temporal_metadata(id),
6931            metadata: id.default_metadata(),
6932        })
6933    }
6934
6935    /// Load all cached datasets.
6936    pub fn load_all_cached(&self) -> Vec<(DatasetId, Result<LoadedDataset>)> {
6937        LoadableDatasetId::all()
6938            .into_iter()
6939            .filter(|id| self.is_cached(*id))
6940            .map(|id| (id.0, self.load(id)))
6941            .collect()
6942    }
6943
6944    /// Get status of all datasets.
6945    #[must_use]
6946    pub fn status(&self) -> Vec<(DatasetId, bool)> {
6947        LoadableDatasetId::all()
6948            .into_iter()
6949            .map(|id| (id.0, self.is_cached(id)))
6950            .collect()
6951    }
6952}
6953
6954impl Default for DatasetLoader {
6955    fn default() -> Self {
6956        Self::new().expect("Failed to create default DatasetLoader")
6957    }
6958}
6959
6960// =============================================================================
6961// Helper Functions
6962// =============================================================================
6963
6964/// Parse BIO tag into prefix and type.
6965#[cfg(test)]
6966fn parse_bio_tag(tag: &str) -> (&str, &str) {
6967    if tag == "O" {
6968        return ("O", "");
6969    }
6970
6971    // Handle B-PER, I-LOC, etc.
6972    if let Some(pos) = tag.find('-') {
6973        (&tag[..pos], &tag[pos + 1..])
6974    } else {
6975        // No prefix, treat as entity type with implicit B
6976        ("B", tag)
6977    }
6978}
6979
6980/// Map dataset-specific entity types to our EntityType enum.
6981///
6982/// **Prefer `crate::schema::map_to_canonical()` for new code** - it handles
6983/// NORP correctly (as GROUP, not ORG) and preserves GPE/FAC distinctions.
6984///
6985/// # Known Issues (preserved for backwards compatibility)
6986///
6987/// - NORP → Organization (WRONG: should be Group)
6988/// - GPE/FAC/LOC all → Location (loses semantic distinctions)
6989///
6990/// See `src/schema.rs` for the corrected mappings.
6991#[cfg(test)]
6992fn map_entity_type(original: &str) -> EntityType {
6993    // Use the new canonical mapper for consistent semantics
6994    anno::schema::map_to_canonical(original, None)
6995}
6996
6997// =============================================================================
6998// Tests
6999// =============================================================================
7000
7001#[cfg(test)]
7002mod tests {
7003    use super::*;
7004
7005    #[test]
7006    fn test_dataset_id_basics() {
7007        let id = DatasetId::WikiGold;
7008        assert_eq!(id.name(), "WikiGold");
7009    }
7010
7011    #[test]
7012    fn test_convenience_metadata_methods() {
7013        let id = DatasetId::WikiGold;
7014
7015        // WikiGold has known metadata
7016        assert_eq!(id.citation(), Some("Balasuriya et al. (2009)"));
7017        assert_eq!(id.license(), Some("CC-BY-4.0"));
7018        assert_eq!(id.year(), Some(2009));
7019    }
7020
7021    #[test]
7022    fn test_loadable_wrapper_invariants() {
7023        // There should be at least one non-loadable dataset in the catalog.
7024        assert!(
7025            DatasetId::all()
7026                .iter()
7027                .copied()
7028                .any(|d| !LoadableDatasetId::is_loadable_dataset(d)),
7029            "Expected registry to contain some non-loadable datasets"
7030        );
7031
7032        for id in LoadableDatasetId::all() {
7033            let ds: DatasetId = id.into();
7034            assert!(
7035                LoadableDatasetId::is_loadable_dataset(ds),
7036                "LoadableDatasetId must imply is_loadable_dataset()"
7037            );
7038            assert!(LoadableDatasetId::try_from(ds).is_ok());
7039        }
7040    }
7041
7042    #[test]
7043    fn test_parse_plan_is_single_source_of_truth_for_loadability() {
7044        // For *every* registry dataset id:
7045        // - parse_plan(id).is_some()  <=>  TryFrom<DatasetId> succeeds
7046        for &ds in DatasetId::all() {
7047            let plan_exists = LoadableDatasetId::parse_plan(ds).is_some();
7048            let try_ok = LoadableDatasetId::try_from(ds).is_ok();
7049            assert_eq!(
7050                plan_exists, try_ok,
7051                "parse_plan / TryFrom mismatch for {:?}",
7052                ds
7053            );
7054        }
7055
7056        // Also ensure `LoadableDatasetId::all()` only returns ids with plans.
7057        for id in LoadableDatasetId::all() {
7058            let ds: DatasetId = id.into();
7059            assert!(
7060                LoadableDatasetId::parse_plan(ds).is_some(),
7061                "LoadableDatasetId::all() returned {:?} with no parse plan",
7062                ds
7063            );
7064        }
7065    }
7066
7067    #[test]
7068    fn test_registry_hints_do_not_contradict_parse_plan() {
7069        // If a dataset is loadable (i.e., has a parse plan), and the registry provides a strong
7070        // hint, they should agree. (Hints are allowed to be optimistic for non-loadable datasets.)
7071        for &ds in DatasetId::all() {
7072            let Some(plan) = LoadableDatasetId::parse_plan(ds) else {
7073                continue;
7074            };
7075            let Some(hint) = LoadableDatasetId::registry_hint_plan(ds) else {
7076                continue;
7077            };
7078            assert_eq!(hint, plan, "Registry hint mismatch for {:?}", ds);
7079        }
7080    }
7081
7082    #[test]
7083    fn test_huggingface_access_status_requires_hf_id() {
7084        // `access_status: HuggingFace` is our strongest signal that a dataset is automatable
7085        // via the Hub. Keep the registry self-consistent so hinting can stay metadata-driven.
7086        for &ds in DatasetId::all() {
7087            if ds.access_status()
7088                != crate::eval::dataset_registry::DatasetAccessibility::HuggingFace
7089            {
7090                continue;
7091            }
7092            assert!(
7093                ds.hf_id().is_some(),
7094                "Dataset {:?} is marked HuggingFace-accessible but has no hf_id",
7095                ds
7096            );
7097        }
7098    }
7099
7100    #[test]
7101    fn test_huggingface_access_status_is_hintable() {
7102        // If the registry says “HuggingFace”, the loader should be able to produce *some*
7103        // parse-plan hint (not necessarily `HfApiResponse`, since a few datasets are hybrids
7104        // with bespoke parse plans).
7105        for &ds in DatasetId::all() {
7106            if ds.access_status()
7107                != crate::eval::dataset_registry::DatasetAccessibility::HuggingFace
7108            {
7109                continue;
7110            }
7111            assert!(
7112                LoadableDatasetId::registry_hint_plan(ds).is_some(),
7113                "Dataset {:?} is marked HuggingFace-accessible but has no registry hint plan",
7114                ds
7115            );
7116        }
7117    }
7118
7119    #[test]
7120    fn test_parse_bio_tag() {
7121        assert_eq!(parse_bio_tag("O"), ("O", ""));
7122        assert_eq!(parse_bio_tag("B-PER"), ("B", "PER"));
7123        assert_eq!(parse_bio_tag("I-LOC"), ("I", "LOC"));
7124        assert_eq!(parse_bio_tag("B-ORG"), ("B", "ORG"));
7125    }
7126
7127    #[test]
7128    fn test_map_entity_type() {
7129        // Core types
7130        assert_eq!(map_entity_type("PER"), EntityType::Person);
7131        assert_eq!(map_entity_type("PERSON"), EntityType::Person);
7132        assert_eq!(map_entity_type("LOC"), EntityType::Location);
7133        assert_eq!(map_entity_type("ORG"), EntityType::Organization);
7134
7135        // GPE now preserves distinction (Custom, not Location)
7136        assert!(matches!(map_entity_type("GPE"), EntityType::Custom { .. }));
7137
7138        // MISC -> Custom or Other
7139        assert!(matches!(map_entity_type("MISC"), EntityType::Custom { .. }));
7140
7141        // OntoNotes types -> Custom (preserves semantics)
7142        assert!(matches!(
7143            map_entity_type("PRODUCT"),
7144            EntityType::Custom { .. }
7145        ));
7146        assert!(matches!(
7147            map_entity_type("EVENT"),
7148            EntityType::Custom { .. }
7149        ));
7150        assert!(matches!(
7151            map_entity_type("WORK_OF_ART"),
7152            EntityType::Custom { .. }
7153        ));
7154
7155        // Numeric types preserved
7156        assert_eq!(map_entity_type("CARDINAL"), EntityType::Cardinal);
7157    }
7158
7159    #[test]
7160    fn test_dataset_id_display() {
7161        assert_eq!(DatasetId::WikiGold.to_string(), "WikiGold");
7162        assert_eq!(DatasetId::Wnut17.to_string(), "WNUT-17");
7163    }
7164
7165    #[test]
7166    fn test_dataset_id_from_str() {
7167        assert_eq!(
7168            "wikigold".parse::<DatasetId>().unwrap(),
7169            DatasetId::WikiGold
7170        );
7171        assert_eq!("wnut-17".parse::<DatasetId>().unwrap(), DatasetId::Wnut17);
7172        assert_eq!(
7173            "mit_movie".parse::<DatasetId>().unwrap(),
7174            DatasetId::MitMovie
7175        );
7176    }
7177
7178    #[test]
7179    fn test_annotated_sentence_text() {
7180        let sentence = AnnotatedSentence {
7181            tokens: vec![
7182                AnnotatedToken {
7183                    text: "John".into(),
7184                    ner_tag: "B-PER".into(),
7185                },
7186                AnnotatedToken {
7187                    text: "lives".into(),
7188                    ner_tag: "O".into(),
7189                },
7190                AnnotatedToken {
7191                    text: "in".into(),
7192                    ner_tag: "O".into(),
7193                },
7194                AnnotatedToken {
7195                    text: "New".into(),
7196                    ner_tag: "B-LOC".into(),
7197                },
7198                AnnotatedToken {
7199                    text: "York".into(),
7200                    ner_tag: "I-LOC".into(),
7201                },
7202            ],
7203            source_dataset: DatasetId::WikiGold,
7204        };
7205
7206        assert_eq!(sentence.text(), "John lives in New York");
7207    }
7208
7209    #[test]
7210    fn test_annotated_sentence_entities() {
7211        let sentence = AnnotatedSentence {
7212            tokens: vec![
7213                AnnotatedToken {
7214                    text: "John".into(),
7215                    ner_tag: "B-PER".into(),
7216                },
7217                AnnotatedToken {
7218                    text: "Smith".into(),
7219                    ner_tag: "I-PER".into(),
7220                },
7221                AnnotatedToken {
7222                    text: "works".into(),
7223                    ner_tag: "O".into(),
7224                },
7225                AnnotatedToken {
7226                    text: "at".into(),
7227                    ner_tag: "O".into(),
7228                },
7229                AnnotatedToken {
7230                    text: "Google".into(),
7231                    ner_tag: "B-ORG".into(),
7232                },
7233            ],
7234            source_dataset: DatasetId::WikiGold,
7235        };
7236
7237        let entities = sentence.entities();
7238        assert_eq!(entities.len(), 2);
7239        assert_eq!(entities[0].text, "John Smith");
7240        assert_eq!(entities[0].entity_type, EntityType::Person);
7241        assert_eq!(entities[1].text, "Google");
7242        assert_eq!(entities[1].entity_type, EntityType::Organization);
7243    }
7244
7245    #[test]
7246    fn test_parse_conll_format() {
7247        let content = r#"
7248John B-PER
7249Smith I-PER
7250works O
7251at O
7252Google B-ORG
7253. O
7254
7255Apple B-ORG
7256announced O
7257today O
7258. O
7259"#;
7260
7261        let loader = DatasetLoader::new().unwrap();
7262        let dataset = loader.parse_conll(content, DatasetId::WikiGold).unwrap();
7263
7264        assert_eq!(dataset.len(), 2);
7265        assert_eq!(dataset.entity_count(), 3);
7266    }
7267
7268    #[test]
7269    fn test_parse_conll2003_format() {
7270        // CoNLL-2003 has 4 columns: word POS chunk NER
7271        let content = r#"
7272-DOCSTART- -X- -X- O
7273
7274EU NNP B-NP B-ORG
7275rejects VBZ B-VP O
7276German JJ B-NP B-MISC
7277call NN I-NP O
7278. . O O
7279
7280Peter NNP B-NP B-PER
7281Blackburn NNP I-NP I-PER
7282"#;
7283
7284        let loader = DatasetLoader::new().unwrap();
7285        let dataset = loader
7286            .parse_conll(content, DatasetId::CoNLL2003Sample)
7287            .unwrap();
7288
7289        assert_eq!(dataset.len(), 2);
7290
7291        let entities1 = dataset.sentences[0].entities();
7292        assert_eq!(entities1.len(), 2); // EU (ORG), German (MISC)
7293
7294        let entities2 = dataset.sentences[1].entities();
7295        assert_eq!(entities2.len(), 1); // Peter Blackburn (PER)
7296        assert_eq!(entities2[0].text, "Peter Blackburn");
7297    }
7298
7299    #[test]
7300    fn test_historical_datasets_configured() {
7301        // Historical NER datasets should have proper metadata
7302        assert!(!DatasetId::HIPE2022.download_url().is_empty());
7303        assert_eq!(DatasetId::HIPE2022.name(), "HIPE-2022");
7304
7305        assert!(!DatasetId::MedievalCzechCharters.download_url().is_empty());
7306        assert_eq!(
7307            DatasetId::MedievalCzechCharters.name(),
7308            "Medieval Czech Charters"
7309        );
7310
7311        assert!(!DatasetId::TRIDIS.download_url().is_empty());
7312        assert_eq!(DatasetId::TRIDIS.name(), "TRIDIS");
7313
7314        // Should be in all() list
7315        let all = DatasetId::all();
7316        assert!(all.contains(&DatasetId::HIPE2022));
7317        assert!(all.contains(&DatasetId::TRIDIS));
7318    }
7319
7320    #[test]
7321    fn test_queer_nlp_datasets_configured() {
7322        // Queer/gender-inclusive NLP datasets
7323        assert!(!DatasetId::WinoQueer.download_url().is_empty());
7324        assert_eq!(DatasetId::WinoQueer.name(), "WinoQueer");
7325
7326        assert!(!DatasetId::GICoref.download_url().is_empty());
7327        assert_eq!(DatasetId::GICoref.name(), "GICoref");
7328
7329        assert!(!DatasetId::BBQ.download_url().is_empty());
7330        assert_eq!(DatasetId::BBQ.name(), "BBQ");
7331
7332        // Should be in all() list
7333        let all = DatasetId::all();
7334        assert!(all.contains(&DatasetId::WinoQueer));
7335        assert!(all.contains(&DatasetId::GICoref));
7336        assert!(all.contains(&DatasetId::BBQ));
7337    }
7338
7339    #[test]
7340    fn test_joint_re_datasets_configured() {
7341        // Joint NER + Relation Extraction datasets
7342        assert!(
7343            DatasetId::TACRED.requires_license(),
7344            "TACRED is LDC-licensed; download_url may be empty"
7345        );
7346        assert_eq!(DatasetId::TACRED.name(), "TACRED");
7347
7348        assert!(!DatasetId::REBEL.download_url().is_empty());
7349        assert_eq!(DatasetId::REBEL.name(), "REBEL");
7350
7351        // Should be in all() list
7352        let all = DatasetId::all();
7353        assert!(all.contains(&DatasetId::TACRED));
7354        assert!(all.contains(&DatasetId::REBEL));
7355    }
7356
7357    #[test]
7358    fn test_dialogue_coref_datasets_configured() {
7359        // Dialogue/streaming coreference datasets
7360        assert!(!DatasetId::CODICRAC.download_url().is_empty());
7361        assert_eq!(DatasetId::CODICRAC.name(), "CODI-CRAC");
7362
7363        assert!(!DatasetId::AMIMeeting.download_url().is_empty());
7364        assert_eq!(DatasetId::AMIMeeting.name(), "AMI Meeting");
7365
7366        assert!(
7367            DatasetId::ARRAU.requires_license(),
7368            "ARRAU has LDC + research distribution; download_url may be empty"
7369        );
7370        assert!(
7371            DatasetId::ARRAU.name().contains("ARRAU"),
7372            "ARRAU name should contain 'ARRAU'"
7373        );
7374
7375        // Should be in all() list
7376        let all = DatasetId::all();
7377        assert!(all.contains(&DatasetId::CODICRAC));
7378        assert!(all.contains(&DatasetId::AMIMeeting));
7379        assert!(all.contains(&DatasetId::ARRAU));
7380    }
7381
7382    #[test]
7383    fn test_is_historical_classification() {
7384        // Historical datasets
7385        assert!(DatasetId::HIPE2022.is_historical());
7386        assert!(DatasetId::MedievalCzechCharters.is_historical());
7387        assert!(DatasetId::EighteenthCenturyNER.is_historical());
7388        assert!(DatasetId::HistoricalChineseNER.is_historical());
7389
7390        // Non-historical should return false
7391        assert!(!DatasetId::WikiGold.is_historical());
7392        assert!(!DatasetId::CoNLL2003Sample.is_historical());
7393    }
7394
7395    #[test]
7396    fn test_is_bias_evaluation_classification() {
7397        // Bias/fairness datasets
7398        assert!(DatasetId::WinoQueer.is_bias_evaluation());
7399        assert!(DatasetId::BBQ.is_bias_evaluation());
7400        assert!(DatasetId::GICoref.is_bias_evaluation());
7401        assert!(DatasetId::WinoBias.is_bias_evaluation());
7402        assert!(DatasetId::GAP.is_bias_evaluation());
7403
7404        // Non-bias should return false
7405        assert!(!DatasetId::WikiGold.is_bias_evaluation());
7406    }
7407
7408    #[test]
7409    fn test_new_datasets_have_descriptions() {
7410        // All new datasets should have proper descriptions (not the catch-all)
7411        let catch_all = "Dataset not yet fully integrated";
7412
7413        // Historical
7414        assert_ne!(DatasetId::HIPE2022.description(), catch_all);
7415        assert_ne!(DatasetId::TRIDIS.description(), catch_all);
7416
7417        // Queer NLP
7418        assert_ne!(DatasetId::WinoQueer.description(), catch_all);
7419        assert_ne!(DatasetId::BBQ.description(), catch_all);
7420        assert_ne!(DatasetId::GICoref.description(), catch_all);
7421
7422        // Joint NER+RE
7423        assert_ne!(DatasetId::TACRED.description(), catch_all);
7424        assert_ne!(DatasetId::REBEL.description(), catch_all);
7425
7426        // Dialogue
7427        assert_ne!(DatasetId::CODICRAC.description(), catch_all);
7428        assert_ne!(DatasetId::ARRAU.description(), catch_all);
7429    }
7430
7431    #[test]
7432    fn test_coreference_includes_new_datasets() {
7433        // All coreference datasets should be detected
7434        assert!(DatasetId::GICoref.is_coreference());
7435        assert!(DatasetId::CODICRAC.is_coreference());
7436        assert!(DatasetId::ARRAU.is_coreference());
7437        assert!(DatasetId::WinoPron.is_coreference());
7438        assert!(DatasetId::DROC.is_coreference());
7439        assert!(DatasetId::KoCoNovel.is_coreference());
7440    }
7441
7442    #[test]
7443    fn test_chisiec_is_historical_and_relation_extraction() {
7444        // CHisIEC should be both historical and relation extraction
7445        assert!(DatasetId::CHisIEC.is_historical());
7446        assert!(DatasetId::CHisIEC.is_relation_extraction());
7447
7448        // Verify entity types
7449        let types = DatasetId::CHisIEC.entity_types();
7450        assert!(types.contains(&"PER"));
7451        assert!(types.contains(&"LOC"));
7452        assert!(types.contains(&"OFI"));
7453        assert!(types.contains(&"BOOK"));
7454    }
7455
7456    #[test]
7457    fn test_chisiec_from_str() {
7458        // Test various string representations
7459        assert_eq!("chisiec".parse::<DatasetId>().unwrap(), DatasetId::CHisIEC);
7460        assert_eq!(
7461            "ch-is-iec".parse::<DatasetId>().unwrap(),
7462            DatasetId::CHisIEC
7463        );
7464        assert_eq!(
7465            "chinese-historical-ie".parse::<DatasetId>().unwrap(),
7466            DatasetId::CHisIEC
7467        );
7468        assert_eq!(
7469            "ancient-chinese-ner".parse::<DatasetId>().unwrap(),
7470            DatasetId::CHisIEC
7471        );
7472    }
7473
7474    #[test]
7475    fn test_chisiec_parse_ner() {
7476        // Test CHisIEC NER parsing with sample data
7477        let sample_json = r#"[
7478            {
7479                "tokens": "衞鞅奔魏",
7480                "entities": [
7481                    {"type": "PER", "start": 0, "end": 2, "span": "衞鞅"},
7482                    {"type": "LOC", "start": 3, "end": 4, "span": "魏"}
7483                ],
7484                "relations": []
7485            }
7486        ]"#;
7487
7488        let loader = DatasetLoader::new().unwrap();
7489        let result = loader.parse_chisiec(sample_json, DatasetId::CHisIEC);
7490        assert!(result.is_ok());
7491
7492        let dataset = result.unwrap();
7493        assert_eq!(dataset.sentences.len(), 1);
7494
7495        let sentence = &dataset.sentences[0];
7496        // 4 characters: 衞 鞅 奔 魏
7497        assert_eq!(sentence.tokens.len(), 4);
7498
7499        // Check BIO tags
7500        assert_eq!(sentence.tokens[0].ner_tag, "B-PER"); // 衞
7501        assert_eq!(sentence.tokens[1].ner_tag, "I-PER"); // 鞅
7502        assert_eq!(sentence.tokens[2].ner_tag, "O"); // 奔
7503        assert_eq!(sentence.tokens[3].ner_tag, "B-LOC"); // 魏
7504    }
7505
7506    #[test]
7507    fn test_chisiec_parse_relations() {
7508        // Test CHisIEC relation extraction parsing
7509        let sample_json = r#"[
7510            {
7511                "tokens": "嚴公遣玉汝使",
7512                "entities": [
7513                    {"type": "PER", "start": 0, "end": 2, "span": "嚴公"},
7514                    {"type": "PER", "start": 2, "end": 4, "span": "玉汝"}
7515                ],
7516                "relations": [
7517                    {"type": "上下級", "head": 0, "tail": 1, "head_span": "嚴公", "tail_span": "玉汝"}
7518                ]
7519            }
7520        ]"#;
7521
7522        let loader = DatasetLoader::new().unwrap();
7523        let result = loader.parse_chisiec_relations(sample_json);
7524        assert!(result.is_ok());
7525
7526        let docs = result.unwrap();
7527        assert_eq!(docs.len(), 1);
7528
7529        let doc = &docs[0];
7530        assert_eq!(doc.relations.len(), 1);
7531
7532        let rel = &doc.relations[0];
7533        assert_eq!(rel.relation_type, "上下級");
7534        assert_eq!(rel.head_text, "嚴公");
7535        assert_eq!(rel.tail_text, "玉汝");
7536        assert_eq!(rel.head_type, "PER");
7537        assert_eq!(rel.tail_type, "PER");
7538        // Character offsets (嚴公 at 0-2, 玉汝 at 2-4)
7539        assert_eq!(rel.head_span, (0, 2));
7540        assert_eq!(rel.tail_span, (2, 4));
7541    }
7542
7543    #[test]
7544    fn test_chisiec_all_entity_types() {
7545        // Test all 4 CHisIEC entity types: PER, LOC, OFI (官职), BOOK (书籍)
7546        // This is important because OFI (Official position) is domain-specific
7547        let sample_json = r#"[
7548            {
7549                "tokens": "司馬遷為太史令著史記於長安",
7550                "entities": [
7551                    {"type": "PER", "start": 0, "end": 3, "span": "司馬遷"},
7552                    {"type": "OFI", "start": 4, "end": 7, "span": "太史令"},
7553                    {"type": "BOOK", "start": 8, "end": 10, "span": "史記"},
7554                    {"type": "LOC", "start": 11, "end": 13, "span": "長安"}
7555                ],
7556                "relations": []
7557            }
7558        ]"#;
7559
7560        let loader = DatasetLoader::new().unwrap();
7561        let dataset = loader
7562            .parse_chisiec(sample_json, DatasetId::CHisIEC)
7563            .unwrap();
7564
7565        assert_eq!(dataset.sentences.len(), 1);
7566        let sentence = &dataset.sentences[0];
7567
7568        // Verify each entity type is correctly tagged
7569        // 司馬遷 (Sima Qian - historian)
7570        assert_eq!(sentence.tokens[0].ner_tag, "B-PER");
7571        assert_eq!(sentence.tokens[1].ner_tag, "I-PER");
7572        assert_eq!(sentence.tokens[2].ner_tag, "I-PER");
7573
7574        // 為 (was)
7575        assert_eq!(sentence.tokens[3].ner_tag, "O");
7576
7577        // 太史令 (Grand Historian - official position)
7578        assert_eq!(sentence.tokens[4].ner_tag, "B-OFI");
7579        assert_eq!(sentence.tokens[5].ner_tag, "I-OFI");
7580        assert_eq!(sentence.tokens[6].ner_tag, "I-OFI");
7581
7582        // 著 (wrote)
7583        assert_eq!(sentence.tokens[7].ner_tag, "O");
7584
7585        // 史記 (Records of the Grand Historian - book)
7586        assert_eq!(sentence.tokens[8].ner_tag, "B-BOOK");
7587        assert_eq!(sentence.tokens[9].ner_tag, "I-BOOK");
7588
7589        // 於 (at)
7590        assert_eq!(sentence.tokens[10].ner_tag, "O");
7591
7592        // 長安 (Chang'an - capital city)
7593        assert_eq!(sentence.tokens[11].ner_tag, "B-LOC");
7594        assert_eq!(sentence.tokens[12].ner_tag, "I-LOC");
7595    }
7596
7597    #[test]
7598    fn test_chisiec_unicode_character_offsets() {
7599        // Critical: CHisIEC uses CHARACTER offsets, not byte offsets
7600        // This test ensures we handle multi-byte Chinese characters correctly
7601        let sample_json = r#"[
7602            {
7603                "tokens": "曹操",
7604                "entities": [
7605                    {"type": "PER", "start": 0, "end": 2, "span": "曹操"}
7606                ],
7607                "relations": []
7608            }
7609        ]"#;
7610
7611        let loader = DatasetLoader::new().unwrap();
7612        let dataset = loader
7613            .parse_chisiec(sample_json, DatasetId::CHisIEC)
7614            .unwrap();
7615
7616        // 曹操 is 2 characters (but 6 bytes in UTF-8)
7617        let sentence = &dataset.sentences[0];
7618        assert_eq!(sentence.tokens.len(), 2);
7619        assert_eq!(sentence.tokens[0].text, "曹");
7620        assert_eq!(sentence.tokens[1].text, "操");
7621    }
7622
7623    #[test]
7624    fn test_chisiec_multiple_relations_same_document() {
7625        // Test parsing multiple relations in a single document
7626        // Relations: 任職 (holds office), 管理 (manages)
7627        let sample_json = r#"[
7628            {
7629                "tokens": "曹操為丞相管冀州",
7630                "entities": [
7631                    {"type": "PER", "start": 0, "end": 2, "span": "曹操"},
7632                    {"type": "OFI", "start": 3, "end": 5, "span": "丞相"},
7633                    {"type": "LOC", "start": 6, "end": 8, "span": "冀州"}
7634                ],
7635                "relations": [
7636                    {"type": "任職", "head": 0, "tail": 1, "head_span": "曹操", "tail_span": "丞相"},
7637                    {"type": "管理", "head": 0, "tail": 2, "head_span": "曹操", "tail_span": "冀州"}
7638                ]
7639            }
7640        ]"#;
7641
7642        let loader = DatasetLoader::new().unwrap();
7643        let docs = loader.parse_chisiec_relations(sample_json).unwrap();
7644
7645        assert_eq!(docs.len(), 1);
7646        assert_eq!(docs[0].relations.len(), 2);
7647
7648        // First relation: 任職 (office holding)
7649        assert_eq!(docs[0].relations[0].relation_type, "任職");
7650        assert_eq!(docs[0].relations[0].head_type, "PER");
7651        assert_eq!(docs[0].relations[0].tail_type, "OFI");
7652
7653        // Second relation: 管理 (manages)
7654        assert_eq!(docs[0].relations[1].relation_type, "管理");
7655        assert_eq!(docs[0].relations[1].head_type, "PER");
7656        assert_eq!(docs[0].relations[1].tail_type, "LOC");
7657    }
7658
7659    #[test]
7660    fn test_chisiec_distinct_from_historical_chinese_ner() {
7661        // CHisIEC and HistoricalChineseNER are DIFFERENT datasets
7662        // This test documents their distinction
7663
7664        // Both should be classified as historical
7665        assert!(DatasetId::CHisIEC.is_historical());
7666        assert!(DatasetId::HistoricalChineseNER.is_historical());
7667
7668        // But they are different datasets
7669        assert_ne!(DatasetId::CHisIEC, DatasetId::HistoricalChineseNER);
7670
7671        // CHisIEC supports relation extraction; HistoricalChineseNER may not
7672        assert!(DatasetId::CHisIEC.is_relation_extraction());
7673
7674        // Different entity types:
7675        // CHisIEC: PER, LOC, OFI, BOOK (ancient Chinese)
7676        // HistoricalChineseNER: PER, LOC, ORG, DATE, etc. (modern Chinese 1872-1949)
7677        let chisiec_types = DatasetId::CHisIEC.entity_types();
7678        assert!(chisiec_types.contains(&"OFI")); // Official position - unique to CHisIEC
7679        assert!(chisiec_types.contains(&"BOOK")); // Classical texts
7680
7681        // Different names
7682        assert_eq!(DatasetId::CHisIEC.name(), "CHisIEC");
7683        assert_eq!(
7684            DatasetId::HistoricalChineseNER.name(),
7685            "Historical Chinese NER"
7686        );
7687    }
7688
7689    #[test]
7690    fn test_chisiec_empty_entities_handled() {
7691        // Test graceful handling of documents with no entities
7692        let sample_json = r#"[
7693            {
7694                "tokens": "天下太平",
7695                "entities": [],
7696                "relations": []
7697            }
7698        ]"#;
7699
7700        let loader = DatasetLoader::new().unwrap();
7701        let dataset = loader
7702            .parse_chisiec(sample_json, DatasetId::CHisIEC)
7703            .unwrap();
7704
7705        assert_eq!(dataset.sentences.len(), 1);
7706        let sentence = &dataset.sentences[0];
7707
7708        // All tokens should be tagged as O
7709        for token in &sentence.tokens {
7710            assert_eq!(token.ner_tag, "O");
7711        }
7712    }
7713
7714    #[test]
7715    fn test_chisiec_entity_types_in_schema() {
7716        // Verify CHisIEC entity types are properly mapped in the schema
7717        use anno::schema::map_to_canonical;
7718
7719        // PER -> Person
7720        let per_type = map_to_canonical("PER", None);
7721        assert_eq!(per_type, EntityType::Person);
7722
7723        // LOC -> Location
7724        let loc_type = map_to_canonical("LOC", None);
7725        assert_eq!(loc_type, EntityType::Location);
7726
7727        // OFI -> Custom OFFICIAL type (domain-specific for ancient Chinese)
7728        let ofi_type = map_to_canonical("OFI", None);
7729        assert!(matches!(ofi_type, EntityType::Custom { .. }));
7730
7731        // BOOK -> WORK_OF_ART (creative works)
7732        let book_type = map_to_canonical("BOOK", None);
7733        assert!(matches!(book_type, EntityType::Custom { .. }));
7734    }
7735
7736    #[test]
7737    fn test_chisiec_language_and_domain() {
7738        // CHisIEC is Classical Chinese (文言文) - ISO 639-3: lzh
7739        assert_eq!(DatasetId::CHisIEC.language(), "lzh");
7740
7741        // CHisIEC is a historical dataset (24 dynastic histories)
7742        assert_eq!(DatasetId::CHisIEC.domain(), "historical");
7743
7744        // Compare with HistoricalChineseNER which is modern Chinese (1872-1949)
7745        assert_eq!(DatasetId::HistoricalChineseNER.language(), "zh");
7746        assert_eq!(DatasetId::HistoricalChineseNER.domain(), "historical");
7747    }
7748
7749    // =========================================================================
7750    // African Language Dataset Tests
7751    // =========================================================================
7752
7753    #[test]
7754    fn test_african_datasets_configured() {
7755        // MasakhaNER datasets should have download URLs
7756        assert!(!DatasetId::MasakhaNER.download_url().is_empty());
7757        assert!(!DatasetId::MasakhaNER2.download_url().is_empty());
7758        assert!(!DatasetId::AfriSenti.download_url().is_empty());
7759        assert!(!DatasetId::AfriQA.download_url().is_empty());
7760        assert!(!DatasetId::MasakhaNEWS.download_url().is_empty());
7761        assert!(!DatasetId::MasakhaPOS.download_url().is_empty());
7762
7763        // Names should be set
7764        assert_eq!(DatasetId::MasakhaNER.name(), "MasakhaNER");
7765        assert_eq!(DatasetId::MasakhaNER2.name(), "MasakhaNER 2.0");
7766        assert_eq!(DatasetId::AfriSenti.name(), "AfriSenti");
7767        assert_eq!(DatasetId::AfriQA.name(), "AfriQA");
7768        assert_eq!(DatasetId::MasakhaNEWS.name(), "MasakhaNEWS");
7769        assert_eq!(DatasetId::MasakhaPOS.name(), "MasakhaPOS");
7770    }
7771
7772    #[test]
7773    fn test_african_datasets_entity_types() {
7774        // MasakhaNER uses PER, ORG, LOC, DATE
7775        let ner_types = DatasetId::MasakhaNER.entity_types();
7776        assert!(ner_types.contains(&"PER"));
7777        assert!(ner_types.contains(&"LOC"));
7778        assert!(ner_types.contains(&"ORG"));
7779        assert!(ner_types.contains(&"DATE"));
7780
7781        // AfriSenti uses sentiment labels
7782        let senti_types = DatasetId::AfriSenti.entity_types();
7783        assert!(senti_types.contains(&"positive"));
7784        assert!(senti_types.contains(&"neutral"));
7785        assert!(senti_types.contains(&"negative"));
7786
7787        // MasakhaNEWS uses topic labels
7788        let news_types = DatasetId::MasakhaNEWS.entity_types();
7789        assert!(news_types.contains(&"politics"));
7790        assert!(news_types.contains(&"sports"));
7791        assert!(news_types.contains(&"business"));
7792
7793        // MasakhaPOS uses Universal Dependencies POS tags
7794        let pos_types = DatasetId::MasakhaPOS.entity_types();
7795        assert!(pos_types.contains(&"NOUN"));
7796        assert!(pos_types.contains(&"VERB"));
7797        assert!(pos_types.contains(&"ADJ"));
7798    }
7799
7800    #[test]
7801    fn test_parse_afrisenti() {
7802        // Test AfriSenti TSV parsing
7803        let sample_tsv = "This movie is great!\tpositive\n\
7804                          Awful experience\tnegative\n\
7805                          It was okay\tneutral";
7806
7807        let loader = DatasetLoader::new().unwrap();
7808        let result = loader.parse_afrisenti(sample_tsv, DatasetId::AfriSenti);
7809        assert!(result.is_ok());
7810
7811        let dataset = result.unwrap();
7812        assert_eq!(dataset.sentences.len(), 3);
7813
7814        // Check first sentence has positive label
7815        assert_eq!(dataset.sentences[0].tokens[0].ner_tag, "B-positive");
7816        // Check second sentence has negative label
7817        assert_eq!(dataset.sentences[1].tokens[0].ner_tag, "B-negative");
7818        // Check third sentence has neutral label
7819        assert_eq!(dataset.sentences[2].tokens[0].ner_tag, "B-neutral");
7820    }
7821
7822    #[test]
7823    fn test_parse_masakhanews() {
7824        // Test MasakhaNEWS TSV parsing
7825        let sample_tsv = "headline\tbody\tcategory\n\
7826                          Breaking: Election Results\tThe results are in...\tpolitics\n\
7827                          Team Wins Championship\tIn an exciting match...\tsports";
7828
7829        let loader = DatasetLoader::new().unwrap();
7830        let result = loader.parse_masakhanews(sample_tsv, DatasetId::MasakhaNEWS);
7831        assert!(result.is_ok());
7832
7833        let dataset = result.unwrap();
7834        // Header line should be skipped
7835        assert_eq!(dataset.sentences.len(), 2);
7836
7837        // Check categories
7838        assert_eq!(dataset.sentences[0].tokens[0].ner_tag, "B-politics");
7839        assert_eq!(dataset.sentences[1].tokens[0].ner_tag, "B-sports");
7840    }
7841
7842    #[test]
7843    fn test_parse_conllu() {
7844        // Test CoNLL-U parsing (MasakhaPOS format)
7845        let sample_conllu = "# sent_id = 1\n\
7846                             # text = John loves Mary\n\
7847                             1\tJohn\tJohn\tPROPN\tNNP\t_\t2\tnsubj\t_\t_\n\
7848                             2\tloves\tlove\tVERB\tVBZ\t_\t0\troot\t_\t_\n\
7849                             3\tMary\tMary\tPROPN\tNNP\t_\t2\tobj\t_\t_\n\
7850                             \n\
7851                             # sent_id = 2\n\
7852                             1\tHe\the\tPRON\tPRP\t_\t2\tnsubj\t_\t_\n\
7853                             2\truns\trun\tVERB\tVBZ\t_\t0\troot\t_\t_\n";
7854
7855        let loader = DatasetLoader::new().unwrap();
7856        let result = loader.parse_conllu(sample_conllu, DatasetId::MasakhaPOS);
7857        assert!(result.is_ok());
7858
7859        let dataset = result.unwrap();
7860        assert_eq!(dataset.sentences.len(), 2);
7861
7862        // First sentence: John loves Mary
7863        assert_eq!(dataset.sentences[0].tokens.len(), 3);
7864        assert_eq!(dataset.sentences[0].tokens[0].text, "John");
7865        assert_eq!(dataset.sentences[0].tokens[0].ner_tag, "B-PROPN");
7866        assert_eq!(dataset.sentences[0].tokens[1].text, "loves");
7867        assert_eq!(dataset.sentences[0].tokens[1].ner_tag, "B-VERB");
7868
7869        // Second sentence: He runs
7870        assert_eq!(dataset.sentences[1].tokens.len(), 2);
7871        assert_eq!(dataset.sentences[1].tokens[0].text, "He");
7872        assert_eq!(dataset.sentences[1].tokens[0].ner_tag, "B-PRON");
7873    }
7874
7875    #[test]
7876    fn test_ancient_language_ud_datasets_are_loadable() {
7877        // Ancient language UD treebanks should be loadable via registry hints
7878        // These have format: "CoNLLU" and categories: [ner, ancient]
7879        let ancient_datasets = [
7880            DatasetId::AncientGreekUD,
7881            DatasetId::LatinUD,
7882            DatasetId::SanskritUD,
7883            DatasetId::OldEnglishUD,
7884            DatasetId::OldNorseUD,
7885        ];
7886
7887        for ds in ancient_datasets {
7888            assert!(
7889                LoadableDatasetId::is_loadable_dataset(ds),
7890                "{:?} should be loadable via registry hint (format={:?})",
7891                ds,
7892                ds.format()
7893            );
7894        }
7895    }
7896
7897    #[test]
7898    fn test_conllu_with_ner_tags_from_ancient_greek() {
7899        // Test CoNLLU parsing with MISC column NER tags (Ancient Greek Perseus format)
7900        // Real format from UD Ancient Greek Perseus
7901        let sample_conllu = "\
7902# sent_id = tlg0012.tlg001.perseus-grc1:1.1
7903# text = μῆνιν ἄειδε θεὰ Πηληϊάδεω Ἀχιλῆος
79041\tμῆνιν\tμῆνις\tNOUN\tn-s---fa-\tCase=Acc|Gender=Fem|Number=Sing\t2\tobj\t_\tO
79052\tἄειδε\tᾄδω\tVERB\tv2sama---\tMood=Imp|Number=Sing|Person=2|Tense=Pres|VerbForm=Fin|Voice=Act\t0\troot\t_\tO
79063\tθεὰ\tθεά\tNOUN\tn-s---fv-\tCase=Voc|Gender=Fem|Number=Sing\t2\tvocative\t_\tO
79074\tΠηληϊάδεω\tΠηληϊάδης\tNOUN\tn-s---mg-\tCase=Gen|Gender=Masc|Number=Sing\t5\tnmod\t_\tB-PER
79085\tἈχιλῆος\tἈχιλλεύς\tPROPN\tn-s---mg-\tCase=Gen|Gender=Masc|Number=Sing\t1\tnmod\t_\tI-PER
7909
7910";
7911
7912        let loader = DatasetLoader::new().unwrap();
7913        let result = loader.parse_conllu(sample_conllu, DatasetId::AncientGreekUD);
7914        assert!(
7915            result.is_ok(),
7916            "Failed to parse Ancient Greek CoNLLU: {:?}",
7917            result.err()
7918        );
7919
7920        let dataset = result.unwrap();
7921        assert_eq!(dataset.sentences.len(), 1);
7922        assert_eq!(dataset.sentences[0].tokens.len(), 5);
7923
7924        // Check Achilles (Ἀχιλῆος) entity
7925        assert_eq!(dataset.sentences[0].tokens[4].text, "Ἀχιλῆος");
7926        // Note: CoNLLU parser may use POS tags if MISC doesn't have NER
7927        // This depends on how the parser handles the MISC column
7928    }
7929
7930    #[test]
7931    fn test_registry_hints_cover_all_conllu_ner_datasets() {
7932        // Datasets with format CoNLLU/CoNLL-U and task NER should get a registry hint
7933        for &ds in DatasetId::all() {
7934            let format = ds.format().unwrap_or("");
7935            let is_conllu = format == "CoNLLU" || format == "CoNLL-U";
7936            let is_ner = ds.is_ner();
7937
7938            if is_conllu && is_ner {
7939                let hint = LoadableDatasetId::registry_hint_plan(ds);
7940                assert!(
7941                    hint.is_some(),
7942                    "{:?} has format={} and is NER but no registry hint",
7943                    ds,
7944                    format
7945                );
7946                if let Some(plan) = hint {
7947                    assert_eq!(
7948                        plan,
7949                        DatasetParsePlan::Conllu,
7950                        "{:?} should use Conllu parse plan",
7951                        ds
7952                    );
7953                }
7954            }
7955        }
7956    }
7957
7958    #[test]
7959    fn test_datasets_with_public_url_and_format_are_hintable() {
7960        // Datasets with a public URL and a parseable format should get hints
7961        let hintable_formats = ["CoNLL", "CoNLLU", "CoNLL-U", "BIO", "IOB2", "JSONL"];
7962
7963        let mut missing_hints = Vec::new();
7964
7965        for &ds in DatasetId::all() {
7966            let url = ds.download_url();
7967            let format = ds.format().unwrap_or("");
7968            let is_ner = ds.is_ner();
7969
7970            // Skip datasets without URLs or non-NER datasets
7971            if url.is_empty() || !is_ner {
7972                continue;
7973            }
7974
7975            // Skip formats we don't auto-detect
7976            if !hintable_formats.contains(&format) {
7977                continue;
7978            }
7979
7980            let hint = LoadableDatasetId::registry_hint_plan(ds);
7981            if hint.is_none() {
7982                missing_hints.push((ds, format));
7983            }
7984        }
7985
7986        // Allow some datasets to not have hints (complex formats, etc.)
7987        // but document them
7988        if !missing_hints.is_empty() {
7989            // These are known to be missing hints (need special parsers)
7990            let known_missing: &[DatasetId] = &[
7991                // Add any datasets that intentionally don't have hints here
7992            ];
7993            for (ds, format) in &missing_hints {
7994                if !known_missing.contains(ds) {
7995                    eprintln!(
7996                        "Warning: {:?} (format={}) has public URL but no registry hint",
7997                        ds, format
7998                    );
7999                }
8000            }
8001        }
8002    }
8003
8004    #[test]
8005    fn test_loadable_count_is_reasonable() {
8006        // Ensure we have a reasonable number of loadable datasets
8007        let loadable_count = LoadableDatasetId::all().len();
8008        let total_count = DatasetId::all().len();
8009
8010        // We should have at least 50% of datasets loadable via either parse_plan or hints
8011        let min_expected = total_count / 2;
8012        assert!(
8013            loadable_count >= min_expected,
8014            "Only {} of {} datasets are loadable (expected at least {})",
8015            loadable_count,
8016            total_count,
8017            min_expected
8018        );
8019    }
8020
8021    #[test]
8022    fn test_datasets_with_urls_have_formats() {
8023        // Datasets with public URLs should ideally have format info for auto-loading
8024        let mut missing_format = Vec::new();
8025
8026        for &ds in DatasetId::all() {
8027            let url = ds.download_url();
8028            let format = ds.format();
8029            let access = ds.access_status();
8030
8031            // Skip datasets that require registration or aren't publicly available
8032            if url.is_empty() {
8033                continue;
8034            }
8035
8036            // Check if format is missing for public datasets
8037            if format.is_none()
8038                && access == crate::eval::dataset_registry::DatasetAccessibility::Public
8039            {
8040                missing_format.push(ds);
8041            }
8042        }
8043
8044        // Log datasets that could benefit from format info
8045        if !missing_format.is_empty() {
8046            eprintln!(
8047                "Datasets with public URLs but no format field ({}):",
8048                missing_format.len()
8049            );
8050            for ds in &missing_format[..missing_format.len().min(10)] {
8051                eprintln!("  - {:?}", ds);
8052            }
8053        }
8054
8055        // We expect most public datasets to have format info
8056        // Allow up to 20% to be missing format (some have unusual formats)
8057        let max_missing = DatasetId::all().len() / 5;
8058        assert!(
8059            missing_format.len() <= max_missing,
8060            "Too many public datasets missing format: {} (max {})",
8061            missing_format.len(),
8062            max_missing
8063        );
8064    }
8065
8066    #[test]
8067    fn test_conll_format_ner_only_datasets_are_parseable() {
8068        // All NER-only datasets with CoNLL/CoNLLU format should have a parse plan
8069        // (Datasets with joint RE/coref tasks may use different column formats)
8070        let mut not_loadable = Vec::new();
8071
8072        for &ds in DatasetId::all() {
8073            let format = ds.format().unwrap_or("");
8074            let is_conll = format == "CoNLL" || format == "CoNLLU" || format == "CoNLL-U";
8075            let is_ner = ds.is_ner();
8076            let is_re = ds.is_relation_extraction();
8077            let is_coref = ds.is_coreference();
8078            let is_event = ds.is_event_coref();
8079
8080            if !is_conll || !is_ner {
8081                continue;
8082            }
8083
8084            // Skip explicitly blocked datasets (they may be present in the registry for metadata,
8085            // but intentionally cannot be downloaded/loaded automatically).
8086            if ds.tasks_or_inferred().contains(&"blocked") {
8087                continue;
8088            }
8089
8090            // Skip joint task datasets (they use CoNLL but with different structure)
8091            if is_re || is_coref || is_event {
8092                continue;
8093            }
8094
8095            // Pure NER CoNLL datasets should be loadable
8096            let is_loadable = LoadableDatasetId::is_loadable_dataset(ds);
8097            if !is_loadable {
8098                not_loadable.push((ds, format));
8099            }
8100        }
8101
8102        if !not_loadable.is_empty() {
8103            eprintln!("Pure NER CoNLL datasets not loadable:");
8104            for (ds, format) in &not_loadable {
8105                eprintln!("  - {:?} (format={})", ds, format);
8106            }
8107        }
8108
8109        // All pure NER CoNLL datasets should be loadable
8110        assert!(
8111            not_loadable.is_empty(),
8112            "{} pure NER CoNLL datasets are not loadable",
8113            not_loadable.len()
8114        );
8115    }
8116
8117    #[test]
8118    fn test_jsonl_ner_datasets_are_parseable() {
8119        // JSONL datasets with NER task should ideally be loadable
8120        let mut jsonl_ner_not_loadable = Vec::new();
8121
8122        for &ds in DatasetId::all() {
8123            let format = ds.format().unwrap_or("");
8124            let is_jsonl = format == "JSONL" || format == "JSON-Lines" || format == "jsonl";
8125            let is_ner = ds.is_ner();
8126
8127            if !is_jsonl || !is_ner {
8128                continue;
8129            }
8130
8131            if !LoadableDatasetId::is_loadable_dataset(ds) {
8132                jsonl_ner_not_loadable.push(ds);
8133            }
8134        }
8135
8136        // Log for debugging
8137        if !jsonl_ner_not_loadable.is_empty() {
8138            eprintln!(
8139                "JSONL NER datasets not loadable ({}):",
8140                jsonl_ner_not_loadable.len()
8141            );
8142            for ds in &jsonl_ner_not_loadable {
8143                eprintln!("  - {:?}", ds);
8144            }
8145        }
8146    }
8147
8148    #[test]
8149    fn test_parse_afriqa() {
8150        // Test AfriQA JSON parsing
8151        let sample_json = r#"[
8152            {
8153                "context": "Lagos is a major city in Nigeria.",
8154                "question": "What is Lagos?",
8155                "answers": {
8156                    "text": ["major city"],
8157                    "answer_start": [11]
8158                }
8159            }
8160        ]"#;
8161
8162        let loader = DatasetLoader::new().unwrap();
8163        let result = loader.parse_afriqa(sample_json, DatasetId::AfriQA);
8164        assert!(result.is_ok());
8165
8166        let dataset = result.unwrap();
8167        assert_eq!(dataset.sentences.len(), 1);
8168
8169        // Check that answer is marked
8170        let tokens = &dataset.sentences[0].tokens;
8171        // The answer "major city" should have B-ANSWER and I-ANSWER tags
8172        let answer_tokens: Vec<_> = tokens
8173            .iter()
8174            .filter(|t| t.ner_tag.contains("ANSWER"))
8175            .collect();
8176        assert!(!answer_tokens.is_empty(), "Should have answer tokens");
8177    }
8178
8179    #[test]
8180    fn test_african_datasets_in_all_list() {
8181        let all = DatasetId::all();
8182        assert!(all.contains(&DatasetId::MasakhaNER));
8183        assert!(all.contains(&DatasetId::MasakhaNER2));
8184        assert!(all.contains(&DatasetId::AfriSenti));
8185        assert!(all.contains(&DatasetId::AfriQA));
8186        assert!(all.contains(&DatasetId::MasakhaNEWS));
8187        assert!(all.contains(&DatasetId::MasakhaPOS));
8188    }
8189
8190    // =========================================================================
8191    // Event Extraction Parser Tests
8192    // =========================================================================
8193
8194    #[test]
8195    fn test_parse_maven_jsonl() {
8196        // Test MAVEN full JSONL format with events array
8197        let sample_jsonl = r#"{"id": "doc1", "content": [{"sentence": "The earthquake struck Tokyo.", "tokens": ["The", "earthquake", "struck", "Tokyo", "."]}], "events": [{"type": "Disaster", "mention": [{"trigger_word": "earthquake", "sent_id": 0, "offset": [1, 2]}]}]}"#;
8198
8199        let loader = DatasetLoader::new().unwrap();
8200        let result = loader.parse_maven(sample_jsonl, DatasetId::MAVEN);
8201        assert!(result.is_ok(), "parse_maven should succeed");
8202
8203        let dataset = result.unwrap();
8204        assert!(!dataset.sentences.is_empty(), "Should have sentences");
8205
8206        // Check event type tag
8207        let has_disaster = dataset
8208            .sentences
8209            .iter()
8210            .any(|s| s.tokens.iter().any(|t| t.ner_tag.contains("Disaster")));
8211        assert!(has_disaster, "Should have Disaster event tag");
8212    }
8213
8214    #[test]
8215    fn test_parse_maven_docid2topic_fallback() {
8216        // Test fallback format (docid2topic.json)
8217        let sample_json = r#"{"doc1": "Natural_Disaster", "doc2": "Political_Event"}"#;
8218
8219        let loader = DatasetLoader::new().unwrap();
8220        let result = loader.parse_maven(sample_json, DatasetId::MAVEN);
8221        assert!(result.is_ok(), "parse_maven fallback should succeed");
8222
8223        let dataset = result.unwrap();
8224        assert_eq!(dataset.sentences.len(), 2, "Should have 2 entries");
8225    }
8226
8227    #[test]
8228    fn test_parse_casie() {
8229        // Test CASIE cybersecurity event format
8230        let sample_jsonl = r#"{"content": "A vulnerability was discovered in Apache.", "cyberevent": {"hopper": [{"events": [{"subtype": "Vulnerability", "nugget": {"text": "vulnerability"}, "argument": [{"text": "Apache", "role": {"type": "Affected_System"}}]}]}]}}"#;
8231
8232        let loader = DatasetLoader::new().unwrap();
8233        let result = loader.parse_casie(sample_jsonl, DatasetId::CASIE);
8234        assert!(result.is_ok(), "parse_casie should succeed");
8235
8236        let dataset = result.unwrap();
8237        assert!(!dataset.sentences.is_empty(), "Should have sentences");
8238
8239        // Check for vulnerability trigger
8240        let has_vuln = dataset
8241            .sentences
8242            .iter()
8243            .any(|s| s.tokens.iter().any(|t| t.ner_tag.contains("Vulnerability")));
8244        assert!(has_vuln, "Should have Vulnerability tag");
8245
8246        // Check for argument
8247        let has_arg = dataset
8248            .sentences
8249            .iter()
8250            .any(|s| s.tokens.iter().any(|t| t.ner_tag.contains("ARG_")));
8251        assert!(has_arg, "Should have argument tag");
8252    }
8253
8254    #[test]
8255    fn test_parse_maven_arg() {
8256        // Test MAVEN-ARG format with arguments
8257        let sample_jsonl = r#"{"id": "doc1", "document": "The company announced layoffs.", "events": [{"type": "Employment", "mention": [{"trigger_word": "layoffs", "offset": [4, 5]}], "argument": {"Employer": [{"content": "company", "offset": [1, 2]}]}}]}"#;
8258
8259        let loader = DatasetLoader::new().unwrap();
8260        let result = loader.parse_maven_arg(sample_jsonl, DatasetId::MAVENArg);
8261        assert!(result.is_ok(), "parse_maven_arg should succeed");
8262
8263        let dataset = result.unwrap();
8264        assert!(!dataset.sentences.is_empty(), "Should have sentences");
8265
8266        // Check for trigger
8267        let has_trigger = dataset
8268            .sentences
8269            .iter()
8270            .any(|s| s.tokens.iter().any(|t| t.ner_tag.contains("Employment")));
8271        assert!(has_trigger, "Should have Employment event tag");
8272
8273        // Check for argument role
8274        let has_employer = dataset
8275            .sentences
8276            .iter()
8277            .any(|s| s.tokens.iter().any(|t| t.ner_tag.contains("ARG_Employer")));
8278        assert!(has_employer, "Should have Employer argument tag");
8279    }
8280
8281    #[test]
8282    fn test_parse_rams() {
8283        // Test RAMS tokenized format
8284        let sample_jsonl = r#"{"doc_key": "doc1", "sentences": [["The", "soldier", "fired", "his", "weapon", "."]], "evt_triggers": [[2, 2, [["conflict.attack", 1.0]]]], "gold_evt_links": [[[0], [1, 1], "attacker"]]}"#;
8285
8286        let loader = DatasetLoader::new().unwrap();
8287        let result = loader.parse_rams(sample_jsonl, DatasetId::RAMS);
8288        assert!(result.is_ok(), "parse_rams should succeed");
8289
8290        let dataset = result.unwrap();
8291        assert!(!dataset.sentences.is_empty(), "Should have sentences");
8292
8293        // Check for event trigger
8294        let has_event = dataset
8295            .sentences
8296            .iter()
8297            .any(|s| s.tokens.iter().any(|t| t.ner_tag.starts_with("B-")));
8298        assert!(has_event, "Should have event tags");
8299    }
8300
8301    #[test]
8302    fn test_parse_trec() {
8303        // Test TREC question classification format
8304        let sample =
8305            "NUM:dist How far is it from Denver to Aspen ?\nLOC:city What county is Modesto in ?\n";
8306
8307        let loader = DatasetLoader::new().unwrap();
8308        let result = loader.parse_trec(sample, DatasetId::TREC);
8309        assert!(result.is_ok(), "parse_trec should succeed");
8310
8311        let dataset = result.unwrap();
8312        assert_eq!(dataset.sentences.len(), 2, "Should have 2 questions");
8313
8314        // Check coarse labels
8315        assert!(dataset.sentences[0].tokens[0].ner_tag.contains("NUM"));
8316        assert!(dataset.sentences[1].tokens[0].ner_tag.contains("LOC"));
8317    }
8318
8319    #[test]
8320    fn test_parse_litbank_ner_improved() {
8321        // Test improved LitBank NER parser with word-level tokenization
8322        let sample_ann = "T1\tPER 0 5\tAlice\nT2\tPER 10 14\tBob\nT3\tORG 20 28\tMicrosoft";
8323
8324        let loader = DatasetLoader::new().unwrap();
8325        let result = loader.parse_litbank(sample_ann, DatasetId::LitBank);
8326        assert!(result.is_ok(), "parse_litbank should succeed");
8327
8328        let dataset = result.unwrap();
8329        assert!(!dataset.sentences.is_empty(), "Should have sentences");
8330
8331        // Check that entities are tokenized into words (not just single tokens)
8332        let sentence = &dataset.sentences[0];
8333        let entity_tokens: Vec<_> = sentence
8334            .tokens
8335            .iter()
8336            .filter(|t| t.ner_tag.starts_with("B-") || t.ner_tag.starts_with("I-"))
8337            .collect();
8338
8339        // Should have at least the entity tokens (Alice, Bob, Microsoft)
8340        assert!(
8341            entity_tokens.len() >= 3,
8342            "Should have at least 3 entity tokens, got {}",
8343            entity_tokens.len()
8344        );
8345
8346        // Check BIO tagging is correct (B- for first word, I- for subsequent words)
8347        let mut found_b_tag = false;
8348        for token in &sentence.tokens {
8349            if token.ner_tag.starts_with("B-") {
8350                found_b_tag = true;
8351                // First word of entity should be B-
8352                assert!(
8353                    token.ner_tag.starts_with("B-"),
8354                    "First word of entity should have B- tag"
8355                );
8356            }
8357        }
8358        assert!(found_b_tag, "Should have at least one B- tag");
8359    }
8360
8361    #[test]
8362    fn test_parse_tweettopic() {
8363        // Test TweetTopic JSONL format
8364        let sample_jsonl = r#"{"text": "Amazing game last night!", "label": 4, "label_name": "sports_&_gaming"}
8365{"text": "New AI breakthrough announced", "label": 5, "label_name": "science_&_technology"}"#;
8366
8367        let loader = DatasetLoader::new().unwrap();
8368        let result = loader.parse_tweettopic(sample_jsonl, DatasetId::TweetTopic);
8369        assert!(result.is_ok(), "parse_tweettopic should succeed");
8370
8371        let dataset = result.unwrap();
8372        assert_eq!(dataset.sentences.len(), 2, "Should have 2 tweets");
8373
8374        // Check label names are used
8375        assert!(dataset.sentences[0].tokens[0]
8376            .ner_tag
8377            .contains("sports_&_gaming"));
8378        assert!(dataset.sentences[1].tokens[0]
8379            .ner_tag
8380            .contains("science_&_technology"));
8381    }
8382
8383    #[test]
8384    fn test_african_dataset_from_str() {
8385        // Test string parsing for African datasets
8386        assert_eq!(
8387            "masakhaner".parse::<DatasetId>().unwrap(),
8388            DatasetId::MasakhaNER
8389        );
8390        assert_eq!(
8391            "masakhaner2".parse::<DatasetId>().unwrap(),
8392            DatasetId::MasakhaNER2
8393        );
8394        assert_eq!(
8395            "afrisenti".parse::<DatasetId>().unwrap(),
8396            DatasetId::AfriSenti
8397        );
8398        assert_eq!("afriqa".parse::<DatasetId>().unwrap(), DatasetId::AfriQA);
8399        assert_eq!(
8400            "masakhanews".parse::<DatasetId>().unwrap(),
8401            DatasetId::MasakhaNEWS
8402        );
8403        assert_eq!(
8404            "masakhapos".parse::<DatasetId>().unwrap(),
8405            DatasetId::MasakhaPOS
8406        );
8407
8408        // Alternative spellings
8409        assert_eq!(
8410            "masakhane-ner".parse::<DatasetId>().unwrap(),
8411            DatasetId::MasakhaNER
8412        );
8413        assert_eq!(
8414            "afri-senti".parse::<DatasetId>().unwrap(),
8415            DatasetId::AfriSenti
8416        );
8417        assert_eq!(
8418            "masakhane-news".parse::<DatasetId>().unwrap(),
8419            DatasetId::MasakhaNEWS
8420        );
8421    }
8422
8423    #[test]
8424    fn test_afrisenti_parse_with_tonal_diacritics() {
8425        // Test Yoruba text with tonal diacritics (common in AfriSenti)
8426        let yoruba_tsv = "Ó dára púpọ̀!\tpositive\n\
8427                          Kò dára rárá\tnegative\n\
8428                          Ẹ ṣé, mo dupẹ́\tpositive";
8429
8430        let loader = DatasetLoader::new().unwrap();
8431        let result = loader.parse_afrisenti(yoruba_tsv, DatasetId::AfriSenti);
8432        assert!(result.is_ok());
8433
8434        let dataset = result.unwrap();
8435        assert_eq!(dataset.sentences.len(), 3);
8436
8437        // Verify the text is preserved correctly with diacritics
8438        assert!(dataset.sentences[0].tokens[0].text.contains("dára"));
8439        assert!(dataset.sentences[1].tokens[0].text.contains("rárá"));
8440        assert!(dataset.sentences[2].tokens[0].text.contains("dupẹ́"));
8441    }
8442
8443    #[test]
8444    fn test_masakhaner_parse_with_ethiopic_script() {
8445        // Test Amharic text in MasakhaNER CoNLL format (Ethiopic script)
8446        let amharic_conll = "ዶክተር B-PER\n\
8447                             አቢይ I-PER\n\
8448                             አህመድ I-PER\n\
8449                             ኢትዮጵያ B-LOC\n\
8450                             ውስጥ O\n\
8451                             ተወለዱ O\n";
8452
8453        let loader = DatasetLoader::new().unwrap();
8454        let result = loader.parse_conll(amharic_conll, DatasetId::MasakhaNER);
8455        assert!(result.is_ok());
8456
8457        let dataset = result.unwrap();
8458        assert_eq!(dataset.sentences.len(), 1);
8459
8460        let tokens = &dataset.sentences[0].tokens;
8461        assert_eq!(tokens.len(), 6);
8462
8463        // Verify Ethiopic script is preserved
8464        assert_eq!(tokens[0].text, "ዶክተር");
8465        assert_eq!(tokens[0].ner_tag, "B-PER");
8466        assert_eq!(tokens[3].text, "ኢትዮጵያ");
8467        assert_eq!(tokens[3].ner_tag, "B-LOC");
8468    }
8469
8470    #[test]
8471    fn test_conllu_parse_with_nguni_clicks() {
8472        // Test isiXhosa/isiZulu text with click consonants (MasakhaPOS)
8473        let xhosa_conllu = "# sent_id = xho_test_1\n\
8474                           # text = UMongameli uCyril Ramaphosa\n\
8475                           1\tUMongameli\tumongameli\tNOUN\tN\t_\t0\troot\t_\t_\n\
8476                           2\tuCyril\tuCyril\tPROPN\tNNP\t_\t1\tappos\t_\t_\n\
8477                           3\tRamaphosa\tRamaphosa\tPROPN\tNNP\t_\t2\tflat:name\t_\t_\n\
8478                           \n\
8479                           # sent_id = xho_test_2\n\
8480                           # text = Ndiqala ukuthetha isiXhosa\n\
8481                           1\tNdiqala\tqala\tVERB\tV\t_\t0\troot\t_\t_\n\
8482                           2\tukuthetha\tthetha\tVERB\tV\t_\t1\txcomp\t_\t_\n\
8483                           3\tisiXhosa\tisiXhosa\tNOUN\tN\t_\t2\tobj\t_\t_\n";
8484
8485        let loader = DatasetLoader::new().unwrap();
8486        let result = loader.parse_conllu(xhosa_conllu, DatasetId::MasakhaPOS);
8487        assert!(result.is_ok());
8488
8489        let dataset = result.unwrap();
8490        assert_eq!(dataset.sentences.len(), 2);
8491
8492        // Verify text with click-like consonants is preserved
8493        assert_eq!(dataset.sentences[0].tokens[0].text, "UMongameli");
8494        assert_eq!(dataset.sentences[1].tokens[2].text, "isiXhosa");
8495    }
8496
8497    #[test]
8498    fn test_masakhanews_parse_with_arabic_variants() {
8499        // MasakhaNEWS includes Algerian/Moroccan Arabic variants
8500        let news_tsv = "headline\tbody\tcategory\n\
8501                        الأخبار العاجلة\tتفاصيل الخبر...\tpolitics\n\
8502                        رياضة محلية\tمباراة اليوم...\tsports";
8503
8504        let loader = DatasetLoader::new().unwrap();
8505        let result = loader.parse_masakhanews(news_tsv, DatasetId::MasakhaNEWS);
8506        assert!(result.is_ok());
8507
8508        let dataset = result.unwrap();
8509        // Header skipped, 2 data rows
8510        assert_eq!(dataset.sentences.len(), 2);
8511
8512        // Check Arabic text is preserved
8513        assert!(dataset.sentences[0].tokens[0].text.contains("الأخبار"));
8514        assert_eq!(dataset.sentences[0].tokens[0].ner_tag, "B-politics");
8515        assert_eq!(dataset.sentences[1].tokens[0].ner_tag, "B-sports");
8516    }
8517
8518    #[test]
8519    fn test_afriqa_multilingual_qa() {
8520        // AfriQA has questions in target language, context may be in English
8521        let qa_json = r#"[
8522            {
8523                "context": "Yorùbá is a tonal language spoken in Nigeria.",
8524                "question": "Kí ni Yorùbá?",
8525                "answers": {
8526                    "text": ["tonal language"],
8527                    "answer_start": [13]
8528                },
8529                "language": "yo"
8530            }
8531        ]"#;
8532
8533        let loader = DatasetLoader::new().unwrap();
8534        let result = loader.parse_afriqa(qa_json, DatasetId::AfriQA);
8535        assert!(result.is_ok());
8536
8537        let dataset = result.unwrap();
8538        assert_eq!(dataset.sentences.len(), 1);
8539    }
8540
8541    // NOTE: We intentionally do not embed language-code→URL expansion helpers here.
8542    // If we add them, they should live in the registry (metadata) or a dedicated dataset
8543    // URL builder module, not in the loader tests.
8544
8545    // =========================================================================
8546    // Comprehensive Parser Smoke Tests (one per DatasetParsePlan)
8547    // =========================================================================
8548
8549    #[test]
8550    fn test_parse_content_rejects_empty_for_all_loadable_datasets() {
8551        // Global invariant: no dataset should parse from empty content.
8552        let loader = DatasetLoader::new().unwrap();
8553        for loadable in LoadableDatasetId::all() {
8554            let id: DatasetId = loadable.into();
8555            let err = loader
8556                .parse_content_impl("   \n\t", id)
8557                .expect_err("empty content must error");
8558            let msg = format!("{err}");
8559            assert!(
8560                msg.to_lowercase().contains("empty"),
8561                "Expected an 'empty' error message for {:?}, got: {}",
8562                id,
8563                msg
8564            );
8565        }
8566    }
8567
8568    #[test]
8569    fn test_parse_docred_smoke() {
8570        let sample = r#"{"doc_key":"d1","sentence":["John","met","Mary","in","Paris","."],"ner":[[0,0,"PER"],[2,2,"PER"],[4,4,"LOC"]],"relations":[]}"#;
8571        let loader = DatasetLoader::new().unwrap();
8572        let ds = loader.parse_docred(sample, DatasetId::DocRED).unwrap();
8573        assert!(!ds.sentences.is_empty());
8574        assert!(ds.sentences[0]
8575            .tokens
8576            .iter()
8577            .any(|t| t.ner_tag.starts_with("B-")));
8578    }
8579
8580    #[test]
8581    fn test_parse_cadec_jsonl_smoke() {
8582        // Character offsets in `entities` are interpreted over reconstructed text from tokens.
8583        // "I took aspirin" → aspirin starts at char 7, ends at 14 (exclusive).
8584        let sample = r#"{"tokens":["I","took","aspirin"],"entities":[{"text":"aspirin","label":"DRUG","start":7,"end":14}]}"#;
8585        let loader = DatasetLoader::new().unwrap();
8586        let ds = loader.parse_cadec_jsonl(sample, DatasetId::CADEC).unwrap();
8587        assert!(!ds.sentences.is_empty());
8588        assert!(ds.sentences[0]
8589            .tokens
8590            .iter()
8591            .any(|t| t.ner_tag == "B-DRUG" || t.ner_tag == "I-DRUG"));
8592    }
8593
8594    #[test]
8595    fn test_parse_cadec_hf_api_smoke() {
8596        let sample = r#"{"rows":[{"row":{"text":"I took aspirin","ade":"aspirin"}}]}"#;
8597        let loader = DatasetLoader::new().unwrap();
8598        let ds = loader.parse_cadec_hf_api(sample, DatasetId::CADEC).unwrap();
8599        assert!(!ds.sentences.is_empty());
8600        assert!(ds.sentences[0]
8601            .tokens
8602            .iter()
8603            .any(|t| t.ner_tag.contains("adverse_drug_event")));
8604    }
8605
8606    #[test]
8607    fn test_parse_bc5cdr_smoke() {
8608        let sample = "Aspirin\tNN\tO\tB-CHEMICAL\nhelps\tVBZ\tO\tO\n\n";
8609        let loader = DatasetLoader::new().unwrap();
8610        let ds = loader.parse_bc5cdr(sample, DatasetId::BC5CDR).unwrap();
8611        assert_eq!(ds.sentences.len(), 1);
8612        assert_eq!(ds.sentences[0].tokens[0].ner_tag, "B-CHEMICAL");
8613    }
8614
8615    #[test]
8616    fn test_parse_ncbi_disease_smoke() {
8617        let sample = "Cancer\tNN\tO\tB-Disease\nprogresses\tVBZ\tO\tO\n\n";
8618        let loader = DatasetLoader::new().unwrap();
8619        let ds = loader
8620            .parse_ncbi_disease(sample, DatasetId::NCBIDisease)
8621            .unwrap();
8622        assert_eq!(ds.sentences.len(), 1);
8623        assert!(ds.sentences[0].tokens[0].ner_tag.starts_with("B-"));
8624    }
8625
8626    #[test]
8627    fn test_parse_gap_smoke() {
8628        let sample =
8629            "ID\tText\tPronoun\tPronoun-offset\tA\tA-offset\tA-coref\tB\tB-offset\tB-coref\tURL\n\
8630g1\tJohn met Mary. He waved.\tHe\t14\tJohn\t0\tTRUE\tMary\t9\tFALSE\thttp://example\n";
8631        let loader = DatasetLoader::new().unwrap();
8632        let ds = loader.parse_gap(sample, DatasetId::GAP).unwrap();
8633        assert_eq!(ds.sentences.len(), 1);
8634        assert!(!ds.sentences[0].tokens.is_empty());
8635    }
8636
8637    #[test]
8638    fn test_parse_preco_jsonl_smoke() {
8639        let sample = r#"{"sentences":[["John","went","home","."],["He","slept","."]]}"#;
8640        let loader = DatasetLoader::new().unwrap();
8641        let ds = loader.parse_preco_jsonl(sample, DatasetId::PreCo).unwrap();
8642        assert_eq!(ds.sentences.len(), 2);
8643        assert_eq!(ds.sentences[0].tokens[0].text, "John");
8644    }
8645
8646    #[test]
8647    fn test_parse_wikiann_json_array_smoke() {
8648        let sample =
8649            r#"[{"tokens":["John","went","to","Paris"],"ner_tags":["B-PER","O","O","B-LOC"]}]"#;
8650        let loader = DatasetLoader::new().unwrap();
8651        let ds = loader.parse_wikiann_json(sample, DatasetId::UNER).unwrap();
8652        assert_eq!(ds.sentences.len(), 1);
8653        assert_eq!(ds.sentences[0].tokens[0].ner_tag, "B-PER");
8654    }
8655
8656    #[test]
8657    fn test_parse_hf_api_response_smoke() {
8658        let sample = r#"{
8659  "features":[{"name":"tokens"},{"name":"ner_tags","type":{"feature":{"names":["O","B-PER","I-PER"]}}}],
8660  "rows":[{"row_idx":0,"row":{"tokens":["John"],"ner_tags":[1]}}]
8661}"#;
8662        let loader = DatasetLoader::new().unwrap();
8663        let ds = loader
8664            .parse_hf_api_response(sample, DatasetId::UniversalNER)
8665            .unwrap();
8666        assert_eq!(ds.sentences.len(), 1);
8667        assert_eq!(ds.sentences[0].tokens[0].ner_tag, "B-PER");
8668    }
8669
8670    #[test]
8671    fn test_parse_hf_api_response_temporal_standoff_smoke() {
8672        let sample = r#"{
8673  "features":[{"name":"text"},{"name":"time_expressions"}],
8674  "rows":[{"row_idx":0,"row":{
8675    "text":"A 10/30/89 .",
8676    "time_expressions":[{"text":"10/30/89","start_char":2,"end_char":10,"tid":"t1","type":"DATE","value":"1989-10-30"}],
8677    "event_expressions":[],
8678    "signal_expressions":[]
8679  }}]
8680}"#;
8681        let loader = DatasetLoader::new().unwrap();
8682        let ds = loader
8683            .parse_hf_api_response(sample, DatasetId::TimexRecognitionSentenceOriginal)
8684            .unwrap();
8685        assert_eq!(ds.sentences.len(), 1);
8686        assert_eq!(ds.sentences[0].tokens.len(), 3);
8687        assert_eq!(ds.sentences[0].tokens[0].text, "A");
8688        assert_eq!(ds.sentences[0].tokens[0].ner_tag, "O");
8689        assert_eq!(ds.sentences[0].tokens[1].text, "10/30/89");
8690        assert_eq!(ds.sentences[0].tokens[1].ner_tag, "B-TIMEX");
8691        assert_eq!(ds.sentences[0].tokens[2].text, ".");
8692        assert_eq!(ds.sentences[0].tokens[2].ner_tag, "O");
8693    }
8694
8695    #[test]
8696    fn test_parse_hf_api_response_pairwise_discourse_smoke() {
8697        let sample = r#"{
8698  "features":[{"name":"unit1_txt"},{"name":"unit2_txt"},{"name":"label"}],
8699  "rows":[{"row_idx":0,"row":{
8700    "unit1_txt":"Because it rained",
8701    "unit2_txt":"the game was canceled",
8702    "label":"Cause"
8703  }}]
8704}"#;
8705        let loader = DatasetLoader::new().unwrap();
8706        let ds = loader
8707            .parse_hf_api_response(sample, DatasetId::DisrptEngDepScidtbRels)
8708            .unwrap();
8709        assert_eq!(ds.sentences.len(), 1);
8710        assert_eq!(ds.sentences[0].tokens.len(), 1);
8711        assert_eq!(
8712            ds.sentences[0].tokens[0].text,
8713            "Because it rained [SEP] the game was canceled"
8714        );
8715        assert_eq!(ds.sentences[0].tokens[0].ner_tag, "B-Cause");
8716    }
8717
8718    #[test]
8719    fn test_parse_hf_api_response_disrpt_conllu_seg_smoke() {
8720        let sample = r#"{
8721  "features":[{"name":"form"},{"name":"misc"}],
8722  "rows":[{"row_idx":0,"row":{
8723    "form":["We","propose","a","method","."],
8724    "misc":["Seg=B-seg","Seg=O","Seg=O","Seg=B-seg","Seg=O"]
8725  }}]
8726}"#;
8727        let loader = DatasetLoader::new().unwrap();
8728        let ds = loader
8729            .parse_hf_api_response(sample, DatasetId::DisrptEngDepScidtbConlluSeg)
8730            .unwrap();
8731        assert_eq!(ds.sentences.len(), 1);
8732        let tags: Vec<&str> = ds.sentences[0]
8733            .tokens
8734            .iter()
8735            .map(|t| t.ner_tag.as_str())
8736            .collect();
8737        assert_eq!(tags, vec!["B-SEG", "I-SEG", "I-SEG", "B-SEG", "I-SEG"]);
8738    }
8739
8740    #[test]
8741    fn test_parse_agnews_smoke() {
8742        let sample = r#"{"text":"Stocks rally on earnings","label":2}"#;
8743        let loader = DatasetLoader::new().unwrap();
8744        let ds = loader.parse_agnews(sample, DatasetId::AGNews).unwrap();
8745        assert_eq!(ds.sentences.len(), 1);
8746        assert!(ds.sentences[0].tokens[0].ner_tag.starts_with("B-"));
8747    }
8748
8749    #[test]
8750    fn test_parse_dbpedia14_smoke() {
8751        let sample = r#"{"content":"The Beatles released Abbey Road","label":5}"#;
8752        let loader = DatasetLoader::new().unwrap();
8753        let ds = loader
8754            .parse_dbpedia14(sample, DatasetId::DBPedia14)
8755            .unwrap();
8756        assert_eq!(ds.sentences.len(), 1);
8757        assert!(ds.sentences[0].tokens[0].ner_tag.starts_with("B-"));
8758    }
8759
8760    #[test]
8761    fn test_parse_yahoo_answers_smoke() {
8762        let sample = r#"{"question_title":"Why is the sky blue?","topic":1}"#;
8763        let loader = DatasetLoader::new().unwrap();
8764        let ds = loader
8765            .parse_yahoo_answers(sample, DatasetId::YahooAnswers)
8766            .unwrap();
8767        assert_eq!(ds.sentences.len(), 1);
8768        assert!(ds.sentences[0].tokens[0].ner_tag.starts_with("B-"));
8769    }
8770
8771    // =========================================================================
8772    // Dataset Registry Integration Tests
8773    // =========================================================================
8774
8775    #[test]
8776    fn test_sec_filings_has_raw_url() {
8777        // Verify SEC-filings dataset has a raw GitHub URL for direct download
8778        let url = DatasetId::SECFilingsNER.download_url();
8779        assert!(
8780            url.contains("raw.githubusercontent.com"),
8781            "SEC-filings should have raw GitHub URL, got: {}",
8782            url
8783        );
8784        assert!(
8785            url.ends_with(".txt"),
8786            "SEC-filings should point to a .txt file, got: {}",
8787            url
8788        );
8789    }
8790
8791    #[test]
8792    fn test_twiconv_has_format() {
8793        // Verify TwiConv dataset has format field
8794        let format = DatasetId::TwiConv.format();
8795        assert!(format.is_some(), "TwiConv should have format field");
8796        // TwiConv uses CoNLL format for coreference data
8797        assert_eq!(format.unwrap(), "CoNLL", "TwiConv should be CoNLL format");
8798    }
8799
8800    #[test]
8801    fn test_mudoco_has_format() {
8802        // Verify MuDoCo dataset has format field
8803        let format = DatasetId::MuDoCo.format();
8804        assert!(format.is_some(), "MuDoCo should have format field");
8805        assert_eq!(format.unwrap(), "JSON", "MuDoCo should be JSON format");
8806    }
8807
8808    #[test]
8809    fn test_all_public_ud_datasets_have_conllu_format() {
8810        // All Universal Dependencies datasets should have CoNLLU format
8811        let ud_datasets = vec![
8812            DatasetId::AncientGreekUD,
8813            DatasetId::LatinUD,
8814            DatasetId::SanskritUD,
8815            DatasetId::OldEnglishUD,
8816            DatasetId::UDEsperantoCairo,
8817        ];
8818
8819        for ds in ud_datasets {
8820            let format = ds.format();
8821            assert!(format.is_some(), "{:?} should have format field", ds);
8822            assert_eq!(
8823                format.unwrap(),
8824                "CoNLLU",
8825                "{:?} should be CoNLLU format",
8826                ds
8827            );
8828        }
8829    }
8830
8831    #[test]
8832    fn test_datasets_with_public_urls_are_accessible() {
8833        // Verify that key public datasets have valid URLs (format check only)
8834        let test_cases = vec![
8835            (DatasetId::AncientGreekUD, "universaldependencies"),
8836            (DatasetId::LatinUD, "universaldependencies"),
8837            (DatasetId::SECFilingsNER, "entity-recognition-datasets"),
8838            (DatasetId::TwiConv, "twiconv"), // lowercase for case-insensitive match
8839        ];
8840
8841        for (ds, expected_substring) in test_cases {
8842            let url = ds.download_url();
8843            assert!(!url.is_empty(), "{:?} should have a download URL", ds);
8844            assert!(
8845                url.to_lowercase()
8846                    .contains(&expected_substring.to_lowercase()),
8847                "{:?} URL should contain '{}', got: {}",
8848                ds,
8849                expected_substring,
8850                url
8851            );
8852        }
8853    }
8854
8855    #[test]
8856    fn test_loadable_datasets_count_is_stable() {
8857        // Track the number of loadable datasets to detect regressions
8858        // This should only increase as we add more loaders
8859        let loadable = LoadableDatasetId::all();
8860        let count = loadable.len();
8861
8862        // As of 2025-12-15, after recent fixes we have 295 loadable datasets
8863        assert!(
8864            count >= 295,
8865            "Expected at least 295 loadable datasets, got {}. \
8866             This may indicate a regression in the loading system.",
8867            count
8868        );
8869    }
8870
8871    #[test]
8872    fn test_conll_format_variants_all_detected() {
8873        // Ensure CoNLL format variants for pure NER datasets are properly detected
8874        // We exclude RE/coref datasets as they have special parsing needs
8875        for &ds in DatasetId::all() {
8876            let format = ds.format();
8877            if let Some(fmt) = format {
8878                let is_conll_variant =
8879                    fmt == "CoNLL" || fmt == "CoNLLU" || fmt == "CoNLL-U" || fmt == "CoNLL03";
8880
8881                // Only check pure NER datasets (no coref, no RE)
8882                let is_pure_ner = ds.supports_ner() && !ds.supports_coref() && !ds.supports_re();
8883
8884                if is_conll_variant && is_pure_ner {
8885                    // Should be loadable via hint system
8886                    let hint = LoadableDatasetId::registry_hint_plan(ds);
8887                    assert!(
8888                        hint.is_some() || LoadableDatasetId::is_loadable_dataset(ds),
8889                        "{:?} with format {} and pure NER task should be loadable",
8890                        ds,
8891                        fmt
8892                    );
8893                }
8894            }
8895        }
8896    }
8897
8898    #[test]
8899    fn test_parse_csv_ner_smoke() {
8900        // Test CSV NER format (E-NER/EDGAR-NER style: Token,Tag)
8901        let sample = "\
8902-DOCSTART-,O
8903,O
8904Check,O
8905the,O
8906appropriate,O
8907box,O
8908,O
8909Nuveen,I-BUSINESS
8910New,I-BUSINESS
8911York,I-BUSINESS
8912Fund,I-BUSINESS
8913,O
8914
8915The,O
8916SEC,I-GOVERNMENT
8917filed,O
8918charges,O
8919.
8920
8921John,I-PERSON
8922Smith,I-PERSON
8923is,O
8924the,O
8925CEO,O
8926.
8927";
8928        let loader = DatasetLoader::new().unwrap();
8929        let ds = loader.parse_csv_ner(sample, DatasetId::ENer).unwrap();
8930
8931        // Should have 3 sentences (separated by empty lines and -DOCSTART-)
8932        assert_eq!(
8933            ds.sentences.len(),
8934            3,
8935            "Expected 3 sentences, got {:?}",
8936            ds.sentences.len()
8937        );
8938
8939        // First sentence should have BUSINESS entities
8940        let first_sentence = &ds.sentences[0];
8941        assert!(
8942            first_sentence
8943                .tokens
8944                .iter()
8945                .any(|t| t.ner_tag == "I-BUSINESS"),
8946            "First sentence should contain I-BUSINESS tags"
8947        );
8948
8949        // Second sentence should have GOVERNMENT entity
8950        let second_sentence = &ds.sentences[1];
8951        assert!(
8952            second_sentence
8953                .tokens
8954                .iter()
8955                .any(|t| t.ner_tag == "I-GOVERNMENT"),
8956            "Second sentence should contain I-GOVERNMENT tag"
8957        );
8958
8959        // Third sentence should have PERSON entities
8960        let third_sentence = &ds.sentences[2];
8961        assert!(
8962            third_sentence
8963                .tokens
8964                .iter()
8965                .any(|t| t.ner_tag == "I-PERSON"),
8966            "Third sentence should contain I-PERSON tags"
8967        );
8968
8969        // Check specific token/tag pairs
8970        let nuveen_token = first_sentence.tokens.iter().find(|t| t.text == "Nuveen");
8971        assert!(nuveen_token.is_some(), "Should have Nuveen token");
8972        assert_eq!(nuveen_token.unwrap().ner_tag, "I-BUSINESS");
8973
8974        let john_token = third_sentence.tokens.iter().find(|t| t.text == "John");
8975        assert!(john_token.is_some(), "Should have John token");
8976        assert_eq!(john_token.unwrap().ner_tag, "I-PERSON");
8977    }
8978
8979    #[test]
8980    fn test_csv_ner_format_is_detected() {
8981        // Ensure CSV format datasets with NER tasks are properly detected as loadable
8982        let ener_hint = LoadableDatasetId::registry_hint_plan(DatasetId::ENer);
8983        assert_eq!(
8984            ener_hint,
8985            Some(DatasetParsePlan::CsvNer),
8986            "ENer should use CsvNer parse plan"
8987        );
8988
8989        // Verify ENer is loadable
8990        assert!(
8991            LoadableDatasetId::is_loadable_dataset(DatasetId::ENer),
8992            "ENer should be loadable"
8993        );
8994    }
8995
8996    // =========================================================================
8997    // Tests for newly added dataset loaders (2025-12)
8998    // =========================================================================
8999
9000    #[test]
9001    fn test_newly_added_conll_datasets_are_loadable() {
9002        let new_conll = [
9003            DatasetId::QxoRef,
9004            DatasetId::GICoref,
9005            DatasetId::WNUT16,
9006            DatasetId::NoiseBench,
9007            DatasetId::CrossWeigh,
9008            DatasetId::ZELDA,
9009            DatasetId::GENIANested,
9010        ];
9011
9012        for id in new_conll {
9013            assert!(
9014                LoadableDatasetId::is_loadable_dataset(id),
9015                "{:?} should be loadable with Conll parse plan",
9016                id
9017            );
9018            assert_eq!(
9019                LoadableDatasetId::parse_plan(id),
9020                Some(DatasetParsePlan::Conll),
9021                "{:?} should use Conll plan",
9022                id
9023            );
9024        }
9025    }
9026
9027    #[test]
9028    fn test_newly_added_jsonl_datasets_are_loadable() {
9029        let new_jsonl = [
9030            DatasetId::REBEL,
9031            DatasetId::BBQ,
9032            DatasetId::RealToxicityPrompts,
9033            DatasetId::BookCoref,
9034            DatasetId::BookCorefSplit,
9035            DatasetId::WIESP2022NER,
9036            DatasetId::FewRel,
9037            DatasetId::PIIMasking200k,
9038            DatasetId::B2NERD,
9039            DatasetId::OpenNER,
9040            DatasetId::FictionNER750M,
9041        ];
9042
9043        for id in new_jsonl {
9044            assert!(
9045                LoadableDatasetId::is_loadable_dataset(id),
9046                "{:?} should be loadable with JsonlNer parse plan",
9047                id
9048            );
9049            assert_eq!(
9050                LoadableDatasetId::parse_plan(id),
9051                Some(DatasetParsePlan::JsonlNer),
9052                "{:?} should use JsonlNer plan",
9053                id
9054            );
9055        }
9056    }
9057
9058    #[test]
9059    fn test_genia_nested_conll_parse() {
9060        // GENIA nested NER uses multi-layered BIO tags
9061        let nested_conll = "IL-2\tB-protein\n\
9062                            gene\tI-protein\n\
9063                            expression\tO\n\
9064                            \n\
9065                            T\tB-cell_type\n\
9066                            cells\tI-cell_type\n";
9067
9068        let loader = DatasetLoader::new().unwrap();
9069        let result = loader.parse_conll(nested_conll, DatasetId::GENIANested);
9070        assert!(result.is_ok());
9071
9072        let dataset = result.unwrap();
9073        assert_eq!(dataset.sentences.len(), 2);
9074        assert_eq!(dataset.sentences[0].tokens[0].ner_tag, "B-protein");
9075        assert_eq!(dataset.sentences[1].tokens[0].ner_tag, "B-cell_type");
9076    }
9077
9078    #[test]
9079    fn test_gicoref_gender_inclusive_parse() {
9080        // GICoref uses neopronouns and singular they
9081        let gicoref_conll = "Alex\tB-PER\n\
9082                             uses\tO\n\
9083                             they\tB-PER\n\
9084                             pronouns\tO\n\
9085                             \n\
9086                             Jordan\tB-PER\n\
9087                             introduced\tO\n\
9088                             themself\tB-PER\n";
9089
9090        let loader = DatasetLoader::new().unwrap();
9091        let result = loader.parse_conll(gicoref_conll, DatasetId::GICoref);
9092        assert!(result.is_ok());
9093
9094        let dataset = result.unwrap();
9095        assert_eq!(dataset.sentences.len(), 2);
9096
9097        // Verify pronoun tokens are correctly tagged
9098        assert_eq!(dataset.sentences[0].tokens[2].text, "they");
9099        assert_eq!(dataset.sentences[0].tokens[2].ner_tag, "B-PER");
9100        assert_eq!(dataset.sentences[1].tokens[2].text, "themself");
9101        assert_eq!(dataset.sentences[1].tokens[2].ner_tag, "B-PER");
9102    }
9103
9104    #[test]
9105    fn test_fewrel_jsonl_parse() {
9106        // FewRel has relation extraction in JSONL format
9107        // The parser expects integer tags mapping to MultiNERD labels:
9108        // 0=O, 1=B-PER, 3=B-ORG
9109        let fewrel_sample =
9110            r#"{"tokens":["John","works","at","Google","."],"ner_tags":[1,0,0,3,0]}"#;
9111
9112        let loader = DatasetLoader::new().unwrap();
9113        let result = loader.parse_jsonl_ner(fewrel_sample, DatasetId::FewRel);
9114        assert!(result.is_ok());
9115
9116        let dataset = result.unwrap();
9117        assert_eq!(dataset.sentences.len(), 1);
9118        assert_eq!(dataset.sentences[0].tokens.len(), 5);
9119        assert_eq!(dataset.sentences[0].tokens[0].ner_tag, "B-PER");
9120        assert_eq!(dataset.sentences[0].tokens[3].ner_tag, "B-ORG");
9121    }
9122
9123    #[test]
9124    fn test_b2nerd_business_entities_parse() {
9125        // B2NERD focuses on business/financial entity types
9126        // Using standard MultiNERD indices: 3=B-ORG (closest to COMPANY), 0=O
9127        let b2nerd_sample =
9128            r#"{"tokens":["Apple","Inc",".","reports","Q4","earnings"],"ner_tags":[3,4,4,0,0,0]}"#;
9129
9130        let loader = DatasetLoader::new().unwrap();
9131        let result = loader.parse_jsonl_ner(b2nerd_sample, DatasetId::B2NERD);
9132        assert!(result.is_ok());
9133
9134        let dataset = result.unwrap();
9135        assert_eq!(dataset.sentences.len(), 1);
9136        assert_eq!(dataset.sentences[0].tokens[0].ner_tag, "B-ORG");
9137        assert_eq!(dataset.sentences[0].tokens[1].ner_tag, "I-ORG");
9138    }
9139
9140    #[test]
9141    fn test_ud_classical_languages_are_loadable() {
9142        // Universal Dependencies treebanks for classical/ancient languages
9143        let ud_datasets = [
9144            DatasetId::AncientGreekUD,
9145            DatasetId::LatinUD,
9146            DatasetId::SanskritUD,
9147            DatasetId::OldEnglishUD,
9148            DatasetId::OldNorseUD,
9149            DatasetId::UDEsperantoCairo,
9150        ];
9151
9152        for id in ud_datasets {
9153            assert!(
9154                LoadableDatasetId::is_loadable_dataset(id),
9155                "{:?} should be loadable with Conllu parse plan",
9156                id
9157            );
9158            assert_eq!(
9159                LoadableDatasetId::parse_plan(id),
9160                Some(DatasetParsePlan::Conllu),
9161                "{:?} should use Conllu plan",
9162                id
9163            );
9164        }
9165    }
9166
9167    #[test]
9168    fn test_hipe2022_tsv_is_loadable() {
9169        assert!(
9170            LoadableDatasetId::is_loadable_dataset(DatasetId::HIPE2022),
9171            "HIPE2022 should be loadable"
9172        );
9173        assert_eq!(
9174            LoadableDatasetId::parse_plan(DatasetId::HIPE2022),
9175            Some(DatasetParsePlan::TsvNer),
9176            "HIPE2022 should use TsvNer plan"
9177        );
9178    }
9179
9180    #[test]
9181    fn test_ener_csv_is_loadable() {
9182        assert!(
9183            LoadableDatasetId::is_loadable_dataset(DatasetId::ENer),
9184            "ENer should be loadable"
9185        );
9186        assert_eq!(
9187            LoadableDatasetId::parse_plan(DatasetId::ENer),
9188            Some(DatasetParsePlan::CsvNer),
9189            "ENer should use CsvNer plan"
9190        );
9191    }
9192
9193    #[test]
9194    fn test_loadable_count_increased() {
9195        // Regression test: ensure we have at least 295 loadable datasets
9196        // (Updated 2025-12 after adding CoNLL/JSONL/CoNLLU batches)
9197        let loadable_count = LoadableDatasetId::all().len();
9198        assert!(
9199            loadable_count >= 295,
9200            "Expected at least 295 loadable datasets, got {}",
9201            loadable_count
9202        );
9203    }
9204
9205    // =========================================================================
9206    // Domain-Specific Parser Tests
9207    // =========================================================================
9208
9209    #[test]
9210    fn test_biomedical_conll_with_chemical_entities() {
9211        // CHEMDNER-style biomedical NER with chemical entity types
9212        let chemdner_conll = "Aspirin\tB-CHEMICAL\n\
9213                              inhibits\tO\n\
9214                              COX-2\tB-GENE\n\
9215                              expression\tO\n\
9216                              \n\
9217                              Metformin\tB-CHEMICAL\n\
9218                              treats\tO\n\
9219                              diabetes\tB-DISEASE\n";
9220
9221        let loader = DatasetLoader::new().unwrap();
9222        let result = loader.parse_conll(chemdner_conll, DatasetId::CHEMDNER);
9223        assert!(result.is_ok());
9224
9225        let dataset = result.unwrap();
9226        assert_eq!(dataset.sentences.len(), 2);
9227        assert_eq!(dataset.sentences[0].tokens[0].ner_tag, "B-CHEMICAL");
9228        assert_eq!(dataset.sentences[0].tokens[2].ner_tag, "B-GENE");
9229        assert_eq!(dataset.sentences[1].tokens[2].ner_tag, "B-DISEASE");
9230    }
9231
9232    #[test]
9233    fn test_historical_ner_with_archaic_spelling() {
9234        // Historical NER should handle archaic spellings and diacritics
9235        let historical_conll = "Præsident\tB-PER\n\
9236                                 Washington\tI-PER\n\
9237                                 addresseth\tO\n\
9238                                 ye\tO\n\
9239                                 Congreſs\tB-ORG\n";
9240
9241        let loader = DatasetLoader::new().unwrap();
9242        let result = loader.parse_conll(historical_conll, DatasetId::EighteenthCenturyNER);
9243        assert!(result.is_ok());
9244
9245        let dataset = result.unwrap();
9246        assert_eq!(dataset.sentences.len(), 1);
9247        // Verify archaic characters are preserved
9248        assert!(dataset.sentences[0].tokens[0].text.contains('æ'));
9249        assert!(dataset.sentences[0].tokens[4].text.contains('ſ'));
9250    }
9251
9252    #[test]
9253    fn test_multilingual_code_switching_ner() {
9254        // LinCE/CALCS datasets have code-switched text (e.g., Spanish-English)
9255        let codeswitched_conll = "My\tO\n\
9256                                   abuela\tB-PER\n\
9257                                   lives\tO\n\
9258                                   in\tO\n\
9259                                   Ciudad\tB-LOC\n\
9260                                   de\tI-LOC\n\
9261                                   México\tI-LOC\n";
9262
9263        let loader = DatasetLoader::new().unwrap();
9264        let result = loader.parse_conll(codeswitched_conll, DatasetId::LinCE);
9265        assert!(result.is_ok());
9266
9267        let dataset = result.unwrap();
9268        assert_eq!(dataset.sentences.len(), 1);
9269        assert_eq!(dataset.sentences[0].tokens[1].text, "abuela");
9270        assert_eq!(dataset.sentences[0].tokens[1].ner_tag, "B-PER");
9271        // Multi-word location
9272        assert_eq!(dataset.sentences[0].tokens[4].ner_tag, "B-LOC");
9273        assert_eq!(dataset.sentences[0].tokens[6].ner_tag, "I-LOC");
9274    }
9275
9276    #[test]
9277    fn test_indigenous_language_ner() {
9278        // Test parsing of indigenous language NER (Guarani/Shipibo-Konibo)
9279        let guarani_conll = "Paraguái\tB-LOC\n\
9280                              ha\tO\n\
9281                              yvypora\tO\n\
9282                              oiko\tO\n\
9283                              Asunción\tB-LOC\n\
9284                              pe\tO\n";
9285
9286        let loader = DatasetLoader::new().unwrap();
9287        let result = loader.parse_conll(guarani_conll, DatasetId::GuaraniNER);
9288        assert!(result.is_ok());
9289
9290        let dataset = result.unwrap();
9291        assert_eq!(dataset.sentences.len(), 1);
9292        assert_eq!(dataset.sentences[0].tokens[0].text, "Paraguái");
9293        assert_eq!(dataset.sentences[0].tokens[0].ner_tag, "B-LOC");
9294    }
9295
9296    #[test]
9297    fn test_ancient_greek_conllu_with_polytonic() {
9298        // Ancient Greek with polytonic diacritics
9299        let greek_conllu = "# sent_id = grc_test_1\n\
9300                            # text = ἐπιστήμη καὶ δικαιοσύνη\n\
9301                            1\tἐπιστήμη\tἐπιστήμη\tNOUN\tN\tCase=Nom|Gender=Fem|Number=Sing\t0\troot\t_\tSpaceAfter=Yes\n\
9302                            2\tκαὶ\tκαί\tCCONJ\tC\t_\t3\tcc\t_\tSpaceAfter=Yes\n\
9303                            3\tδικαιοσύνη\tδικαιοσύνη\tNOUN\tN\tCase=Nom|Gender=Fem|Number=Sing\t1\tconj\t_\tSpaceAfter=No\n";
9304
9305        let loader = DatasetLoader::new().unwrap();
9306        let result = loader.parse_conllu(greek_conllu, DatasetId::AncientGreekUD);
9307        assert!(result.is_ok());
9308
9309        let dataset = result.unwrap();
9310        assert_eq!(dataset.sentences.len(), 1);
9311        // Verify polytonic Greek is preserved
9312        assert_eq!(dataset.sentences[0].tokens[0].text, "ἐπιστήμη");
9313        assert_eq!(dataset.sentences[0].tokens[2].text, "δικαιοσύνη");
9314    }
9315
9316    #[test]
9317    fn test_latin_conllu_with_macrons() {
9318        // Latin with optional macrons for vowel length
9319        let latin_conllu = "# sent_id = lat_test_1\n\
9320                            # text = Rōma āterna est\n\
9321                            1\tRōma\tRoma\tPROPN\tNNP\tCase=Nom|Gender=Fem|Number=Sing\t3\tnsubj\t_\tSpaceAfter=Yes\n\
9322                            2\tāterna\taeternus\tADJ\tA\tCase=Nom|Gender=Fem|Number=Sing\t1\tamod\t_\tSpaceAfter=Yes\n\
9323                            3\test\tsum\tAUX\tV\tMood=Ind|Number=Sing|Person=3|Tense=Pres|VerbForm=Fin|Voice=Act\t0\troot\t_\tSpaceAfter=No\n";
9324
9325        let loader = DatasetLoader::new().unwrap();
9326        let result = loader.parse_conllu(latin_conllu, DatasetId::LatinUD);
9327        assert!(result.is_ok());
9328
9329        let dataset = result.unwrap();
9330        assert_eq!(dataset.sentences.len(), 1);
9331        // Verify macrons are preserved
9332        assert!(dataset.sentences[0].tokens[0].text.contains('ō'));
9333        assert!(dataset.sentences[0].tokens[1].text.contains('ā'));
9334    }
9335
9336    #[test]
9337    fn test_sanskrit_conllu_with_devanagari() {
9338        // Sanskrit in Devanagari script
9339        let sanskrit_conllu = "# sent_id = sa_test_1\n\
9340                               # text = रामः सीतां पश्यति\n\
9341                               1\tरामः\tराम\tNOUN\tN\tCase=Nom|Gender=Masc|Number=Sing\t3\tnsubj\t_\tSpaceAfter=Yes\n\
9342                               2\tसीतां\tसीता\tPROPN\tNNP\tCase=Acc|Gender=Fem|Number=Sing\t3\tobj\t_\tSpaceAfter=Yes\n\
9343                               3\tपश्यति\tदृश्\tVERB\tV\tMood=Ind|Number=Sing|Person=3|Tense=Pres|VerbForm=Fin|Voice=Act\t0\troot\t_\tSpaceAfter=No\n";
9344
9345        let loader = DatasetLoader::new().unwrap();
9346        let result = loader.parse_conllu(sanskrit_conllu, DatasetId::SanskritUD);
9347        assert!(result.is_ok());
9348
9349        let dataset = result.unwrap();
9350        assert_eq!(dataset.sentences.len(), 1);
9351        assert_eq!(dataset.sentences[0].tokens.len(), 3);
9352        // Verify Devanagari is preserved
9353        assert_eq!(dataset.sentences[0].tokens[0].text, "रामः");
9354        assert_eq!(dataset.sentences[0].tokens[1].text, "सीतां");
9355    }
9356
9357    #[test]
9358    fn test_klingon_conllu_is_loadable() {
9359        // Klingon (tlh) is a constructed language in TaggedPBCKlingon
9360        assert!(
9361            LoadableDatasetId::is_loadable_dataset(DatasetId::TaggedPBCKlingon),
9362            "Klingon dataset should be loadable"
9363        );
9364        assert_eq!(
9365            LoadableDatasetId::parse_plan(DatasetId::TaggedPBCKlingon),
9366            Some(DatasetParsePlan::Conllu),
9367            "Klingon should use Conllu plan"
9368        );
9369    }
9370
9371    #[test]
9372    fn test_financial_ner_entities() {
9373        // FinanceNER with financial entity types
9374        let finance_conll = "Tesla\tB-COMPANY\n\
9375                             stock\tO\n\
9376                             rose\tO\n\
9377                             5%\tB-PERCENTAGE\n\
9378                             after\tO\n\
9379                             Q4\tB-PERIOD\n\
9380                             earnings\tO\n";
9381
9382        let loader = DatasetLoader::new().unwrap();
9383        let result = loader.parse_conll(finance_conll, DatasetId::FinanceNER);
9384        assert!(result.is_ok());
9385
9386        let dataset = result.unwrap();
9387        assert_eq!(dataset.sentences.len(), 1);
9388        assert_eq!(dataset.sentences[0].tokens[0].ner_tag, "B-COMPANY");
9389        assert_eq!(dataset.sentences[0].tokens[3].ner_tag, "B-PERCENTAGE");
9390    }
9391
9392    #[test]
9393    fn test_recipe_ner_food_entities() {
9394        // RecipeNER with culinary entity types
9395        let recipe_conll = "Add\tO\n\
9396                            2\tB-QUANTITY\n\
9397                            cups\tI-QUANTITY\n\
9398                            of\tO\n\
9399                            flour\tB-INGREDIENT\n\
9400                            and\tO\n\
9401                            1\tB-QUANTITY\n\
9402                            tsp\tI-QUANTITY\n\
9403                            salt\tB-INGREDIENT\n";
9404
9405        let loader = DatasetLoader::new().unwrap();
9406        let result = loader.parse_conll(recipe_conll, DatasetId::RecipeNER);
9407        assert!(result.is_ok());
9408
9409        let dataset = result.unwrap();
9410        assert_eq!(dataset.sentences.len(), 1);
9411        assert_eq!(dataset.sentences[0].tokens[4].ner_tag, "B-INGREDIENT");
9412        assert_eq!(dataset.sentences[0].tokens[8].ner_tag, "B-INGREDIENT");
9413    }
9414
9415    #[test]
9416    fn test_astronomy_ner_entities() {
9417        // AstroNER with astronomical entity types
9418        let astro_conll = "The\tO\n\
9419                           Andromeda\tB-GALAXY\n\
9420                           Galaxy\tI-GALAXY\n\
9421                           is\tO\n\
9422                           2.5\tB-DISTANCE\n\
9423                           million\tI-DISTANCE\n\
9424                           light-years\tI-DISTANCE\n\
9425                           away\tO\n";
9426
9427        let loader = DatasetLoader::new().unwrap();
9428        let result = loader.parse_conll(astro_conll, DatasetId::AstroNER);
9429        assert!(result.is_ok());
9430
9431        let dataset = result.unwrap();
9432        assert_eq!(dataset.sentences.len(), 1);
9433        assert_eq!(dataset.sentences[0].tokens[1].ner_tag, "B-GALAXY");
9434        assert_eq!(dataset.sentences[0].tokens[4].ner_tag, "B-DISTANCE");
9435    }
9436
9437    #[test]
9438    fn test_nested_ner_datasets_are_loadable() {
9439        // Nested NER datasets (entities within entities)
9440        let nested_datasets = [DatasetId::GENIANested, DatasetId::ChineseNestedNER];
9441
9442        for id in nested_datasets {
9443            assert!(
9444                LoadableDatasetId::is_loadable_dataset(id),
9445                "{:?} should be loadable",
9446                id
9447            );
9448        }
9449    }
9450
9451    #[test]
9452    fn test_discontinuous_ner_datasets_are_loadable() {
9453        // Discontinuous NER datasets (non-contiguous entity spans)
9454        let discontinuous_datasets = [
9455            DatasetId::GermEvalDiscontinuous,
9456            DatasetId::PubMedDiscontinuous,
9457        ];
9458
9459        for id in discontinuous_datasets {
9460            assert!(
9461                LoadableDatasetId::is_loadable_dataset(id),
9462                "{:?} should be loadable",
9463                id
9464            );
9465        }
9466    }
9467
9468    #[test]
9469    fn test_social_media_ner_datasets_are_loadable() {
9470        // Social media NER datasets (noisy text, hashtags, mentions)
9471        let social_datasets = [
9472            DatasetId::WNUT16,
9473            DatasetId::TwiConv,
9474            DatasetId::NERsocialFood,
9475        ];
9476
9477        for id in social_datasets {
9478            assert!(
9479                LoadableDatasetId::is_loadable_dataset(id),
9480                "{:?} should be loadable",
9481                id
9482            );
9483            assert_eq!(
9484                LoadableDatasetId::parse_plan(id),
9485                Some(DatasetParsePlan::Conll),
9486                "{:?} should use Conll plan",
9487                id
9488            );
9489        }
9490    }
9491
9492    #[test]
9493    fn test_literary_ner_datasets_are_loadable() {
9494        // Literary NER datasets (fiction, novels)
9495        let literary_datasets = [
9496            DatasetId::CharacterCodex,
9497            DatasetId::FictionNER750M,
9498            DatasetId::BookCoref,
9499        ];
9500
9501        for id in literary_datasets {
9502            assert!(
9503                LoadableDatasetId::is_loadable_dataset(id),
9504                "{:?} should be loadable",
9505                id
9506            );
9507        }
9508    }
9509
9510    #[test]
9511    fn test_jsonl_ner_with_empty_tokens_handled() {
9512        // Edge case: JSONL with some empty tokens
9513        let jsonl_with_empty = r#"{"tokens":["Hello","","world"],"ner_tags":[0,0,0]}"#;
9514
9515        let loader = DatasetLoader::new().unwrap();
9516        let result = loader.parse_jsonl_ner(jsonl_with_empty, DatasetId::MultiWOZNER);
9517        assert!(result.is_ok());
9518
9519        let dataset = result.unwrap();
9520        assert_eq!(dataset.sentences.len(), 1);
9521        // Empty token should still be preserved in parsing
9522        assert_eq!(dataset.sentences[0].tokens.len(), 3);
9523    }
9524
9525    #[test]
9526    fn test_conll_with_long_entity_spans() {
9527        // Test parsing of very long entity spans (e.g., legal document titles)
9528        let long_span_conll = "The\tB-DOCUMENT\n\
9529                               United\tI-DOCUMENT\n\
9530                               States\tI-DOCUMENT\n\
9531                               Constitution\tI-DOCUMENT\n\
9532                               Article\tI-DOCUMENT\n\
9533                               I\tI-DOCUMENT\n\
9534                               Section\tI-DOCUMENT\n\
9535                               8\tI-DOCUMENT\n\
9536                               Clause\tI-DOCUMENT\n\
9537                               3\tI-DOCUMENT\n";
9538
9539        let loader = DatasetLoader::new().unwrap();
9540        let result = loader.parse_conll(long_span_conll, DatasetId::LegNER);
9541        assert!(result.is_ok());
9542
9543        let dataset = result.unwrap();
9544        assert_eq!(dataset.sentences.len(), 1);
9545        assert_eq!(dataset.sentences[0].tokens.len(), 10);
9546        // All tokens should be part of the same entity
9547        assert_eq!(dataset.sentences[0].tokens[0].ner_tag, "B-DOCUMENT");
9548        assert_eq!(dataset.sentences[0].tokens[9].ner_tag, "I-DOCUMENT");
9549    }
9550
9551    #[test]
9552    fn test_all_added_conll_datasets_are_loadable() {
9553        // Comprehensive check for all newly added CoNLL datasets
9554        let added_conll = [
9555            DatasetId::HistNERo,
9556            DatasetId::DutchArchaeology,
9557            DatasetId::FINER,
9558            DatasetId::CALCS2018,
9559            DatasetId::MedievalCharterNER,
9560            DatasetId::RockNER,
9561            DatasetId::AIDACoNLL,
9562            DatasetId::NNE,
9563            DatasetId::IndicNER,
9564            DatasetId::NorNE,
9565            DatasetId::TASTEset,
9566            DatasetId::TechNER,
9567            DatasetId::FinTechPatent,
9568            DatasetId::WaterAgriNER,
9569            DatasetId::RussianCulturalNER,
9570            DatasetId::BASHI,
9571            DatasetId::ENER,
9572        ];
9573
9574        for id in added_conll {
9575            assert!(
9576                LoadableDatasetId::is_loadable_dataset(id),
9577                "{:?} should be loadable",
9578                id
9579            );
9580        }
9581    }
9582
9583    #[test]
9584    fn test_all_added_jsonl_datasets_are_loadable() {
9585        // Comprehensive check for all newly added JSONL datasets
9586        let added_jsonl = [
9587            DatasetId::MultiWOZNER,
9588            DatasetId::HinglishNER,
9589            DatasetId::AgCNER,
9590            DatasetId::LongDocNER,
9591            DatasetId::MultiBioNERLong,
9592            DatasetId::ReasoningNER,
9593            DatasetId::BioNERLLaMA,
9594            DatasetId::LexGLUENER,
9595            DatasetId::FinBenNER,
9596            DatasetId::FiNER139,
9597            DatasetId::SciNER,
9598            DatasetId::AIONER,
9599            DatasetId::WIESPAstro,
9600            DatasetId::CEREC,
9601            DatasetId::DELICATE,
9602            DatasetId::CSN,
9603        ];
9604
9605        for id in added_jsonl {
9606            assert!(
9607                LoadableDatasetId::is_loadable_dataset(id),
9608                "{:?} should be loadable",
9609                id
9610            );
9611        }
9612    }
9613
9614    #[test]
9615    fn test_all_added_ud_datasets_are_loadable() {
9616        // Comprehensive check for all newly added UD datasets
9617        let added_ud = [
9618            DatasetId::CopticScriptorium,
9619            DatasetId::TaggedPBCEsperanto,
9620            DatasetId::TaggedPBCKlingon,
9621            DatasetId::AkkadianUD,
9622            DatasetId::AncientHebrewUD,
9623            DatasetId::ClassicalChineseUD,
9624            DatasetId::CopticUD,
9625            DatasetId::GothicUD,
9626            DatasetId::HittiteUD,
9627            DatasetId::OldChurchSlavonicUD,
9628            DatasetId::LatinITTB,
9629            DatasetId::LatinPROIEL,
9630            DatasetId::EsperantoUD,
9631            DatasetId::NavajoMorph,
9632        ];
9633
9634        for id in added_ud {
9635            assert!(
9636                LoadableDatasetId::is_loadable_dataset(id),
9637                "{:?} should be loadable",
9638                id
9639            );
9640            assert_eq!(
9641                LoadableDatasetId::parse_plan(id),
9642                Some(DatasetParsePlan::Conllu),
9643                "{:?} should use Conllu plan",
9644                id
9645            );
9646        }
9647    }
9648
9649    // =========================================================================
9650    // Edge Case and Robustness Tests
9651    // =========================================================================
9652
9653    #[test]
9654    fn test_conll_handles_windows_line_endings() {
9655        // Test CRLF line endings (Windows format)
9656        let windows_conll = "John\tB-PER\r\nSmith\tI-PER\r\n\r\nLondon\tB-LOC\r\n";
9657
9658        let loader = DatasetLoader::new().unwrap();
9659        let result = loader.parse_conll(windows_conll, DatasetId::WikiGold);
9660        assert!(result.is_ok());
9661
9662        let dataset = result.unwrap();
9663        assert_eq!(dataset.sentences.len(), 2);
9664    }
9665
9666    #[test]
9667    fn test_conll_handles_extra_whitespace() {
9668        // Test lines with trailing/leading whitespace
9669        let whitespace_conll = "  John  \t  B-PER  \n  meets  \t  O  \n\n";
9670
9671        let loader = DatasetLoader::new().unwrap();
9672        let result = loader.parse_conll(whitespace_conll, DatasetId::WikiGold);
9673        // Should handle gracefully (may skip malformed lines)
9674        assert!(result.is_ok());
9675    }
9676
9677    #[test]
9678    fn test_conllu_handles_multiword_tokens() {
9679        // CoNLL-U multi-word tokens (e.g., "don't" -> "do" + "n't")
9680        let mwt_conllu = "# sent_id = test\n\
9681                          # text = I can't go\n\
9682                          1\tI\tI\tPRON\tPRP\t_\t3\tnsubj\t_\tSpaceAfter=Yes\n\
9683                          2-3\tcan't\t_\t_\t_\t_\t_\t_\t_\tSpaceAfter=Yes\n\
9684                          2\tca\tcan\tAUX\tMD\t_\t3\taux\t_\t_\n\
9685                          3\tn't\tnot\tPART\tRB\t_\t0\troot\t_\t_\n\
9686                          4\tgo\tgo\tVERB\tVB\t_\t3\txcomp\t_\tSpaceAfter=No\n";
9687
9688        let loader = DatasetLoader::new().unwrap();
9689        let result = loader.parse_conllu(mwt_conllu, DatasetId::LatinUD);
9690        assert!(result.is_ok());
9691
9692        let dataset = result.unwrap();
9693        assert_eq!(dataset.sentences.len(), 1);
9694        // MWT range tokens (2-3) should be skipped, only atomic tokens kept
9695        assert!(dataset.sentences[0].tokens.len() >= 3);
9696    }
9697
9698    #[test]
9699    fn test_conllu_handles_empty_nodes() {
9700        // CoNLL-U empty nodes (e.g., 2.1 for elided elements)
9701        let empty_node_conllu = "# sent_id = test\n\
9702                                 # text = I saw and heard\n\
9703                                 1\tI\tI\tPRON\tPRP\t_\t2\tnsubj\t_\tSpaceAfter=Yes\n\
9704                                 2\tsaw\tsee\tVERB\tVBD\t_\t0\troot\t_\tSpaceAfter=Yes\n\
9705                                 2.1\tI\tI\tPRON\tPRP\t_\t4\tnsubj\t_\t_\n\
9706                                 3\tand\tand\tCCONJ\tCC\t_\t4\tcc\t_\tSpaceAfter=Yes\n\
9707                                 4\theard\thear\tVERB\tVBD\t_\t2\tconj\t_\tSpaceAfter=No\n";
9708
9709        let loader = DatasetLoader::new().unwrap();
9710        let result = loader.parse_conllu(empty_node_conllu, DatasetId::LatinUD);
9711        assert!(result.is_ok());
9712
9713        let dataset = result.unwrap();
9714        assert_eq!(dataset.sentences.len(), 1);
9715        // Empty nodes (2.1) should be skipped
9716        assert_eq!(dataset.sentences[0].tokens.len(), 4);
9717    }
9718
9719    #[test]
9720    fn test_conll_with_bio_tag_normalization() {
9721        // Some corpora use I- at start of entity (should be B-)
9722        let malformed_bio = "Paris\tI-LOC\n\
9723                             is\tO\n\
9724                             beautiful\tO\n";
9725
9726        let loader = DatasetLoader::new().unwrap();
9727        let result = loader.parse_conll(malformed_bio, DatasetId::WikiGold);
9728        assert!(result.is_ok());
9729
9730        let dataset = result.unwrap();
9731        assert_eq!(dataset.sentences.len(), 1);
9732        // Parser should normalize or preserve the tag
9733        let first_tag = &dataset.sentences[0].tokens[0].ner_tag;
9734        assert!(first_tag == "I-LOC" || first_tag == "B-LOC");
9735    }
9736
9737    #[test]
9738    fn test_conll_with_unicode_normalization() {
9739        // Test that precomposed (U+00E9) and decomposed (U+0065 U+0301) forms
9740        // both parse and produce equivalent tokens.
9741        let composed = "Caf\u{00e9}\tB-LOC\n";
9742        let decomposed = "Cafe\u{0301}\tB-LOC\n";
9743        assert_ne!(
9744            composed, decomposed,
9745            "test inputs must actually differ in bytes"
9746        );
9747
9748        let loader = DatasetLoader::new().unwrap();
9749        let d1 = loader.parse_conll(composed, DatasetId::WikiGold).unwrap();
9750        let d2 = loader.parse_conll(decomposed, DatasetId::WikiGold).unwrap();
9751
9752        assert_eq!(d1.sentences.len(), d2.sentences.len());
9753        assert_eq!(d1.sentences[0].tokens.len(), d2.sentences[0].tokens.len());
9754        assert_eq!(
9755            d1.sentences[0].tokens[0].ner_tag,
9756            d2.sentences[0].tokens[0].ner_tag
9757        );
9758    }
9759
9760    #[test]
9761    fn test_cadec_hf_api_unicode_prefix_case_insensitive_span_search_is_safe() {
9762        // Regression: span finding must not rely on Unicode lowercasing index alignment.
9763        let loader = DatasetLoader::new().unwrap();
9764        // Include `features` and compact JSON so `is_hf_api_response()` recognizes it.
9765        let content = r#"{"features":[{"name":"text"},{"name":"ade"},{"name":"term_PT"}],"rows":[{"row":{"text":"Müller reported HEADACHE after taking aspirin.","ade":"headache","term_PT":"Headache"}}]}"#;
9766
9767        let ds = loader
9768            .parse_content_str(content, DatasetId::CADEC)
9769            .expect("parse CADEC HF API");
9770        assert_eq!(ds.id, DatasetId::CADEC);
9771        assert!(!ds.sentences.is_empty());
9772
9773        let sent = &ds.sentences[0];
9774        // We should tag the ADE as B-/I- adverse_drug_event in BIO space.
9775        assert!(
9776            sent.tokens
9777                .iter()
9778                .any(|t| t.ner_tag == "B-adverse_drug_event" || t.ner_tag == "I-adverse_drug_event"),
9779            "Expected ADE tags in tokens: {:?}",
9780            sent.tokens
9781        );
9782    }
9783
9784    #[test]
9785    fn test_jsonl_with_unicode_tokens() {
9786        // JSONL with various Unicode characters
9787        let unicode_jsonl = r#"{"tokens":["北京","🎉","Москва","القاهرة"],"ner_tags":[5,0,5,5]}"#;
9788
9789        let loader = DatasetLoader::new().unwrap();
9790        let result = loader.parse_jsonl_ner(unicode_jsonl, DatasetId::MultiWOZNER);
9791        assert!(result.is_ok());
9792
9793        let dataset = result.unwrap();
9794        assert_eq!(dataset.sentences.len(), 1);
9795        assert_eq!(dataset.sentences[0].tokens.len(), 4);
9796        assert_eq!(dataset.sentences[0].tokens[0].text, "北京");
9797        assert_eq!(dataset.sentences[0].tokens[1].text, "🎉");
9798        assert_eq!(dataset.sentences[0].tokens[2].text, "Москва");
9799    }
9800
9801    #[test]
9802    fn test_parse_plan_consistency_with_is_loadable() {
9803        // Invariant: parse_plan returns Some iff is_loadable_dataset returns true
9804        for &id in DatasetId::all() {
9805            let has_plan = LoadableDatasetId::parse_plan(id).is_some();
9806            let is_loadable = LoadableDatasetId::is_loadable_dataset(id);
9807            assert_eq!(
9808                has_plan, is_loadable,
9809                "Mismatch for {:?}: parse_plan={}, is_loadable={}",
9810                id, has_plan, is_loadable
9811            );
9812        }
9813    }
9814
9815    #[test]
9816    fn test_dataset_coverage_by_plan() {
9817        // Verify we have good coverage across different parse plans
9818        let mut conll_count = 0;
9819        let mut jsonl_count = 0;
9820        let mut conllu_count = 0;
9821        let mut _other_count = 0;
9822
9823        for id in LoadableDatasetId::all() {
9824            let ds: DatasetId = id.into();
9825            match LoadableDatasetId::parse_plan(ds) {
9826                Some(DatasetParsePlan::Conll) => conll_count += 1,
9827                Some(DatasetParsePlan::JsonlNer) => jsonl_count += 1,
9828                Some(DatasetParsePlan::Conllu) => conllu_count += 1,
9829                Some(_) => _other_count += 1,
9830                None => {}
9831            }
9832        }
9833
9834        // Should have substantial coverage for each format
9835        assert!(
9836            conll_count >= 50,
9837            "Expected at least 50 CoNLL datasets loadable, got {}",
9838            conll_count
9839        );
9840        assert!(
9841            jsonl_count >= 20,
9842            "Expected at least 20 JSONL datasets loadable, got {}",
9843            jsonl_count
9844        );
9845        assert!(
9846            conllu_count >= 10,
9847            "Expected at least 10 CoNLLU datasets loadable, got {}",
9848            conllu_count
9849        );
9850    }
9851
9852    #[test]
9853    fn test_loadable_datasets_have_valid_metadata() {
9854        // All loadable datasets should have basic metadata
9855        for id in LoadableDatasetId::all() {
9856            let ds: DatasetId = id.into();
9857
9858            // Name should not be empty
9859            assert!(!ds.name().is_empty(), "{:?} has empty name", ds);
9860
9861            // Description should exist for most
9862            // (not asserting as some may legitimately have no description)
9863        }
9864    }
9865}