Readers¶

@ReaderRegistry.register("arabic_cross_dialectal_json")
class ArabicCrossDialectalJSONReader(DetokenizingLineReader):
    """Reader for the ArabicCrossDialectal JSON format with detokenization support."""

    def _read_lines(self, lines: FileSource, detokenizer: Detokenizer) -> Iterable[NERDocument]:
        """
        Reads lines in the ArabicCrossDialectal JSON format and parses them into NERDocuments.

        Args:
            lines: FileSource containing JSON data with tokens and NER
                annotations.
            detokenizer: Detokenizer for reconstructing the source text.

        Yields:
            NERDocument instances with detokenized text and entity annotations.
        """

        for sample in json.loads(lines.read()):
            tokens = sample["tokens"]
            text, spans = detokenizer.detokenize(tokens)

            yield NERDocument(
                [
                    LabeledText(
                        text,
                        default_tagset(
                            [
                                Annotation(_flatten_label(label), (Span(spans[start].start, spans[end].stop),))
                                for start, end, label in sample["ner"]
                            ]
                        ),
                    )
                ]
            )

@ReaderRegistry.register("bioc_xml")
@dataclass(slots=True)
class BioCXMLReader(LineReader):
    """Reader for NER datasets in the BioC XML format."""

    def _read_lines(self, lines: FileSource) -> Iterable[NERDocument]:
        """
        Parses BioC XML and yields NERDocuments with annotations, including any discontinuous spans.

        Args:
            lines: FileSource containing BioC XML.
        Yields:
            NERDocument instances with passage-level annotations.
        """

        document = etree.fromstringlist(lines.lines())
        document_ids = set()
        for document in document.xpath("/collection/document"):
            [document_id] = document.xpath("./id/text()")
            if document_id in document_ids:
                logger.warning(f"Skipping duplicate document ID {document_id} in file")
                continue

            document_ids.add(document_id)

            passages = []
            for passage in document.xpath("./passage"):
                [text] = passage.xpath("./text/text()")
                text = text.strip()
                [passage_offset] = passage.xpath("./offset/text()")
                passage_offset = int(passage_offset)
                labels = []

                for annotation in passage.xpath("./annotation"):
                    # Each annotation may have one or multiple locations for discontinuous spans
                    offsets = tuple(
                        Span.from_run_length(int(location.get("offset")) - passage_offset, int(location.get("length")))
                        for location in annotation.xpath("./location")
                    )
                    [label] = annotation.xpath('./infon[@key="type"]/text()')
                    [annotation_text] = annotation.xpath("./text/text()")

                    found_subspans = " ".join(text[subspan.start : subspan.stop] for subspan in offsets)
                    # If the given annotation text and the given start offset of the span in the text don't match, attempt to match them up and fix them automatically
                    # Note: This approach does not handle cases where the length of the span is also incorrect
                    if found_subspans != annotation_text:
                        fixed_offsets = []
                        expected_tokens = re.split("( )", annotation_text)
                        token_offset = 0
                        for subspan in offsets:
                            length = 0
                            for token in expected_tokens:
                                length += len(token)
                                token_offset += 1
                                if length >= len(subspan):
                                    break

                            found = text[subspan.start : subspan.stop]
                            expected = "".join(expected_tokens[:token_offset]).strip()
                            if found == expected:
                                fixed_offsets.append(subspan)
                            else:
                                fixed_offsets.append(_fix_mismatched_offset(text, expected, found, subspan))
                            expected_tokens = expected_tokens[token_offset:]

                    labels.append(Annotation(label, offsets))

                # NOTE: Uses newline delimiters since passages are usually the title and abstract of an article
                passages.append(LabeledText(text, default_tagset(labels), "passage", space_after="\n"))

            # No space after the last passage in the document
            passages[-1].space_after = ""
            yield NERDocument(passages)

