mi-crow 1.0.0__py3-none-any.whl → 1.0.0.post1__py3-none-any.whl

mi_crow/datasets/base_dataset.py

@@ -114,11 +114,14 @@ class BaseDataset(ABC):
             OSError: If file system operations fail
             RuntimeError: If dataset operations fail
         """
+        if len(ds) == 0:
+            return ds
+        
         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))
+                return load_from_disk(str(self._dataset_dir), keep_in_memory=not use_memory_mapping)
             except OSError as e:
                 raise OSError(f"Failed to save/load dataset at {self._dataset_dir}. Error: {e}") from e
             except Exception as e:
@@ -522,6 +525,73 @@ class BaseDataset(ABC):
 
         return cls(ds, store=store, loading_strategy=loading_strategy)
 
+    @classmethod
+    def from_disk(
+        cls,
+        store: Store,
+        *,
+        loading_strategy: LoadingStrategy = LoadingStrategy.MEMORY,
+        **kwargs: Any,
+    ) -> "BaseDataset":
+        """
+        Load dataset from already-saved Arrow files on disk.
+
+        Use this when you've previously saved a dataset and want to reload it
+        without re-downloading from HuggingFace or re-applying transformations.
+
+        Args:
+            store: Store instance pointing to where the dataset was saved
+                   (dataset will be loaded from store.base_path/store.dataset_prefix/)
+            loading_strategy: Loading strategy (MEMORY or DISK only, not STREAMING)
+            **kwargs: Additional arguments (for subclass compatibility)
+
+        Returns:
+            BaseDataset instance loaded from disk
+
+        Raises:
+            ValueError: If store is None or loading_strategy is STREAMING
+            FileNotFoundError: If dataset directory doesn't exist
+            RuntimeError: If dataset loading fails
+
+        Example:
+            # First: save dataset
+            dataset_store = LocalStore("store/my_dataset")
+            dataset = ClassificationDataset.from_huggingface(..., store=dataset_store)
+            # Dataset saved to: store/my_dataset/datasets/*.arrow
+
+            # Later: reload from disk
+            dataset_store = LocalStore("store/my_dataset")
+            dataset = ClassificationDataset.from_disk(store=dataset_store)
+        """
+        if store is None:
+            raise ValueError("store cannot be None")
+
+        if loading_strategy == LoadingStrategy.STREAMING:
+            raise ValueError("STREAMING loading strategy not supported for from_disk(). Use MEMORY or DISK.")
+
+        dataset_dir = Path(store.base_path) / store.dataset_prefix
+
+        if not dataset_dir.exists():
+            raise FileNotFoundError(
+                f"Dataset directory not found: {dataset_dir}. "
+                f"Make sure you've previously saved a dataset to this store location."
+            )
+
+        # Verify it's a valid Arrow dataset directory
+        arrow_files = list(dataset_dir.glob("*.arrow"))
+        if not arrow_files:
+            raise FileNotFoundError(
+                f"No Arrow files found in {dataset_dir}. Directory exists but doesn't contain a valid dataset."
+            )
+
+        try:
+            use_memory_mapping = loading_strategy == LoadingStrategy.DISK
+            ds = load_from_disk(str(dataset_dir), keep_in_memory=not use_memory_mapping)
+        except Exception as e:
+            raise RuntimeError(f"Failed to load dataset from {dataset_dir}. Error: {e}") from e
+
+        return cls(ds, store=store, loading_strategy=loading_strategy)
+
     @classmethod
     def from_csv(
         cls,

mi_crow/datasets/classification_dataset.py

@@ -3,7 +3,7 @@ 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 datasets import Dataset, IterableDataset, load_dataset, load_from_disk
 
 from mi_crow.datasets.base_dataset import BaseDataset
 from mi_crow.datasets.loading_strategy import IndexLike, LoadingStrategy
@@ -117,8 +117,7 @@ class ClassificationDataset(BaseDataset):
 
         item = {"text": text}
         for cat_field in self._category_fields:
-            category = row.get(cat_field)
-            if category is None:
+            if cat_field not in row:
                 raise ValueError(
                     f"Category field '{cat_field}' not found in dataset row. Available fields: {list(row.keys())}"
                 )
@@ -157,9 +156,7 @@ class ClassificationDataset(BaseDataset):
             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."
-            )
+            raise NotImplementedError("Indexing not supported for STREAMING datasets. Use iter_items or iter_batches.")
 
         dataset_len = len(self)
         if dataset_len == 0:
@@ -446,6 +443,89 @@ class ClassificationDataset(BaseDataset):
             category_field=category_field,
         )
 
+    @classmethod
+    def from_disk(
+        cls,
+        store: Store,
+        *,
+        loading_strategy: LoadingStrategy = LoadingStrategy.MEMORY,
+        text_field: str = "text",
+        category_field: Union[str, List[str]] = "category",
+    ) -> "ClassificationDataset":
+        """
+        Load classification dataset from already-saved Arrow files on disk.
+
+        Use this when you've previously saved a dataset and want to reload it
+        without re-downloading from HuggingFace or re-applying transformations.
+
+        Args:
+            store: Store instance pointing to where the dataset was saved
+            loading_strategy: Loading strategy (MEMORY or DISK only)
+            text_field: Name of the column containing text
+            category_field: Name(s) of the column(s) containing category/label
+
+        Returns:
+            ClassificationDataset instance loaded from disk
+
+        Raises:
+            FileNotFoundError: If dataset directory doesn't exist or contains no Arrow files
+            ValueError: If required fields are not in the loaded dataset
+
+        Example:
+            # First: save dataset
+            dataset_store = LocalStore("store/wgmix_test")
+            dataset = ClassificationDataset.from_huggingface(
+                "allenai/wildguardmix",
+                store=dataset_store,
+                limit=100
+            )
+            # Dataset saved to: store/wgmix_test/datasets/*.arrow
+
+            # Later: reload from disk
+            dataset_store = LocalStore("store/wgmix_test")
+            dataset = ClassificationDataset.from_disk(
+                store=dataset_store,
+                text_field="prompt",
+                category_field="prompt_harm_label"
+            )
+        """
+
+        if store is None:
+            raise ValueError("store cannot be None")
+
+        if loading_strategy == LoadingStrategy.STREAMING:
+            raise ValueError("STREAMING loading strategy not supported for from_disk(). Use MEMORY or DISK.")
+
+        dataset_dir = Path(store.base_path) / store.dataset_prefix
+
+        if not dataset_dir.exists():
+            raise FileNotFoundError(
+                f"Dataset directory not found: {dataset_dir}. "
+                f"Make sure you've previously saved a dataset to this store location."
+            )
+
+        # Verify it's a valid Arrow dataset directory
+        arrow_files = list(dataset_dir.glob("*.arrow"))
+        if not arrow_files:
+            raise FileNotFoundError(
+                f"No Arrow files found in {dataset_dir}. Directory exists but doesn't contain a valid dataset."
+            )
+
+        try:
+            use_memory_mapping = loading_strategy == LoadingStrategy.DISK
+            ds = load_from_disk(str(dataset_dir), keep_in_memory=not use_memory_mapping)
+        except Exception as e:
+            raise RuntimeError(f"Failed to load dataset from {dataset_dir}. Error: {e}") from e
+
+        # Create ClassificationDataset with the loaded dataset and field names
+        return cls(
+            ds,
+            store=store,
+            loading_strategy=loading_strategy,
+            text_field=text_field,
+            category_field=category_field,
+        )
+
     @classmethod
     def from_csv(
         cls,
@@ -483,24 +563,37 @@ class ClassificationDataset(BaseDataset):
             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)
+        if store is None:
+            raise ValueError("store cannot be None")
 
-        dataset = super().from_csv(
+        use_streaming = 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.")
+
+        # Load CSV using parent's static method
+        ds = cls._load_csv_source(
             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,
+            streaming=use_streaming,
             **kwargs,
         )
+
+        # Apply postprocessing if not streaming
+        if not use_streaming and (stratify_by or drop_na):
+            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,
+                stratify_by=stratify_by,
+                stratify_seed=stratify_seed,
+                drop_na_columns=drop_na_columns,
+            )
+
         return cls(
-            dataset._ds,
+            ds,
             store=store,
             loading_strategy=loading_strategy,
             text_field=text_field,
@@ -542,23 +635,36 @@ class ClassificationDataset(BaseDataset):
             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)
+        if store is None:
+            raise ValueError("store cannot be None")
+
+        use_streaming = 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.")
 
-        dataset = super().from_json(
+        # Load JSON using parent's static method
+        ds = cls._load_json_source(
             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,
+            streaming=use_streaming,
             **kwargs,
         )
+
+        # Apply postprocessing if not streaming
+        if not use_streaming and (stratify_by or drop_na):
+            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,
+                stratify_by=stratify_by,
+                stratify_seed=stratify_seed,
+                drop_na_columns=drop_na_columns,
+            )
+
         return cls(
-            dataset._ds,
+            ds,
             store=store,
             loading_strategy=loading_strategy,
             text_field=text_field,

mi_crow/datasets/text_dataset.py

@@ -1,9 +1,10 @@
 from __future__ import annotations
 
+import random
 from pathlib import Path
 from typing import Any, Dict, Iterator, List, Optional, Sequence, Union
 
-from datasets import Dataset, IterableDataset, load_dataset
+from datasets import Dataset, IterableDataset, load_dataset, load_from_disk
 
 from mi_crow.datasets.base_dataset import BaseDataset
 from mi_crow.datasets.loading_strategy import IndexLike, LoadingStrategy
@@ -20,7 +21,7 @@ class TextDataset(BaseDataset):
         self,
         ds: Dataset | IterableDataset,
         store: Store,
-        loading_strategy: LoadingStrategy = LoadingStrategy.MEMORY,
+        loading_strategy: LoadingStrategy = LoadingStrategy.DISK,
         text_field: str = "text",
     ):
         """
