CONLL¶

@CoNLLDialectRegistry.register("conll")
@dataclass(slots=True)
class CoNLL:
    """
    Simple CoNLL-style format with form and NER tag columns.

    Attributes:
        form: The token text.
        ner: The NER tag in BIO format.
    """

    form: str
    ner: BIOField

@CoNLLDialectRegistry.register("conll2003")
@dataclass(slots=True)
class CoNLL2003:
    """
    Standard CoNLL2003 format with -DOCSTART- handling.

    Attributes:
        docstart: Class variable indicating the document start marker.
        ignore_docstart: Class variable indicating whether to ignore
            document start markers.
        form: The token form/word.
        pos: Part-of-speech tag.
        syntactic_chunk: Syntactic chunk tag in BIO format.
        ner: The Named Entity Recognition tag in BIO format.
    """

    docstart: ClassVar[str] = "-DOCSTART-"
    ignore_docstart: ClassVar[bool] = False

    form: str
    pos: MaybeHyphen
    syntactic_chunk: Annotated[
        BIO | None,
        BeforeValidator(BIO.from_optional_string),
        maybe_empty("-"),
        PlainSerializer(str, when_used="unless-none"),
    ]
    ner: BIOField

@CoNLLDialectRegistry.register("conll2003_ignore_docstart")
@dataclass(slots=True)
class CoNLL2003IgnoreDocstart(CoNLL2003):
    """
    CoNLL2003 format that ignores -DOCSTART- markers.

    Attributes:
        ignore_docstart: Class variable indicating whether to ignore
            document start markers.
    """

    ignore_docstart: ClassVar[bool] = True

@CoNLLDialectRegistry.register("conll2003_pioner")
@dataclass(slots=True)
class CoNLL2003Pioner:
    """
    CoNLL2003-style format with underscore placeholders and no -DOCSTART- for the pioNER dataset.

    Attributes:
        form: The token text.
        pos: Part-of-speech tag.
        syntactic_chunk: Syntactic chunk tag in BIO format.
        ner: The NER tag in BIO format.
    """

    form: str
    pos: MaybeUnderscore
    syntactic_chunk: Annotated[
        BIO | None,
        BeforeValidator(BIO.from_optional_string),
        maybe_empty(),
        PlainSerializer(str, when_used="unless-none"),
    ]
    ner: BIOField

@CoNLLDialectRegistry.register("conll2003_two_column")
@dataclass(slots=True)
class CoNLL2003TwoColumn:
    """
    Two-column CoNLL2003-style format with only word form and NER columns.

    Attributes:
        docstart: Class variable indicating the document start marker.
        ignore_docstart: Class variable indicating whether to ignore
            document start markers.
        form: The token text.
        ner: The NER tag in BIO format.
    """

    docstart: ClassVar[str] = "-DOCSTART-"
    ignore_docstart: ClassVar[bool] = False

    form: str
    ner: BIOField

@CoNLLDialectRegistry.register("conll_bio_first")
@dataclass(slots=True)
class CoNLLBioFirst:
    """
    Simple CoNLL-style format in which NER tags are stored in the first column.

    Attributes:
        ner: The NER tag in BIO format.
        form: The token text.
    """

    ner: BIOField
    form: str

class CoNLLColumnarRegistry[T: ColumnarCoNLL](Registry[T]):
    """Registry for CoNLL dialects."""

class CoNLLDialectRegistry[T: CommonRowType](Registry[T]):
    """Registry for CoNLL dialects."""

@CoNLLDialectRegistry.register("conll_herodotos")
@dataclass(slots=True)
class CoNLLHerodotos:
    """
    Two-column CoNLL-style format used by the Herodotos-Project-NER dataset. It is similar to `CoNLLBioFirst` but with inverted IOB tags (e.g., PERS-B instead of B-PERS) and "0" instead of "O".

    Attributes:
        ner: The Named Entity Recognition tag in inverted BIO format.
        form: The token form/word.
    """

    ner: Annotated[BIO, BeforeValidator(_herodotos_inverted_bio_from_string), PlainSerializer(str)]
    form: str