@ReaderRegistry.register("byte_offset_jsonl")
@dataclass(slots=True)
class ByteOffsetJSONLReader(LineReader):
    """
    Reader for the ByteOffsetJSONL format.

    Attributes:
        arguments: Configuration for the ByteOffsetJSONL format reader.
    """

    arguments: ByteOffsetJSONLArguments

    def _read_lines(self, lines: FileSource) -> Iterable[NERDocument]:
        """
        Reads ByteOffsetJSONL files with byte-level annotations.

        Args:
            lines: FileSource containing JSONL files.

        Yields:
            NERDocument instances with byte-aligned annotations.

        Raises:
            ValueError: If an annotation offset mismatch is detected.
        """

        text_key = self.arguments.text_key
        offsets_key = self.arguments.offsets_key
        target_key = self.arguments.annotated_span_target_key

        for line in lines.lines():
            sample = json.loads(line)
            # Treat as byte string
            text = sample[text_key]
            text_bytes = text.encode("utf-8")
            annotations = []

            if target_key:
                targets = sample[target_key]
                # Targets field of the form "PER: ... $$ ORG: ..."
                expected_spans = (
                    [target.split(":", 1)[1].strip() for target in targets.split("$$")] if targets else None
                )
            else:
                expected_spans = None

            # Track the location within the string in alignment with the byte string
            string_offset = 0
            # Track the position after the previous span in the byte string
            previous_byte_end = 0
            for i, offset in enumerate(sample[offsets_key]):
                start = offset["start_byte"]
                end = offset["limit_byte"]

                # Add the number of characters of text between spans
                string_offset += len(text_bytes[previous_byte_end:start].decode("utf-8"))
                # Track the start and end offsets of the span in the string
                string_start = string_offset
                string_offset += len(text_bytes[start:end].decode("utf-8"))

                previous_byte_end = end

                if (
                    expected_spans is not None
                    and (current_span := text[string_start:string_offset]) != expected_spans[i]
                ):
                    raise ValueError(
                        f"Annotation offset mismatch. Expected {expected_spans[i]!r} but got {current_span!r} at offset ({start}, {end})/({string_start}, {string_offset})"
                    )

                annotations.append(Annotation(offset["label"], (Span(string_start, string_offset),)))

            yield NERDocument([LabeledText(text, default_tagset(annotations))])

class CoNLLPreprocessorRegistry(Registry[Callable[[Iterable[str]], Iterator[str]]]):
    """Registry for CoNLL preprocessor functions."""

