Formats¶

@dataclass(slots=True)
class Dataset:
    """
    Represents a full dataset made up of one or more subsets.

    Attributes:
        path: Path to the main local dataset directory.
        metadata: Metadata for the dataset.
        subsets: List of subsets in the dataset.
    """

    path: Path
    metadata: DatasetMetadata
    subsets: dict[tuple[str, ...], Subset] = field(init=False)

    def __post_init__(self) -> None:
        """Initializes the subsets from the dataset metadata."""
        self.subsets = {
            tuple(subset_info.hierarchy): Subset(subset_info.metadata, subset_info.hierarchy, self.path)
            for subset_info in self.metadata._iter_subsets()
        }

    def __iter__(self) -> Iterator[Subset]:
        """
        Iterate over all available subsets.

        Yields:
            All subsets of the dataset.
        """

        yield from self.subsets.values()

    @classmethod
    def load(cls, benchmark_path: Path, dataset_name: str) -> Self:
        """
        Load a dataset from the given benchmark path and dataset name.

        Args:
            benchmark_path: The root directory of the locally processed
                benchmark.
            dataset_name: The name of the dataset to load.

        Returns:
            A `Dataset` instance representing the given dataset.
        """

        dataset_directory = _local_processed_directory(benchmark_path) / dataset_name
        return cls(dataset_directory, DatasetMetadata.from_json(dataset_directory / METADATA_FILENAME))

@dataclass(slots=True)
class DatasetMetadata:
    """
    Metadata of a dataset.

    Attributes:
        name: Name of the dataset.
        subsets: Dictionary mapping subset names to SubsetMetadata
            instances or further SubsetHierarchy levels.
        languages: Set of languages included in the dataset.
        main_splits: Main SplitMetadata instance if available.
    """

    name: str
    subsets: SubsetHierarchy
    languages: set[str]
    main_splits: SubsetMetadata | None = None

    def _iter_subsets(self) -> Iterator[_SubsetWithMetadata]:
        """
        Iterates over all subsets in the hierarchy recursively.

        Yields:
            SubsetInformation for each subset.
        """

        if self.main_splits is not None:
            yield _SubsetWithMetadata(self.main_splits, [])

        # BFS through the subset hierarchy
        subset_queue = [([name], hierarchy) for name, hierarchy in self.subsets.items()]
        while subset_queue:
            path, current = subset_queue.pop()
            if isinstance(current, SubsetMetadata):
                yield _SubsetWithMetadata(current, path)
                continue

            for name, hierarchy in current.items():
                subset_queue.append((path + [name], hierarchy))

    @classmethod
    def from_json(cls, manifest_path: Path) -> Self:
        """
        Creates an instance of `DatasetMetadata` from a JSON file.

        Args:
            manifest_path: Path to the JSON metadata file.

        Returns:
            An instance of `DatasetMetadata`.
        """

        metadata_parser = TypeAdapter(cls)
        with manifest_path.open("rb") as file:
            return metadata_parser.validate_json(file.read())

    def dump(self, directory: Path) -> None:
        """
        Write the dataset metadata to a JSON file in the specified directory.

        Args:
            directory: Directory into which the metadata should be
                written.
        """

        # Write metadata for the fully processed dataset
        metadata_serializer = TypeAdapter(self.__class__)
        with (directory / METADATA_FILENAME).open("wb") as file:
            file.write(metadata_serializer.dump_json(self))

