1#[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
82pub use super::dataset_registry::DatasetId;
132
133#[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 TsvNer,
172 CsvNer,
174}
175
176impl LoadableDatasetId {
177 #[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 #[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 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 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 | DatasetId::QxoRef
316 | DatasetId::GICoref
317 | DatasetId::WNUT16
318 | DatasetId::NoiseBench
319 | DatasetId::CrossWeigh
320 | DatasetId::ZELDA
321 | DatasetId::GENIANested => DatasetParsePlan::Conll,
322
323 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 | 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 DatasetId::UNER | DatasetId::MSNER => DatasetParsePlan::WikiannJson,
387
388 DatasetId::TweetNER7 => DatasetParsePlan::TweetNer7,
390
391 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 => DatasetParsePlan::DocredJson,
418
419 DatasetId::CHisIEC => DatasetParsePlan::ChisiecJson,
421
422 DatasetId::CADEC | DatasetId::ShARe13 | DatasetId::ShARe14 => {
424 DatasetParsePlan::CadecHybrid
425 }
426
427 DatasetId::BC5CDR => DatasetParsePlan::Bc5cdr,
429 DatasetId::NCBIDisease => DatasetParsePlan::NcbiDisease,
430
431 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 DatasetId::AfriSenti => DatasetParsePlan::AfriSenti,
441 DatasetId::AfriQA => DatasetParsePlan::AfriQa,
442 DatasetId::MasakhaNEWS => DatasetParsePlan::MasakhaNews,
443 DatasetId::MasakhaPOS => DatasetParsePlan::Conllu,
444
445 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 DatasetId::MAVEN => DatasetParsePlan::Maven,
454 DatasetId::MAVENArg => DatasetParsePlan::MavenArg,
455 DatasetId::CASIE => DatasetParsePlan::Casie,
456 DatasetId::RAMS => DatasetParsePlan::Rams,
457
458 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 #[must_use]
502 fn registry_hint_plan(id: DatasetId) -> Option<DatasetParsePlan> {
503 if id == DatasetId::CHisIEC {
505 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 return Some(DatasetParsePlan::GoogleReCorpus);
519 }
520 if id == DatasetId::TweetNER7 {
521 return Some(DatasetParsePlan::TweetNer7);
522 }
523
524 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 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 | 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 | DatasetId::OntoNotes50
604 | DatasetId::GUM
605 | DatasetId::CoNLL04RE
607 ) {
608 return Some(DatasetParsePlan::Conll);
609 }
610
611 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 | 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 | DatasetId::SCINERNested
647 | DatasetId::AgriNER
648 | DatasetId::MOFDataset
649 | DatasetId::SolidStateDoping
650 ) {
651 return Some(DatasetParsePlan::JsonlNer);
652 }
653
654 if matches!(
656 id,
657 DatasetId::AncientGreekUD
658 | DatasetId::LatinUD
659 | DatasetId::SanskritUD
660 | DatasetId::OldEnglishUD
661 | DatasetId::OldNorseUD
662 | DatasetId::UDEsperantoCairo
663 | 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 if matches!(id, DatasetId::HIPE2022) {
684 return Some(DatasetParsePlan::TsvNer);
685 }
686
687 if matches!(id, DatasetId::ENer) {
689 return Some(DatasetParsePlan::CsvNer);
690 }
691
692 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 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 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 "TSV" | "ConllCoref" | "CoNLLCoref"
740 if is_coref && !is_ner && !is_re && !is_event =>
741 {
742 Some(DatasetParsePlan::GapTsv)
744 }
745 "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 Some(DatasetParsePlan::PrecoJsonl)
758 }
759
760 "JSON" | "JSONL" if is_re && !is_coref && !is_event => {
763 Some(DatasetParsePlan::DocredJson)
764 }
765
766 "JSONL" if is_event && !is_ner && !is_coref && !is_re => {
768 Some(DatasetParsePlan::Maven)
769 }
770
771 "TSV" if is_ner && !is_coref && !is_re && !is_event => {
773 Some(DatasetParsePlan::TsvNer)
774 }
775
776 "CSV" if is_ner && !is_coref && !is_re && !is_event => {
778 Some(DatasetParsePlan::CsvNer)
779 }
780
781 _ => None,
783 };
784
785 if plan.is_some() {
786 return plan;
787 }
788 }
789
790 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
847impl DatasetId {
858 #[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()), 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#[derive(Debug, Clone, Serialize, Deserialize)]
884pub struct CacheManifestEntry {
885 pub dataset_id: String,
887 pub source_url: String,
891 #[serde(default)]
897 pub resolved_url: Option<String>,
898 pub sha256: String,
900 pub file_size: u64,
902 pub downloaded_at: String,
904 pub sentence_count: usize,
906 pub entity_count: usize,
908 pub anno_version: String,
910}
911
912#[derive(Debug, Clone, Default, Serialize, Deserialize)]
920pub struct CacheManifest {
921 pub version: u32,
923 pub entries: HashMap<String, CacheManifestEntry>,
925}
926
927impl CacheManifest {
928 pub const CURRENT_VERSION: u32 = 1;
930
931 #[must_use]
933 pub fn new() -> Self {
934 Self {
935 version: Self::CURRENT_VERSION,
936 entries: HashMap::new(),
937 }
938 }
939
940 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 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 pub fn update_entry(&mut self, entry: CacheManifestEntry) {
968 self.entries.insert(entry.dataset_id.clone(), entry);
969 }
970
971 #[must_use]
973 pub fn get(&self, dataset_id: &str) -> Option<&CacheManifestEntry> {
974 self.entries.get(dataset_id)
975 }
976
977 #[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 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 }
995}
996
997#[derive(Debug, Clone, Default, Serialize, Deserialize)]
1006pub struct TemporalMetadata {
1007 pub kb_version: Option<String>,
1009 pub temporal_cutoff: Option<String>, pub entity_creation_dates: Option<HashMap<String, String>>, }
1014
1015#[derive(Debug, Clone, Serialize, Deserialize)]
1021pub struct AnnotatedToken {
1022 pub text: String,
1024 pub ner_tag: String,
1026}
1027
1028#[derive(Debug, Clone, Serialize, Deserialize)]
1030pub struct AnnotatedSentence {
1031 pub tokens: Vec<AnnotatedToken>,
1033 pub source_dataset: DatasetId,
1035}
1036
1037impl AnnotatedSentence {
1038 #[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 #[must_use]
1050 pub fn entities(&self) -> Vec<super::datasets::GoldEntity> {
1051 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; }
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 if tag.starts_with("B-") {
1070 let entity_type = tag.trim_start_matches("B-");
1071 let mut end = i + 1;
1072 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 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 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#[derive(Debug, Clone, Default, Serialize, Deserialize)]
1148pub struct DatasetMetadata {
1149 pub description: Option<String>,
1151 pub license: Option<String>,
1153 pub citation: Option<String>,
1155 pub split: Option<String>,
1157 pub language: Option<String>,
1159 pub domain: Option<String>,
1161 pub original_source: Option<String>,
1163 pub version: Option<String>,
1165 pub annotators: Option<u32>,
1167 pub guidelines_url: Option<String>,
1169}
1170
1171#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
1173pub enum DataSource {
1174 S3Cache,
1176 LocalCache,
1178 OriginalUrl,
1180 #[default]
1182 Skipped,
1183 Embedded,
1185}
1186
1187impl DataSource {
1188 #[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 #[must_use]
1202 pub fn is_cached(&self) -> bool {
1203 matches!(self, Self::S3Cache | Self::LocalCache | Self::Embedded)
1204 }
1205
1206 #[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#[derive(Debug, Clone, Serialize, Deserialize)]
1221pub struct LoadedDataset {
1222 pub id: DatasetId,
1224 pub sentences: Vec<AnnotatedSentence>,
1226 pub loaded_at: String, pub source_url: String,
1230 #[serde(default)]
1232 pub data_source: DataSource,
1233 #[serde(default)]
1235 pub temporal_metadata: Option<TemporalMetadata>,
1236 #[serde(default)]
1238 pub metadata: DatasetMetadata,
1239}
1240
1241impl LoadedDataset {
1242 #[must_use]
1244 pub fn len(&self) -> usize {
1245 self.sentences.len()
1246 }
1247
1248 #[must_use]
1250 pub fn is_empty(&self) -> bool {
1251 self.sentences.is_empty()
1252 }
1253
1254 #[must_use]
1256 pub fn entity_count(&self) -> usize {
1257 self.sentences.iter().map(|s| s.entities().len()).sum()
1258 }
1259
1260 #[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 #[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 #[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#[derive(Debug, Clone, Serialize, Deserialize)]
1300pub struct DatasetStats {
1301 pub name: String,
1303 pub sentences: usize,
1305 pub tokens: usize,
1307 pub entities: usize,
1309 pub entities_by_type: HashMap<String, usize>,
1311}
1312
1313#[derive(Debug, Clone)]
1315pub struct RelationDocument {
1316 pub text: String,
1318 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
1338pub struct DatasetLoader {
1371 cache_dir: PathBuf,
1372 s3_bucket: Option<String>,
1374 manifest: std::sync::RwLock<CacheManifest>,
1378}
1379
1380impl DatasetLoader {
1381 #[cfg(feature = "eval")]
1386 const DEFAULT_MAX_DOWNLOAD_BYTES: u64 = 50 * 1024 * 1024; #[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 } 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 let (_, query) = url.split_once('?')?;
1429 for part in query.split('&') {
1430 let (k, v) = part.split_once('=')?;
1431 if k == "dataset" {
1432 let decoded = v.replace("%2F", "/").replace("%2f", "/");
1434 return Some(decoded);
1435 }
1436 }
1437 None
1438 }
1439
1440 pub fn new() -> Result<Self> {
1447 anno::env::load_dotenv();
1450
1451 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 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 #[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 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 #[must_use]
1515 pub fn cache_dir(&self) -> &std::path::Path {
1516 &self.cache_dir
1517 }
1518
1519 #[must_use]
1521 pub fn s3_enabled(&self) -> bool {
1522 self.s3_bucket.is_some()
1523 }
1524
1525 #[must_use]
1527 pub fn s3_bucket(&self) -> Option<&str> {
1528 self.s3_bucket.as_deref()
1529 }
1530
1531 #[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 #[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 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 if entry.sentence_count == 0 {
1601 return false;
1602 }
1603 }
1604 }
1605
1606 true
1607 }
1608
1609 #[must_use]
1611 pub fn cache_path(&self, id: LoadableDatasetId) -> PathBuf {
1612 self.cache_path_for(id.0)
1613 }
1614
1615 #[must_use]
1617 pub fn is_cached(&self, id: LoadableDatasetId) -> bool {
1618 self.is_cached_for(id.0)
1619 }
1620
1621 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 #[cfg(feature = "eval")]
1654 pub fn load_or_download(&self, id: LoadableDatasetId) -> Result<LoadedDataset> {
1655 let dataset_id = id.0;
1656 if self.is_cached(id) {
1658 return self.load(id);
1659 }
1660
1661 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 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 if let Some(entry) = manifest_entry {
1676 let _ = self.update_manifest(entry);
1677 }
1678 return Ok(dataset);
1679 }
1680 }
1681
1682 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 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 let mut dataset = self.parse_content_impl(&content, dataset_id)?;
1696 dataset.data_source = DataSource::OriginalUrl;
1697
1698 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); 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 #[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 #[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 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 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 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 if content.is_none() {
1795 let s3_key = format!("datasets/{}", id.cache_filename());
1796 let s3_uri = format!("s3://{}/{}", bucket, s3_key);
1797
1798 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 let manifest = self.download_manifest_entry_from_s3(bucket, id).ok();
1822
1823 Ok((content, manifest))
1824 }
1825
1826 #[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 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 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 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 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 #[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 #[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 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 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 let preferred_hf_config: Option<String> = if id.language().eq_ignore_ascii_case("mul") {
2076 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 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 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 if url.contains("huggingface.co/datasets/") {
2122 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 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 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 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 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 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 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 #[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 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 #[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 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 if pass == 0 {
2312 if let Some(p) = prefer {
2313 if config != p {
2314 continue;
2315 }
2316 } else {
2317 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 #[cfg(feature = "eval")]
2348 fn download_hf_dataset_file_from_hub(&self, dataset: &str) -> Result<(String, String)> {
2349 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 let preferred = [
2394 "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 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 let file_url = format!(
2478 "https://huggingface.co/datasets/{}/resolve/main/{}",
2479 dataset, filename
2480 );
2481
2482 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 #[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 #[cfg(feature = "eval")]
2525 fn download_hf_dataset_paginated(&self, id: DatasetId, base_url: &str) -> Result<String> {
2526 #[cfg(feature = "onnx")] {
2529 if let Ok(content) = self.try_hf_hub_download(id) {
2530 return Ok(content);
2531 }
2532 }
2533
2534 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 let url = if base_url.contains("offset=") {
2551 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 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 if features.is_none() {
2577 features = parsed.get("features").cloned();
2578 }
2579
2580 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 if let Some(rows) = parsed.get("rows").and_then(|v| v.as_array()) {
2590 if rows.is_empty() {
2591 break; }
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 if let Some(total) = total_rows {
2602 if all_rows.len() >= total {
2603 break;
2604 }
2605 } else if rows.len() < PAGE_SIZE {
2606 break;
2608 }
2609
2610 offset += PAGE_SIZE;
2611 } else {
2612 break;
2614 }
2615 }
2616 Err(e) => {
2617 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 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 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 #[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 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 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 anno::env::load_dotenv();
2691
2692 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 #[cfg(not(all(feature = "eval", feature = "onnx")))]
2724 #[allow(dead_code)] fn try_hf_hub_download(&self, _id: DatasetId) -> Result<String> {
2726 Err(Error::InvalidInput("hf-hub not available".to_string()))
2727 }
2728
2729 #[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 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 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 if (500..600).contains(&status) && attempt < MAX_RETRIES {
2787 let wait_ms = 1000 * (1 << attempt); 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 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 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); 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 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 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 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 #[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; 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 #[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 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 fn get_temporal_metadata(id: DatasetId) -> Option<TemporalMetadata> {
2985 match id {
2986 DatasetId::TweetNER7 => {
2987 Some(TemporalMetadata {
2989 kb_version: None, temporal_cutoff: Some("2017-01-01".to_string()), entity_creation_dates: None, })
2993 }
2994 DatasetId::BroadTwitterCorpus => {
2995 Some(TemporalMetadata {
2997 kb_version: None,
2998 temporal_cutoff: Some("2018-01-01".to_string()), 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 Some(TemporalMetadata {
3010 kb_version: None,
3011 temporal_cutoff: None,
3012 entity_creation_dates: None,
3013 })
3014 }
3015 _ => None, }
3017 }
3018
3019 pub fn parse_content_str(&self, content: &str, id: DatasetId) -> Result<LoadedDataset> {
3027 self.parse_content_impl(content, id)
3028 }
3029
3030 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 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 fn is_hf_api_response(&self, content: &str) -> bool {
3096 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 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 let is_mit_format = matches!(id, DatasetId::MitMovie | DatasetId::MitRestaurant);
3112
3113 for line in content.lines() {
3114 let line = line.trim();
3115
3116 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 if line.starts_with("-DOCSTART-") {
3129 continue;
3130 }
3131
3132 let (text, ner_tag) = if is_mit_format {
3134 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 let parts: Vec<&str> = line.split_whitespace().collect();
3144 if parts.is_empty() {
3145 continue;
3146 }
3147
3148 if parts.len() >= 4 {
3149 (parts[0].to_string(), parts[3].to_string())
3151 } else if parts.len() >= 2 {
3152 (parts[0].to_string(), parts[parts.len() - 1].to_string())
3154 } else {
3155 (parts[0].to_string(), "O".to_string())
3157 }
3158 };
3159
3160 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 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 fn parse_jsonl_ner(&self, content: &str, id: DatasetId) -> Result<LoadedDataset> {
3216 let mut sentences = Vec::new();
3217
3218 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 let parsed: serde_json::Value = match serde_json::from_str(line) {
3234 Ok(v) => v,
3235 Err(_) => continue, };
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; }
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 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 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 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 let ner_tag = if let Some(tag_idx) = tag.as_u64() {
3332 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 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 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 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 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 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 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 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 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 fn parse_tweetner7(&self, content: &str, id: DatasetId) -> Result<LoadedDataset> {
3649 let tag_labels = [
3655 "B-corporation", "B-creative_work", "B-event", "B-group", "B-location", "B-person", "B-product", "I-corporation", "I-creative_work", "I-event", "I-group", "I-location", "I-person", "I-product", "O", ];
3671
3672 let mut sentences = Vec::new();
3673
3674 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, };
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 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 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 if line.starts_with("TOKEN\t") || line.starts_with("TOKEN ") {
3764 continue;
3765 }
3766
3767 if line.starts_with('#') {
3769 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 let parts: Vec<&str> = line.split('\t').collect();
3781 if parts.len() < 2 {
3782 continue; }
3784
3785 let token_text = parts[0].to_string();
3786 let ner_label = parts.get(1).unwrap_or(&"O");
3787
3788 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 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 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 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 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 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 let (token_text, ner_tag) = if let Some(rest) = line.strip_prefix(',') {
3877 if let Some(idx) = rest.find(',') {
3879 if idx == 0 {
3881 (",".to_string(), rest[1..].to_string())
3883 } else {
3884 (String::new(), rest[..idx].to_string())
3886 }
3887 } else {
3888 (String::new(), rest.to_string())
3890 }
3891 } else if let Some(idx) = line.rfind(',') {
3892 let token = line[..idx].to_string();
3894 let tag = line[idx + 1..].to_string();
3895 (token, tag)
3896 } else {
3897 continue;
3899 };
3900
3901 if token_text.is_empty() && ner_tag.is_empty() {
3903 continue;
3904 }
3905
3906 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 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 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 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 fn parse_docred(&self, content: &str, id: DatasetId) -> Result<LoadedDataset> {
4008 let mut sentences = Vec::new();
4009
4010 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 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 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 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 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 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 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 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; };
4224 let ade_end_byte = ade_start_byte + ade_text.len();
4225
4226 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 let ner_tag = if word_start >= ade_start_byte && word_end <= ade_end_byte {
4238 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 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 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 let parsed: serde_json::Value = match serde_json::from_str(line) {
4307 Ok(v) => v,
4308 Err(_) => continue, };
4310
4311 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 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; let end = char_offset - 1;
4327 token_offsets.push((text, start, end));
4328 }
4329
4330 for (text, _, _) in &token_offsets {
4332 annotated_tokens.push(AnnotatedToken {
4333 text: text.clone(),
4334 ner_tag: "O".to_string(),
4335 });
4336 }
4337
4338 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 if let Some(spans) = entity.get("spans").and_then(|v| v.as_array()) {
4350 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 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 let start = start_val as usize;
4387 let end = end_val as usize;
4388
4389 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 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 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 if line.starts_with("-DOCSTART-") {
4483 continue;
4484 }
4485
4486 if line.is_empty() {
4487 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 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 let normalized_tag = if ner_tag.contains("Entity")
4505 || ner_tag.contains("CHEMICAL")
4506 || ner_tag.contains("DISEASE")
4507 {
4508 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 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 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 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 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 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 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 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 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 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 if parsed_count < 3 {
4690 log::warn!("PreCo JSONL parse error on line {}: {}", line_count, e);
4691 }
4692 continue; }
4694 };
4695
4696 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(), })
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 fn parse_litbank(&self, content: &str, id: DatasetId) -> Result<LoadedDataset> {
4747 let now = chrono::Utc::now().to_rfc3339();
4750 let mut sentences = Vec::new();
4751 let mut entities: Vec<(usize, usize, String, String)> = Vec::new(); 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 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 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 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 continue;
4793 }
4794
4795 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 entities.sort_by_key(|(start, _, _, _)| *start);
4819
4820 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 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 let _text: String = text_chars
4847 .into_iter()
4848 .collect::<String>()
4849 .trim()
4850 .to_string();
4851
4852 let mut tokens: Vec<AnnotatedToken> = Vec::new();
4855
4856 entities.sort_by_key(|(start, _, _, _)| *start);
4858
4859 let mut last_end = 0usize;
4860 for (start, end, entity_text, label) in &entities {
4861 if *start > last_end {
4863 let gap_size = *start - last_end;
4866 if gap_size > 0 {
4867 let estimated_words = (gap_size / 6).max(1);
4869 for _ in 0..estimated_words.min(10) {
4870 tokens.push(AnnotatedToken {
4872 text: "[...]".to_string(),
4873 ner_tag: "O".to_string(),
4874 });
4875 }
4876 }
4877 }
4878
4879 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 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 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 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 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 if content.trim().starts_with('[') {
5019 let docs = super::coref_loader::parse_preco_json(&content)?;
5021 Ok(docs.into_iter().map(|d| d.to_coref_document()).collect())
5022 } else {
5023 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 if serde_json::from_str::<serde_json::Value>(line).is_ok() {
5032 json_objects.push(line);
5033 }
5034 }
5035 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 self.parse_litbank_coref(&content)
5044 }
5045 DatasetId::ECBPlus => {
5046 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 super::coref_loader::parse_ecb_plus_coref(&content)
5057 }
5058 DatasetId::WikiCoref => {
5059 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 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 #[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 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 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 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 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 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 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 let mut chains: Vec<Vec<Mention>> = Vec::new();
5204 let mut mention_to_chain: HashMap<String, usize> = HashMap::new();
5205
5206 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 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 let to_merge = std::mem::take(&mut chains[idx2]);
5221 chains[idx1].extend(to_merge);
5222 for m in &chains[idx1] {
5224 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; }
5267
5268 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 let doc = CorefDocument::new(&text, coref_chains);
5284 Ok(vec![doc])
5285 }
5286
5287 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 self.parse_docred_relations(&content)
5327 }
5328 DatasetId::CHisIEC => {
5329 self.parse_chisiec_relations(&content)
5331 }
5332 DatasetId::CADEC => {
5333 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 #[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 fn parse_docred_relations(&self, content: &str) -> Result<Vec<RelationDocument>> {
5362 use super::relation::RelationGold;
5363
5364 let mut documents = Vec::new();
5365
5366 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 let tokens_arr = match doc.get("sentence").and_then(|v| v.as_array()) {
5380 Some(t) => t,
5381 None => continue,
5382 };
5383
5384 let text: String = tokens_arr
5386 .iter()
5387 .filter_map(|t| t.as_str())
5388 .collect::<Vec<_>>()
5389 .join(" ");
5390
5391 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 char_pos += tok_str.len();
5400 if i < tokens_arr.len() - 1 {
5401 char_pos += 1; }
5403 } else {
5404 token_to_char.push(char_pos);
5405 }
5406 }
5407
5408 let ner_spans = doc.get("ner").and_then(|v| v.as_array());
5410
5411 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 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 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 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 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 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 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 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 fn parse_chisiec(&self, content: &str, id: DatasetId) -> Result<LoadedDataset> {
5586 let mut sentences = Vec::new();
5587
5588 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 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 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 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 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 fn parse_chisiec_relations(&self, content: &str) -> Result<Vec<RelationDocument>> {
5727 use super::relation::RelationGold;
5728
5729 let mut documents = Vec::new();
5730
5731 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 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 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 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 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 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 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 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 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 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 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 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 let context = doc.get("context").and_then(|v| v.as_str()).unwrap_or("");
5899
5900 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 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 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 let prefix: String = context.chars().take(start).collect();
5926 let word_idx = prefix.split_whitespace().count();
5927
5928 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 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 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 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 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 if let Some(space_idx) = line.find(' ') {
6049 let label = &line[..space_idx];
6050 let question = line[space_idx + 1..].trim();
6051
6052 let coarse_label = label.split(':').next().unwrap_or(label);
6054
6055 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 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 for line in content.lines() {
6099 let line = line.trim();
6100 if line.is_empty() || line.starts_with('#') {
6101 continue;
6102 }
6103
6104 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 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 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 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 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 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 if let Some(events) = doc.get("events").and_then(|e| e.as_array()) {
6316 is_jsonl = true;
6317
6318 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 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 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 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 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 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 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 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 let subtype = event
6463 .get("subtype")
6464 .and_then(|s| s.as_str())
6465 .unwrap_or("Event");
6466
6467 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 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 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 !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 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 let _doc_text = doc.get("document").and_then(|d| d.as_str()).unwrap_or("");
6579
6580 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 if line.starts_with('#') {
6876 continue;
6877 }
6878
6879 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 let fields: Vec<&str> = line.split('\t').collect();
6892 if fields.len() >= 4 {
6893 let id_field = fields[0];
6895 if id_field.contains('-') || id_field.contains('.') {
6896 continue;
6897 }
6898
6899 let form = fields[1]; let upos = fields[3]; current_tokens.push(AnnotatedToken {
6903 text: form.to_string(),
6904 ner_tag: format!("B-{}", upos), });
6906 }
6907 }
6908
6909 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 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 #[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#[cfg(test)]
6966fn parse_bio_tag(tag: &str) -> (&str, &str) {
6967 if tag == "O" {
6968 return ("O", "");
6969 }
6970
6971 if let Some(pos) = tag.find('-') {
6973 (&tag[..pos], &tag[pos + 1..])
6974 } else {
6975 ("B", tag)
6977 }
6978}
6979
6980#[cfg(test)]
6992fn map_entity_type(original: &str) -> EntityType {
6993 anno::schema::map_to_canonical(original, None)
6995}
6996
6997#[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 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 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 &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 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 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 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 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 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 assert!(matches!(map_entity_type("GPE"), EntityType::Custom { .. }));
7137
7138 assert!(matches!(map_entity_type("MISC"), EntityType::Custom { .. }));
7140
7141 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 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 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); let entities2 = dataset.sentences[1].entities();
7295 assert_eq!(entities2.len(), 1); assert_eq!(entities2[0].text, "Peter Blackburn");
7297 }
7298
7299 #[test]
7300 fn test_historical_datasets_configured() {
7301 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 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 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 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 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 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 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 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 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 assert!(!DatasetId::WikiGold.is_historical());
7392 assert!(!DatasetId::CoNLL2003Sample.is_historical());
7393 }
7394
7395 #[test]
7396 fn test_is_bias_evaluation_classification() {
7397 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 assert!(!DatasetId::WikiGold.is_bias_evaluation());
7406 }
7407
7408 #[test]
7409 fn test_new_datasets_have_descriptions() {
7410 let catch_all = "Dataset not yet fully integrated";
7412
7413 assert_ne!(DatasetId::HIPE2022.description(), catch_all);
7415 assert_ne!(DatasetId::TRIDIS.description(), catch_all);
7416
7417 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 assert_ne!(DatasetId::TACRED.description(), catch_all);
7424 assert_ne!(DatasetId::REBEL.description(), catch_all);
7425
7426 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 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 assert!(DatasetId::CHisIEC.is_historical());
7446 assert!(DatasetId::CHisIEC.is_relation_extraction());
7447
7448 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 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 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 assert_eq!(sentence.tokens.len(), 4);
7498
7499 assert_eq!(sentence.tokens[0].ner_tag, "B-PER"); assert_eq!(sentence.tokens[1].ner_tag, "I-PER"); assert_eq!(sentence.tokens[2].ner_tag, "O"); assert_eq!(sentence.tokens[3].ner_tag, "B-LOC"); }
7505
7506 #[test]
7507 fn test_chisiec_parse_relations() {
7508 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 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 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 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 assert_eq!(sentence.tokens[3].ner_tag, "O");
7576
7577 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 assert_eq!(sentence.tokens[7].ner_tag, "O");
7584
7585 assert_eq!(sentence.tokens[8].ner_tag, "B-BOOK");
7587 assert_eq!(sentence.tokens[9].ner_tag, "I-BOOK");
7588
7589 assert_eq!(sentence.tokens[10].ner_tag, "O");
7591
7592 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 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 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 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 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 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 assert!(DatasetId::CHisIEC.is_historical());
7666 assert!(DatasetId::HistoricalChineseNER.is_historical());
7667
7668 assert_ne!(DatasetId::CHisIEC, DatasetId::HistoricalChineseNER);
7670
7671 assert!(DatasetId::CHisIEC.is_relation_extraction());
7673
7674 let chisiec_types = DatasetId::CHisIEC.entity_types();
7678 assert!(chisiec_types.contains(&"OFI")); assert!(chisiec_types.contains(&"BOOK")); 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 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 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 use anno::schema::map_to_canonical;
7718
7719 let per_type = map_to_canonical("PER", None);
7721 assert_eq!(per_type, EntityType::Person);
7722
7723 let loc_type = map_to_canonical("LOC", None);
7725 assert_eq!(loc_type, EntityType::Location);
7726
7727 let ofi_type = map_to_canonical("OFI", None);
7729 assert!(matches!(ofi_type, EntityType::Custom { .. }));
7730
7731 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 assert_eq!(DatasetId::CHisIEC.language(), "lzh");
7740
7741 assert_eq!(DatasetId::CHisIEC.domain(), "historical");
7743
7744 assert_eq!(DatasetId::HistoricalChineseNER.language(), "zh");
7746 assert_eq!(DatasetId::HistoricalChineseNER.domain(), "historical");
7747 }
7748
7749 #[test]
7754 fn test_african_datasets_configured() {
7755 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 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 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 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 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 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 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 assert_eq!(dataset.sentences[0].tokens[0].ner_tag, "B-positive");
7816 assert_eq!(dataset.sentences[1].tokens[0].ner_tag, "B-negative");
7818 assert_eq!(dataset.sentences[2].tokens[0].ner_tag, "B-neutral");
7820 }
7821
7822 #[test]
7823 fn test_parse_masakhanews() {
7824 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 assert_eq!(dataset.sentences.len(), 2);
7836
7837 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 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 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 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 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 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 assert_eq!(dataset.sentences[0].tokens[4].text, "Ἀχιλῆος");
7926 }
7929
7930 #[test]
7931 fn test_registry_hints_cover_all_conllu_ner_datasets() {
7932 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 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 if url.is_empty() || !is_ner {
7972 continue;
7973 }
7974
7975 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 if !missing_hints.is_empty() {
7989 let known_missing: &[DatasetId] = &[
7991 ];
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 let loadable_count = LoadableDatasetId::all().len();
8008 let total_count = DatasetId::all().len();
8009
8010 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 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 if url.is_empty() {
8033 continue;
8034 }
8035
8036 if format.is_none()
8038 && access == crate::eval::dataset_registry::DatasetAccessibility::Public
8039 {
8040 missing_format.push(ds);
8041 }
8042 }
8043
8044 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 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 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 if ds.tasks_or_inferred().contains(&"blocked") {
8087 continue;
8088 }
8089
8090 if is_re || is_coref || is_event {
8092 continue;
8093 }
8094
8095 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 ¬_loadable {
8105 eprintln!(" - {:?} (format={})", ds, format);
8106 }
8107 }
8108
8109 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 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 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 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 let tokens = &dataset.sentences[0].tokens;
8171 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 #[test]
8195 fn test_parse_maven_jsonl() {
8196 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 assert!(
8341 entity_tokens.len() >= 3,
8342 "Should have at least 3 entity tokens, got {}",
8343 entity_tokens.len()
8344 );
8345
8346 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 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 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 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 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 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 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 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 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 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 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 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 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 assert_eq!(dataset.sentences.len(), 2);
8511
8512 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 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 #[test]
8550 fn test_parse_content_rejects_empty_for_all_loadable_datasets() {
8551 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 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 #[test]
8776 fn test_sec_filings_has_raw_url() {
8777 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 let format = DatasetId::TwiConv.format();
8795 assert!(format.is_some(), "TwiConv should have format field");
8796 assert_eq!(format.unwrap(), "CoNLL", "TwiConv should be CoNLL format");
8798 }
8799
8800 #[test]
8801 fn test_mudoco_has_format() {
8802 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 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 let test_cases = vec![
8835 (DatasetId::AncientGreekUD, "universaldependencies"),
8836 (DatasetId::LatinUD, "universaldependencies"),
8837 (DatasetId::SECFilingsNER, "entity-recognition-datasets"),
8838 (DatasetId::TwiConv, "twiconv"), ];
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 let loadable = LoadableDatasetId::all();
8860 let count = loadable.len();
8861
8862 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 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 let is_pure_ner = ds.supports_ner() && !ds.supports_coref() && !ds.supports_re();
8883
8884 if is_conll_variant && is_pure_ner {
8885 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 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 assert_eq!(
8933 ds.sentences.len(),
8934 3,
8935 "Expected 3 sentences, got {:?}",
8936 ds.sentences.len()
8937 );
8938
8939 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 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 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 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 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 assert!(
8991 LoadableDatasetId::is_loadable_dataset(DatasetId::ENer),
8992 "ENer should be loadable"
8993 );
8994 }
8995
8996 #[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 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 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 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 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 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 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 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 #[test]
9210 fn test_biomedical_conll_with_chemical_entities() {
9211 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 assert_eq!(dataset.sentences[0].tokens.len(), 3);
9523 }
9524
9525 #[test]
9526 fn test_conll_with_long_entity_spans() {
9527 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 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 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 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 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 #[test]
9654 fn test_conll_handles_windows_line_endings() {
9655 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 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 assert!(result.is_ok());
9675 }
9676
9677 #[test]
9678 fn test_conllu_handles_multiword_tokens() {
9679 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 assert!(dataset.sentences[0].tokens.len() >= 3);
9696 }
9697
9698 #[test]
9699 fn test_conllu_handles_empty_nodes() {
9700 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 assert_eq!(dataset.sentences[0].tokens.len(), 4);
9717 }
9718
9719 #[test]
9720 fn test_conll_with_bio_tag_normalization() {
9721 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 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 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 let loader = DatasetLoader::new().unwrap();
9764 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 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 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 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 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 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 for id in LoadableDatasetId::all() {
9856 let ds: DatasetId = id.into();
9857
9858 assert!(!ds.name().is_empty(), "{:?} has empty name", ds);
9860
9861 }
9864 }
9865}