Data Pipes¶

class ChecksumValidationError(Exception):
    """Raised when a checksum differs from an expected value."""

@dataclass
class ConvertModule:
    """
    Module for converting dataset files to the MELD parquet format.

    Attributes:
        config: ConvertStep or DatasetConvertStep configuration for the
            conversion process.
    """

    config: ConvertStep | DatasetConvertStep

    def __post_init__(self) -> None:
        """
        Initializes the ConvertModule with a reader based on the configuration.
        """
        if isinstance(self.config, DatasetConvertStep):
            self._reader = DatasetReader(self.config)
        else:
            self._reader = reader_from_config(self.config.reader)

    def _label_set(self) -> _LabelSet:
        """
        Creates a label set configured based on the reader's tagsets.

        Returns:
            _LabelSet with label sets for label inference or explicit
            label maps.
        """

        # Skip inference of possible labels in the dataset if a label_map has been explicitly defined
        if isinstance(self._reader, _ReaderWithTagset) and self._reader.tagsets is not None:
            # Only infer labels for tagsets without an explicit label map
            infer_labels = {tagset.name for tagset in self._reader.tagsets if tagset.label_map is None}
            bio_labels = {
                tagset.name: set() if tagset.label_map is None else set(map(str, tagset.label_map.values()))
                for tagset in self._reader.tagsets
            }

            labels = {
                tagset.name: set()
                if tagset.label_map is None
                else {label.entity_type for label in tagset.label_map.values() if label.entity_type is not None}
                for tagset in self._reader.tagsets
            }
            return _LabelSet(labels, bio_labels, infer_labels)

        # Assumes other readers only support a single, default tagset
        return _LabelSet({DEFAULT_TAGSET: set()}, {DEFAULT_TAGSET: set()}, {DEFAULT_TAGSET})

    def __call__(self, state: DataPipeState) -> DataPipeState:
        """
        Converts all splits to the MELD parquet format, infers tag sets if necessary, and writes output files.

        Args:
            state: The DataPipeState containing splits to convert.

        Returns:
            Updated DataPipeState with metadata of the processed splits.

        Raises:
            ValueError: If an empty document is encountered during
                processing.
            ValueError: If overlapping subsets are encountered.
            ValueError: If an empty split is encountered.
            ValueError: If the output format is unsupported.
            ValueError: If reading a split fails.
        """

        if not state.splits:
            raise ValueError("ConvertModule received no data splits to process")

        # NOTE: Could support additional output formats in the future
        if state.output_format != "parquet":
            raise ValueError(f"Unsupported output format: {state.output_format}")

        main_splits = None
        subsets = {}
        languages = set()
        label_sets: dict[str, _LabelSet] = {}

        state.meld_directory.mkdir(exist_ok=True)

        for split in tqdm(state.splits, unit=" splits", desc="Converting splits"):
            # Normalize split language
            try:
                language_tag = Language.get(split.language)
                split.language = language_tag.to_alpha3() + (
                    "" if language_tag.script is None else ("-" + language_tag.script)
                )
            except LanguageTagError:
                logger.warning(f"Keeping invalid language code {split.language!r}")

            languages.add(split.language)

            # Filename with format based extension
            split_filename = f"{split.name}.{state.output_format}"
            subset_slug = ""
            if not split.subset:
                path = state.meld_directory / split_filename
                if main_splits is None:
                    main_splits = SubsetMetadata._default_with_language(split.language)
                current_subset_metadata = main_splits
                current_split_metadata = main_splits.splits[split.name] = SplitMetadata(
                    path.relative_to(state.meld_directory)
                )
            else:
                subset_slug = split.slug()

                path = state.meld_directory / subset_slug
                path.mkdir(exist_ok=True)
                path /= split_filename

                component_dict: SubsetHierarchy = subsets
                for component in split.subset[:-1]:
                    name = component.name
                    if name not in component_dict:
                        component_dict[name] = {}

                    new_dict = component_dict[name]
                    if not isinstance(new_dict, dict):
                        raise TypeError(f"Overlapping subsets encountered: {new_dict}")

                    component_dict = new_dict

                last_subset = split.subset[-1].name
                if (current_subset_metadata := component_dict.get(last_subset)) is None:
                    current_subset_metadata = SubsetMetadata._default_with_language(split.language)
                    component_dict[last_subset] = current_subset_metadata
                elif not isinstance(current_subset_metadata, SubsetMetadata):
                    raise TypeError(f"Overlapping subsets encountered: {last_subset}")

                current_split_metadata = current_subset_metadata.splits[split.name] = SplitMetadata(
                    path.relative_to(state.meld_directory)
                )

            documents = iter(self._reader.read_split(split))
            # Infers whether documents are pre-tokenized from the first document
            try:
                first_document = next(documents)
                current_subset_metadata.pre_tokenized = has_tokens = first_document.bio is not None
                current_subset_metadata.tagsets = tagsets = sorted(first_document.spans[0].labels.keys())
            except StopIteration:
                raise ValueError(f"Encountered empty split {split}")

            if (subset_labels := label_sets.get(subset_slug)) is None:
                label_sets[subset_slug] = subset_labels = self._label_set()

            sentence_span_filename = f"{state.dataset_name}--{subset_slug}--{split.name}.parquet"
            if state.sentence_span_dir is None:
                sentence_span_file = resources.files("meld") / "package_data/sentence_spans" / sentence_span_filename
            else:
                sentence_span_file = state.sentence_span_dir / sentence_span_filename

            # Only sentence tokenize if the dataset isn't already segmented into sentences
            sentence_boundaries = state.split_data_pipe_metadata(
                [subset.name for subset in split.subset], split.name
            ).sentence_boundaries

            splitter = (
                contextlib.nullcontext()
                if sentence_boundaries == "full"
                else SentenceSplitter(
                    sentence_boundaries,
                    sentence_span_file,
                    read_spans=not state.sentence_span_dir,
                )
            )

            document_count = 0
            sequence_count = 0
            filtered_count = 0

            with (
                NERParquetWriter.open(
                    state.dataset_name,
                    subset_slug,
                    split.name,
                    path,
                    has_tokens,
                    tagsets,
                ) as writer,
                splitter as splitter,
            ):
                for document in itertools.chain((first_document,), documents):
                    # Sort annotations for reproducibility first
                    _sort_annotations(document)
                    if splitter is not None:
                        document = splitter.sentence_tokenize(document)

                    # Sanity check
                    if any(not span.text for span in document.spans):
                        if self.config.filter_empty_documents:
                            filtered_count += 1
                            continue
                        raise ValueError(f"Empty text span found in {document!r}")

                    if subset_labels.tagsets_to_infer:
                        subset_labels.add_labels(document)

                    writer.write_document(document)
                    document_count += 1
                    sequence_count += len(document.spans)

            if filtered_count > 0:
                logger.warning(f"Removed {filtered_count} empty documents in split {subset_slug}.{split.name}")

            current_split_metadata.document_count = document_count
            current_split_metadata.sequence_count = sequence_count
            # Sorts labels for consistency and sets "bio_labels" to None if no BIO labels were found for any tagset
            current_subset_metadata.labels = {
                name: sorted(tagset_labels) for name, tagset_labels in subset_labels.labels.items()
            }
            current_subset_metadata.bio_labels = (
                None
                if not any(subset_labels.bio_labels.values())
                else {name: sorted(tagset_labels) for name, tagset_labels in subset_labels.bio_labels.items()}
            )

        state.metadata = _ProcessedDatasetMetadata(
            subsets,
            main_splits,
            languages,
        )
        return state

