cureta 0.1.0__py3-none-any.whl

cureta/__init__.py

@@ -0,0 +1,136 @@
+"""Cureta: Unified Data Curation Framework.
+
+.. attribute:: __version__
+
+    Package version string resolved from installed package metadata
+    (e.g. ``"0.1.0"``).  Falls back to ``"0.0.0+unknown"`` when the
+    package is not installed (e.g. in an editable development checkout
+    without ``pip install -e .`` having been run).
+
+A pipeline-based library that standardises data processing into a sequence
+of **Taggers** which enrich a standard **Document** object with extracted
+features and metadata.
+
+Execution surfaces
+------------------
+run_pipeline          -- Canonical entry point; runs end-to-end, returns PipelineResult.
+iter_pipeline_chunks  -- Generator variant; yields ChunkResult after each chunk.
+quick_tag             -- REPL one-shot; runs taggers on a single string, returns dict.
+stream_tagged_batches -- Streaming without disk writes; yields pd.DataFrame batches.
+
+Result types
+------------
+PipelineResult  -- Structured summary of a completed run.
+ChunkResult     -- Per-chunk timing and output path.
+
+Config / composition types
+--------------------------
+PipelineConfig      -- Pydantic model for pipeline-level configuration.
+PipelineStageConfig -- Pydantic model for a single pipeline stage.
+TaggerConfig        -- Pydantic model for per-tagger configuration.
+ClusterConfig       -- Pydantic model for SkyPilot cluster configuration.
+ExecutionMode       -- Enum for SIMD / SDMI / AUTO execution strategies.
+ExecutionConfig     -- Pydantic model for per-resource-type execution mode.
+OutputFormat        -- Enum for output file format (parquet / jsonl).
+
+Authoring types
+---------------
+Document    -- The canonical data envelope.
+BaseTagger  -- Abstract base for all taggers.
+CPUTagger   -- Base class for CPU-bound taggers.
+GPUTagger   -- Base class for GPU-bound taggers.
+
+Composition
+-----------
+Pipeline  -- Orchestrator; use Pipeline.from_config() or Pipeline.from_taggers()
+             to build pipelines programmatically without YAML.
+"""
+
+from importlib.metadata import PackageNotFoundError
+from importlib.metadata import version as _pkg_version
+
+try:
+    __version__: str = _pkg_version("cureta")
+except PackageNotFoundError:
+    __version__ = "0.0.0+unknown"
+
+from cureta import tag_ops  # noqa: F401 - re-exported for convenience
+from cureta.base_tagger import BaseTagger, CPUTagger, GPUTagger
+from cureta.config import (
+    ClusterConfig,
+    ExecutionConfig,
+    ExecutionMode,
+    OutputFormat,
+    PipelineConfig,
+    PipelineStageConfig,
+    TaggerConfig,
+)
+from cureta.document import Document
+from cureta.manifest import read_manifest, write_manifest
+from cureta.pipeline import Pipeline
+from cureta.results import ChunkResult, PipelineResult
+from cureta.runner import (
+    iter_pipeline_chunks,
+    quick_tag,
+    run_pipeline,
+    stream_tagged_batches,
+)
+__all__ = [
+    # Package version
+    "__version__",
+    # Execution surfaces (canonical)
+    "run_pipeline",
+    "iter_pipeline_chunks",
+    "quick_tag",
+    "stream_tagged_batches",
+    # Result types
+    "PipelineResult",
+    "ChunkResult",
+    # Config types
+    "PipelineConfig",
+    "PipelineStageConfig",
+    "TaggerConfig",
+    "ClusterConfig",
+    "ExecutionMode",
+    "ExecutionConfig",
+    "OutputFormat",
+    # Composition
+    "Pipeline",
+    # Authoring
+    "Document",
+    "BaseTagger",
+    "CPUTagger",
+    "GPUTagger",
+    # Helper modules
+    "tag_ops",
+    # Manifest / Provenance
+    "write_manifest",
+    "read_manifest",
+    # Optional / lazy-loaded
+    "IndicLIDTagger",
+    "DistributedRunner",
+    "SkyPilotCluster",
+]
+
+
+def __getattr__(name: str):  # noqa: ANN202
+    """Lazy import for optional-dependency modules.
+
+    DistributedRunner requires ``ray[data]``.  IndicLIDTagger and the other
+    GPU taggers require ``torch``/``transformers``.  Deferring these imports
+    avoids breaking users who have not installed the optional extras.
+    """
+    if name == "IndicLIDTagger":
+        from cureta.taggers.indic_lid_tagger import IndicLIDTagger
+
+        return IndicLIDTagger
+    if name == "DistributedRunner":
+        from cureta.infra.distributed_runner import DistributedRunner
+
+        return DistributedRunner
+    if name == "SkyPilotCluster":
+        from cureta.infra.skypilot_cluster import SkyPilotCluster
+
+        return SkyPilotCluster
+    msg = f"module {__name__!r} has no attribute {name!r}"
+    raise AttributeError(msg)