@ReaderRegistry.register("conll")
@dataclass(slots=True)
class CoNLLReader:
    """
    Reader for CoNLL-formatted NER datasets.

    Attributes:
        arguments: CoNLLArguments configuration for reading.
        detokenizer_type: Type of detokenization to apply for text
            reconstruction.
        tagsets: Optional list of NamedTagSet instances for tag
            configuration.
    """

    arguments: CoNLLArguments
    detokenizer_type: DetokenizerType
    tagsets: list[NamedTagSet] | None = None

    def __post_init__(self) -> None:
        """Initializes tagsets from a label map if provided."""
        # Invert label map for index to BIO conversion
        self.tagsets = (
            None
            if self.arguments.label_map is None
            else [
                NamedTagSet(tagset, label_map={index: label for label, index in label_map.items()})
                for tagset, label_map in self.arguments.label_map.items()
            ]
        )

    def _conll_to_document(self, sentences: list[Sentence], detokenizer: Detokenizer) -> NERDocument:
        """
        Converts a list of CoNLL sentences to a NERDocument.

        Args:
            sentences: List of Sentence instances.
            detokenizer: Detokenizer for text reconstruction.

        Returns:
            A NERDocument with detokenized text and labeled tokens.
        """

        first_meta = sentences[0].meta
        return detokenizer.tokens_to_document(
            [
                _to_labeled_tokens(
                    sentence,
                    self.tagsets,
                    self.arguments.bioes_to_bio,
                )
                for sentence in sentences
            ],
            None
            if first_meta is None or "text" not in first_meta
            else [_get_sentence_text(sentence.meta) for sentence in sentences],
        )

    def _nersuite_to_document(self, sentences: list[Sentence[NerSuiteCoNLL]], _detokenizer: Detokenizer) -> NERDocument:
        """
        Converts NERSuite-style sentences to a NERDocument with proper offset handling.

        Args:
            sentences: List of NERSuite Sentence instances.

        Returns:
            A NERDocument with text, offsets, and labeled tokens.
        """

        labeled_text = []
        labeled_tokens = []
        last_end = 0

        for sentence in sentences:
            pre_tokenized = _to_labeled_tokens(
                sentence,
                self.tagsets,
                self.arguments.bioes_to_bio,
            )
            labeled_tokens.append(pre_tokenized)
            text, spans = _detokenize_nersuite_offsets(sentence.rows)

            if labeled_text:
                labeled_text[-1].space_after = (sentence.rows[0].start - last_end) * " "

            last_end = sentence.rows[-1].end

            labeled_text.append(
                LabeledText(text, {tagset: bio_to_spans(tags, spans) for tagset, tags in pre_tokenized.labels.items()})
            )

        return NERDocument(labeled_text, labeled_tokens)

    def _read_file_shard(
        self, source: FileSource, detokenizer: Detokenizer, metadata: dict[str, Any], shards_are_documents: bool = False
    ) -> Iterable[NERDocument]:
        """
        Reads and processes a file shard into `NERDocument`s.

        Args:
            source: FileSource containing the input file.
            detokenizer: Detokenizer for text reconstruction.
            metadata: Processing metadata for the file.
            shards_are_documents: Whether each shard represents a
                complete document.
        Yields:
            NERDocument instances.

        Raises:
            KeyError: If the CoNLL dialect is unsupported.
        """

        try:
            dialect = CoNLLDialectRegistry.get(self.arguments.dialect)
        except KeyError:
            raise KeyError(f"Unsupported dialect for CoNLL files: {self.arguments.dialect}")

        parse_arguments = metadata.get("parse_arguments") or {}

        # Select dialect specific CoNLL row to NERDocument processor
        process_document = self._nersuite_to_document if dialect == NerSuiteCoNLL else self._conll_to_document

        lines = source.lines()
        if self.arguments.preprocessor is not None:
            lines = CoNLLPreprocessorRegistry.get(self.arguments.preprocessor)(lines)

        sentences: list[Sentence] = []
        for document in conll.parse(
            lines, dialect, self.arguments.delimiter, self.arguments.enforce_blank_lines, **parse_arguments
        ):
            if shards_are_documents:
                # Assume no nested documents if the "shards_are_documents" option is enabled
                try:
                    [sentence] = document
                    sentences.append(sentence)
                except ValueError:
                    raise ValueError(
                        'Expected exactly one sentence in document when "shards_are_documents" is True, '
                        f"but got {len(document)} sentences."
                    )
            else:
                # Type of document at this point can currently not be inferred
                yield process_document(document, detokenizer)  # pyright: ignore

        if sentences:
            yield process_document(sentences, detokenizer)

    def _read_dataset_shard(
        self, source: DatasetSource, detokenizer: Detokenizer, shards_are_documents: bool = False
    ) -> Iterable[NERDocument]:
        """
        Reads and processes a HuggingFace dataset shard into NERDocuments.

        Args:
            source: DatasetSource containing the dataset.
            detokenizer: Detokenizer for text reconstruction.
            shards_are_documents: Whether each shard represents a
                complete document.
        Yields:
            NERDocument instances.

        Raises:
            KeyError: If the dialect is unsupported.
            ValueError: If a preprocessor is specified which is not
                supported by CoNLL-style columnar datasets.
            ValueError: If multiple tagsets are provided which are
                currently unsupported.
        """

        try:
            dialect = CoNLLColumnarRegistry.get(self.arguments.dialect)
        except KeyError:
            raise KeyError(f"Unsupported dialect for CoNLL-style columnar data: {self.arguments.dialect}")

        if self.arguments.preprocessor is not None:
            raise ValueError(
                f"Preprocessor {self.arguments.preprocessor} not supported for CoNLL-style columnar datasets"
            )

        if self.tagsets is None:
            tag_column = source.default_tag_column
        else:
            try:
                [tagset] = self.tagsets
                tag_column = tagset.column or source.default_tag_column
            except ValueError:
                raise ValueError("Multiple tagsets are currently unsupported for CoNLL-style datasets")

        dataset = source.dataset.rename_columns({source.text_column: "form", tag_column: "ner"})
        previous_document_index = None

        # Pick source specific label map if possible and fall back to the global mapping
        sentences = []
        for sentence in conll.parse_columns(dataset.to_iterable_dataset(), dialect):
            if (
                not shards_are_documents
                and previous_document_index != (doc_idx := sentence.rows[0].doc_idx)
                and sentences
            ):
                previous_document_index = doc_idx
                yield detokenizer.tokens_to_document(sentences)
                sentences = []

            sentences.append(_to_labeled_tokens(sentence, self.tagsets, self.arguments.bioes_to_bio))

        if sentences:
            yield detokenizer.tokens_to_document(sentences)

    def read_split(self, split: SplitData) -> Iterable[NERDocument]:
        """
        Reads a split and yields NERDocuments for each shard.

        Args:
            split: Split instance with data sources.
        Yields:
            NERDocument instances from reading each shard.

        Raises:
            ValueError: If a shard has an unsupported type.
        """

        for shard in split.data:
            detokenizer = Detokenizer(split.language, self.detokenizer_type)
            if isinstance(shard, FileSource):
                yield from self._read_file_shard(
                    shard, detokenizer, split.metadata, self.arguments.shards_are_documents
                )
            elif isinstance(shard, DatasetSource):
                yield from self._read_dataset_shard(shard, detokenizer, self.arguments.shards_are_documents)
            else:
                raise TypeError(f"Encountered shard with unsupported type: {shard}")