class DataPipe:
    """
    Orchestrates dataset processing through a series of modular steps.

    Manages the execution order of data pipe steps, validates dependencies,
    and maintains state throughout the pipeline.

    Args:
        step_definitions: Iterable of DataPipeStep and MetadataStep
            configurations.
        predefined_dataset_metadata: List of sub-metadata providing
            additional information about the dataset to process.
        loader: Optional loader module to initialize the first
            stage.
        assume_completed: Set of step IDs to consider already
            completed for validation purposes in nested data pipes.

    Raises:
        DataPipeError: If a duplicate step is encountered of a type
            for which duplicates are not allowed.
        DataPipeError: If a step's dependencies are not satisfied.
    """

    _stages: list[list[DataPipeModule]]

    def __init__(
        self,
        step_definitions: Iterable[DataPipeStep | MetadataStep],
        predefined_dataset_metadata: list[SubMetadata] | None = None,
        loader: Loader | None = None,
        assume_completed: Iterable[str] | None = None,
    ) -> None:
        if predefined_dataset_metadata is None:
            predefined_dataset_metadata = []
        self._predefined_dataset_metadata = predefined_dataset_metadata

        self._stage_definitions = []
        completed_steps = set() if loader is None else {loader.loader}
        if assume_completed:
            completed_steps.update(assume_completed)
        initial_steps = completed_steps.copy()
        last_step = loader

        for i, step in enumerate(step_definitions):
            step_type = step.step
            if step_type in completed_steps and step_type not in _DUPLICATES_ALLOWED:
                raise DataPipeError(f"Duplicate {step_type!r} step at position {i}")
            if step_type == last_step:
                self._stage_definitions[-1].append(step)
                continue
            dependencies = _STEP_DEPENDENCIES[step_type]
            if not (dependencies is None or any(dependency in completed_steps for dependency in dependencies)):
                raise DataPipeError(
                    f"Missing dependency for {step_type!r} step at position {i}."
                    f" Needs at least one completed step of {_STEP_DEPENDENCIES[step_type]}"
                )
            self._stage_definitions.append([step])
            completed_steps.add(step_type)
            last_step = step_type

        match loader:
            case HuggingfaceLoader(arguments=arguments):
                self._stages = [[HuggingfaceModule(arguments)]]
            case GitLoader(arguments=arguments):
                self._stages = [[GitModule(arguments)]]
            case None | GenericLoader():
                self._stages = []

        self._stages.extend(
            [_config_to_stage(step, predefined_dataset_metadata, initial_steps) for step in parallel_steps]
            for parallel_steps in self._stage_definitions
        )

    def __call__(self, state: DataPipeState) -> DataPipeState:
        """
        Executes the data pipe with the given state.

        Args:
            state: Initial DataPipeState to process.

        Returns:
            Updated DataPipeState after all stages have run.
        """

        # Create a scoped state to support nested data pipes with their own contexts
        sub_state = state.scoped_state()
        for stage in self._stages:
            for step in stage:
                sub_state = step(sub_state)

        # Merge changes into the shared state
        state.update(sub_state)
        return state

    def run(
        self,
        dataset_name: str,
        intermediate_directory: Path,
        meld_directory: Path,
        sentence_span_dir: Path | None = None,
    ) -> DataPipeState:
        """
        Creates a new DataPipeState and executes the data pipe.

        Args:
            dataset_name: Name of the dataset being processed.
            intermediate_directory: Directory for intermediate
                downloaded, extracted or cached files.
            meld_directory: Directory for final MELD output.
            sentence_span_dir: Optional directory for sentence span
                files used for applying sentence segmentation to
                document-level datasets.

        Returns:
            Final DataPipeState after data pipe execution.
        """

        return self(
            DataPipeState(
                dataset_name,
                intermediate_directory,
                meld_directory,
                self._predefined_dataset_metadata,
                sentence_span_dir=sentence_span_dir,
            )
        )