@CoNLLDialectRegistry.register("conll_jnlpba")
@dataclass(slots=True)
class CoNLLJNLPBA:
    """
    CoNLL format for JNLPBA dataset with MEDLINE document markers.

    Attributes:
        docstart: Class variable indicating the document start marker.
        ignore_docstart: Class variable indicating whether to ignore
            document start markers.
        form: The token text.
        ner: The NER tag in BIO format.
    """

    docstart: ClassVar[str] = "###MEDLINE:"
    ignore_docstart: ClassVar[bool] = False

    form: str
    ner: BIOField

@dataclass(slots=True)
class CoNLLMetadata:
    """
    Stores CoNLL format metadata, handling both regular comments and metadata encoded as key-value pairs.

    Attributes:
        meta: Dictionary of key-value metadata.
        comments: List of other comment strings.
    """

    meta: dict[str, str] = field(default_factory=dict)
    comments: list[str] = field(default_factory=list)

    def get_meta(self, metadata_key: str) -> str | None:
        """
        Retrieve metadata by key.

        Args:
            metadata_key: The key to look up.

        Returns:
            The metadata value, or None if not found.
        """

        return self.meta.get(metadata_key)

    def __getitem__(self, metadata_key: str) -> str:
        """
        Retrieve metadata by key.

        Args:
            metadata_key: The key to look up.

        Returns:
            The metadata value.

        Raises:
            KeyError: If the key is not found.
        """

        return self.meta[metadata_key]

    def __contains__(self, metadata_key: str) -> bool:
        """
        Check if a metadata key exists.

        Args:
            metadata_key: The key to check.

        Returns:
            `True` if the key exists in the key-value store.
        """

        return metadata_key in self.meta

    @classmethod
    def with_key_value(cls, comments: Iterable[tuple[str, str] | str]) -> Self:
        """
        Create CoNLLMetadata from an iterable of plain comments and parsed key-value pairs.

        Args:
            comments: Iterable of tuples (key, value) or plain comment
                strings.

        Returns:
            A new CoNLLMetadata instance.
        """

        meta = {}
        general = []
        for comment in comments:
            if isinstance(comment, tuple):
                meta[comment[0]] = comment[1]
            else:
                general.append(comment)

        return cls(meta, general)

@CoNLLDialectRegistry.register("conll_stackoverflow_ner")
@dataclass(slots=True)
class CoNLLStackOverflowNER:
    """
    CoNLL format for the StackOverflow-NER dataset.

    Attributes:
        docstart: Class variable indicating the document start marker.
        ignore_docstart: Class variable indicating whether to ignore
            document start markers.
        form: The token text.
        ner: The NER tag in BIO format.
        form_2: Secondary form field.
        markdown: Markdown syntax metadata in BIO format.
    """

    docstart: ClassVar[str] = "-DOCSTART-"
    ignore_docstart: ClassVar[bool] = False

    form: str
    ner: BIOField
    form_2: str
    markdown: BIOField

@dataclass(slots=True)
class CoNLLU:
    """
    CoNLL-U format with 10 standard columns. For reference, see: https://universaldependencies.org/format.html

    Attributes:
        id: Token ID (int, Decimal for empty nodes, or range for multi-
            word tokens).
        form: Token text.
        lemma: Lemma or stem.
        upos: Universal part-of-speech tag.
        xpos: Language-specific part-of-speech tag.
        feats: Dictionary of morphological features.
        head: Head token ID.
        deprel: Dependency relation.
        deps: Enhanced dependency graph in the form of a list of head-deprel pairs.
        misc: Dictionary of other miscellaneous information.
    """

    id: Annotated[int | Decimal | range, PlainValidator(conllu_id)]
    # Handle literal underscores as just underscores
    # Making this decision is left to the implementation according to https://universaldependencies.org/format.html
    form: str
    lemma: str
    upos: MaybeUnderscore
    xpos: MaybeUnderscore
    feats: Annotated[dict[str, str] | None, key_value_list(), maybe_empty()]
    head: Annotated[int | None, maybe_empty()]
    deprel: MaybeUnderscore
    deps: Annotated[dict[int, str] | None, key_value_list(key_value_sep=":"), maybe_empty()]
    # Handle SpaceAfter=No
    misc: Annotated[dict[str, str] | None, key_value_list(), maybe_empty()]

    @staticmethod
    def parse_comments(comments: list[str]) -> CoNLLMetadata:
        """
        Parse CoNLL-U style comments into CoNLLMetadata.

        Args:
            comments: List of comment strings starting with "#".

        Returns:
            Parsed CoNLLMetadata.
        """

        return CoNLLMetadata.with_key_value(parse_key_value(comment) for comment in comments)

    @staticmethod
    def is_document_start(metadata: CoNLLMetadata) -> bool:
        """
        Check if the metadata indicates the start of a new document.

        Args:
            metadata: The CoNLLMetadata to check.

        Returns:
            `True` if a "newdoc id" is present, indicating the start of
            a new document.
        """

        return metadata.get_meta("newdoc id") is not None