@dataclass(slots=True)
class DatasetReader:
    """
    Reader for Huggingface Dataset sources with NER annotations.

    Reads datasets in various formats and converts them to NERDocuments,
    handling detokenization, BIO tags, and label mapping.

    Attributes:
        config: Dataset conversion configuration.
        tagsets: Optional list of NamedTagSet instances for tag
            configuration.
    """

    config: DatasetConvertStep
    tagsets: list[NamedTagSet] | None = None

    def __post_init__(self) -> None:
        """Initializes tagsets from label maps if provided."""
        self.tagsets = (
            None
            if self.config.tagsets is None
            else [
                # Strings are just column names
                NamedTagSet(name, tagset)
                if isinstance(tagset, str)
                # Invert label map for index to BIO conversion
                else NamedTagSet(
                    name,
                    tagset.column,
                    None if tagset.label_map is None else {index: label for label, index in tagset.label_map.items()},
                )
                for name, tagset in self.config.tagsets.items()
            ]
        )

    def read_split(self, split: SplitData) -> Iterable[NERDocument]:
        """
        Reads a split and yields NERDocuments from a Huggingface dataset.

        Args:
            split: Split instance with DatasetSource containing the
                dataset.

        Yields:
            NERDocument instances from detokenizing and processing dataset rows.

        Raises:
            ValueError: If an input shard is not a HuggingFace dataset.
        """

        detokenizer = Detokenizer(split.language, self.config.detokenizer_type)
        sequence_type = self.config.sequence_type

        for shard in split.data:
            if not isinstance(shard, DatasetSource):
                raise TypeError(f"DatasetReader can only process datasets, but found non-dataset shard: {shard}")

            dataset = shard.dataset

            tagsets = [NamedTagSet()] if self.tagsets is None else self.tagsets

            for row in dataset.select_columns(
                [shard.text_column, *{tagset.column or shard.default_tag_column for tagset in tagsets}]
            ).to_iterable_dataset():
                mapped_tagsets = {}
                for tagset in tagsets:
                    tag_column = tagset.column or shard.default_tag_column
                    label_map = tagset.label_map
                    tags = row[tag_column]
                    if label_map is None:
                        if self.config.bio_type == "iob_type_only":
                            # BIO tag format with an implicit "I" position. Used by Polyglot-NER
                            bio_tags = [BIO("O") if iob == "O" else BIO("I", iob) for iob in tags]
                        else:
                            bio_tags = list(map(BIO.from_string, tags))
                        mapped_tagsets[tagset.name] = bio_tags
                    else:
                        mapped_tagsets[tagset.name] = [label_map[tag] for tag in tags]

                # NOTE: Assumes all text columns are sentence level
                yield detokenizer.tokens_to_document(
                    [LabeledTokens(row[shard.text_column], mapped_tagsets, sequence_type)]
                )