@@ -115,9 +116,7 @@ class TextDataset(BaseDataset):
             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."
-            )
+            raise NotImplementedError("Indexing not supported for STREAMING datasets. Use iter_items or iter_batches.")
 
         dataset_len = len(self)
         if dataset_len == 0:
@@ -217,6 +216,48 @@ class TextDataset(BaseDataset):
             return list(self.iter_items())
         return list(self._ds["text"])
 
+    def random_sample(self, n: int, seed: Optional[int] = None) -> "TextDataset":
+        """Create a new TextDataset with n randomly sampled items.
+
+        Args:
+            n: Number of items to sample
+            seed: Optional random seed for reproducibility
+
+        Returns:
+            New TextDataset instance with sampled items
+
+        Raises:
+            NotImplementedError: If loading_strategy is STREAMING
+            ValueError: If n <= 0
+        """
+        if self._loading_strategy == LoadingStrategy.STREAMING:
+            raise NotImplementedError(
+                "random_sample() not supported for STREAMING datasets. Use iter_items() and sample manually."
+            )
+        
+        if n <= 0:
+            raise ValueError(f"n must be > 0, got: {n}")
+        
+        dataset_len = len(self)
+        if n >= dataset_len:
+            if seed is not None:
+                random.seed(seed)
+            indices = list(range(dataset_len))
+            random.shuffle(indices)
+            sampled_ds = self._ds.select(indices)
+        else:
+            if seed is not None:
+                random.seed(seed)
+            indices = random.sample(range(dataset_len), n)
+            sampled_ds = self._ds.select(indices)
+        
+        return TextDataset(
+            sampled_ds,
+            store=self._store,
+            loading_strategy=self._loading_strategy,
+            text_field=self._text_field,
+        )
+
     @classmethod
     def from_huggingface(
         cls,
@@ -300,6 +341,81 @@ class TextDataset(BaseDataset):
 
         return cls(ds, store=store, loading_strategy=loading_strategy, text_field=text_field)
 
+    @classmethod
+    def from_disk(
+        cls,
+        store: Store,
+        *,
+        loading_strategy: LoadingStrategy = LoadingStrategy.MEMORY,
+        text_field: str = "text",
+    ) -> "TextDataset":
+        """
+        Load text dataset from already-saved Arrow files on disk.
+
+        Use this when you've previously saved a dataset and want to reload it
+        without re-downloading from HuggingFace or re-applying transformations.
+
+        Args:
+            store: Store instance pointing to where the dataset was saved
+            loading_strategy: Loading strategy (MEMORY or DISK only)
+            text_field: Name of the column containing text
+
+        Returns:
+            TextDataset instance loaded from disk
+
+        Raises:
+            FileNotFoundError: If dataset directory doesn't exist or contains no Arrow files
+
+        Example:
+            # First: save dataset
+            dataset_store = LocalStore("store/my_texts")
+            dataset = TextDataset.from_huggingface(
+                "wikipedia",
+                store=dataset_store,
+                limit=1000
+            )
+            # Dataset saved to: store/my_texts/datasets/*.arrow
+
+            # Later: reload from disk
+            dataset_store = LocalStore("store/my_texts")
+            dataset = TextDataset.from_disk(store=dataset_store)
+        """
+
+        if store is None:
+            raise ValueError("store cannot be None")
+
+        if loading_strategy == LoadingStrategy.STREAMING:
+            raise ValueError("STREAMING loading strategy not supported for from_disk(). Use MEMORY or DISK.")
+
+        dataset_dir = Path(store.base_path) / store.dataset_prefix
+
+        if not dataset_dir.exists():
+            raise FileNotFoundError(
+                f"Dataset directory not found: {dataset_dir}. "
+                f"Make sure you've previously saved a dataset to this store location."
+            )
+
+        # Verify it's a valid Arrow dataset directory
+        arrow_files = list(dataset_dir.glob("*.arrow"))
+        if not arrow_files:
+            raise FileNotFoundError(
+                f"No Arrow files found in {dataset_dir}. Directory exists but doesn't contain a valid dataset."
+            )
+
+        try:
+            use_memory_mapping = loading_strategy == LoadingStrategy.DISK
+            ds = load_from_disk(str(dataset_dir), keep_in_memory=not use_memory_mapping)
+        except Exception as e:
+            raise RuntimeError(f"Failed to load dataset from {dataset_dir}. Error: {e}") from e
+
+        # Create TextDataset with the loaded dataset and field name
+        return cls(
+            ds,
+            store=store,
+            loading_strategy=loading_strategy,
+            text_field=text_field,
+        )
+
     @classmethod
     def from_csv(
         cls,
@@ -335,20 +451,33 @@ class TextDataset(BaseDataset):
             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(
+        if store is None:
+            raise ValueError("store cannot be None")
+
+        use_streaming = 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.")
+
+        # Load CSV using parent's static method
+        ds = cls._load_csv_source(
             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,
+            streaming=use_streaming,
             **kwargs,
         )
+
+        # Apply postprocessing if not streaming
+        if not use_streaming and (stratify_by or drop_na):
+            drop_na_columns = [text_field] if drop_na else None
+            ds = cls._postprocess_non_streaming_dataset(
+                ds,
+                stratify_by=stratify_by,
+                stratify_seed=stratify_seed,
+                drop_na_columns=drop_na_columns,
+            )
+
         return cls(
-            dataset._ds,
+            ds,
             store=store,
             loading_strategy=loading_strategy,
             text_field=text_field,
@@ -387,20 +516,32 @@ class TextDataset(BaseDataset):
             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(
+        if store is None:
+            raise ValueError("store cannot be None")
+
+        use_streaming = 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.")
+
+        # Load JSON using parent's static method
+        ds = cls._load_json_source(
             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,
+            streaming=use_streaming,
             **kwargs,
         )
-        # Re-initialize with text_field
+
+        # Apply postprocessing if not streaming
+        if not use_streaming and (stratify_by or drop_na):
+            drop_na_columns = [text_field] if drop_na else None
+            ds = cls._postprocess_non_streaming_dataset(
+                ds,
+                stratify_by=stratify_by,
+                stratify_seed=stratify_seed,
+                drop_na_columns=drop_na_columns,
+            )
+
         return cls(
-            dataset._ds,
+            ds,
             store=store,
             loading_strategy=loading_strategy,
             text_field=text_field,

mi_crow/hooks/controller.py

@@ -7,7 +7,11 @@ import torch
 import torch.nn as nn
 
 from mi_crow.hooks.hook import Hook, HookType, HOOK_FUNCTION_INPUT, HOOK_FUNCTION_OUTPUT
-from mi_crow.hooks.utils import extract_tensor_from_input, extract_tensor_from_output
+from mi_crow.hooks.utils import (
+    extract_tensor_from_input,
+    extract_tensor_from_output,
+    apply_modification_to_output
+)
 from mi_crow.utils import get_logger
 
 if TYPE_CHECKING:
@@ -86,13 +90,14 @@ class Controller(Hook):
         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)
+        modified_tensor = self.modify_activations(module, input_tensor, output_tensor)
+        
+        if modified_tensor is not None and isinstance(modified_tensor, torch.Tensor):
+            target_device = None
+            if self.context is not None and hasattr(self.context, 'device') and self.context.device:
+                target_device = torch.device(self.context.device)
+            apply_modification_to_output(output, modified_tensor, target_device=target_device)
 
     def _hook_fn(
             self,