@CoNLLDialectRegistry.register("conllu_plus")
@dataclass(slots=True)
class CoNLLUPlus:
    """
    CoNLL-U Plus format with NER tags used by the NYTK-NerKor dataset. For reference, see: https://universaldependencies.org/ext-format.html

    Attributes:
        form: The token text.
        lemma: Lemma or stem.
        upos: Universal part-of-speech tag.
        xpos: Language-specific part-of-speech tag.
        feats: Dictionary of morphological features.
        ner: The NER tag in BIO format.
        emmorph_lemma: Lemma derived from the Hungarian emMorph
            morphological analyzer
    """

    form: str
    lemma: str
    upos: MaybeUnderscore
    xpos: MaybeUnderscore
    feats: Annotated[dict[str, str] | None, key_value_list(), maybe_empty()]
    ner: BIOField
    emmorph_lemma: str

    @staticmethod
    def parse_comments(comments: list[str]) -> CoNLLMetadata:
        """
        Parse CoNLL-U style comments into CoNLLMetadata.

        Args:
            comments: List of comment strings starting with "#".

        Returns:
            Parsed CoNLLMetadata.
        """

        return CoNLLMetadata.with_key_value(parse_key_value(comment) for comment in comments)

    @staticmethod
    def is_document_start(_: CoNLLMetadata) -> bool:
        """
        Check if the metadata indicates the start of a new document.

        Returns:
            Always returns True for CoNLLUPlus.
        """

        return True

@CoNLLDialectRegistry.register("conll_weibo_ner")
@dataclass(slots=True)
class CoNLLWeiboNER:
    """
    CoNLL-style format from the WeiboNER dataset in which word position indices are stored as part of the word form column.
    Word position indices are removed during deserialization.

    Attributes:
        form: The token form/word with word position indices removed.
        ner: The NER tag in BIO format.
    """

    form: Annotated[str, BeforeValidator(_strip_word_position)]
    ner: BIOField

@CoNLLDialectRegistry.register("conll_with_pos")
@dataclass(slots=True)
class CoNLLWithPOS:
    """
    CoNLL-style format with POS and NER tags.

    Attributes:
        form: The token text.
        pos: The part-of-speech tag.
        ner: The NER tag in BIO format.
    """

    form: str
    pos: str
    ner: BIOField

@runtime_checkable
class ColumnarCoNLL(Protocol):
    """
    Protocol for CoNLL-style formats that support column-based segmentation."""

    @staticmethod
    def segment_columns(rows: Iterable[dict[str, Any]]) -> Iterable[list[dict[str, Any]]]:
        """
        Segment sentence indexed rows into sentence-level groups.

        Args:
            rows: Iterable of indexed rows.

        Returns:
            Iterable of sentence-level row groups.
        """
        ...

@runtime_checkable
class CommentedCoNLL(Protocol):
    """
    Protocol for CoNLL-style formats that support comment-based metadata."""

    @staticmethod
    def parse_comments(comments: list[str]) -> CoNLLMetadata:
        """
        Parse CoNLL metadata from comments.

        Args:
            comments: List of comment strings.

        Returns:
            Parsed CoNLLMetadata.
        """
        ...

    @staticmethod
    def is_document_start(metadata: CoNLLMetadata) -> bool:
        """
        Check if metadata indicates the start of a new document.

        Args:
            metadata: The CoNLLMetadata to check.

        Returns:
            True at the start of a new document.
        """
        ...

@runtime_checkable
class DocumentSeparatedCoNLL(Protocol):
    """
    Protocol for CoNLL-style formats separated by document markers.

    Attributes:
        docstart: Class variable indicating the document start marker.
        ignore_docstart: Class variable indicating whether to ignore
            document start markers.
    """

    docstart: ClassVar[str]
    ignore_docstart: ClassVar[bool]