@dataclass(slots=True)
class DatasetSource:
    """
    Represents a Huggingface Dataset source for NER data.

    Attributes:
        dataset: The Huggingface Dataset instance.
        text_column: Name of the column containing text tokens.
        default_tag_column: Name of the default NER tag column.
    """

    dataset: Dataset
    text_column: str
    default_tag_column: str

@ReaderRegistry.register("dataset_spans")
@dataclass(slots=True)
class DatasetSpanReader:
    """
    Reader for Huggingface Dataset sources with pre-aligned span annotations.

    Reads datasets where spans are already aligned with text content.
    """

    def read_split(self, split: SplitData) -> Iterable[NERDocument]:
        """
        Reads a split and yields NERDocuments from a Huggingface dataset with pre-aligned spans.

        Args:
            split: Split instance with DatasetSource instances.

        Yields:
            NERDocument instances with text and aligned span annotations.

        Raises:
            ValueError: If an input shard is not a dataset.
            ValueError: If span mismatch is detected.
        """

        for shard in split.data:
            if not isinstance(shard, DatasetSource):
                raise TypeError(f"DatasetSpanReader can only process datasets, but found non-dataset shard: {shard}")

            dataset = shard.dataset

            for row in dataset.select_columns([shard.text_column, shard.default_tag_column]).to_iterable_dataset():
                annotations = []
                text = row[shard.text_column]
                for annotation in row[shard.default_tag_column]:
                    expected = annotation["name"]
                    span = Span(*annotation["span"])
                    actual = text[span.start : span.stop]

                    if actual != expected:
                        raise ValueError(f"Span mismatch {actual!r} : {expected!r}")

                    annotations.append(Annotation(annotation["type"], (span,)))

                yield NERDocument([LabeledText(text, default_tagset(annotations))])

class DetokenizingLineReader(metaclass=ABCMeta):
    """Abstract base class for readers that require detokenization of tokenized text."""

    @abstractmethod
    def _read_lines(self, lines: FileSource, detokenizer: Detokenizer) -> Iterable[NERDocument]:
        """
        Reads and detokenizes text from a file.

        Args:
            lines: FileSource containing tokenized text.
            detokenizer: Detokenizer for text reconstruction.

        Yields:
            NERDocument instances with detokenized text.
        """

    def read_split(self, split: SplitData) -> Iterable[NERDocument]:
        """
        Reads splits and yields documents after detokenization.

        Args:
            split: Split instance with data sources.

        Yields:
            NERDocument instances.

        Raises:
            ValueError: If a shard is not a text file.
        """

        detokenizer = Detokenizer(split.language)
        for shard in split.data:
            if not isinstance(shard, FileSource):
                raise TypeError(
                    f"{self.__class__.__name__} only supports text files, but found shard of type {type(shard)}"
                )

            yield from self._read_lines(shard, detokenizer)

