mi-crow 0.1.1.post12__py3-none-any.whl

amber/datasets/classification_dataset.py

@@ -0,0 +1,566 @@
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any, Dict, Iterator, List, Optional, Sequence, Union
+
+from datasets import Dataset, IterableDataset, load_dataset
+
+from amber.datasets.base_dataset import BaseDataset
+from amber.datasets.loading_strategy import IndexLike, LoadingStrategy
+from amber.store.store import Store
+
+
+class ClassificationDataset(BaseDataset):
+    """
+    Classification dataset with text and category/label columns.
+    Each item is a dict with 'text' and label column(s) as keys.
+    Supports single or multiple label columns.
+    """
+
+    def __init__(
+        self,
+        ds: Dataset | IterableDataset,
+        store: Store,
+        loading_strategy: LoadingStrategy = LoadingStrategy.MEMORY,
+        text_field: str = "text",
+        category_field: Union[str, List[str]] = "category",
+    ):
+        """
+        Initialize classification dataset.
+
+        Args:
+            ds: HuggingFace Dataset or IterableDataset
+            store: Store instance
+            loading_strategy: Loading strategy
+            text_field: Name of the column containing text
+            category_field: Name(s) of the column(s) containing category/label.
+                          Can be a single string or a list of strings for multiple labels.
+
+        Raises:
+            ValueError: If text_field or category_field is empty, or fields not found in dataset
+        """
+        self._validate_text_field(text_field)
+
+        # Normalize category_field to list
+        if isinstance(category_field, str):
+            self._category_fields = [category_field]
+        else:
+            self._category_fields = list(category_field)
+
+        self._validate_category_fields(self._category_fields)
+
+        # Validate dataset
+        is_iterable = isinstance(ds, IterableDataset)
+        if not is_iterable:
+            if text_field not in ds.column_names:
+                raise ValueError(f"Dataset must have a '{text_field}' column; got columns: {ds.column_names}")
+            for cat_field in self._category_fields:
+                if cat_field not in ds.column_names:
+                    raise ValueError(f"Dataset must have a '{cat_field}' column; got columns: {ds.column_names}")
+            # Set format with all required columns
+            format_columns = [text_field] + self._category_fields
+            ds.set_format("python", columns=format_columns)
+
+        self._text_field = text_field
+        self._category_field = category_field  # Keep original for backward compatibility
+        super().__init__(ds, store=store, loading_strategy=loading_strategy)
+
+    def _validate_text_field(self, text_field: str) -> None:
+        """Validate text_field parameter.
+
+        Args:
+            text_field: Text field name to validate
+
+        Raises:
+            ValueError: If text_field is empty or not a string
+        """
+        if not text_field or not isinstance(text_field, str) or not text_field.strip():
+            raise ValueError(f"text_field must be a non-empty string, got: {text_field!r}")
+
+    def _validate_category_fields(self, category_fields: List[str]) -> None:
+        """Validate category_fields parameter.
+
+        Args:
+            category_fields: List of category field names to validate
+
+        Raises:
+            ValueError: If category_fields is empty or contains invalid values
+        """
+        if not category_fields:
+            raise ValueError("category_field cannot be empty")
+
+        for cat_field in category_fields:
+            if not cat_field or not isinstance(cat_field, str) or not cat_field.strip():
+                raise ValueError(f"All category fields must be non-empty strings, got invalid field: {cat_field!r}")
+
+    def _extract_item_from_row(self, row: Dict[str, Any]) -> Dict[str, Any]:
+        """Extract item (text + categories) from a dataset row.
+
+        Args:
+            row: Dataset row dictionary
+
+        Returns:
+            Dictionary with 'text' and category fields as keys
+
+        Raises:
+            ValueError: If required fields are not found in row
+        """
+        if self._text_field in row:
+            text = row[self._text_field]
+        elif "text" in row:
+            text = row["text"]
+        else:
+            raise ValueError(
+                f"Text field '{self._text_field}' or 'text' not found in dataset row. "
+                f"Available fields: {list(row.keys())}"
+            )
+
+        item = {"text": text}
+        for cat_field in self._category_fields:
+            category = row.get(cat_field)
+            if category is None:
+                raise ValueError(
+                    f"Category field '{cat_field}' not found in dataset row. Available fields: {list(row.keys())}"
+                )
+            category = row.get(cat_field)  # Potentially None
+            item[cat_field] = category
+
+        return item
+
+    def __len__(self) -> int:
+        """
+        Return the number of items in the dataset.
+
+        Raises:
+            NotImplementedError: If loading_strategy is STREAMING
+        """
+        if self._loading_strategy == LoadingStrategy.STREAMING:
+            raise NotImplementedError("len() not supported for STREAMING datasets")
+        return self._ds.num_rows
+
+    def __getitem__(self, idx: IndexLike) -> Union[Dict[str, Any], List[Dict[str, Any]]]:
+        """
+        Get item(s) by index. Returns dict with 'text' and label column(s) as keys.
+
+        For single label: {"text": "...", "category": "..."}
+        For multiple labels: {"text": "...", "label1": "...", "label2": "..."}
+
+        Args:
+            idx: Index (int), slice, or sequence of indices
+
+        Returns:
+            Single item dict or list of item dicts
+
+        Raises:
+            NotImplementedError: If loading_strategy is STREAMING
+            IndexError: If index is out of bounds
+            ValueError: If dataset is empty
+        """
+        if self._loading_strategy == LoadingStrategy.STREAMING:
+            raise NotImplementedError(
+                "Indexing not supported for STREAMING datasets. Use iter_items or iter_batches."
+            )
+
+        dataset_len = len(self)
+        if dataset_len == 0:
+            raise ValueError("Cannot index into empty dataset")
+
+        if isinstance(idx, int):
+            if idx < 0:
+                idx = dataset_len + idx
+            if idx < 0 or idx >= dataset_len:
+                raise IndexError(f"Index {idx} out of bounds for dataset of length {dataset_len}")
+            row = self._ds[idx]
+            return self._extract_item_from_row(row)
+
+        if isinstance(idx, slice):
+            start, stop, step = idx.indices(dataset_len)
+            if step != 1:
+                indices = list(range(start, stop, step))
+                selected = self._ds.select(indices)
+            else:
+                selected = self._ds.select(range(start, stop))
+            return [self._extract_item_from_row(row) for row in selected]
+
+        if isinstance(idx, Sequence):
+            # Validate all indices are in bounds
+            invalid_indices = [i for i in idx if not (0 <= i < dataset_len)]
+            if invalid_indices:
+                raise IndexError(f"Indices out of bounds: {invalid_indices} (dataset length: {dataset_len})")
+            selected = self._ds.select(list(idx))
+            return [self._extract_item_from_row(row) for row in selected]
+
+        raise TypeError(f"Invalid index type: {type(idx)}")
+
+    def iter_items(self) -> Iterator[Dict[str, Any]]:
+        """
+        Iterate over items one by one. Yields dict with 'text' and label column(s) as keys.
+
+        For single label: {"text": "...", "category_column_1": "..."}
+        For multiple labels: {"text": "...", "category_column_1": "...", "category_column_2": "..."}
+
+        Yields:
+            Item dictionaries with text and category fields
+
+        Raises:
+            ValueError: If required fields are not found in any row
+        """
+        for row in self._ds:
+            yield self._extract_item_from_row(row)
+
+    def iter_batches(self, batch_size: int) -> Iterator[List[Dict[str, Any]]]:
+        """
+        Iterate over items in batches. Each batch is a list of dicts with 'text' and label column(s) as keys.
+
+        For single label: [{"text": "...", "category_column_1": "..."}, ...]
+        For multiple labels: [{"text": "...", "category_column_1": "...", "category_column_2": "..."}, ...]
+
+        Args:
+            batch_size: Number of items per batch
+
+        Yields:
+            Lists of item dictionaries (batches)
+
+        Raises:
+            ValueError: If batch_size <= 0 or required fields are not found in any row
+        """
+        if batch_size <= 0:
+            raise ValueError(f"batch_size must be > 0, got: {batch_size}")
+
+        if self._loading_strategy == LoadingStrategy.STREAMING:
+            batch = []
+            for row in self._ds:
+                batch.append(self._extract_item_from_row(row))
+                if len(batch) >= batch_size:
+                    yield batch
+                    batch = []
+            if batch:
+                yield batch
+        else:
+            # Use select to get batches with proper format
+            for i in range(0, len(self), batch_size):
+                end = min(i + batch_size, len(self))
+                batch_list = self[i:end]
+                yield batch_list
+
+    def get_categories(self) -> Union[List[Any], Dict[str, List[Any]]]:  # noqa: C901
+        """
+        Get unique categories in the dataset, excluding None values.
+
+        Returns:
+            - For single label column: List of unique category values
+            - For multiple label columns: Dict mapping column name to list of unique categories
+
+        Raises:
+            NotImplementedError: If loading_strategy is STREAMING and dataset is large
+        """
+        if len(self._category_fields) == 1:
+            # Single label: return list for backward compatibility
+            cat_field = self._category_fields[0]
+            if self._loading_strategy == LoadingStrategy.STREAMING:
+                categories = set()
+                for item in self.iter_items():
+                    cat = item[cat_field]
+                    if cat is not None:
+                        categories.add(cat)
+                return sorted(list(categories))  # noqa: C414
+            categories = [cat for cat in set(self._ds[cat_field]) if cat is not None]
+            return sorted(categories)
+        else:
+            # Multiple labels: return dict
+            result = {}
+            if self._loading_strategy == LoadingStrategy.STREAMING:
+                # Collect categories from all items
+                category_sets = {field: set() for field in self._category_fields}
+                for item in self.iter_items():
+                    for field in self._category_fields:
+                        cat = item[field]
+                        if cat is not None:
+                            category_sets[field].add(cat)
+                for field in self._category_fields:
+                    result[field] = sorted(list(category_sets[field]))  # noqa: C414
+            else:
+                # Use direct column access
+                for field in self._category_fields:
+                    categories = [cat for cat in set(self._ds[field]) if cat is not None]
+                    result[field] = sorted(categories)
+            return result
+
+    def extract_texts_from_batch(self, batch: List[Dict[str, Any]]) -> List[Optional[str]]:
+        """Extract text strings from a batch of classification items.
+
+        Args:
+            batch: List of dicts with 'text' and category fields
+
+        Returns:
+            List of text strings from the batch
+
+        Raises:
+            ValueError: If 'text' key is not found in any batch item
+        """
+        texts = []
+        for item in batch:
+            if "text" not in item:
+                raise ValueError(f"'text' key not found in batch item. Available keys: {list(item.keys())}")
+            texts.append(item["text"])
+        return texts
+
+    def get_all_texts(self) -> List[Optional[str]]:
+        """Get all texts from the dataset.
+
+        Returns:
+            List of all text strings
+
+        Raises:
+            NotImplementedError: If loading_strategy is STREAMING and dataset is very large
+        """
+        if self._loading_strategy == LoadingStrategy.STREAMING:
+            return [item["text"] for item in self.iter_items()]
+        return list(self._ds[self._text_field])
+
+    def get_categories_for_texts(self, texts: List[Optional[str]]) -> Union[List[Any], List[Dict[str, Any]]]:
+        """
+        Get categories for given texts (if texts match dataset texts).
+
+        Args:
+            texts: List of text strings to look up
+
+        Returns:
+            - For single label column: List of category values (one per text)
+            - For multiple label columns: List of dicts with label columns as keys
+
+        Raises:
+            NotImplementedError: If loading_strategy is STREAMING
+            ValueError: If texts list is empty
+        """
+        if self._loading_strategy == LoadingStrategy.STREAMING:
+            raise NotImplementedError("get_categories_for_texts not supported for STREAMING datasets")
+
+        if not texts:
+            raise ValueError("texts list cannot be empty")
+
+        if len(self._category_fields) == 1:
+            # Single label: return list for backward compatibility
+            cat_field = self._category_fields[0]
+            text_to_category = {row[self._text_field]: row[cat_field] for row in self._ds}
+            return [text_to_category.get(text) for text in texts]
+        else:
+            # Multiple labels: return list of dicts
+            text_to_categories = {
+                row[self._text_field]: {field: row[field] for field in self._category_fields} for row in self._ds
+            }
+            return [text_to_categories.get(text) for text in texts]
+
+    @classmethod
+    def from_huggingface(
+        cls,
+        repo_id: str,
+        store: Store,
+        *,
+        split: str = "train",
+        loading_strategy: LoadingStrategy = LoadingStrategy.MEMORY,
+        revision: Optional[str] = None,
+        text_field: str = "text",
+        category_field: Union[str, List[str]] = "category",
+        filters: Optional[Dict[str, Any]] = None,
+        limit: Optional[int] = None,
+        stratify_by: Optional[str] = None,
+        stratify_seed: Optional[int] = None,
+        streaming: Optional[bool] = None,
+        drop_na: bool = False,
+        **kwargs,
+    ) -> "ClassificationDataset":
+        """
+        Load classification dataset from HuggingFace Hub.
+
+        Args:
+            repo_id: HuggingFace dataset repository ID
+            store: Store instance
+            split: Dataset split
+            loading_strategy: Loading strategy
+            revision: Optional git revision
+            text_field: Name of the column containing text
+            category_field: Name(s) of the column(s) containing category/label
+            filters: Optional filters to apply (dict of column: value)
+            limit: Optional limit on number of rows
+            stratify_by: Optional column used for stratified sampling (non-streaming only)
+            stratify_seed: Optional RNG seed for stratified sampling
+            streaming: Optional override for streaming
+            drop_na: Whether to drop rows with None/empty text or categories
+            **kwargs: Additional arguments for load_dataset
+
+        Returns:
+            ClassificationDataset instance
+
+        Raises:
+            ValueError: If parameters are invalid
+            RuntimeError: If dataset loading fails
+        """
+        use_streaming = streaming if streaming is not None else (loading_strategy == LoadingStrategy.STREAMING)
+
+        if (stratify_by or drop_na) and use_streaming:
+            raise NotImplementedError(
+                "Stratification and drop_na are not supported for streaming datasets. Use MEMORY or DISK."
+            )
+
+        try:
+            ds = load_dataset(
+                path=repo_id,
+                split=split,
+                revision=revision,
+                streaming=use_streaming,
+                **kwargs,
+            )
+
+            if use_streaming:
+                if filters or limit:
+                    raise NotImplementedError(
+                        "filters and limit are not supported when streaming datasets. Choose MEMORY or DISK."
+                    )
+            else:
+                drop_na_columns = None
+                if drop_na:
+                    cat_fields = [category_field] if isinstance(category_field, str) else category_field
+                    drop_na_columns = [text_field] + list(cat_fields)
+
+                ds = cls._postprocess_non_streaming_dataset(
+                    ds,
+                    filters=filters,
+                    limit=limit,
+                    stratify_by=stratify_by,
+                    stratify_seed=stratify_seed,
+                    drop_na_columns=drop_na_columns,
+                )
+        except Exception as e:
+            raise RuntimeError(
+                f"Failed to load classification dataset from HuggingFace Hub: "
+                f"repo_id={repo_id!r}, split={split!r}, text_field={text_field!r}, "
+                f"category_field={category_field!r}. Error: {e}"
+            ) from e
+
+        return cls(
+            ds,
+            store=store,
+            loading_strategy=loading_strategy,
+            text_field=text_field,
+            category_field=category_field,
+        )
+
+    @classmethod
+    def from_csv(
+        cls,
+        source: Union[str, Path],
+        store: Store,
+        *,
+        loading_strategy: LoadingStrategy = LoadingStrategy.MEMORY,
+        text_field: str = "text",
+        category_field: Union[str, List[str]] = "category",
+        delimiter: str = ",",
+        stratify_by: Optional[str] = None,
+        stratify_seed: Optional[int] = None,
+        drop_na: bool = False,
+        **kwargs,
+    ) -> "ClassificationDataset":
+        """
+        Load classification dataset from CSV file.
+
+        Args:
+            source: Path to CSV file
+            store: Store instance
+            loading_strategy: Loading strategy
+            text_field: Name of the column containing text
+            category_field: Name(s) of the column(s) containing category/label
+            delimiter: CSV delimiter (default: comma)
+            stratify_by: Optional column used for stratified sampling
+            stratify_seed: Optional RNG seed for stratified sampling
+            drop_na: Whether to drop rows with None/empty text or categories
+            **kwargs: Additional arguments for load_dataset
+
+        Returns:
+            ClassificationDataset instance
+
+        Raises:
+            FileNotFoundError: If CSV file doesn't exist
+            RuntimeError: If dataset loading fails
+        """
+        drop_na_columns = None
+        if drop_na:
+            cat_fields = [category_field] if isinstance(category_field, str) else category_field
+            drop_na_columns = [text_field] + list(cat_fields)
+
+        dataset = super().from_csv(
+            source,
+            store=store,
+            loading_strategy=loading_strategy,
+            text_field=text_field,
+            delimiter=delimiter,
+            stratify_by=stratify_by,
+            stratify_seed=stratify_seed,
+            drop_na_columns=drop_na_columns,
+            **kwargs,
+        )
+        return cls(
+            dataset._ds,
+            store=store,
+            loading_strategy=loading_strategy,
+            text_field=text_field,
+            category_field=category_field,
+        )
+
+    @classmethod
+    def from_json(
+        cls,
+        source: Union[str, Path],
+        store: Store,
+        *,
+        loading_strategy: LoadingStrategy = LoadingStrategy.MEMORY,
+        text_field: str = "text",
+        category_field: Union[str, List[str]] = "category",
+        stratify_by: Optional[str] = None,
+        stratify_seed: Optional[int] = None,
+        drop_na: bool = False,
+        **kwargs,
+    ) -> "ClassificationDataset":
+        """
+        Load classification dataset from JSON/JSONL file.
+
+        Args:
+            source: Path to JSON or JSONL file
+            store: Store instance
+            loading_strategy: Loading strategy
+            text_field: Name of the field containing text
+            category_field: Name(s) of the field(s) containing category/label
+            stratify_by: Optional column used for stratified sampling
+            stratify_seed: Optional RNG seed for stratified sampling
+            drop_na: Whether to drop rows with None/empty text or categories
+            **kwargs: Additional arguments for load_dataset
+
+        Returns:
+            ClassificationDataset instance
+
+        Raises:
+            FileNotFoundError: If JSON file doesn't exist
+            RuntimeError: If dataset loading fails
+        """
+        drop_na_columns = None
+        if drop_na:
+            cat_fields = [category_field] if isinstance(category_field, str) else category_field
+            drop_na_columns = [text_field] + list(cat_fields)
+
+        dataset = super().from_json(
+            source,
+            store=store,
+            loading_strategy=loading_strategy,
+            text_field=text_field,
+            stratify_by=stratify_by,
+            stratify_seed=stratify_seed,
+            drop_na_columns=drop_na_columns,
+            **kwargs,
+        )
+        return cls(
+            dataset._ds,
+            store=store,
+            loading_strategy=loading_strategy,
+            text_field=text_field,
+            category_field=category_field,
+        )

amber/datasets/loading_strategy.py

@@ -0,0 +1,29 @@
+from __future__ import annotations
+
+from enum import Enum
+from typing import Union, Sequence, TypeAlias
+
+
+class LoadingStrategy(Enum):
+    """
+    Strategy for loading dataset data.
+    
+    Choose the best strategy for your use case:
+    
+    - MEMORY: Load entire dataset into memory (fastest random access, highest memory usage)
+      Best for: Small datasets that fit in memory, when you need fast random access
+    
+    - DISK: Save to disk, read dynamically via memory-mapped Arrow files
+      (supports len/getitem, lower memory usage)
+      Best for: Large datasets that don't fit in memory, when you need random access
+    
+    - STREAMING: True streaming mode using IterableDataset (lowest memory, no len/getitem support)
+      Best for: Very large datasets, when you only need sequential iteration
+    """
+    MEMORY = "memory"  # Load all into memory (fastest random access, highest memory usage)
+    DISK = "disk"  # Save to disk, read dynamically via memory-mapped Arrow files (supports len/getitem, lower memory usage)
+    STREAMING = "streaming"  # True streaming mode using IterableDataset (lowest memory, no len/getitem support)
+
+
+IndexLike: TypeAlias = Union[int, slice, Sequence[int]]
+