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

amber/datasets/text_dataset.py

@@ -0,0 +1,488 @@
+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 TextDataset(BaseDataset):
+    """
+    Text-only dataset with support for multiple sources and loading strategies.
+    Each item is a string (text snippet).
+    """
+
+    def __init__(
+        self,
+        ds: Dataset | IterableDataset,
+        store: Store,
+        loading_strategy: LoadingStrategy = LoadingStrategy.MEMORY,
+        text_field: str = "text",
+    ):
+        """
+        Initialize text dataset.
+
+        Args:
+            ds: HuggingFace Dataset or IterableDataset
+            store: Store instance
+            loading_strategy: Loading strategy
+            text_field: Name of the column containing text
+
+        Raises:
+            ValueError: If text_field is empty or not found in dataset
+        """
+        self._validate_text_field(text_field)
+
+        # Validate and prepare 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}")
+            # Keep only text column for memory efficiency
+            columns_to_remove = [c for c in ds.column_names if c != text_field]
+            if columns_to_remove:
+                ds = ds.remove_columns(columns_to_remove)
+            if text_field != "text":
+                ds = ds.rename_column(text_field, "text")
+            ds.set_format("python", columns=["text"])
+
+        self._text_field = text_field
+        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 _extract_text_from_row(self, row: Dict[str, Any]) -> Optional[str]:
+        """Extract text from a dataset row.
+
+        Args:
+            row: Dataset row dictionary
+
+        Returns:
+            Text string from the row
+
+        Raises:
+            ValueError: If text field is 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())}"
+            )
+        return text
+
+    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[Optional[str], List[Optional[str]]]:
+        """
+        Get text item(s) by index.
+
+        Args:
+            idx: Index (int), slice, or sequence of indices
+
+        Returns:
+            Single text string or list of text strings
+
+        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}")
+            return self._ds[idx]["text"]
+
+        if isinstance(idx, slice):
+            start, stop, step = idx.indices(dataset_len)
+            if step != 1:
+                indices = list(range(start, stop, step))
+                out = self._ds.select(indices)["text"]
+            else:
+                out = self._ds.select(range(start, stop))["text"]
+            return list(out)
+
+        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})")
+            out = self._ds.select(list(idx))["text"]
+            return list(out)
+
+        raise TypeError(f"Invalid index type: {type(idx)}")
+
+    def iter_items(self) -> Iterator[Optional[str]]:
+        """
+        Iterate over text items one by one.
+
+        Yields:
+            Text strings from the dataset
+
+        Raises:
+            ValueError: If text field is not found in any row
+        """
+        for row in self._ds:
+            yield self._extract_text_from_row(row)
+
+    def iter_batches(self, batch_size: int) -> Iterator[List[Optional[str]]]:
+        """
+        Iterate over text items in batches.
+
+        Args:
+            batch_size: Number of items per batch
+
+        Yields:
+            Lists of text strings (batches)
+
+        Raises:
+            ValueError: If batch_size <= 0 or text field is 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_text_from_row(row))
+                if len(batch) >= batch_size:
+                    yield batch
+                    batch = []
+            if batch:
+                yield batch
+        else:
+            for batch in self._ds.iter(batch_size=batch_size):
+                yield list(batch["text"])
+
+    def extract_texts_from_batch(self, batch: List[Optional[str]]) -> List[Optional[str]]:
+        """Extract text strings from a batch.
+
+        For TextDataset, batch items are already strings, so return as-is.
+
+        Args:
+            batch: List of text strings
+
+        Returns:
+            List of text strings (same as input)
+        """
+        return batch
+
+    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
+        """
+        if self._loading_strategy == LoadingStrategy.STREAMING:
+            return list(self.iter_items())
+        return list(self._ds["text"])
+
+    @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",
+        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,
+    ) -> "TextDataset":
+        """
+        Load text 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
+            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 deterministic stratification
+            streaming: Optional override for streaming
+            drop_na: Whether to drop rows with None/empty text
+            **kwargs: Additional arguments for load_dataset
+
+        Returns:
+            TextDataset 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 = [text_field] if drop_na else None
+                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 text dataset from HuggingFace Hub: "
+                f"repo_id={repo_id!r}, split={split!r}, text_field={text_field!r}. "
+                f"Error: {e}"
+            ) from e
+
+        return cls(ds, store=store, loading_strategy=loading_strategy, text_field=text_field)
+
+    @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: bool = False,
+        **kwargs,
+    ) -> "TextDataset":
+        """
+        Load text 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 to use for stratified sampling
+            stratify_seed: Optional RNG seed for stratified sampling
+            drop_na: Whether to drop rows with None/empty text
+            **kwargs: Additional arguments for load_dataset
+
+        Returns:
+            TextDataset instance
+
+        Raises:
+            FileNotFoundError: If CSV file doesn't exist
+            RuntimeError: If dataset loading fails
+        """
+        drop_na_columns = [text_field] if drop_na else None
+        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,
+        )
+
+    @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: bool = False,
+        **kwargs,
+    ) -> "TextDataset":
+        """
+        Load text 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
+            stratify_by: Optional column to use for stratified sampling
+            stratify_seed: Optional RNG seed for stratified sampling
+            drop_na: Whether to drop rows with None/empty text
+            **kwargs: Additional arguments for load_dataset
+
+        Returns:
+            TextDataset instance
+
+        Raises:
+            FileNotFoundError: If JSON file doesn't exist
+            RuntimeError: If dataset loading fails
+        """
+        drop_na_columns = [text_field] if drop_na else None
+        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,
+        )
+        # Re-initialize with text_field
+        return cls(
+            dataset._ds,
+            store=store,
+            loading_strategy=loading_strategy,
+            text_field=text_field,
+        )
+
+    @classmethod
+    def from_local(
+        cls,
+        source: Union[str, Path],
+        store: Store,
+        *,
+        loading_strategy: LoadingStrategy = LoadingStrategy.MEMORY,
+        text_field: str = "text",
+        recursive: bool = True,
+    ) -> "TextDataset":
+        """
+        Load from a local directory or file(s).
+
+        Supported:
+          - Directory of .txt files (each file becomes one example)
+          - JSONL/JSON/CSV/TSV files with a text column
+
+        Args:
+            source: Path to directory or file
+            store: Store instance
+            loading_strategy: Loading strategy
+            text_field: Name of the column/field containing text
+            recursive: Whether to recursively search directories for .txt files
+
+        Returns:
+            TextDataset instance
+
+        Raises:
+            FileNotFoundError: If source path doesn't exist
+            ValueError: If source is invalid or unsupported file type
+            RuntimeError: If file operations fail
+        """
+        p = Path(source)
+        if not p.exists():
+            raise FileNotFoundError(f"Source path does not exist: {source}")
+
+        if p.is_dir():
+            txts: List[str] = []
+            pattern = "**/*.txt" if recursive else "*.txt"
+            try:
+                for fp in sorted(p.glob(pattern)):
+                    txts.append(fp.read_text(encoding="utf-8", errors="ignore"))
+            except OSError as e:
+                raise RuntimeError(f"Failed to read text files from directory {source}. Error: {e}") from e
+
+            if not txts:
+                raise ValueError(f"No .txt files found in directory: {source} (recursive={recursive})")
+
+            ds = Dataset.from_dict({"text": txts})
+        else:
+            suffix = p.suffix.lower()
+            if suffix in {".jsonl", ".json"}:
+                return cls.from_json(
+                    source,
+                    store=store,
+                    loading_strategy=loading_strategy,
+                    text_field=text_field,
+                )
+            elif suffix in {".csv"}:
+                return cls.from_csv(
+                    source,
+                    store=store,
+                    loading_strategy=loading_strategy,
+                    text_field=text_field,
+                )
+            elif suffix in {".tsv"}:
+                return cls.from_csv(
+                    source,
+                    store=store,
+                    loading_strategy=loading_strategy,
+                    text_field=text_field,
+                    delimiter="\t",
+                )
+            else:
+                raise ValueError(
+                    f"Unsupported file type: {suffix} for source: {source}. "
+                    f"Use directory of .txt, or JSON/JSONL/CSV/TSV."
+                )
+
+        return cls(ds, store=store, loading_strategy=loading_strategy, text_field=text_field)

amber/hooks/__init__.py

@@ -0,0 +1,20 @@
+from amber.hooks.hook import Hook, HookType, HookError
+from amber.hooks.detector import Detector
+from amber.hooks.controller import Controller
+from amber.hooks.implementations.layer_activation_detector import LayerActivationDetector
+from amber.hooks.implementations.model_input_detector import ModelInputDetector
+from amber.hooks.implementations.model_output_detector import ModelOutputDetector
+from amber.hooks.implementations.function_controller import FunctionController
+
+__all__ = [
+    "Hook",
+    "HookType",
+    "HookError",
+    "Detector",
+    "Controller",
+    "LayerActivationDetector",
+    "ModelInputDetector",
+    "ModelOutputDetector",
+    "FunctionController",
+]
+

amber/hooks/controller.py

@@ -0,0 +1,171 @@
+from __future__ import annotations
+
+import abc
+from typing import TYPE_CHECKING
+
+import torch
+import torch.nn as nn
+
+from amber.hooks.hook import Hook, HookType, HOOK_FUNCTION_INPUT, HOOK_FUNCTION_OUTPUT
+from amber.hooks.utils import extract_tensor_from_input, extract_tensor_from_output
+from amber.utils import get_logger
+
+if TYPE_CHECKING:
+    pass
+
+logger = get_logger(__name__)
+
+
+class Controller(Hook):
+    """
+    Abstract base class for controller hooks that modify activations during inference.
+    
+    Controllers can modify inputs (pre_forward) or outputs (forward) of layers.
+    They are designed to actively change the behavior of the model during inference.
+    """
+
+    def __init__(
+            self,
+            hook_type: HookType | str = HookType.FORWARD,
+            hook_id: str | None = None,
+            layer_signature: str | int | None = None
+    ):
+        """
+        Initialize a controller hook.
+        
+        Args:
+            hook_type: Type of hook (HookType.FORWARD or HookType.PRE_FORWARD)
+            hook_id: Unique identifier
+            layer_signature: Layer to attach to (optional, for compatibility)
+        """
+        super().__init__(layer_signature=layer_signature, hook_type=hook_type, hook_id=hook_id)
+
+    def _handle_pre_forward(
+            self,
+            module: torch.nn.Module,
+            input: HOOK_FUNCTION_INPUT
+    ) -> HOOK_FUNCTION_INPUT | None:
+        """Handle pre-forward hook execution.
+        
+        Args:
+            module: The PyTorch module being hooked
+            input: Tuple of input tensors to the module
+            
+        Returns:
+            Modified input tuple or None to keep original
+        """
+        input_tensor = extract_tensor_from_input(input)
+
+        if input_tensor is None:
+            return None
+
+        modified_tensor = self.modify_activations(module, input_tensor, input_tensor)
+
+        if modified_tensor is not None and isinstance(modified_tensor, torch.Tensor):
+            result = list(input)
+            if len(result) > 0:
+                result[0] = modified_tensor
+            return tuple(result)
+        return None
+
+    def _handle_forward(
+            self,
+            module: torch.nn.Module,
+            input: HOOK_FUNCTION_INPUT,
+            output: HOOK_FUNCTION_OUTPUT
+    ) -> None:
+        """Handle forward hook execution.
+
+        Args:
+            module: The PyTorch module being hooked
+            input: Tuple of input tensors to the module
+            output: Output tensor(s) from the module
+        """
+        output_tensor = extract_tensor_from_output(output)
+
+        if output_tensor is None:
+            return
+
+        # Extract input tensor if available for modify_activations
+        input_tensor = extract_tensor_from_input(input)
+
+        # Note: forward hooks can't modify output in PyTorch, but we call modify_activations
+        # for consistency. The actual modification happens via the hook mechanism.
+        # We still call it so controllers can capture/process activations.
+        self.modify_activations(module, input_tensor, output_tensor)
+
+    def _hook_fn(
+            self,
+            module: torch.nn.Module,
+            input: HOOK_FUNCTION_INPUT,
+            output: HOOK_FUNCTION_OUTPUT
+    ) -> None | HOOK_FUNCTION_INPUT:
+        """
+        Internal hook function that modifies activations.
+        
+        If the instance also inherits from Detector, first processes activations
+        as a Detector (saves metadata), then modifies activations as a Controller.
+        
+        Args:
+            module: The PyTorch module being hooked
+            input: Tuple of input tensors to the module
+            output: Output tensor(s) from the module
+            
+        Returns:
+            For pre_forward hooks: modified inputs (tuple) or None to keep original
+            For forward hooks: None (forward hooks cannot modify output in PyTorch)
+            
+        Raises:
+            RuntimeError: If modify_activations raises an exception
+        """
+        if not self._enabled:
+            return None
+
+        # Check if this instance also inherits from Detector
+        if self._is_both_controller_and_detector():
+            # First, process activations as a Detector (save metadata)
+            try:
+                self.process_activations(module, input, output)
+            except Exception as e:
+                logger.warning(
+                    f"Error in {self.__class__.__name__} detector process_activations: {e}",
+                    exc_info=True
+                )
+
+        try:
+            if self.hook_type == HookType.PRE_FORWARD:
+                return self._handle_pre_forward(module, input)
+            else:
+                self._handle_forward(module, input, output)
+                return None
+        except Exception as e:
+            raise RuntimeError(
+                f"Error in controller {self.id} modify_activations: {e}"
+            ) from e
+
+    @abc.abstractmethod
+    def modify_activations(
+            self,
+            module: nn.Module,
+            inputs: torch.Tensor | None,
+            output: torch.Tensor | None
+    ) -> torch.Tensor | None:
+        """
+        Modify activations from the hooked layer.
+        
+        For pre_forward hooks: receives input tensor, should return modified input tensor.
+        For forward hooks: receives input and output tensors, should return modified output tensor.
+        
+        Args:
+            module: The PyTorch module being hooked
+            inputs: Input tensor (None for forward hooks if not available)
+            output: Output tensor (None for pre_forward hooks)
+            
+        Returns:
+            Modified input tensor (for pre_forward) or modified output tensor (for forward).
+            Return None to keep original tensor unchanged.
+            
+        Raises:
+            Exception: Subclasses may raise exceptions for invalid inputs or modification errors
+        """
+        raise NotImplementedError("modify_activations must be implemented by subclasses")