@ReaderRegistry.register("ebm_nlp_standoff")
@dataclass
class EBMNLPStandoffReader:
    """
    Reader for the EBM-NLP standoff format with broad and granular annotation layers.

    Attributes:
        tagsets: List of NamedTagSet instances for broad and granular
            tags.
        arguments: Configuration for the standoff format reader.
    """

    arguments: EBMNLPStandoffArguments

    tagsets: ClassVar[list[NamedTagSet]] = [NamedTagSet("broad"), NamedTagSet("granular")]

    def __post_init__(self) -> None:
        """Constructs label maps for broad and granular annotation layers."""
        self._broad_label_map = {
            name: {index: label for label, index in label_map.items()}
            for name, label_map in self.arguments.broad_label_map.items()
        }
        self._label_map = {
            name: {index: label for label, index in label_map.items()}
            for name, label_map in self.arguments.label_map.items()
        }

    def _read_document(
        self,
        text: FileSource,
        token_source: FileSource,
        annotations: dict[str, _AnnotatedSource],
    ) -> NERDocument:
        """
        Reads a complete EBM-NLP document with both annotation layers.

        Args:
            text: FileSource containing the original text.
            token_source: FileSource containing pre-tokenized text.
            annotations: Dictionary mapping annotation types to sources.

        Returns:
            NERDocument with both broad and granular annotations.

        Raises:
            ValueError: If an IOB label sequence length doesn't match
                the number of tokens.
        """

        original_text = text.read().strip()
        tokens = list(map(str.strip, token_source.lines()))
        token_spans = align_tokens_with_text(tokens, original_text)
        token_count = len(token_spans)

        broad_span_annotations = []
        span_annotations = []
        for name, source in annotations.items():
            if source.annotation_type == "broad":
                label_map = self._broad_label_map[name]
            else:
                label_map = self._label_map[name]

            bio_labels = [BIO.from_string(label_map[int(label_id)]) for label_id in source.source.lines()]
            if len(bio_labels) != token_count:
                raise ValueError(
                    f"IOB Label sequence must match the number of tokens in the document. Expected: {token_count}, actual: {len(bio_labels)}"
                )

            spans = bio_to_spans(bio_labels, token_spans)

            if source.annotation_type == "broad":
                broad_span_annotations.extend(spans)
            else:
                span_annotations.extend(spans)

        return NERDocument(
            [LabeledText(original_text, {"broad": broad_span_annotations, "granular": span_annotations}, "passage")]
        )

    def read_split(self, split: SplitData) -> Iterable[NERDocument]:
        """
        Reads a split and yields NERDocuments for the EBM-NLP dataset with broad and granular annotation layers.

        Args:
            split: Split instance with dictionary sources containing
                text, token, and annotation files.

        Yields:
            NERDocument instances from reading documents with both annotation layers.

        Raises:
            ValueError: If the shard is not a dictionary.
        """

        for shard in split.data:
            if not isinstance(shard, dict):
                raise TypeError(
                    f"{self.__class__.__name__} only supports dictionary sources, but found shard of type {type(shard)}"
                )

            # NOTE: EBM-NLP 2.0 does not include annotation files for some documents where all annotations are "O"
            # This requires constructing a sparse mapping between annotations and documents to ensure
            # that all are assigned to the correct split
            annotations = defaultdict(dict)
            for annotation in shard["broad_annotations"]:
                annotations[annotation.path.stem.removesuffix(".AGGREGATED")][_get_annotation_type(annotation)] = (
                    _AnnotatedSource("broad", annotation)
                )

            for annotation in shard["annotations"]:
                annotations[annotation.path.stem.removesuffix(".AGGREGATED")][_get_annotation_type(annotation)] = (
                    _AnnotatedSource("granular", annotation)
                )

            text_map = {source.path.stem: source for source in shard["text"]}
            tokens_map = {source.path.stem: source for source in shard["tokens"]}

            for document_id, annotation_sources in annotations.items():
                yield self._read_document(
                    text_map[document_id],
                    tokens_map[document_id],
                    annotation_sources,
                )