class NERParquetWriter:
    """
    Writer for NER data in MELD parquet format.

    Args:
        dataset: Name of the dataset.
        subset: Name of the subset.
        split: Name of the split (e.g., train, test).
        writer: The PyArrow ParquetWriter instance to write to.
        schema: The arrow schema to use for writing.
    """

    _BUFFER_SIZE = 64

    def __init__(self, dataset: str, subset: str, split: str, writer: ParquetWriter, schema: Schema) -> None:
        self._sequence_id_prefix = f"{dataset} {subset} {split}".encode()
        self._writer = writer
        self._annotation_serializer = TypeAdapter(list[Annotation])
        self._tokens_serializer = TypeAdapter(list[BIOField])
        self._schema = schema
        # Parse tagsets from schema metadata, assuming tagset columns always appear last and in the given order
        self._tagsets = json.loads(schema.metadata[b"tagsets"])
        # Automatically detect whether the schema contains tokens and iob labels
        self._is_tokenized = schema.get_field_index("tokens") != -1

    def _initialize_buffer(self) -> None:
        """Initialize internal column buffers for batched writing."""
        self._sentence_buffer = arrays = [[] for _ in range(len(self._schema))]
        self._record_buffer = {
            "sequence_id": arrays[0],
            "document_index": arrays[1],
            "document_position": arrays[2],
            "sequence_type": arrays[3],
            "text": arrays[4],
            "space_after": arrays[5],
        }
        tagsets_start = 6

        if self._is_tokenized:
            self._record_buffer["tokens"] = arrays[tagsets_start]
            tagsets_start += 1

        for i, tagset in enumerate(self._tagsets, tagsets_start):
            self._record_buffer[tagset] = arrays[i]

        if self._is_tokenized:
            for i, tagset in enumerate(self._tagsets, tagsets_start + len(self._tagsets)):
                self._record_buffer[f"{tagset}_iob"] = arrays[i]

    def __enter__(self) -> Self:
        self._initialize_buffer()
        self._document_index = 0
        return self

    def write_document(self, document: NERDocument) -> None:
        """
        Write a document to the underlying parquet file.

        Args:
            document: The NER document to write.
        """

        tokenized = document.bio if document.bio is not None else itertools.repeat(None)
        arrays = self._sentence_buffer

        for document_position, (labeled_text, labeled_tokens) in enumerate(zip(document.spans, tokenized)):
            # Generate reproducible UUID v5 IDs for each sentence or passage in the data
            # NOTE: Assumes documents are written in a deterministic order!
            sequence_id = uuid.uuid5(
                _UUID_NAMESPACE,
                self._sequence_id_prefix
                + self._document_index.to_bytes(64, "little")
                + document_position.to_bytes(64, "little"),
            ).bytes

            arrays[0].append(sequence_id)
            arrays[1].append(self._document_index)
            arrays[2].append(document_position)
            arrays[3].append(SEQUENCE_TYPES[labeled_text.sequence_type])
            arrays[4].append(labeled_text.text)
            arrays[5].append(labeled_text.space_after)
            tagsets_start = 6

            if labeled_tokens is not None:
                arrays[tagsets_start].append(labeled_tokens.tokens)
                tagsets_start += 1

            for i, tagset in enumerate(self._tagsets, tagsets_start):
                arrays[i].append(self._annotation_serializer.dump_python(labeled_text.labels[tagset]))

            tagsets_start += len(self._tagsets)

            if labeled_tokens is not None:
                for i, tagset in enumerate(self._tagsets, tagsets_start):
                    arrays[i].append(self._tokens_serializer.dump_python(labeled_tokens.labels[tagset]))

        # Batches can be different lengths but must contain at least 64 sentences if possible
        if len(self._sentence_buffer[0]) > self._BUFFER_SIZE:
            self._writer.write_batch(RecordBatch.from_pydict(self._record_buffer, self._schema))
            self._initialize_buffer()

        self._document_index += 1

    def __exit__(
        self, exc_type: type[BaseException] | None, exc_value: BaseException | None, traceback: TracebackType | None
    ) -> None:
        # Write final batch
        if self._sentence_buffer[0]:
            self._writer.write_batch(RecordBatch.from_pydict(self._record_buffer, self._schema))

    @classmethod
    @contextmanager
    def open(
        cls,
        dataset: str,
        subset: str,
        split: str,
        path: str | Path | IO | NativeFile,
        with_tokens: bool,
        tagsets: list[str],
    ) -> Iterator[Self]:
        """
        Open a NERParquetWriter as a context manager.

        Args:
            dataset: Name of the dataset.
            subset: Name of the subset.
            split: Name of the split.
            path: Path to to which the parquet file will be written.
            with_tokens: Whether to include token columns in the output.
            tagsets: List of tagsets to include.

        Yields:
            NERParquetWriter instance.
        """

        schema = FULL_SCHEMA if with_tokens else TEXT_ONLY_SCHEMA
        schema = FULL_SCHEMA if with_tokens else TEXT_ONLY_SCHEMA
        for tagset in tagsets:
            schema = schema.append(tagset_field(tagset))

        schema = schema.with_metadata({"tagsets": json.dumps(tagsets), **schema.metadata})

        if with_tokens:
            for tagset in tagsets:
                schema = schema.append(iob_field(tagset))

        with (
            ParquetWriter(path, schema, compression="zstd") as parquet_writer,
            cls(dataset, subset, split, parquet_writer, schema) as writer,
        ):
            yield writer