cureta/base_tagger.py

@@ -0,0 +1,270 @@
+"""Tagger architecture for the Cureta framework.
+
+Provides the abstract base class :class:`BaseTagger` and two intermediate
+abstract subclasses -- :class:`CPUTagger` and :class:`GPUTagger` -- that every
+concrete tagger must inherit from.
+
+Each tagger exposes:
+
+* :meth:`process_document` -- Tier-1 single-document logic (simple API).
+* :meth:`run` -- Tier-2 batch interface for vectorised processing.
+
+Implement exactly **one** of the two. The framework calls ``run()``; the
+default ``run()`` delegates to ``process_document()`` row by row.
+
+* :attr:`resources` -- ``dict`` with ``num_cpus`` / ``num_gpus``.
+* :attr:`use_actors` -- whether Ray should use ``ActorPoolStrategy``.
+
+Hierarchy
+---------
+::
+
+    BaseTagger  (ABC)
+    |- CPUTagger  (ABC)  -- taggers that run on CPU only
+    |- GPUTagger  (ABC)  -- taggers that require GPU resources
+"""
+
+from __future__ import annotations
+
+import logging
+import re
+import uuid
+from abc import ABC
+from typing import Any
+
+import pandas as pd
+
+from cureta.config import ExecutionMode, TaggerConfig
+from cureta.document import Document
+
+logger = logging.getLogger(__name__)
+
+
+class BaseTagger(ABC):
+    """Abstract base class for all feature-extraction taggers.
+
+    Parameters
+    ----------
+    config : TaggerConfig
+        Validated configuration for this tagger instance.
+
+    Raises
+    ------
+    ValueError
+        If :meth:`validate_config` determines the configuration is invalid.
+
+    Notes
+    -----
+    **Column preservation guarantee:** Any column present in the input
+    batch will be present, unchanged, in the output batch. Taggers may
+    modify ``document_id``, ``raw_content``, ``metadata``, and ``tags``.
+    They must not modify or drop any other column. Tier-2 taggers that
+    override :meth:`run` are responsible for upholding this contract.
+
+    **Two-tier API:** Override :meth:`process_document` for simple per-document
+    logic (Tier-1, ~90% of taggers). Override :meth:`run` for vectorised or
+    GPU batch operations (Tier-2). The framework dispatches to whichever is
+    overridden; if neither is overridden, pipeline construction raises
+    :class:`TypeError`.
+
+    **Preferred execution mode:** subclasses may set the class attribute
+    ``preferred_mode`` to hint to the runner which execution strategy works
+    best for them.  A mismatch between the tagger's preference and the
+    pipeline's configured mode triggers a startup warning.  ``None`` means
+    no preference — either mode is acceptable.
+
+    """
+
+    preferred_mode: ExecutionMode | None = None
+
+    def __init__(
+        self,
+        config: TaggerConfig | None = None,
+        **parameters: Any,
+    ) -> None:
+        """Initialise the tagger with its config or keyword parameters.
+
+        When *config* is ``None`` the tagger constructs a :class:`TaggerConfig`
+        automatically from the class name and any *parameters* kwargs, enabling
+        direct instantiation without YAML::
+
+            tagger = NumWordsTagger(lang="hi")
+        """
+        if config is None:
+            config = TaggerConfig(
+                tagger_name=self._default_name(),
+                tagger_class=f"{type(self).__module__}.{type(self).__name__}",
+                parameters=parameters,
+            )
+        self.config = config
+        self.validate_config()
+        logger.info("Initialised tagger '%s'", config.tagger_name)
+
+    @classmethod
+    def _default_name(cls) -> str:
+        """Derive a snake_case registry name from the class name.
+
+        ``NumWordsTagger`` → ``num_words``, ``CuretaIDTagger`` → ``cureta_id``.
+        """
+        name = cls.__name__.removesuffix("Tagger")
+        # Insert _ between a lowercase letter and the next uppercase letter.
+        s1 = re.sub(r"([a-z])([A-Z])", r"\1_\2", name)
+        # Insert _ between a run of uppercase letters and a trailing uppercase+lowercase.
+        s2 = re.sub(r"([A-Z]+)([A-Z][a-z])", r"\1_\2", s1)
+        return s2.lower()
+
+    # ------------------------------------------------------------------
+    # Abstract interface
+    # ------------------------------------------------------------------
+
+    def process_document(self, document: Document) -> Document:
+        """Tier-1 API: process a single document.
+
+        Override this for simple, per-document logic. The framework handles
+        batching, column preservation, and tag merging.
+
+        If you override ``run()`` instead (Tier-2), this method is unused and
+        you do not need to implement it.
+
+        Raises
+        ------
+        NotImplementedError
+            If neither this method nor ``run()`` is overridden.
+
+        """
+        raise NotImplementedError(
+            f"Tagger '{self.name}' must override either `process_document` "
+            f"(Tier-1) or `run` (Tier-2). See docs/writing_taggers.md."
+        )
+
+    # ------------------------------------------------------------------
+    # Batch interface (used by Ray map_batches)
+    # ------------------------------------------------------------------
+
+    def run(self, batch: pd.DataFrame) -> pd.DataFrame:
+        """Process a batch of documents as a pandas DataFrame.
+
+        The default implementation iterates rows, deserialises each into
+        a :class:`Document`, calls :meth:`process_document`, and writes
+        the enriched fields back.  Subclasses may override this for
+        vectorised / batch-optimised processing.
+
+        ``load_model()`` is called lazily on the first invocation so
+        that heavy resources are initialised inside Ray workers rather
+        than on the driver (which would cause serialisation failures).
+
+        Extra columns present in the input batch are preserved unchanged
+        in the output. Only ``document_id``, ``raw_content``,
+        ``metadata``, and ``tags`` may be modified by the tagger.
+        """
+        self.load_model()
+        out_rows: list[dict[str, Any]] = []
+        for _, row in batch.iterrows():
+            row_dict = row.to_dict()
+            doc = Document(
+                document_id=row_dict.get("document_id") or str(uuid.uuid4()),
+                raw_content=str(row_dict.get("raw_content", "")),
+                metadata=dict(row_dict.get("metadata") or {}),
+                tags=dict(row_dict.get("tags") or {}),
+            )
+            enriched = self.process_document(doc)
+            row_dict["document_id"] = enriched.document_id
+            row_dict["raw_content"] = enriched.raw_content
+            row_dict["metadata"] = enriched.metadata
+            row_dict["tags"] = enriched.tags
+            out_rows.append(row_dict)
+        return pd.DataFrame(out_rows)
+
+    # ------------------------------------------------------------------
+    # Config validation
+    # ------------------------------------------------------------------
+
+    def validate_config(self) -> None:
+        """Verify that the tagger's parameters are sane.
+
+        The default implementation performs no checks.  Override to
+        enforce domain-specific constraints and raise :class:`ValueError`.
+        """
+        return None
+
+    # ------------------------------------------------------------------
+    # Resource management
+    # ------------------------------------------------------------------
+
+    def load_model(self) -> None:
+        """Load heavy model weights into memory.
+
+        This hook is called by the Ray runner immediately before the tagger
+        is used to process a shard (SDMI) or when the actor starts (SIMD).
+        Implementations should override this to perform deferred initialization.
+        """
+        return None
+
+    def unload_model(self) -> None:
+        """Unload model weights and clear memory.
+
+        This hook is called by the Ray runner in SDMI mode after a shard is
+        fully processed by this tagger to ensure memory is cleared before
+        loading the next tagger.
+        """
+        return None
+
+    # ------------------------------------------------------------------
+    # Resource requirements
+    # ------------------------------------------------------------------
+
+    @property
+    def resources(self) -> dict[str, int]:
+        """Return Ray resource requirements for ``map_batches``.
+
+        Returns ``{"num_cpus": ..., "num_gpus": ...}``.
+        """
+        return {"num_cpus": 1, "num_gpus": 0}
+
+    @property
+    def use_actors(self) -> bool:
+        """Whether Ray should use ``ActorPoolStrategy`` for this tagger.
+
+        Defaults to ``False``.  GPU taggers override to ``True`` so that
+        heavy models stay loaded across batches.
+        """
+        return False
+
+    # ------------------------------------------------------------------
+    # Utilities
+    # ------------------------------------------------------------------
+
+    @property
+    def name(self) -> str:
+        """Return the tagger's canonical name from its config."""
+        return self.config.tagger_name
+
+    def __repr__(self) -> str:
+        """Return a string representation of the tagger."""
+        return f"{self.__class__.__name__}(name={self.name!r})"
+
+
+class CPUTagger(BaseTagger, ABC):
+    """Base for taggers that execute exclusively on CPU."""
+
+    resource_type: str = "cpu"
+
+
+class GPUTagger(BaseTagger, ABC):
+    """Base for taggers that require GPU resources.
+
+    Uses ``ActorPoolStrategy`` by default so that large models remain
+    loaded in GPU memory across batches.
+    """
+
+    resource_type: str = "gpu"
+
+    @property
+    def resources(self) -> dict[str, int]:
+        """Request 1 CPU + 1 GPU."""
+        return {"num_cpus": 1, "num_gpus": 1}
+
+    @property
+    def use_actors(self) -> bool:
+        """GPU taggers use actors to keep models warm."""
+        return True