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

amber/__init__.py

@@ -0,0 +1,15 @@
+"""Amber: helper package for the Engineer Thesis project.
+
+This module is intentionally minimal. It exists to define the top-level package
+and to enable code coverage to include the package. Importing it should succeed
+without side effects.
+"""
+
+# A tiny bit of executable code to make the package measurable by coverage.
+PACKAGE_NAME = "amber"
+__version__ = "0.0.0"
+
+
+def ping() -> str:
+    """Return a simple response to verify the package is wired correctly."""
+    return "pong"

amber/datasets/__init__.py

@@ -0,0 +1,11 @@
+from amber.datasets.base_dataset import BaseDataset
+from amber.datasets.text_dataset import TextDataset
+from amber.datasets.classification_dataset import ClassificationDataset
+from amber.datasets.loading_strategy import LoadingStrategy
+
+__all__ = [
+    "BaseDataset",
+    "TextDataset",
+    "ClassificationDataset",
+    "LoadingStrategy",
+]

amber/datasets/base_dataset.py

@@ -0,0 +1,640 @@
+from __future__ import annotations
+
+import math
+import random
+from abc import ABC, abstractmethod
+from collections import defaultdict
+from pathlib import Path
+from typing import Any, Dict, Iterator, List, Optional, Union
+
+from datasets import Dataset, IterableDataset, load_dataset, load_from_disk
+
+from amber.datasets.loading_strategy import IndexLike, LoadingStrategy
+from amber.store.store import Store
+
+
+class BaseDataset(ABC):
+    """
+    Abstract base class for datasets with support for multiple sources,
+    loading strategies, and Store integration.
+
+    Loading Strategies:
+    - MEMORY: Load entire dataset into memory (fastest random access, highest memory usage)
+    - DISK: Save to disk, read dynamically via memory-mapped Arrow files
+      (supports len/getitem, lower memory usage)
+    - STREAMING: True streaming mode using IterableDataset
+      (lowest memory, no len/getitem support, no stratification and limit support)
+    """
+
+    def __init__(
+        self,
+        ds: Dataset | IterableDataset,
+        store: Store,
+        loading_strategy: LoadingStrategy = LoadingStrategy.MEMORY,
+    ):
+        """
+        Initialize dataset.
+
+        Args:
+            ds: HuggingFace Dataset or IterableDataset
+            store: Store instance for caching/persistence
+            loading_strategy: How to load data (MEMORY, DISK, or STREAMING)
+
+        Raises:
+            ValueError: If store is None, loading_strategy is invalid, or dataset operations fail
+            OSError: If file system operations fail
+        """
+        self._validate_initialization_params(store, loading_strategy)
+
+        self._store = store
+        self._loading_strategy = loading_strategy
+        self._dataset_dir: Path = Path(store.base_path) / store.dataset_prefix
+
+        is_iterable_input = isinstance(ds, IterableDataset)
+
+        if loading_strategy == LoadingStrategy.MEMORY:
+            # MEMORY: Convert to Dataset if needed, save to disk, load fully into memory
+            self._is_iterable = False
+            if is_iterable_input:
+                ds = Dataset.from_generator(lambda: iter(ds))
+            self._ds = self._save_and_load_dataset(ds, use_memory_mapping=False)
+        elif loading_strategy == LoadingStrategy.DISK:
+            # DISK: Save to disk, use memory-mapped Arrow files (supports len/getitem)
+            self._is_iterable = False
+            if is_iterable_input:
+                ds = Dataset.from_generator(lambda: iter(ds))
+            self._ds = self._save_and_load_dataset(ds, use_memory_mapping=True)
+        elif loading_strategy == LoadingStrategy.STREAMING:
+            # STREAMING: Convert to IterableDataset, don't save to disk (no len/getitem)
+            if not is_iterable_input:
+                ds = ds.to_iterable_dataset()
+            self._is_iterable = True
+            self._ds = ds
+            # Don't save to disk for iterable-only mode
+        else:
+            raise ValueError(
+                f"Unknown loading strategy: {loading_strategy}. Must be one of: {[s.value for s in LoadingStrategy]}"
+            )
+
+    def _validate_initialization_params(self, store: Store, loading_strategy: LoadingStrategy) -> None:
+        """Validate initialization parameters.
+
+        Args:
+            store: Store instance to validate
+            loading_strategy: Loading strategy to validate
+
+        Raises:
+            ValueError: If store is None or loading_strategy is invalid
+        """
+        if store is None:
+            raise ValueError("store cannot be None")
+
+        if not isinstance(loading_strategy, LoadingStrategy):
+            raise ValueError(f"loading_strategy must be a LoadingStrategy enum value, got: {type(loading_strategy)}")
+
+    def _has_valid_dataset_dir(self) -> bool:
+        """Check if dataset directory path is valid (non-empty base_path).
+
+        Returns:
+            True if base_path is not empty, False otherwise
+        """
+        return bool(self._store.base_path and str(self._store.base_path).strip())
+
+    def _save_and_load_dataset(self, ds: Dataset, use_memory_mapping: bool = True) -> Dataset:
+        """Save dataset to disk and load it back (with optional memory mapping).
+
+        Args:
+            ds: Dataset to save and load
+            use_memory_mapping: Whether to use memory mapping (True for DISK)
+
+        Returns:
+            Loaded dataset
+
+        Raises:
+            OSError: If file system operations fail
+            RuntimeError: If dataset operations fail
+        """
+        if self._has_valid_dataset_dir():
+            try:
+                self._dataset_dir.mkdir(parents=True, exist_ok=True)
+                ds.save_to_disk(str(self._dataset_dir))
+                return load_from_disk(str(self._dataset_dir))
+            except OSError as e:
+                raise OSError(f"Failed to save/load dataset at {self._dataset_dir}. Error: {e}") from e
+            except Exception as e:
+                raise RuntimeError(f"Failed to process dataset at {self._dataset_dir}. Error: {e}") from e
+        else:
+            return ds
+
+    @classmethod
+    def _postprocess_non_streaming_dataset(
+        cls,
+        ds: Dataset,
+        *,
+        filters: Optional[Dict[str, Any]] = None,
+        limit: Optional[int] = None,
+        stratify_by: Optional[str] = None,
+        stratify_seed: Optional[int] = None,
+        drop_na_columns: Optional[List[str]] = None,
+    ) -> Dataset:
+        """Apply filters, stratified sampling, and limits to an in-memory dataset."""
+
+        if drop_na_columns:
+            ds = cls._drop_na(ds, drop_na_columns)
+
+        if filters:
+            ds = cls._apply_filters(ds, filters)
+
+        limit_applied = False
+        if stratify_by:
+            sample_size = limit if limit is not None else len(ds)
+            if sample_size is not None and sample_size <= 0:
+                raise ValueError(f"limit must be > 0 when stratifying, got: {sample_size}")
+            ds = cls._stratified_sample(
+                ds,
+                stratify_by=stratify_by,
+                sample_size=sample_size,
+                seed=stratify_seed,
+            )
+            limit_applied = True
+
+        if limit is not None and not limit_applied:
+            if limit <= 0:
+                raise ValueError(f"limit must be > 0, got: {limit}")
+            ds = ds.select(range(min(limit, len(ds))))
+
+        return ds
+
+    @staticmethod
+    def _drop_na(ds: Dataset, columns: List[str]) -> Dataset:
+        """Drop rows where any of the specified columns are None or empty string."""
+
+        def _is_valid(example: Dict[str, Any]) -> bool:
+            for col in columns:
+                val = example.get(col)
+                if val is None:
+                    return False
+                if isinstance(val, str) and not val.strip():
+                    return False
+            return True
+
+        return ds.filter(_is_valid)
+
+    @staticmethod
+    def _apply_filters(ds: Dataset, filters: Dict[str, Any]) -> Dataset:
+        """Apply exact-match filters to a Dataset."""
+
+        def _predicate(example: Dict[str, Any]) -> bool:
+            return all(example.get(key) == value for key, value in filters.items())
+
+        return ds.filter(_predicate)
+
+    @staticmethod
+    def _stratified_sample(  # noqa: C901
+        ds: Dataset,
+        *,
+        stratify_by: str,
+        sample_size: Optional[int],
+        seed: Optional[int],
+    ) -> Dataset:
+        """Return a stratified sample of the dataset with the requested size."""
+
+        if stratify_by not in ds.column_names:
+            raise ValueError(f"Column '{stratify_by}' not found in dataset columns: {ds.column_names}")
+
+        total_rows = len(ds)
+        if total_rows == 0:
+            return ds
+
+        if sample_size is None:
+            sample_size = total_rows
+
+        sample_size = min(sample_size, total_rows)
+        if sample_size <= 0:
+            raise ValueError("sample_size must be greater than 0 for stratification")
+
+        column_values = ds[stratify_by]
+        label_to_indices: Dict[Any, List[int]] = defaultdict(list)
+        for idx, label in enumerate(column_values):
+            label_to_indices[label].append(idx)
+
+        label_counts = {label: len(indices) for label, indices in label_to_indices.items()}
+        allocations: Dict[Any, int] = {}
+        fractional_parts: List[tuple[float, int, Any]] = []
+
+        allocated_total = 0
+        for order, (label, count) in enumerate(label_counts.items()):
+            exact_allocation = (count / total_rows) * sample_size
+            base_allocation = min(count, int(math.floor(exact_allocation)))
+            allocations[label] = base_allocation
+            allocated_total += base_allocation
+            fractional_parts.append((exact_allocation - base_allocation, order, label))
+
+        remaining = sample_size - allocated_total
+        fractional_parts.sort(key=lambda item: (-item[0], item[1]))
+        for _, _, label in fractional_parts:
+            if remaining <= 0:
+                break
+            available = label_counts[label] - allocations[label]
+            if available <= 0:
+                continue
+            take = min(available, remaining)
+            allocations[label] += take
+            remaining -= take
+
+        rng = random.Random(seed)
+        selected_indices: List[int] = []
+        for label, count in allocations.items():
+            if count <= 0:
+                continue
+            indices = label_to_indices[label]
+            if count >= len(indices):
+                chosen = list(indices)
+            else:
+                chosen = rng.sample(indices, count)
+            selected_indices.extend(chosen)
+
+        rng.shuffle(selected_indices)
+        return ds.select(selected_indices)
+
+    @staticmethod
+    def _load_csv_source(
+        source: Union[str, Path],
+        *,
+        delimiter: str,
+        streaming: bool,
+        **kwargs,
+    ) -> Dataset | IterableDataset:
+        """Load a CSV dataset from disk using HuggingFace datasets."""
+
+        p = Path(source)
+        if not p.exists():
+            raise FileNotFoundError(f"CSV file not found: {source}")
+        if not p.is_file():
+            raise ValueError(f"Source must be a file, got: {source}")
+
+        try:
+            return load_dataset(
+                "csv",
+                data_files=str(p),
+                split="train",
+                delimiter=delimiter,
+                streaming=streaming,
+                **kwargs,
+            )
+        except Exception as e:
+            raise RuntimeError(f"Failed to load CSV dataset from {source}. Error: {e}") from e
+
+    @staticmethod
+    def _load_json_source(
+        source: Union[str, Path],
+        *,
+        streaming: bool,
+        **kwargs,
+    ) -> Dataset | IterableDataset:
+        """Load a JSON/JSONL dataset from disk using HuggingFace datasets."""
+
+        p = Path(source)
+        if not p.exists():
+            raise FileNotFoundError(f"JSON file not found: {source}")
+        if not p.is_file():
+            raise ValueError(f"Source must be a file, got: {source}")
+
+        try:
+            return load_dataset(
+                "json",
+                data_files=str(p),
+                split="train",
+                streaming=streaming,
+                **kwargs,
+            )
+        except Exception as e:
+            raise RuntimeError(f"Failed to load JSON dataset from {source}. Error: {e}") from e
+
+    def get_batch(self, start: int, batch_size: int) -> List[Any]:
+        """
+        Get a contiguous batch of items.
+
+        Args:
+            start: Starting index
+            batch_size: Number of items to retrieve
+
+        Returns:
+            List of items
+
+        Raises:
+            NotImplementedError: If loading_strategy is STREAMING
+        """
+        if self._loading_strategy == LoadingStrategy.STREAMING:
+            raise NotImplementedError("get_batch not supported for STREAMING datasets. Use iter_batches instead.")
+        if batch_size <= 0:
+            return []
+        end = min(start + batch_size, len(self))
+        if start >= end:
+            return []
+        return self[start:end]
+
+    def head(self, n: int = 5) -> List[Any]:
+        """
+        Get first n items.
+
+        Works for all loading strategies.
+
+        Args:
+            n: Number of items to retrieve (default: 5)
+
+        Returns:
+            List of first n items
+        """
+        if self._loading_strategy == LoadingStrategy.STREAMING:
+            items = []
+            for i, item in enumerate(self.iter_items()):
+                if i >= n:
+                    break
+                items.append(item)
+            return items
+        return self[:n]
+
+    def sample(self, n: int = 5) -> List[Any]:
+        """
+        Get n random items from the dataset.
+
+        Works for MEMORY and DISK strategies only.
+
+        Args:
+            n: Number of items to sample
+
+        Returns:
+            List of n randomly sampled items
+
+        Raises:
+            NotImplementedError: If loading_strategy is STREAMING
+        """
+        if self._loading_strategy == LoadingStrategy.STREAMING:
+            raise NotImplementedError(
+                "sample() not supported for STREAMING datasets. Use iter_items() and sample manually."
+            )
+
+        dataset_len = len(self)
+        if n <= 0:
+            return []
+        if n >= dataset_len:
+            # Return all items in random order
+            indices = list(range(dataset_len))
+            random.shuffle(indices)
+            return [self[i] for i in indices]
+
+        # Sample n random indices
+        indices = random.sample(range(dataset_len), n)
+        # Use __getitem__ with list of indices
+        return self[indices]
+
+    @property
+    def is_streaming(self) -> bool:
+        """Whether this dataset is streaming (DISK or STREAMING)."""
+        return self._loading_strategy in (LoadingStrategy.DISK, LoadingStrategy.STREAMING)
+
+    @abstractmethod
+    def __len__(self) -> int:
+        """Return the number of items in the dataset."""
+        pass
+
+    @abstractmethod
+    def __getitem__(self, idx: IndexLike) -> Any:
+        """Get item(s) by index."""
+        pass
+
+    @abstractmethod
+    def iter_items(self) -> Iterator[Any]:
+        """Iterate over items one by one."""
+        pass
+
+    @abstractmethod
+    def iter_batches(self, batch_size: int) -> Iterator[List[Any]]:
+        """Iterate over items in batches."""
+        pass
+
+    @abstractmethod
+    def extract_texts_from_batch(self, batch: List[Any]) -> List[str]:
+        """Extract text strings from a batch.
+
+        Args:
+            batch: A batch as returned by iter_batches()
+
+        Returns:
+            List of text strings ready for model inference
+        """
+        pass
+
+    @abstractmethod
+    def get_all_texts(self) -> List[str]:
+        """Get all texts from the dataset.
+
+        Returns:
+            List of all text strings in the dataset
+
+        Raises:
+            NotImplementedError: If not supported for streaming datasets
+        """
+        pass
+
+    # --- Factory methods ---
+
+    @classmethod
+    def from_huggingface(
+        cls,
+        repo_id: str,
+        store: Store,
+        *,
+        split: str = "train",
+        loading_strategy: LoadingStrategy = LoadingStrategy.MEMORY,
+        revision: Optional[str] = None,
+        streaming: Optional[bool] = None,
+        filters: Optional[Dict[str, Any]] = None,
+        limit: Optional[int] = None,
+        stratify_by: Optional[str] = None,
+        stratify_seed: Optional[int] = None,
+        **kwargs,
+    ) -> "BaseDataset":
+        """
+        Load dataset from HuggingFace Hub.
+
+        Args:
+            repo_id: HuggingFace dataset repository ID
+            store: Store instance
+            split: Dataset split (e.g., "train", "validation")
+            loading_strategy: Loading strategy (MEMORY, DISK, or STREAMING)
+            revision: Optional git revision/branch/tag
+            streaming: Optional override for streaming (if None, uses loading_strategy)
+            filters: Optional dict of column->value pairs used for exact-match filtering
+            limit: Optional maximum number of rows to keep (applied after filtering/stratification)
+            stratify_by: Optional column to use for stratified sampling (non-streaming only)
+            stratify_seed: Optional RNG seed for deterministic stratification
+            **kwargs: Additional arguments passed to load_dataset
+
+        Returns:
+            BaseDataset instance
+
+        Raises:
+            ValueError: If repo_id is empty or store is None
+            RuntimeError: If dataset loading fails
+        """
+        if not repo_id or not isinstance(repo_id, str) or not repo_id.strip():
+            raise ValueError(f"repo_id must be a non-empty string, got: {repo_id!r}")
+
+        if store is None:
+            raise ValueError("store cannot be None")
+
+        # Determine if we should use streaming for HuggingFace load_dataset
+        use_streaming = streaming if streaming is not None else (loading_strategy == LoadingStrategy.STREAMING)
+
+        if stratify_by and loading_strategy == LoadingStrategy.STREAMING:
+            raise NotImplementedError("Stratification is not supported for STREAMING datasets.")
+
+        try:
+            ds = load_dataset(
+                path=repo_id,
+                split=split,
+                revision=revision,
+                streaming=use_streaming,
+                **kwargs,
+            )
+        except Exception as e:
+            raise RuntimeError(
+                f"Failed to load dataset from HuggingFace Hub: repo_id={repo_id!r}, "
+                f"split={split!r}, revision={revision!r}. Error: {e}"
+            ) from e
+
+        if use_streaming:
+            if filters or limit or stratify_by:
+                raise NotImplementedError(
+                    "filters, limit, and stratification are not supported when streaming datasets. "
+                    "Choose MEMORY or DISK loading strategy instead."
+                )
+        else:
+            ds = cls._postprocess_non_streaming_dataset(
+                ds,
+                filters=filters,
+                limit=limit,
+                stratify_by=stratify_by,
+                stratify_seed=stratify_seed,
+            )
+
+        return cls(ds, store=store, loading_strategy=loading_strategy)
+
+    @classmethod
+    def from_csv(
+        cls,
+        source: Union[str, Path],
+        store: Store,
+        *,
+        loading_strategy: LoadingStrategy = LoadingStrategy.MEMORY,
+        text_field: str = "text",
+        delimiter: str = ",",
+        stratify_by: Optional[str] = None,
+        stratify_seed: Optional[int] = None,
+        drop_na_columns: Optional[List[str]] = None,
+        **kwargs,
+    ) -> "BaseDataset":
+        """
+        Load 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
+            delimiter: CSV delimiter (default: comma)
+            stratify_by: Optional column used for stratified sampling (non-streaming only)
+            stratify_seed: Optional RNG seed for stratified sampling
+            drop_na_columns: Optional list of columns to check for None/empty values
+            **kwargs: Additional arguments passed to load_dataset
+
+        Returns:
+            BaseDataset instance
+
+        Raises:
+            FileNotFoundError: If CSV file doesn't exist
+            ValueError: If store is None or source is invalid
+            RuntimeError: If dataset loading fails
+        """
+        if store is None:
+            raise ValueError("store cannot be None")
+
+        use_streaming = loading_strategy == LoadingStrategy.STREAMING
+        if (stratify_by or drop_na_columns) and use_streaming:
+            raise NotImplementedError("Stratification and drop_na are not supported for STREAMING datasets.")
+
+        ds = cls._load_csv_source(
+            source,
+            delimiter=delimiter,
+            streaming=use_streaming,
+            **kwargs,
+        )
+
+        if not use_streaming and (stratify_by or drop_na_columns):
+            ds = cls._postprocess_non_streaming_dataset(
+                ds,
+                stratify_by=stratify_by,
+                stratify_seed=stratify_seed,
+                drop_na_columns=drop_na_columns,
+            )
+
+        return cls(ds, store=store, loading_strategy=loading_strategy)
+
+    @classmethod
+    def from_json(
+        cls,
+        source: Union[str, Path],
+        store: Store,
+        *,
+        loading_strategy: LoadingStrategy = LoadingStrategy.MEMORY,
+        text_field: str = "text",
+        stratify_by: Optional[str] = None,
+        stratify_seed: Optional[int] = None,
+        drop_na_columns: Optional[List[str]] = None,
+        **kwargs,
+    ) -> "BaseDataset":
+        """
+        Load dataset from JSON or 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 (for JSON objects)
+            stratify_by: Optional column used for stratified sampling (non-streaming only)
+            stratify_seed: Optional RNG seed for stratified sampling
+            drop_na_columns: Optional list of columns to check for None/empty values
+            **kwargs: Additional arguments passed to load_dataset
+
+        Returns:
+            BaseDataset instance
+
+        Raises:
+            FileNotFoundError: If JSON file doesn't exist
+            ValueError: If store is None or source is invalid
+            RuntimeError: If dataset loading fails
+        """
+        if store is None:
+            raise ValueError("store cannot be None")
+
+        use_streaming = loading_strategy == LoadingStrategy.STREAMING
+        if (stratify_by or drop_na_columns) and use_streaming:
+            raise NotImplementedError("Stratification and drop_na are not supported for STREAMING datasets.")
+
+        ds = cls._load_json_source(
+            source,
+            streaming=use_streaming,
+            **kwargs,
+        )
+
+        if not use_streaming and (stratify_by or drop_na_columns):
+            ds = cls._postprocess_non_streaming_dataset(
+                ds,
+                stratify_by=stratify_by,
+                stratify_seed=stratify_seed,
+                drop_na_columns=drop_na_columns,
+            )
+
+        return cls(ds, store=store, loading_strategy=loading_strategy)