@CoNLLColumnarRegistry.register("flat")
@dataclass(slots=True)
class FlatIndices:
    """
    CoNLL-style columnar format for data without blank lines, indicating documents and sentences via index columns.

    Attributes:
        ner: NER label.
        form: The token text.
        doc_idx: Document index.
        sent_idx: Sentence index.
    """

    ner: int
    form: str
    doc_idx: int
    sent_idx: int

    @staticmethod
    def segment_columns(rows: Iterable[dict[str, Any]]) -> Iterable[list[dict[str, Any]]]:
        """
        Segment flat index rows into sentence-level groups.

        Args:
            rows: Iterable of flat index rows.

        Returns:
            Iterable of sentence-level row groups.
        """

        current_index = 0
        sentence_rows = []
        for row in rows:
            index = row["sent_idx"]
            if index != current_index:
                current_index = index
                yield sentence_rows
                sentence_rows = []

            sentence_rows.append(row)

        if sentence_rows:
            yield sentence_rows

@CoNLLDialectRegistry.register("ner_suite")
@dataclass(slots=True)
class NerSuiteCoNLL:
    """
    NERSuite CoNLL-style format with character indices mapping tokens to the source text. For reference, see: https://nersuite.nlplab.org/advanced_usage.html

    Attributes:
        ner: The NER tag in BIO format.
        start: Start character offset of the token in the source text.
        end: End character offset of the token in the source text.
        form: The token text.
        lemma: Lemma or stem.
        pos: Part-of-speech tag.
        chunk: Chunk tag in BIO format.
    """

    ner: BIOField
    start: int
    end: int
    form: str
    lemma: str
    pos: MaybeUnderscore
    chunk: BIOField

class RowParser[T]:
    """
    Parser for CoNLL-style rows supporting arbitrary CoNLL dialects.

    **Type Parameters:**

    T
        The CoNLL dialect class used for parsing rows.

    Args:
        dialect: The CoNLL dialect class to parse rows with.
    """

    def __init__(self, dialect: type[T]) -> None:
        self.row_parser = TypeAdapter(dialect)
        fields = dataclasses.fields(dialect)  # type: ignore
        self.fields = [field.name for field in fields]

    def validate_row(self, row: list[str]) -> T:
        """
        Validate and parse a row using the parser's CoNLL dialect.

        Args:
            row: List of fields in the row.

        Returns:
            Parsed Sentence using the parser's CoNLL dialect.
        """

        return self.row_parser.validate_python({field: value for field, value in zip(self.fields, row)})

@dataclass(slots=True)
class Sentence[T: RowType]:
    """
    A parsed CoNLL-style sentence with optional metadata.

    **Type Parameters:**

    T
        The type of rows in the sentence.

    Attributes:
        rows: List of parsed CoNLL-style rows.
        meta: Optional CoNLLMetadata for the sentence.
    """

    rows: list[T]
    meta: CoNLLMetadata | None = None

@CoNLLDialectRegistry.register("conllu_uner")
@dataclass(slots=True)
class UNERCoNLLU:
    """
    Universal NER CoNLL-U-style format with UNER tags and (optionally) original NER tags.

    Attributes:
        id: Token ID.
        form: The token text.
        ner: The NER tag in BIO format.
        original_ner: Original NER tags for converted datasets.
        annotator: Optional name of the annotator.
    """

    id: int | Decimal
    form: str
    ner: BIOField
    original_ner: Annotated[list[BIO] | None, BeforeValidator(_maybe_multi_label), maybe_empty("-")]
    annotator: MaybeHyphen

    @staticmethod
    def parse_comments(comments: list[str]) -> CoNLLMetadata:
        """
        Parse CoNLL-U style comments into CoNLLMetadata.

        Args:
            comments: List of comment strings starting with "#".

        Returns:
            Parsed CoNLLMetadata.
        """

        return CoNLLMetadata.with_key_value(parse_key_value(comment) for comment in comments)

    @staticmethod
    def is_document_start(metadata: CoNLLMetadata) -> bool:
        """
        Check if the metadata indicates the start of a new document.

        Args:
            metadata: The CoNLLMetadata to check.

        Returns:
            `True` if a "newdoc id" is present, indicating the start of
            a new document.
        """

        return metadata.get_meta("newdoc id") is not None