anno_eval/eval/
bio_adapter.rs

1//! BIO tag sequence adapter.
2//!
3//! This module provides utilities for converting between BIO-tagged sequences
4//! and entity spans.
5//!
6//! # Supported Schemes
7//!
8//! - IOB1: Inside-Outside-Begin (I appears first, B only when needed)
9//! - IOB2: Inside-Outside-Begin (B always starts entity) - **most common**
10//! - IOE1: Inside-Outside-End (E appears last, I continues)
11//! - IOE2: Inside-Outside-End (E always ends entity)
12//! - IOBES/BILOU: Begin-Inside-Last-Outside-Unit
13//!
14//! # Example
15//!
16//! ```rust
17//! use anno_eval::eval::bio_adapter::{BioScheme, bio_to_entities, entities_to_bio};
18//! use anno::{Entity, EntityType};
19//!
20//! // Convert BIO tags to entities
21//! let tokens = ["John", "Smith", "works", "at", "Apple"];
22//! let tags = ["B-PER", "I-PER", "O", "O", "B-ORG"];
23//! let entities = bio_to_entities(&tokens, &tags, BioScheme::IOB2).unwrap();
24//!
25//! assert_eq!(entities.len(), 2);
26//! assert_eq!(entities[0].text, "John Smith");
27//! assert_eq!(entities[1].text, "Apple");
28//! ```
29
30use anno::{Entity, EntityType, Result};
31use std::fmt;
32
33/// BIO tagging scheme.
34#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
35pub enum BioScheme {
36    /// IOB1: I appears first, B only when two entities of same type are adjacent
37    IOB1,
38    /// IOB2: B always starts an entity (most common, seqeval default)
39    #[default]
40    IOB2,
41    /// IOE1: E appears last, I continues
42    IOE1,
43    /// IOE2: E always ends an entity
44    IOE2,
45    /// IOBES/BILOU: Begin-Inside-Last-Outside-Unit (single-token entities use U/S)
46    IOBES,
47}
48
49impl fmt::Display for BioScheme {
50    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
51        match self {
52            BioScheme::IOB1 => write!(f, "IOB1"),
53            BioScheme::IOB2 => write!(f, "IOB2"),
54            BioScheme::IOE1 => write!(f, "IOE1"),
55            BioScheme::IOE2 => write!(f, "IOE2"),
56            BioScheme::IOBES => write!(f, "IOBES/BILOU"),
57        }
58    }
59}
60
61/// A parsed BIO tag.
62#[derive(Debug, Clone)]
63struct ParsedTag {
64    prefix: char,
65    entity_type: Option<String>,
66}
67
68impl ParsedTag {
69    fn parse(tag: &str) -> Self {
70        if tag == "O" || tag == "o" {
71            return Self {
72                prefix: 'O',
73                entity_type: None,
74            };
75        }
76
77        // Handle B-PER, I-LOC, etc.
78        if tag.len() >= 2 && (tag.chars().nth(1) == Some('-') || tag.chars().nth(1) == Some('_')) {
79            let prefix = tag.chars().next().unwrap_or('O').to_ascii_uppercase();
80            let entity_type = tag[2..].to_string();
81            return Self {
82                prefix,
83                entity_type: Some(entity_type),
84            };
85        }
86
87        // Fallback: treat as O
88        Self {
89            prefix: 'O',
90            entity_type: None,
91        }
92    }
93
94    fn is_outside(&self) -> bool {
95        self.prefix == 'O'
96    }
97
98    fn is_begin(&self) -> bool {
99        self.prefix == 'B'
100    }
101
102    fn is_inside(&self) -> bool {
103        self.prefix == 'I'
104    }
105
106    fn is_end(&self) -> bool {
107        self.prefix == 'E' || self.prefix == 'L'
108    }
109
110    fn is_single(&self) -> bool {
111        self.prefix == 'S' || self.prefix == 'U'
112    }
113}
114
115/// Convert BIO-tagged tokens to entity spans.
116///
117/// # Arguments
118///
119/// * `tokens` - Slice of token strings
120/// * `tags` - Slice of BIO tags (same length as tokens)
121/// * `scheme` - BIO tagging scheme to use
122///
123/// # Returns
124///
125/// Vector of Entity spans with character offsets computed from tokens.
126///
127/// # Example
128///
129/// ```rust
130/// use anno_eval::eval::bio_adapter::{BioScheme, bio_to_entities};
131///
132/// let tokens = ["The", "United", "Nations", "met", "today"];
133/// let tags = ["O", "B-ORG", "I-ORG", "O", "O"];
134///
135/// let entities = bio_to_entities(&tokens, &tags, BioScheme::IOB2).unwrap();
136/// assert_eq!(entities.len(), 1);
137/// assert_eq!(entities[0].text, "United Nations");
138/// ```
139pub fn bio_to_entities<S: AsRef<str>>(
140    tokens: &[S],
141    tags: &[S],
142    scheme: BioScheme,
143) -> Result<Vec<Entity>> {
144    if tokens.len() != tags.len() {
145        return Err(crate::Error::invalid_input(format!(
146            "Token count ({}) != tag count ({})",
147            tokens.len(),
148            tags.len()
149        )));
150    }
151
152    // Compute character offsets for each token
153    let mut offsets = Vec::with_capacity(tokens.len());
154    let mut current_offset = 0;
155    for token in tokens {
156        let token_str = token.as_ref();
157        offsets.push((current_offset, current_offset + token_str.len()));
158        current_offset += token_str.len() + 1; // +1 for space
159    }
160
161    let mut entities = Vec::new();
162    let mut current_entity: Option<(usize, String)> = None; // (start_idx, type)
163
164    for (i, tag_str) in tags.iter().enumerate() {
165        let tag = ParsedTag::parse(tag_str.as_ref());
166
167        match scheme {
168            BioScheme::IOB2 => {
169                if tag.is_begin() || tag.is_single() {
170                    // Finish previous entity if any
171                    if let Some((start_idx, entity_type)) = current_entity.take() {
172                        entities.push(build_entity(
173                            tokens,
174                            &offsets,
175                            start_idx,
176                            i - 1,
177                            &entity_type,
178                        ));
179                    }
180                    // Start new entity
181                    current_entity = tag.entity_type.clone().map(|t| (i, t));
182
183                    // Single-token entity in IOBES mode
184                    if tag.is_single() {
185                        if let Some((start_idx, entity_type)) = current_entity.take() {
186                            entities.push(build_entity(
187                                tokens,
188                                &offsets,
189                                start_idx,
190                                i,
191                                &entity_type,
192                            ));
193                        }
194                    }
195                } else if tag.is_inside() {
196                    // Continue entity if types match
197                    if let Some((_, ref current_type)) = current_entity {
198                        if tag.entity_type.as_ref() != Some(current_type) {
199                            // Type mismatch - close current and start new
200                            if let Some((start_idx, entity_type)) = current_entity.take() {
201                                entities.push(build_entity(
202                                    tokens,
203                                    &offsets,
204                                    start_idx,
205                                    i - 1,
206                                    &entity_type,
207                                ));
208                            }
209                            current_entity = tag.entity_type.clone().map(|t| (i, t));
210                        }
211                    } else {
212                        // I without B - start new entity (lenient)
213                        current_entity = tag.entity_type.clone().map(|t| (i, t));
214                    }
215                } else if tag.is_end() {
216                    // Close entity
217                    if let Some((start_idx, entity_type)) = current_entity.take() {
218                        entities.push(build_entity(tokens, &offsets, start_idx, i, &entity_type));
219                    }
220                } else if tag.is_outside() {
221                    // Close any open entity
222                    if let Some((start_idx, entity_type)) = current_entity.take() {
223                        entities.push(build_entity(
224                            tokens,
225                            &offsets,
226                            start_idx,
227                            i - 1,
228                            &entity_type,
229                        ));
230                    }
231                }
232            }
233            BioScheme::IOB1 => {
234                // IOB1: B only appears between adjacent same-type entities
235                if tag.is_begin() {
236                    if let Some((start_idx, entity_type)) = current_entity.take() {
237                        entities.push(build_entity(
238                            tokens,
239                            &offsets,
240                            start_idx,
241                            i - 1,
242                            &entity_type,
243                        ));
244                    }
245                    current_entity = tag.entity_type.clone().map(|t| (i, t));
246                } else if tag.is_inside() {
247                    if current_entity.is_none()
248                        || current_entity.as_ref().map(|(_, t)| t) != tag.entity_type.as_ref()
249                    {
250                        // New entity starts with I in IOB1
251                        if let Some((start_idx, entity_type)) = current_entity.take() {
252                            entities.push(build_entity(
253                                tokens,
254                                &offsets,
255                                start_idx,
256                                i - 1,
257                                &entity_type,
258                            ));
259                        }
260                        current_entity = tag.entity_type.clone().map(|t| (i, t));
261                    }
262                } else if tag.is_outside() {
263                    if let Some((start_idx, entity_type)) = current_entity.take() {
264                        entities.push(build_entity(
265                            tokens,
266                            &offsets,
267                            start_idx,
268                            i - 1,
269                            &entity_type,
270                        ));
271                    }
272                }
273            }
274            BioScheme::IOBES => {
275                if tag.is_begin() {
276                    if let Some((start_idx, entity_type)) = current_entity.take() {
277                        entities.push(build_entity(
278                            tokens,
279                            &offsets,
280                            start_idx,
281                            i - 1,
282                            &entity_type,
283                        ));
284                    }
285                    current_entity = tag.entity_type.clone().map(|t| (i, t));
286                } else if tag.is_inside() {
287                    // Continue
288                } else if tag.is_end() {
289                    if let Some((start_idx, entity_type)) = current_entity.take() {
290                        entities.push(build_entity(tokens, &offsets, start_idx, i, &entity_type));
291                    }
292                } else if tag.is_single() {
293                    if let Some((start_idx, entity_type)) = current_entity.take() {
294                        entities.push(build_entity(
295                            tokens,
296                            &offsets,
297                            start_idx,
298                            i - 1,
299                            &entity_type,
300                        ));
301                    }
302                    if let Some(t) = tag.entity_type.clone() {
303                        entities.push(build_entity(tokens, &offsets, i, i, &t));
304                    }
305                } else if tag.is_outside() {
306                    if let Some((start_idx, entity_type)) = current_entity.take() {
307                        entities.push(build_entity(
308                            tokens,
309                            &offsets,
310                            start_idx,
311                            i - 1,
312                            &entity_type,
313                        ));
314                    }
315                }
316            }
317            // IOE1/IOE2 similar logic but ending-focused
318            BioScheme::IOE1 | BioScheme::IOE2 => {
319                if tag.is_inside() || tag.is_begin() {
320                    if current_entity.is_none() {
321                        current_entity = tag.entity_type.clone().map(|t| (i, t));
322                    }
323                } else if tag.is_end() {
324                    if current_entity.is_none() {
325                        current_entity = tag.entity_type.clone().map(|t| (i, t));
326                    }
327                    if let Some((start_idx, entity_type)) = current_entity.take() {
328                        entities.push(build_entity(tokens, &offsets, start_idx, i, &entity_type));
329                    }
330                } else if tag.is_outside() {
331                    if let Some((start_idx, entity_type)) = current_entity.take() {
332                        entities.push(build_entity(
333                            tokens,
334                            &offsets,
335                            start_idx,
336                            i - 1,
337                            &entity_type,
338                        ));
339                    }
340                }
341            }
342        }
343    }
344
345    // Close any remaining entity
346    if let Some((start_idx, entity_type)) = current_entity {
347        entities.push(build_entity(
348            tokens,
349            &offsets,
350            start_idx,
351            tokens.len() - 1,
352            &entity_type,
353        ));
354    }
355
356    Ok(entities)
357}
358
359/// Build an Entity from token range.
360fn build_entity<S: AsRef<str>>(
361    tokens: &[S],
362    offsets: &[(usize, usize)],
363    start_idx: usize,
364    end_idx: usize,
365    entity_type: &str,
366) -> Entity {
367    let text: String = tokens[start_idx..=end_idx]
368        .iter()
369        .map(|t| t.as_ref())
370        .collect::<Vec<_>>()
371        .join(" ");
372
373    let char_start = offsets[start_idx].0;
374    let char_end = offsets[end_idx].1;
375
376    Entity::new(
377        &text,
378        string_to_entity_type(entity_type),
379        char_start,
380        char_end,
381        1.0,
382    )
383}
384
385/// Convert entity type string to EntityType.
386fn string_to_entity_type(s: &str) -> EntityType {
387    match s.to_uppercase().as_str() {
388        "PER" | "PERSON" => EntityType::Person,
389        "ORG" | "ORGANIZATION" => EntityType::Organization,
390        "LOC" | "LOCATION" | "GPE" => EntityType::Location,
391        "MISC" | "MISCELLANEOUS" => EntityType::custom("MISC", anno::EntityCategory::Misc),
392        "DATE" => EntityType::Date,
393        "TIME" => EntityType::Time,
394        "MONEY" | "CURRENCY" => EntityType::Money,
395        "PERCENT" | "PERCENTAGE" => EntityType::Percent,
396        other => EntityType::custom(other, anno::EntityCategory::Misc),
397    }
398}
399
400/// Convert entities back to BIO tags.
401///
402/// # Arguments
403///
404/// * `text` - The original text
405/// * `tokens` - Token boundaries as (start, end) character offsets
406/// * `entities` - Entities to convert
407/// * `scheme` - BIO scheme to use
408///
409/// # Returns
410///
411/// Vector of BIO tags, one per token.
412pub fn entities_to_bio(
413    tokens: &[(usize, usize)],
414    entities: &[Entity],
415    scheme: BioScheme,
416) -> Vec<String> {
417    let mut tags = vec!["O".to_string(); tokens.len()];
418
419    for entity in entities {
420        let type_label = entity.entity_type.as_label().to_uppercase();
421
422        // Find tokens that overlap with this entity
423        let mut entity_tokens: Vec<usize> = Vec::new();
424        for (i, &(tok_start, tok_end)) in tokens.iter().enumerate() {
425            if tok_start < entity.end() && tok_end > entity.start() {
426                entity_tokens.push(i);
427            }
428        }
429
430        if entity_tokens.is_empty() {
431            continue;
432        }
433
434        match scheme {
435            BioScheme::IOB2 => {
436                for (j, &tok_idx) in entity_tokens.iter().enumerate() {
437                    tags[tok_idx] = if j == 0 {
438                        format!("B-{}", type_label)
439                    } else {
440                        format!("I-{}", type_label)
441                    };
442                }
443            }
444            BioScheme::IOB1 => {
445                // B only if previous token was same type
446                for (j, &tok_idx) in entity_tokens.iter().enumerate() {
447                    let needs_b = j == 0
448                        && tok_idx > 0
449                        && tags[tok_idx - 1].ends_with(&format!("-{}", type_label));
450                    tags[tok_idx] = if needs_b {
451                        format!("B-{}", type_label)
452                    } else {
453                        format!("I-{}", type_label)
454                    };
455                }
456            }
457            BioScheme::IOBES => {
458                let len = entity_tokens.len();
459                for (j, &tok_idx) in entity_tokens.iter().enumerate() {
460                    tags[tok_idx] = if len == 1 {
461                        format!("S-{}", type_label)
462                    } else if j == 0 {
463                        format!("B-{}", type_label)
464                    } else if j == len - 1 {
465                        format!("E-{}", type_label)
466                    } else {
467                        format!("I-{}", type_label)
468                    };
469                }
470            }
471            BioScheme::IOE2 => {
472                let len = entity_tokens.len();
473                for (j, &tok_idx) in entity_tokens.iter().enumerate() {
474                    tags[tok_idx] = if j == len - 1 {
475                        format!("E-{}", type_label)
476                    } else {
477                        format!("I-{}", type_label)
478                    };
479                }
480            }
481            BioScheme::IOE1 => {
482                let len = entity_tokens.len();
483                for (j, &tok_idx) in entity_tokens.iter().enumerate() {
484                    // E only if next token is same type
485                    let needs_e = j == len - 1
486                        && tok_idx + 1 < tokens.len()
487                        && tags
488                            .get(tok_idx + 1)
489                            .map(|t| t.ends_with(&format!("-{}", type_label)))
490                            .unwrap_or(false);
491                    tags[tok_idx] = if needs_e {
492                        format!("E-{}", type_label)
493                    } else {
494                        format!("I-{}", type_label)
495                    };
496                }
497            }
498        }
499    }
500
501    tags
502}
503
504/// Validate BIO tag sequence for a given scheme.
505///
506/// Returns errors for invalid transitions (e.g., O -> I in strict IOB2).
507///
508/// Invalid transitions are a common issue in NER model outputs, particularly
509/// when not using CRF layers for constraint enforcement during training.
510pub fn validate_bio_sequence<S: AsRef<str>>(tags: &[S], scheme: BioScheme) -> Vec<String> {
511    let mut errors = Vec::new();
512    let mut prev_tag = ParsedTag {
513        prefix: 'O',
514        entity_type: None,
515    };
516
517    for (i, tag_str) in tags.iter().enumerate() {
518        let tag = ParsedTag::parse(tag_str.as_ref());
519
520        match scheme {
521            // I must follow B or I of same type
522            BioScheme::IOB2 if tag.is_inside() => {
523                if prev_tag.is_outside() {
524                    errors.push(format!(
525                        "Position {}: I-{} follows O (should be B-{})",
526                        i,
527                        tag.entity_type.as_deref().unwrap_or("?"),
528                        tag.entity_type.as_deref().unwrap_or("?")
529                    ));
530                } else if tag.entity_type != prev_tag.entity_type {
531                    errors.push(format!(
532                        "Position {}: I-{} follows {}-{} (type mismatch)",
533                        i,
534                        tag.entity_type.as_deref().unwrap_or("?"),
535                        prev_tag.prefix,
536                        prev_tag.entity_type.as_deref().unwrap_or("?")
537                    ));
538                }
539            }
540            BioScheme::IOBES => {
541                // E/L must follow B or I of same type
542                if tag.is_end() && !prev_tag.is_begin() && !prev_tag.is_inside() {
543                    errors.push(format!(
544                        "Position {}: E-{} without preceding B or I",
545                        i,
546                        tag.entity_type.as_deref().unwrap_or("?")
547                    ));
548                }
549                // I must follow B or I
550                if tag.is_inside() && !prev_tag.is_begin() && !prev_tag.is_inside() {
551                    errors.push(format!(
552                        "Position {}: I-{} without preceding B or I",
553                        i,
554                        tag.entity_type.as_deref().unwrap_or("?")
555                    ));
556                }
557            }
558            _ => {} // IOB1, IOE1, IOE2 are more lenient
559        }
560
561        prev_tag = tag;
562    }
563
564    errors
565}
566
567/// Repair strategy for invalid BIO sequences.
568#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
569pub enum RepairStrategy {
570    /// Convert invalid I tags to B tags (most common approach).
571    #[default]
572    PromoteToBegin,
573    /// Discard invalid transitions by converting to O.
574    Discard,
575    /// Keep invalid tags as-is (lenient parsing).
576    Lenient,
577}
578
579/// Repair invalid BIO tag sequences.
580///
581/// Production NER systems often produce invalid sequences (e.g., O followed by I).
582/// This function repairs such sequences according to the chosen strategy.
583///
584/// # Strategies
585///
586/// - `PromoteToBegin`: Convert orphan I tags to B tags (recommended)
587/// - `Discard`: Convert invalid tags to O
588/// - `Lenient`: Keep as-is (caller handles in parsing)
589///
590/// # Example
591///
592/// ```rust
593/// use anno_eval::eval::bio_adapter::{repair_bio_sequence, RepairStrategy, BioScheme};
594///
595/// let invalid = vec!["O", "I-PER", "I-PER", "O"];  // Invalid: O->I
596/// let repaired = repair_bio_sequence(&invalid, BioScheme::IOB2, RepairStrategy::PromoteToBegin);
597/// assert_eq!(repaired, vec!["O", "B-PER", "I-PER", "O"]);  // Fixed
598/// ```
599pub fn repair_bio_sequence<S: AsRef<str>>(
600    tags: &[S],
601    scheme: BioScheme,
602    strategy: RepairStrategy,
603) -> Vec<String> {
604    if strategy == RepairStrategy::Lenient {
605        return tags.iter().map(|t| t.as_ref().to_string()).collect();
606    }
607
608    let mut result: Vec<String> = Vec::with_capacity(tags.len());
609    let mut prev_tag = ParsedTag {
610        prefix: 'O',
611        entity_type: None,
612    };
613
614    for tag_str in tags {
615        let tag = ParsedTag::parse(tag_str.as_ref());
616        let mut repaired = tag_str.as_ref().to_string();
617
618        match scheme {
619            BioScheme::IOB2 if tag.is_inside() => {
620                let needs_repair = prev_tag.is_outside() || tag.entity_type != prev_tag.entity_type;
621
622                if needs_repair {
623                    match strategy {
624                        RepairStrategy::PromoteToBegin => {
625                            if let Some(ref t) = tag.entity_type {
626                                repaired = format!("B-{}", t);
627                            }
628                        }
629                        RepairStrategy::Discard => {
630                            repaired = "O".to_string();
631                        }
632                        RepairStrategy::Lenient => {}
633                    }
634                }
635            }
636            BioScheme::IOBES
637                if (tag.is_inside() || tag.is_end())
638                    && !prev_tag.is_begin()
639                    && !prev_tag.is_inside() =>
640            {
641                match strategy {
642                    RepairStrategy::PromoteToBegin => {
643                        if let Some(ref t) = tag.entity_type {
644                            // If single invalid I or E, make it S (single)
645                            repaired = format!("S-{}", t);
646                        }
647                    }
648                    RepairStrategy::Discard => {
649                        repaired = "O".to_string();
650                    }
651                    RepairStrategy::Lenient => {}
652                }
653            }
654            _ => {} // Other schemes more lenient
655        }
656
657        prev_tag = ParsedTag::parse(&repaired);
658        result.push(repaired);
659    }
660
661    result
662}
663
664#[cfg(test)]
665mod tests {
666    use super::*;
667
668    #[test]
669    fn test_iob2_basic() {
670        let tokens = ["John", "Smith", "works", "at", "Apple"];
671        let tags = ["B-PER", "I-PER", "O", "O", "B-ORG"];
672
673        let entities =
674            bio_to_entities(&tokens, &tags, BioScheme::IOB2).expect("valid BIO tags should parse");
675
676        assert_eq!(entities.len(), 2);
677        assert_eq!(entities[0].text, "John Smith");
678        assert_eq!(entities[0].entity_type, EntityType::Person);
679        assert_eq!(entities[1].text, "Apple");
680        assert_eq!(entities[1].entity_type, EntityType::Organization);
681    }
682
683    #[test]
684    fn test_iob2_adjacent_same_type() {
685        let tokens = ["John", "and", "Mary"];
686        let tags = ["B-PER", "O", "B-PER"];
687
688        let entities =
689            bio_to_entities(&tokens, &tags, BioScheme::IOB2).expect("valid BIO tags should parse");
690
691        assert_eq!(entities.len(), 2);
692        assert_eq!(entities[0].text, "John");
693        assert_eq!(entities[1].text, "Mary");
694    }
695
696    #[test]
697    fn test_iob2_multi_token_org() {
698        let tokens = ["The", "United", "Nations", "Security", "Council"];
699        let tags = ["O", "B-ORG", "I-ORG", "I-ORG", "I-ORG"];
700
701        let entities =
702            bio_to_entities(&tokens, &tags, BioScheme::IOB2).expect("valid BIO tags should parse");
703
704        assert_eq!(entities.len(), 1);
705        assert_eq!(entities[0].text, "United Nations Security Council");
706        assert_eq!(entities[0].entity_type, EntityType::Organization);
707    }
708
709    #[test]
710    fn test_iobes_single_token() {
711        let tokens = ["John", "works", "here"];
712        let tags = ["S-PER", "O", "O"];
713
714        let entities = bio_to_entities(&tokens, &tags, BioScheme::IOBES).unwrap();
715
716        assert_eq!(entities.len(), 1);
717        assert_eq!(entities[0].text, "John");
718    }
719
720    #[test]
721    fn test_iobes_bie_sequence() {
722        let tokens = ["New", "York", "City"];
723        let tags = ["B-LOC", "I-LOC", "E-LOC"];
724
725        let entities = bio_to_entities(&tokens, &tags, BioScheme::IOBES).unwrap();
726
727        assert_eq!(entities.len(), 1);
728        assert_eq!(entities[0].text, "New York City");
729    }
730
731    #[test]
732    fn test_validation_iob2() {
733        // Invalid: O -> I
734        let tags = ["O", "I-PER", "I-PER"];
735        let errors = validate_bio_sequence(&tags, BioScheme::IOB2);
736        assert!(!errors.is_empty());
737        assert!(errors[0].contains("follows O"));
738
739        // Valid: B -> I
740        let tags = ["B-PER", "I-PER", "O"];
741        let errors = validate_bio_sequence(&tags, BioScheme::IOB2);
742        assert!(errors.is_empty());
743    }
744
745    #[test]
746    fn test_validation_type_mismatch() {
747        // Invalid: I-LOC after B-PER
748        let tags = ["B-PER", "I-LOC"];
749        let errors = validate_bio_sequence(&tags, BioScheme::IOB2);
750        assert!(!errors.is_empty());
751        assert!(errors[0].contains("type mismatch"));
752    }
753
754    #[test]
755    fn test_repair_promote_to_begin() {
756        let invalid = vec!["O", "I-PER", "I-PER", "O"];
757        let repaired =
758            repair_bio_sequence(&invalid, BioScheme::IOB2, RepairStrategy::PromoteToBegin);
759        assert_eq!(repaired, vec!["O", "B-PER", "I-PER", "O"]);
760    }
761
762    #[test]
763    fn test_repair_discard() {
764        let invalid = vec!["O", "I-PER", "I-PER", "O"];
765        let repaired = repair_bio_sequence(&invalid, BioScheme::IOB2, RepairStrategy::Discard);
766        // First I-PER becomes O (orphan), second I-PER also becomes O (no valid predecessor)
767        assert_eq!(repaired, vec!["O", "O", "O", "O"]);
768    }
769
770    #[test]
771    fn test_repair_lenient() {
772        let invalid = vec!["O", "I-PER", "I-PER", "O"];
773        let repaired = repair_bio_sequence(&invalid, BioScheme::IOB2, RepairStrategy::Lenient);
774        assert_eq!(repaired, vec!["O", "I-PER", "I-PER", "O"]);
775    }
776
777    #[test]
778    fn test_repair_type_change() {
779        // B-PER followed by I-LOC - type mismatch
780        let invalid = vec!["B-PER", "I-LOC", "O"];
781        let repaired =
782            repair_bio_sequence(&invalid, BioScheme::IOB2, RepairStrategy::PromoteToBegin);
783        assert_eq!(repaired, vec!["B-PER", "B-LOC", "O"]);
784    }
785
786    #[test]
787    fn test_roundtrip() {
788        let tokens = ["The", "United", "Nations", "met", "in", "New", "York"];
789        let tags = ["O", "B-ORG", "I-ORG", "O", "O", "B-LOC", "I-LOC"];
790
791        let entities =
792            bio_to_entities(&tokens, &tags, BioScheme::IOB2).expect("valid BIO tags should parse");
793
794        // Create token offsets for roundtrip
795        let mut offsets = Vec::new();
796        let mut pos = 0;
797        for t in &tokens {
798            offsets.push((pos, pos + t.len()));
799            pos += t.len() + 1;
800        }
801
802        let recovered_tags = entities_to_bio(&offsets, &entities, BioScheme::IOB2);
803
804        assert_eq!(recovered_tags, tags);
805    }
806
807    #[test]
808    fn test_empty_input() {
809        let tokens: [&str; 0] = [];
810        let tags: [&str; 0] = [];
811
812        let entities =
813            bio_to_entities(&tokens, &tags, BioScheme::IOB2).expect("valid BIO tags should parse");
814        assert!(entities.is_empty());
815    }
816
817    #[test]
818    fn test_all_outside() {
819        let tokens = ["The", "cat", "sat"];
820        let tags = ["O", "O", "O"];
821
822        let entities =
823            bio_to_entities(&tokens, &tags, BioScheme::IOB2).expect("valid BIO tags should parse");
824        assert!(entities.is_empty());
825    }
826
827    #[test]
828    fn test_mismatched_lengths() {
829        let tokens = ["John", "Smith"];
830        let tags = ["B-PER"];
831
832        let result = bio_to_entities(&tokens, &tags, BioScheme::IOB2);
833        assert!(result.is_err());
834    }
835
836    #[test]
837    fn test_character_offsets() {
838        let tokens = ["John", "Smith"];
839        let tags = ["B-PER", "I-PER"];
840
841        let entities =
842            bio_to_entities(&tokens, &tags, BioScheme::IOB2).expect("valid BIO tags should parse");
843
844        assert_eq!(entities.len(), 1);
845        assert_eq!(entities[0].start(), 0);
846        // "John" (4) + space (1) + "Smith" (5) = 10, but text is "John Smith"
847        // start of "John" = 0, end of "Smith" = 4 + 1 + 5 = 10
848        assert_eq!(entities[0].end(), 10);
849    }
850
851    #[test]
852    fn test_iob1_scheme() {
853        // In IOB1, B is only used when two same-type entities are adjacent
854        let tokens = ["John", "Mary", "works"];
855        // Both start with I in IOB1 (no adjacency issue)
856        let tags = ["I-PER", "I-PER", "O"];
857
858        let entities =
859            bio_to_entities(&tokens, &tags, BioScheme::IOB1).expect("valid IOB1 tags should parse");
860
861        // In IOB1 with same types, I-I continues the entity
862        assert_eq!(entities.len(), 1);
863        assert_eq!(entities[0].text, "John Mary");
864    }
865
866    #[test]
867    fn test_custom_entity_types() {
868        let tokens = ["CRISPR", "is", "a", "technology"];
869        let tags = ["B-TECH", "O", "O", "O"];
870
871        let entities =
872            bio_to_entities(&tokens, &tags, BioScheme::IOB2).expect("valid BIO tags should parse");
873
874        assert_eq!(entities.len(), 1);
875        assert!(matches!(entities[0].entity_type, EntityType::Custom { .. }));
876    }
877
878    // =============================================================================
879    // IOE Scheme Tests
880    // =============================================================================
881
882    #[test]
883    fn test_ioe2_basic() {
884        // IOE2: E always ends an entity
885        let tokens = ["New", "York", "City"];
886        let tags = ["I-LOC", "I-LOC", "E-LOC"];
887
888        let entities = bio_to_entities(&tokens, &tags, BioScheme::IOE2).unwrap();
889
890        assert_eq!(entities.len(), 1);
891        assert_eq!(entities[0].text, "New York City");
892        assert_eq!(entities[0].entity_type, EntityType::Location);
893    }
894
895    #[test]
896    fn test_ioe2_multiple_entities() {
897        let tokens = ["John", "works", "at", "Apple", "Inc"];
898        let tags = ["E-PER", "O", "O", "I-ORG", "E-ORG"];
899
900        let entities = bio_to_entities(&tokens, &tags, BioScheme::IOE2).unwrap();
901
902        assert_eq!(entities.len(), 2);
903        assert_eq!(entities[0].text, "John");
904        assert_eq!(entities[1].text, "Apple Inc");
905    }
906
907    #[test]
908    fn test_ioe1_basic() {
909        // IOE1: E only appears when needed (similar to IOB1)
910        let tokens = ["New", "York"];
911        let tags = ["I-LOC", "I-LOC"];
912
913        let entities =
914            bio_to_entities(&tokens, &tags, BioScheme::IOE1).expect("valid IOE1 tags should parse");
915
916        assert_eq!(entities.len(), 1);
917        assert_eq!(entities[0].text, "New York");
918    }
919
920    #[test]
921    fn test_entities_to_bio_ioe2() {
922        let _tokens = ["The", "Big", "Apple"];
923        let entities = vec![Entity::new("Big Apple", EntityType::Location, 4, 14, 0.9)];
924
925        // Create token offsets: "The" (0-3), "Big" (4-7), "Apple" (8-13)
926        let offsets = vec![(0, 3), (4, 7), (8, 13)];
927
928        let tags = entities_to_bio(&offsets, &entities, BioScheme::IOE2);
929
930        assert_eq!(tags[0], "O");
931        // EntityType::Location.as_label() returns "LOC"
932        assert_eq!(tags[1], "I-LOC");
933        assert_eq!(tags[2], "E-LOC");
934    }
935
936    // =============================================================================
937    // Repair for Different Schemes
938    // =============================================================================
939
940    #[test]
941    fn test_repair_iobes_orphan_inside() {
942        let invalid = vec!["O", "I-PER", "O"];
943        let repaired =
944            repair_bio_sequence(&invalid, BioScheme::IOBES, RepairStrategy::PromoteToBegin);
945        // Orphan I should become S (single) in IOBES
946        assert_eq!(repaired, vec!["O", "S-PER", "O"]);
947    }
948
949    #[test]
950    fn test_repair_iobes_orphan_end() {
951        let invalid = vec!["O", "E-PER", "O"];
952        let repaired =
953            repair_bio_sequence(&invalid, BioScheme::IOBES, RepairStrategy::PromoteToBegin);
954        // Orphan E should become S (single) in IOBES
955        assert_eq!(repaired, vec!["O", "S-PER", "O"]);
956    }
957
958    // =============================================================================
959    // Roundtrip Tests for All Schemes
960    // =============================================================================
961
962    #[test]
963    fn test_roundtrip_iobes() {
964        let tokens = ["The", "United", "Nations"];
965        let tags = ["O", "B-ORG", "E-ORG"];
966
967        let entities = bio_to_entities(&tokens, &tags, BioScheme::IOBES).unwrap();
968
969        let mut offsets = Vec::new();
970        let mut pos = 0;
971        for t in &tokens {
972            offsets.push((pos, pos + t.len()));
973            pos += t.len() + 1;
974        }
975
976        let recovered = entities_to_bio(&offsets, &entities, BioScheme::IOBES);
977        assert_eq!(recovered, tags);
978    }
979
980    #[test]
981    fn test_roundtrip_ioe2() {
982        let tokens = ["Visit", "New", "York"];
983        let tags = ["O", "I-LOC", "E-LOC"];
984
985        let entities = bio_to_entities(&tokens, &tags, BioScheme::IOE2).unwrap();
986
987        let mut offsets = Vec::new();
988        let mut pos = 0;
989        for t in &tokens {
990            offsets.push((pos, pos + t.len()));
991            pos += t.len() + 1;
992        }
993
994        let recovered = entities_to_bio(&offsets, &entities, BioScheme::IOE2);
995        assert_eq!(recovered, tags);
996    }
997}
anno_eval/eval/bio_adapter.rs

anno_eval/eval/
bio_adapter.rs