class DataPipeError(Exception):
    """Exception raised for errors in a data pipe."""

@dataclass(slots=True)
class DataPipeState:
    """
    State of a data pipe.

    Attributes:
        dataset_name: Name of the dataset being processed.
        intermediate_directory: Directory for intermediate files.
        meld_directory: Directory for final MELD output.
        data_pipe_metadata: List of metadata configurations to apply.
        file_streams: Open file streams for processing.
        generated_files: Set of files generated during processing.
        splits: List of split configurations.
        metadata: Current dataset metadata.
        output_format: Format for output files.
        sentence_span_dir: Optional directory for sentence span files to
            segment document-level datasets.
    """

    dataset_name: str
    intermediate_directory: Path
    meld_directory: Path
    data_pipe_metadata: list[SubMetadata]
    file_streams: dict[Path, AbstractContextManager[Iterable[bytes]]] = field(default_factory=dict)
    generated_files: set[Path] = field(default_factory=set)
    splits: list[SplitData] = field(default_factory=list)
    metadata: _ProcessedDatasetMetadata = field(default_factory=_ProcessedDatasetMetadata)
    output_format: str = "parquet"
    sentence_span_dir: Path | None = None

    def scoped_state(self) -> Self:
        """
        Creates a new state with shared resources but without accumulated splits or metadata.

        Returns:
            New DataPipeState instance with shared file streams and
            generated files.
        """

        # State with shared file streams and generated files but without previous splits or metadata
        return self.__class__(
            self.dataset_name,
            self.intermediate_directory,
            self.meld_directory,
            self.data_pipe_metadata,
            self.file_streams,
            self.generated_files,
            output_format=self.output_format,
            sentence_span_dir=self.sentence_span_dir,
        )

    def update(self, other: Self) -> None:
        """
        Updates this state with data from another state.

        Args:
            other: The other DataPipeState to merge with.
        """

        self.splits.extend(other.splits)
        self.metadata.merge(other.metadata)

    def split_data_pipe_metadata(self, subset_hierarchy: list[str], split_name: str) -> Metadata:
        """
        Create metadata for a given subset hierarchy and split.

        Args:
            subset_hierarchy: Path through the subset hierarchy.
            split_name: Name of the split.

        Returns:
            The `Metadata` for the given split.
        """

        applicable_metadata = []

        for sub_metadata in self.data_pipe_metadata:
            # Treat empty selector as a general wildcard "*"
            if not sub_metadata.split:
                applicable_metadata.append(((0, 0), sub_metadata))
                continue

            specificity = max(
                (
                    selector.specificity()
                    for selector in sub_metadata.split
                    if selector.matches(subset_hierarchy, split_name)
                ),
                default=None,
            )
            if specificity is not None:
                applicable_metadata.append((specificity, sub_metadata))

        return _merge_metadata(applicable_metadata)

    def dataset_metadata(self) -> DatasetMetadata:
        """
        Constructs the final DatasetMetadata from the current state.

        Returns:
            Complete DatasetMetadata instance.
        """

        metadata = DatasetMetadata(
            self.dataset_name,
            self.metadata.subsets,
            self.metadata.languages,
            self.metadata.main_splits,
        )

        for subset in metadata._iter_subsets():
            for split_name, split in subset.splits.items():
                split.creation_metadata = self.split_data_pipe_metadata(subset.hierarchy, split_name)

        return metadata

@dataclass(slots=True)
class DownloadModule:
    """
    Module for downloading files from URLs with SHA-256 verification.

    Attributes:
        config: DownloadStep configuration containing URLs and
            checksums.
    """

    config: DownloadStep

    def __call__(self, state: DataPipeState) -> DataPipeState:
        """
        Downloads files from configured URLs and verifies their SHA-256 checksums.

        Args:
            state: The DataPipeState to process.

        Returns:
            Updated DataPipeState with downloaded files added to
            generated_files.
        """

        for url in self.config.urls:
            path = Path(urlparse(url.url).path)
            target_path = state.intermediate_directory / path.name
            # Skip any further checks and downloading of the file other than checksum verification
            # when it already exists in the target directory with the name that could be parsed from the URL
            if target_path.is_file():
                result = target_path
            else:
                state.intermediate_directory.mkdir(exist_ok=True, parents=True)
                # Skip downloading if the file already exists and the filename could not already be correctly inferred from the URL
                result = download.download(
                    url.url, state.intermediate_directory, skip_if_file_exists=True, stream=False
                )

            # Verify checksum of the downloaded file
            _verify_sha256(result, url.sha256)

            state.generated_files.add(result)

        return state