@dataclass(slots=True)
class Split:
    """
    Metadata for a specific data split within a dataset subset.

    Attributes:
        name: Name of this split (e.g. "train", "validation", or "test")
        dataset_name: Name of the dataset this split belongs to
        language: Language of the data in this split
        tagsets: List of tagset names available in this split
        metadata: Metadata about the split including size and creation information
        labels: Dictionary mapping tagset names to their respective label lists
        bio_labels: Optional dictionary mapping tagset names to BIO-formatted label lists
    """

    name: str
    dataset_name: str
    language: str
    tagsets: list[str]
    metadata: SplitMetadata
    labels: dict[str, list[str]]
    bio_labels: dict[str, list[str]] | None = None

@dataclass(slots=True)
class SplitMetadata:
    """
    Metadata for a split in a dataset.

    Attributes:
        path: Path to the split data.
        document_count: Number of documents in the split.
        sequence_count: Number of sequences in the split.
        creation_metadata: Metadata containing information on licensing,
            domains, sources, and the annotation process that apply to
            this split
    """

    path: Path
    document_count: int = 0
    sequence_count: int = 0
    creation_metadata: Metadata = field(default_factory=Metadata)

@dataclass(slots=True)
class Subset(_SubsetWithMetadata):
    """
    Represents a subset of a dataset.

    Attributes:
        metadata: Metadata containing information about the subset.
        hierarchy: A flat representation of the hierarchy, including the
            subset name and names of all its parent subsets
        dataset_path: Path to the main dataset directory.
    """

    dataset_path: Path

    def split_path(self, name: str) -> Path:
        """
        Construct the path to a specific split of the dataset.

        Args:
            name: Name of the split.

        Returns:
            Path to the split.
        """

        return self.dataset_path / self.metadata.splits[name].path

    def open_split(self, name: str) -> ParquetFile:
        """
        Open a given split for reading as a `pyarrow.ParquetFile`.

        Args:
            name: Name of the split to open.

        Returns:
            A `pyarrow.ParquetFile` containing the split data.
        """

        return ParquetFile(self.split_path(name))

    def scan_split(self, name: str) -> LazyFrame:
        """
        Scan the split into a lazy Polars DataFrame.

        Args:
            name: Name of the split to scan.

        Returns:
            A Polars LazyFrame containing the split data.
        """

        return pl.scan_parquet(self.split_path(name))

@dataclass(slots=True)
class SubsetMetadata:
    """
    Metadata for a subset of a dataset.

    Attributes:
        language: Language of the subset.
        pre_tokenized: Whether the subset is pre-tokenized.
        tagsets: Tag sets available in the subset.
        labels: List of entity labels in the dataset.
        splits: Metadata of the splits in the subset by split name.
        bio_labels: List of BIO labels associated with the dataset if available.
    """

    language: str
    pre_tokenized: bool
    tagsets: list[str]
    labels: dict[str, list[str]]
    splits: dict[str, SplitMetadata] = field(default_factory=dict)
    bio_labels: dict[str, list[str]] | None = None

    @classmethod
    def _default_with_language(cls, language: str) -> Self:
        """
        Creates a default SubsetMetadata instance with the given language, no labels, and sets pre_tokenized to False.

        Args:
            language: Language of the subset.

        Returns:
            Default SubsetMetadata instance.
        """

        return cls(language, False, [DEFAULT_TAGSET], {})