anno_eval/eval/dataset/types.rs
1//! Core dataset types for NER evaluation.
2//!
3//! These types are the foundation of the dataset API, providing a uniform
4//! representation for annotated NER examples regardless of source.
5
6use crate::eval::GoldEntity;
7use anno::EntityType;
8use serde::{Deserialize, Serialize};
9
10// ============================================================================
11// Domain and Difficulty Classifications
12// ============================================================================
13
14/// Domain classification for NER examples.
15///
16/// Domains help organize datasets and enable domain-specific filtering/analysis.
17/// A single example belongs to exactly one domain.
18///
19/// # Example
20///
21/// ```rust
22/// use anno_eval::eval::dataset::Domain;
23///
24/// let domain = Domain::News;
25/// assert_eq!(format!("{:?}", domain), "News");
26/// ```
27#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize, Default)]
28#[non_exhaustive]
29pub enum Domain {
30 /// News articles (CoNLL-2003 style)
31 #[default]
32 News,
33 /// Social media text (WNUT style - noisy, informal)
34 SocialMedia,
35 /// Biomedical/clinical text (diseases, drugs, genes)
36 Biomedical,
37 /// Financial/business text (stocks, companies, money)
38 Financial,
39 /// Legal documents (contracts, court cases)
40 Legal,
41 /// Scientific/academic text (papers, abstracts)
42 Scientific,
43 /// Conversational text (dialogue, chat)
44 Conversational,
45 /// Technical documentation (code, manuals)
46 Technical,
47 /// Historical text (archaic language, historical figures)
48 Historical,
49 /// Sports reporting
50 Sports,
51 /// Entertainment news (movies, music, celebrities)
52 Entertainment,
53 /// Political news and discourse
54 Politics,
55 /// E-commerce (products, prices, brands)
56 Ecommerce,
57 /// Academic text (citations, institutions)
58 Academic,
59 /// Email communications
60 Email,
61 /// Weather reports and forecasts
62 Weather,
63 /// Travel content (destinations, hotels)
64 Travel,
65 /// Food and restaurant content
66 Food,
67 /// Real estate listings
68 RealEstate,
69 /// Cybersecurity (CVEs, malware, threat actors)
70 Cybersecurity,
71 /// Multilingual text with native scripts
72 Multilingual,
73}
74
75impl Domain {
76 /// Returns all available domain variants.
77 pub fn all() -> &'static [Domain] {
78 &[
79 Domain::News,
80 Domain::SocialMedia,
81 Domain::Biomedical,
82 Domain::Financial,
83 Domain::Legal,
84 Domain::Scientific,
85 Domain::Conversational,
86 Domain::Technical,
87 Domain::Historical,
88 Domain::Sports,
89 Domain::Entertainment,
90 Domain::Politics,
91 Domain::Ecommerce,
92 Domain::Academic,
93 Domain::Email,
94 Domain::Weather,
95 Domain::Travel,
96 Domain::Food,
97 Domain::RealEstate,
98 Domain::Cybersecurity,
99 Domain::Multilingual,
100 ]
101 }
102
103 /// Human-readable name for the domain.
104 pub fn name(&self) -> &'static str {
105 match self {
106 Domain::News => "News",
107 Domain::SocialMedia => "Social Media",
108 Domain::Biomedical => "Biomedical",
109 Domain::Financial => "Financial",
110 Domain::Legal => "Legal",
111 Domain::Scientific => "Scientific",
112 Domain::Conversational => "Conversational",
113 Domain::Technical => "Technical",
114 Domain::Historical => "Historical",
115 Domain::Sports => "Sports",
116 Domain::Entertainment => "Entertainment",
117 Domain::Politics => "Politics",
118 Domain::Ecommerce => "E-commerce",
119 Domain::Academic => "Academic",
120 Domain::Email => "Email",
121 Domain::Weather => "Weather",
122 Domain::Travel => "Travel",
123 Domain::Food => "Food",
124 Domain::RealEstate => "Real Estate",
125 Domain::Cybersecurity => "Cybersecurity",
126 Domain::Multilingual => "Multilingual",
127 }
128 }
129}
130
131/// Difficulty level for NER examples.
132///
133/// Difficulty is subjective but useful for progressive evaluation:
134/// - Start with Easy to verify basic functionality
135/// - Progress to Medium for realistic performance estimates
136/// - Use Hard/Adversarial for stress testing
137///
138/// # Example
139///
140/// ```rust
141/// use anno_eval::eval::dataset::Difficulty;
142///
143/// let difficulty = Difficulty::Hard;
144/// assert!(difficulty.is_challenging());
145/// ```
146#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize, Default)]
147pub enum Difficulty {
148 /// Simple entities, clear context, unambiguous
149 #[default]
150 Easy,
151 /// Multiple entities, some ambiguity, typical real-world
152 Medium,
153 /// Complex sentences, nested entities, domain jargon
154 Hard,
155 /// Edge cases, adversarial examples, intentionally tricky
156 Adversarial,
157}
158
159impl Difficulty {
160 /// Returns all difficulty levels in order.
161 pub fn all() -> &'static [Difficulty] {
162 &[
163 Difficulty::Easy,
164 Difficulty::Medium,
165 Difficulty::Hard,
166 Difficulty::Adversarial,
167 ]
168 }
169
170 /// Returns true if this is Hard or Adversarial.
171 pub fn is_challenging(&self) -> bool {
172 matches!(self, Difficulty::Hard | Difficulty::Adversarial)
173 }
174
175 /// Numeric difficulty score (0-3) for sorting/filtering.
176 pub fn score(&self) -> u8 {
177 match self {
178 Difficulty::Easy => 0,
179 Difficulty::Medium => 1,
180 Difficulty::Hard => 2,
181 Difficulty::Adversarial => 3,
182 }
183 }
184}
185
186// ============================================================================
187// Annotated Example
188// ============================================================================
189
190/// A single annotated NER example with text, entities, and metadata.
191///
192/// This is the canonical type for representing labeled NER data, whether
193/// from synthetic generation, manual annotation, or loaded from files.
194///
195/// # Example
196///
197/// ```rust
198/// use anno_eval::eval::dataset::{AnnotatedExample, Domain, Difficulty};
199/// use anno_eval::eval::GoldEntity;
200/// use anno::EntityType;
201///
202/// let example = AnnotatedExample::new(
203/// "Microsoft announced earnings",
204/// vec![GoldEntity::new("Microsoft", EntityType::Organization, 0)],
205/// );
206/// assert_eq!(example.text, "Microsoft announced earnings");
207/// assert_eq!(example.entities.len(), 1);
208/// ```
209#[derive(Debug, Clone, Serialize, Deserialize)]
210pub struct AnnotatedExample {
211 /// The input text.
212 pub text: String,
213 /// Gold standard entity annotations.
214 pub entities: Vec<GoldEntity>,
215 /// Domain classification (e.g., News, Biomedical).
216 pub domain: Domain,
217 /// Difficulty level (e.g., Easy, Hard).
218 pub difficulty: Difficulty,
219}
220
221impl AnnotatedExample {
222 /// Create a new annotated example with default domain/difficulty.
223 pub fn new(text: impl Into<String>, entities: Vec<GoldEntity>) -> Self {
224 Self {
225 text: text.into(),
226 entities,
227 domain: Domain::default(),
228 difficulty: Difficulty::default(),
229 }
230 }
231
232 /// Convenience constructor from text and entity tuples (alias for `from_tuples`).
233 ///
234 /// This is the same as `from_tuples` but with the shorter name commonly used
235 /// in test code.
236 #[must_use]
237 pub fn simple(text: impl Into<String>, entities: Vec<(&str, &str)>) -> Self {
238 Self::from_tuples(text, entities)
239 }
240
241 /// Create with explicit domain and difficulty.
242 pub fn with_metadata(
243 text: impl Into<String>,
244 entities: Vec<GoldEntity>,
245 domain: Domain,
246 difficulty: Difficulty,
247 ) -> Self {
248 Self {
249 text: text.into(),
250 entities,
251 domain,
252 difficulty,
253 }
254 }
255
256 /// Convenience constructor for doctests and simple examples.
257 ///
258 /// Entity positions are computed by finding each entity text within the input.
259 ///
260 /// # Panics
261 ///
262 /// Panics if an entity text is not found in the input text.
263 ///
264 /// # Example
265 ///
266 /// ```rust
267 /// use anno_eval::eval::dataset::AnnotatedExample;
268 ///
269 /// let example = AnnotatedExample::from_tuples(
270 /// "John works at Google.",
271 /// vec![("John", "PER"), ("Google", "ORG")],
272 /// );
273 /// assert_eq!(example.entities.len(), 2);
274 /// ```
275 pub fn from_tuples(text: impl Into<String>, entities: Vec<(&str, &str)>) -> Self {
276 let text = text.into();
277 let gold_entities = entities
278 .into_iter()
279 .map(|(entity_text, entity_type_str)| {
280 let start = text.find(entity_text).unwrap_or_else(|| {
281 panic!("Entity '{}' not found in text '{}'", entity_text, text)
282 });
283 let entity_type = EntityType::from_label(entity_type_str);
284 GoldEntity::new(entity_text, entity_type, start)
285 })
286 .collect();
287
288 Self {
289 text,
290 entities: gold_entities,
291 domain: Domain::default(),
292 difficulty: Difficulty::default(),
293 }
294 }
295
296 /// Set the domain and return self (builder pattern).
297 pub fn with_domain(mut self, domain: Domain) -> Self {
298 self.domain = domain;
299 self
300 }
301
302 /// Set the difficulty and return self (builder pattern).
303 pub fn with_difficulty(mut self, difficulty: Difficulty) -> Self {
304 self.difficulty = difficulty;
305 self
306 }
307
308 /// Returns true if this example has no entities (negative example).
309 pub fn is_negative(&self) -> bool {
310 self.entities.is_empty()
311 }
312
313 /// Returns the number of entities.
314 pub fn entity_count(&self) -> usize {
315 self.entities.len()
316 }
317
318 /// Returns unique entity types in this example.
319 pub fn entity_types(&self) -> Vec<&EntityType> {
320 let mut types: Vec<_> = self.entities.iter().map(|e| &e.entity_type).collect();
321 types.sort_by_key(|t| format!("{:?}", t));
322 types.dedup();
323 types
324 }
325
326 /// Convert to (text, entities) tuple for evaluation functions.
327 pub fn as_test_case(&self) -> (&str, &[GoldEntity]) {
328 (&self.text, &self.entities)
329 }
330
331 /// Consume and convert to owned (text, entities) tuple.
332 pub fn into_test_case(self) -> (String, Vec<GoldEntity>) {
333 (self.text, self.entities)
334 }
335}
336
337// ============================================================================
338// Entity Helpers
339// ============================================================================
340
341/// Helper module for creating entities concisely in dataset definitions.
342///
343/// These functions are internal conveniences; external code should use
344/// `GoldEntity::new()` directly.
345pub(crate) mod helpers {
346 use super::*;
347
348 /// Create a standard entity.
349 pub fn entity(text: &str, entity_type: EntityType, start: usize) -> GoldEntity {
350 GoldEntity::new(text, entity_type, start)
351 }
352
353 /// Create a disease entity (biomedical domain).
354 pub fn disease(text: &str, start: usize) -> GoldEntity {
355 GoldEntity::new(
356 text,
357 EntityType::Custom {
358 name: "DISEASE".to_string(),
359 category: anno::EntityCategory::Misc,
360 },
361 start,
362 )
363 }
364
365 /// Create a drug entity (biomedical domain).
366 pub fn drug(text: &str, start: usize) -> GoldEntity {
367 GoldEntity::new(
368 text,
369 EntityType::Custom {
370 name: "DRUG".to_string(),
371 category: anno::EntityCategory::Misc,
372 },
373 start,
374 )
375 }
376
377 /// Create a gene entity (biomedical domain).
378 pub fn gene(text: &str, start: usize) -> GoldEntity {
379 GoldEntity::new(
380 text,
381 EntityType::Custom {
382 name: "GENE".to_string(),
383 category: anno::EntityCategory::Misc,
384 },
385 start,
386 )
387 }
388
389 /// Create a chemical entity (biomedical domain).
390 pub fn chemical(text: &str, start: usize) -> GoldEntity {
391 GoldEntity::new(
392 text,
393 EntityType::Custom {
394 name: "CHEMICAL".to_string(),
395 category: anno::EntityCategory::Misc,
396 },
397 start,
398 )
399 }
400
401 /// Create an email entity.
402 pub fn entity_email(text: &str, start: usize) -> GoldEntity {
403 GoldEntity::new(
404 text,
405 EntityType::Custom {
406 name: "EMAIL".to_string(),
407 category: anno::EntityCategory::Contact,
408 },
409 start,
410 )
411 }
412
413 /// Create a URL entity.
414 pub fn entity_url(text: &str, start: usize) -> GoldEntity {
415 GoldEntity::new(
416 text,
417 EntityType::Custom {
418 name: "URL".to_string(),
419 category: anno::EntityCategory::Misc,
420 },
421 start,
422 )
423 }
424
425 /// Create a phone number entity.
426 pub fn entity_phone(text: &str, start: usize) -> GoldEntity {
427 GoldEntity::new(
428 text,
429 EntityType::Custom {
430 name: "PHONE".to_string(),
431 category: anno::EntityCategory::Contact,
432 },
433 start,
434 )
435 }
436}
437
438#[cfg(test)]
439mod tests {
440 use super::*;
441
442 #[test]
443 fn test_domain_all() {
444 let domains = Domain::all();
445 assert!(domains.len() >= 20);
446 assert!(domains.contains(&Domain::News));
447 assert!(domains.contains(&Domain::Biomedical));
448 }
449
450 #[test]
451 fn test_difficulty_ordering() {
452 assert!(Difficulty::Easy.score() < Difficulty::Medium.score());
453 assert!(Difficulty::Medium.score() < Difficulty::Hard.score());
454 assert!(Difficulty::Hard.score() < Difficulty::Adversarial.score());
455 }
456
457 #[test]
458 fn test_annotated_example_from_tuples() {
459 let example = AnnotatedExample::from_tuples(
460 "John works at Google in NYC.",
461 vec![("John", "PER"), ("Google", "ORG"), ("NYC", "LOC")],
462 );
463 assert_eq!(example.entities.len(), 3);
464 assert_eq!(example.entities[0].text, "John");
465 assert_eq!(example.entities[0].start, 0);
466 assert_eq!(example.entities[1].text, "Google");
467 assert_eq!(example.entities[1].start, 14);
468 }
469
470 #[test]
471 fn test_annotated_example_builder() {
472 let example = AnnotatedExample::new("Test text", vec![])
473 .with_domain(Domain::Biomedical)
474 .with_difficulty(Difficulty::Hard);
475
476 assert_eq!(example.domain, Domain::Biomedical);
477 assert_eq!(example.difficulty, Difficulty::Hard);
478 }
479
480 #[test]
481 fn test_is_negative() {
482 let positive = AnnotatedExample::from_tuples("John is here", vec![("John", "PER")]);
483 let negative = AnnotatedExample::new("No entities here", vec![]);
484
485 assert!(!positive.is_negative());
486 assert!(negative.is_negative());
487 }
488}