DocumentBatch

from nemo_curator.tasks import DocumentBatch

from dataclasses import dataclass
import pandas as pd
import pyarrow as pa
@dataclass
class DocumentBatch(Task[pa.Table | pd.DataFrame]):
    """Task containing a batch of text documents.
    Attributes:
        task_id: Unique identifier for this batch.
        dataset_name: Name of the source dataset.
        data: DataFrame or PyArrow Table containing documents.
    """
    task_id: str
    dataset_name: str
    data: pa.Table | pd.DataFrame

@property
def num_items(self) -> int:
    """Returns the number of documents in this batch."""

def to_pyarrow(self) -> pa.Table:
    """Convert data to PyArrow table."""

def to_pandas(self) -> pd.DataFrame:
    """Convert data to Pandas DataFrame."""

def get_columns(self) -> list[str]:
    """Get column names from the data."""

def validate(self) -> bool:
    """Validate that the batch has required structure.
    Returns:
        True if valid, False if empty or has no columns (logs warning).
    """

import pandas as pd
from nemo_curator.tasks import DocumentBatch
# Create from DataFrame
df = pd.DataFrame({
    "text": ["Document 1 content...", "Document 2 content..."],
    "id": ["doc_001", "doc_002"],
    "url": ["https://example.com/1", "https://example.com/2"],
})
batch = DocumentBatch(
    task_id="batch_001",
    dataset_name="my_dataset",
    data=df,
)
print(f"Batch contains {batch.num_items} documents")

from dataclasses import dataclass
from nemo_curator.stages.base import ProcessingStage
from nemo_curator.tasks import DocumentBatch
@dataclass
class TextFilterStage(ProcessingStage[DocumentBatch, DocumentBatch]):
    """Filter documents based on text length."""
    name: str = "TextFilter"
    min_length: int = 100
    def inputs(self) -> tuple[list[str], list[str]]:
        return ["data"], ["text"]
    def outputs(self) -> tuple[list[str], list[str]]:
        return ["data"], ["text"]
    def process(self, task: DocumentBatch) -> DocumentBatch | None:
        df = task.data
        # Filter by text length
        mask = df["text"].str.len() >= self.min_length
        filtered_df = df[mask]
        if filtered_df.empty:
            return None
        return DocumentBatch(
            task_id=f"{task.task_id}_filtered",
            dataset_name=task.dataset_name,
            data=filtered_df,
            _metadata=task._metadata,
            _stage_perf=task._stage_perf,
        )

def process(self, task: DocumentBatch) -> DocumentBatch:
    df = task.data.copy()
    df["word_count"] = df["text"].str.split().str.len()
    return DocumentBatch(
        task_id=f"{task.task_id}_{self.name}",
        dataset_name=task.dataset_name,
        data=df,
        _metadata=task._metadata,
        _stage_perf=task._stage_perf,
    )

def process(self, task: DocumentBatch) -> DocumentBatch | None:
    df = task.data
    filtered = df[df["score"] > self.threshold]
    if filtered.empty:
        return None  # Filter out entire batch
    return DocumentBatch(
        task_id=f"{task.task_id}_{self.name}",
        dataset_name=task.dataset_name,
        data=filtered,
        _metadata=task._metadata,
        _stage_perf=task._stage_perf,
    )