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

amber/language_model/inference.py

@@ -0,0 +1,525 @@
+"""Inference engine for language models."""
+
+from __future__ import annotations
+
+import datetime
+from typing import Sequence, Any, Dict, List, TYPE_CHECKING
+
+import torch
+from torch import nn
+
+from amber.language_model.utils import get_device_from_model, move_tensors_to_device, extract_logits_from_output
+from amber.utils import get_logger
+
+if TYPE_CHECKING:
+    from amber.language_model.language_model import LanguageModel
+    from amber.hooks.controller import Controller
+    from amber.datasets import BaseDataset
+    from amber.store.store import Store
+
+logger = get_logger(__name__)
+
+
+class _EarlyStopInference(Exception):
+    """Internal exception used to stop model forward pass after a specific layer."""
+
+    def __init__(self, output: Any):
+        super().__init__("Early stop after requested layer")
+        self.output = output
+
+
+class InferenceEngine:
+    """Handles inference operations for LanguageModel."""
+    
+    def __init__(self, language_model: "LanguageModel"):
+        """
+        Initialize inference engine.
+        
+        Args:
+            language_model: LanguageModel instance
+        """
+        self.lm = language_model
+    
+    def _prepare_tokenizer_kwargs(self, tok_kwargs: Dict | None) -> Dict[str, Any]:
+        """
+        Prepare tokenizer keyword arguments with defaults.
+
+        Args:
+            tok_kwargs: Optional tokenizer keyword arguments
+
+        Returns:
+            Dictionary of tokenizer kwargs with defaults applied
+        """
+        if tok_kwargs is None:
+            tok_kwargs = {}
+        return {
+            "padding": True,
+            "truncation": True,
+            "return_tensors": "pt",
+            **tok_kwargs,
+        }
+    
+    def _setup_trackers(self, texts: Sequence[str]) -> None:
+        """
+        Setup input trackers for current texts.
+
+        Args:
+            texts: Sequence of input texts
+        """
+        if self.lm._input_tracker is not None and self.lm._input_tracker.enabled:
+            self.lm._input_tracker.set_current_texts(texts)
+    
+    def _setup_model_input_detectors(self, enc: Dict[str, torch.Tensor]) -> None:
+        """
+        Automatically set inputs from encodings for all registered ModelInputDetector hooks.
+
+        This is necessary because PyTorch's pre_forward hook doesn't receive kwargs,
+        so ModelInputDetector hooks can't automatically capture attention masks when
+        models are called with **kwargs (e.g., model(**encodings)).
+
+        Args:
+            enc: Encoded inputs dictionary
+        """
+        from amber.hooks.implementations.model_input_detector import ModelInputDetector
+        
+        detectors = self.lm.layers.get_detectors()
+        for detector in detectors:
+            if isinstance(detector, ModelInputDetector):
+                detector.set_inputs_from_encodings(enc, module=self.lm.model)
+    
+    def _prepare_controllers(self, with_controllers: bool) -> List["Controller"]:
+        """
+        Prepare controllers for inference, disabling if needed.
+
+        Args:
+            with_controllers: Whether to keep controllers enabled
+
+        Returns:
+            List of controllers that were disabled (to restore later)
+        """
+        controllers_to_restore = []
+        if not with_controllers:
+            controllers = self.lm.layers.get_controllers()
+            for controller in controllers:
+                if controller.enabled:
+                    controller.disable()
+                    controllers_to_restore.append(controller)
+        return controllers_to_restore
+    
+    def _restore_controllers(self, controllers_to_restore: List["Controller"]) -> None:
+        """
+        Restore controllers that were disabled.
+
+        Args:
+            controllers_to_restore: List of controllers to restore
+        """
+        for controller in controllers_to_restore:
+            controller.enable()
+    
+    def _run_model_forward(
+            self,
+            enc: Dict[str, torch.Tensor],
+            autocast: bool,
+            device_type: str,
+            autocast_dtype: torch.dtype | None,
+    ) -> Any:
+        """
+        Run model forward pass with optional autocast.
+        
+        Args:
+            enc: Encoded inputs dictionary
+            autocast: Whether to use automatic mixed precision
+            device_type: Device type string ("cuda", "cpu", etc.)
+            autocast_dtype: Optional dtype for autocast
+            
+        Returns:
+            Model output
+        """
+        try:
+            with torch.inference_mode():
+                if autocast and device_type == "cuda":
+                    amp_dtype = autocast_dtype or torch.float16
+                    with torch.autocast(device_type, dtype=amp_dtype):
+                        return self.lm.model(**enc)
+                return self.lm.model(**enc)
+        except _EarlyStopInference as e:
+            # Early stopping hook raised this to short‑circuit the remaining forward pass.
+            # We return the output captured at the requested layer.
+            return e.output
+    
+    def execute_inference(
+            self,
+            texts: Sequence[str],
+            tok_kwargs: Dict | None = None,
+            autocast: bool = True,
+            autocast_dtype: torch.dtype | None = None,
+            with_controllers: bool = True,
+            stop_after_layer: str | int | None = None,
+    ) -> tuple[Any, Dict[str, torch.Tensor]]:
+        """
+        Execute inference on texts.
+        
+        Args:
+            texts: Sequence of input texts
+            tok_kwargs: Optional tokenizer keyword arguments
+            autocast: Whether to use automatic mixed precision
+            autocast_dtype: Optional dtype for autocast
+            with_controllers: Whether to use controllers during inference
+            stop_after_layer: Optional layer signature (name or index) after which
+                the forward pass should be stopped early
+            
+        Returns:
+            Tuple of (model_output, encodings)
+            
+        Raises:
+            ValueError: If texts is empty or tokenizer is not initialized
+        """
+        if not texts:
+            raise ValueError("Texts list cannot be empty")
+
+        if self.lm.tokenizer is None:
+            raise ValueError("Tokenizer must be initialized before running inference")
+
+        tok_kwargs = self._prepare_tokenizer_kwargs(tok_kwargs)
+        enc = self.lm.tokenize(texts, **tok_kwargs)
+
+        device = get_device_from_model(self.lm.model)
+        device_type = str(device.type)
+        enc = move_tensors_to_device(enc, device)
+
+        self.lm.model.eval()
+
+        self._setup_trackers(texts)
+        self._setup_model_input_detectors(enc)
+
+        controllers_to_restore = self._prepare_controllers(with_controllers)
+
+        hook_handle = None
+        try:
+            if stop_after_layer is not None:
+                # Register a temporary forward hook that stops the forward pass
+                def _early_stop_hook(module: nn.Module, inputs: tuple, output: Any):
+                    raise _EarlyStopInference(output)
+
+                hook_handle = self.lm.layers.register_forward_hook_for_layer(
+                    stop_after_layer, _early_stop_hook
+                )
+
+            output = self._run_model_forward(enc, autocast, device_type, autocast_dtype)
+            return output, enc
+        finally:
+            if hook_handle is not None:
+                try:
+                    hook_handle.remove()
+                except Exception:
+                    pass
+            self._restore_controllers(controllers_to_restore)
+    
+    def extract_logits(self, output: Any) -> torch.Tensor:
+        """
+        Extract logits tensor from model output.
+        
+        Args:
+            output: Model output
+            
+        Returns:
+            Logits tensor
+        """
+        return extract_logits_from_output(output)
+    
+    def _extract_dataset_info(self, dataset: "BaseDataset | None") -> Dict[str, Any]:
+        """
+        Extract dataset information for metadata.
+        
+        Args:
+            dataset: Optional dataset instance
+            
+        Returns:
+            Dictionary with dataset information
+        """
+        if dataset is None:
+            return {}
+        
+        try:
+            ds_id = str(getattr(dataset, "dataset_dir", ""))
+            ds_len = int(len(dataset))
+            return {
+                "dataset_dir": ds_id,
+                "length": ds_len,
+            }
+        except (AttributeError, TypeError, ValueError, RuntimeError):
+            return {
+                "dataset_dir": "",
+                "length": -1,
+            }
+    
+    def _prepare_run_metadata(
+        self,
+        dataset: "BaseDataset | None" = None,
+        run_name: str | None = None,
+        options: Dict[str, Any] | None = None,
+    ) -> tuple[str, Dict[str, Any]]:
+        """
+        Prepare run metadata dictionary.
+        
+        Args:
+            dataset: Optional dataset (for dataset info)
+            run_name: Optional run name (generates if None)
+            options: Optional dict of options to include
+            
+        Returns:
+            Tuple of (run_name, metadata_dict)
+        """
+        if run_name is None:
+            run_name = f"run_{datetime.datetime.now().strftime('%Y%m%d_%H%M%S')}"
+        
+        if options is None:
+            options = {}
+        
+        dataset_info = self._extract_dataset_info(dataset)
+        
+        meta: Dict[str, Any] = {
+            "run_name": run_name,
+            "model": getattr(self.lm.model, "model_name", self.lm.model.__class__.__name__),
+            "options": options.copy(),
+        }
+        
+        if dataset_info:
+            meta["dataset"] = dataset_info
+        
+        return run_name, meta
+    
+    @staticmethod
+    def _save_run_metadata(
+        store: "Store",
+        run_name: str,
+        meta: Dict[str, Any],
+        verbose: bool = False,
+    ) -> None:
+        """
+        Save run metadata to store.
+        
+        Args:
+            store: Store to save to
+            run_name: Run name
+            meta: Metadata dictionary
+            verbose: Whether to log
+        """
+        try:
+            store.put_run_metadata(run_name, meta)
+        except (OSError, IOError, ValueError, RuntimeError) as e:
+            if verbose:
+                logger.warning(f"Failed to save run metadata for {run_name}: {e}")
+    
+    def infer_texts(
+        self,
+        texts: Sequence[str],
+        run_name: str | None = None,
+        batch_size: int | None = None,
+        tok_kwargs: Dict | None = None,
+        autocast: bool = True,
+        autocast_dtype: torch.dtype | None = None,
+        with_controllers: bool = True,
+        clear_detectors_before: bool = False,
+        verbose: bool = False,
+        stop_after_layer: str | int | None = None,
+        save_in_batches: bool = True,
+    ) -> tuple[Any, Dict[str, torch.Tensor]] | tuple[List[Any], List[Dict[str, torch.Tensor]]]:
+        """
+        Run inference on list of strings with optional metadata saving.
+        
+        Args:
+            texts: Sequence of input texts
+            run_name: Optional run name for saving metadata (if None, no metadata saved)
+            batch_size: Optional batch size for processing (if None, processes all at once)
+            tok_kwargs: Optional tokenizer keyword arguments
+            autocast: Whether to use automatic mixed precision
+            autocast_dtype: Optional dtype for autocast
+            with_controllers: Whether to use controllers during inference
+            clear_detectors_before: If True, clears all detector state before running
+            verbose: Whether to log progress
+            stop_after_layer: Optional layer signature (name or index) after which
+                the forward pass should be stopped early
+            save_in_batches: If True, save detector metadata in per‑batch
+                directories. If False, aggregate all detector metadata for
+                the run under a single detectors directory.
+            
+        Returns:
+            If batch_size is None or >= len(texts): Tuple of (model_output, encodings)
+            If batch_size < len(texts): Tuple of (list of outputs, list of encodings)
+            
+        Raises:
+            ValueError: If texts is empty or tokenizer is not initialized
+        """
+        if not texts:
+            raise ValueError("Texts list cannot be empty")
+        
+        if self.lm.tokenizer is None:
+            raise ValueError("Tokenizer must be initialized before running inference")
+        
+        if clear_detectors_before:
+            self.lm.clear_detectors()
+
+        store = self.lm.store
+        if run_name is not None and store is None:
+            raise ValueError("Store must be provided to save metadata")
+        
+        if batch_size is None or batch_size >= len(texts):
+            output, enc = self.execute_inference(
+                texts,
+                tok_kwargs=tok_kwargs,
+                autocast=autocast,
+                autocast_dtype=autocast_dtype,
+                with_controllers=with_controllers,
+                stop_after_layer=stop_after_layer,
+            )
+            
+            if run_name is not None:
+                options = {
+                    "batch_size": len(texts),
+                    "max_length": tok_kwargs.get("max_length") if tok_kwargs else None,
+                }
+                _, meta = self._prepare_run_metadata(dataset=None, run_name=run_name, options=options)
+                self._save_run_metadata(store, run_name, meta, verbose)
+                self.lm.save_detector_metadata(run_name, 0, unified=not save_in_batches)
+            
+            return output, enc
+        
+        all_outputs = []
+        all_encodings = []
+        batch_counter = 0
+        
+        if run_name is not None:
+            options = {
+                "batch_size": batch_size,
+                "max_length": tok_kwargs.get("max_length") if tok_kwargs else None,
+            }
+            _, meta = self._prepare_run_metadata(dataset=None, run_name=run_name, options=options)
+            self._save_run_metadata(store, run_name, meta, verbose)
+        
+        for i in range(0, len(texts), batch_size):
+            batch_texts = texts[i:i + batch_size]
+            output, enc = self.execute_inference(
+                batch_texts,
+                tok_kwargs=tok_kwargs,
+                autocast=autocast,
+                autocast_dtype=autocast_dtype,
+                with_controllers=with_controllers,
+                stop_after_layer=stop_after_layer,
+            )
+            
+            all_outputs.append(output)
+            all_encodings.append(enc)
+            
+            if run_name is not None:
+                self.lm.save_detector_metadata(run_name, batch_counter, unified=not save_in_batches)
+                if verbose:
+                    logger.info(f"Saved batch {batch_counter} for run={run_name}")
+            
+            batch_counter += 1
+        
+        return all_outputs, all_encodings
+    
+    def infer_dataset(
+        self,
+        dataset: "BaseDataset",
+        run_name: str | None = None,
+        batch_size: int = 32,
+        tok_kwargs: Dict | None = None,
+        autocast: bool = True,
+        autocast_dtype: torch.dtype | None = None,
+        with_controllers: bool = True,
+        free_cuda_cache_every: int | None = 0,
+        clear_detectors_before: bool = False,
+        verbose: bool = False,
+        stop_after_layer: str | int | None = None,
+        save_in_batches: bool = True,
+    ) -> str:
+        """
+        Run inference on whole dataset with metadata saving.
+        
+        Args:
+            dataset: Dataset to process
+            run_name: Optional run name (generated if None)
+            batch_size: Batch size for processing
+            tok_kwargs: Optional tokenizer keyword arguments
+            autocast: Whether to use automatic mixed precision
+            autocast_dtype: Optional dtype for autocast
+            with_controllers: Whether to use controllers during inference
+            free_cuda_cache_every: Clear CUDA cache every N batches (0 or None to disable)
+            clear_detectors_before: If True, clears all detector state before running
+            verbose: Whether to log progress
+            stop_after_layer: Optional layer signature (name or index) after which
+                the forward pass should be stopped early
+            
+        Returns:
+            Run name used for saving
+            
+        Raises:
+            ValueError: If model or store is not initialized
+        """
+        if clear_detectors_before:
+            self.lm.clear_detectors()
+
+        model: nn.Module | None = self.lm.model
+        if model is None:
+            raise ValueError("Model must be initialized before running")
+        
+        store = self.lm.store
+        if store is None:
+            raise ValueError("Store must be provided or set on the language model")
+        
+        device = get_device_from_model(model)
+        device_type = str(device.type)
+        
+        options = {
+            "max_length": tok_kwargs.get("max_length") if tok_kwargs else None,
+            "batch_size": int(batch_size),
+        }
+        
+        run_name, meta = self._prepare_run_metadata(dataset=dataset, run_name=run_name, options=options)
+        
+        if verbose:
+            logger.info(
+                f"Starting infer_dataset: run={run_name}, "
+                f"batch_size={batch_size}, device={device_type}"
+            )
+        
+        self._save_run_metadata(store, run_name, meta, verbose)
+        
+        batch_counter = 0
+        
+        with torch.inference_mode():
+            for batch_index, batch in enumerate(dataset.iter_batches(batch_size)):
+                if not batch:
+                    continue
+                
+                texts = dataset.extract_texts_from_batch(batch)
+                
+                self.execute_inference(
+                    texts,
+                    tok_kwargs=tok_kwargs,
+                    autocast=autocast,
+                    autocast_dtype=autocast_dtype,
+                    with_controllers=with_controllers,
+                    stop_after_layer=stop_after_layer,
+                )
+                
+                self.lm.save_detector_metadata(run_name, batch_index, unified=not save_in_batches)
+                
+                batch_counter += 1
+                
+                if device_type == "cuda" and free_cuda_cache_every and free_cuda_cache_every > 0:
+                    if (batch_counter % free_cuda_cache_every) == 0:
+                        torch.cuda.empty_cache()
+                        if verbose:
+                            logger.info("Emptied CUDA cache")
+                
+                if verbose:
+                    logger.info(f"Saved batch {batch_index} for run={run_name}")
+        
+        if verbose:
+            logger.info(f"Completed infer_dataset: run={run_name}, batches_saved={batch_counter}")
+        
+        return run_name
+

amber/language_model/initialization.py

@@ -0,0 +1,126 @@
+"""Model initialization and factory methods."""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+import torch
+from torch import nn
+from transformers import AutoTokenizer, AutoModelForCausalLM, PreTrainedTokenizerBase
+
+from amber.store.store import Store
+from amber.language_model.utils import extract_model_id
+
+if TYPE_CHECKING:
+    from amber.language_model.language_model import LanguageModel
+
+
+def initialize_model_id(
+        model: nn.Module,
+        provided_model_id: str | None = None
+) -> str:
+    """
+    Initialize model ID for LanguageModel.
+    
+    Args:
+        model: PyTorch model module
+        provided_model_id: Optional model ID provided by user
+        
+    Returns:
+        Model ID string
+    """
+    return extract_model_id(model, provided_model_id)
+
+
+def create_from_huggingface(
+        cls: type["LanguageModel"],
+        model_name: str,
+        store: Store,
+        tokenizer_params: dict | None = None,
+        model_params: dict | None = None,
+) -> "LanguageModel":
+    """
+    Load a language model from HuggingFace Hub.
+    
+    Args:
+        cls: LanguageModel class
+        model_name: HuggingFace model identifier
+        store: Store instance for persistence
+        tokenizer_params: Optional tokenizer parameters
+        model_params: Optional model parameters
+        
+    Returns:
+        LanguageModel instance
+        
+    Raises:
+        ValueError: If model_name is invalid
+        RuntimeError: If model loading fails
+    """
+    if not model_name or not isinstance(model_name, str) or not model_name.strip():
+        raise ValueError(f"model_name must be a non-empty string, got: {model_name!r}")
+    
+    if store is None:
+        raise ValueError("store cannot be None")
+    
+    if tokenizer_params is None:
+        tokenizer_params = {}
+    if model_params is None:
+        model_params = {}
+
+    try:
+        tokenizer = AutoTokenizer.from_pretrained(model_name, **tokenizer_params)
+        model = AutoModelForCausalLM.from_pretrained(model_name, **model_params)
+    except Exception as e:
+        raise RuntimeError(
+            f"Failed to load model '{model_name}' from HuggingFace. Error: {e}"
+        ) from e
+
+    return cls(model, tokenizer, store)
+
+
+def create_from_local_torch(
+        cls: type["LanguageModel"],
+        model_path: str,
+        tokenizer_path: str,
+        store: Store
+) -> "LanguageModel":
+    """
+    Load a language model from local HuggingFace paths.
+    
+    Args:
+        cls: LanguageModel class
+        model_path: Path to the model directory or file
+        tokenizer_path: Path to the tokenizer directory or file
+        store: Store instance for persistence
+        
+    Returns:
+        LanguageModel instance
+        
+    Raises:
+        FileNotFoundError: If model or tokenizer paths don't exist
+        RuntimeError: If model loading fails
+    """
+    if store is None:
+        raise ValueError("store cannot be None")
+    
+    model_path_obj = Path(model_path)
+    tokenizer_path_obj = Path(tokenizer_path)
+    
+    if not model_path_obj.exists():
+        raise FileNotFoundError(f"Model path does not exist: {model_path}")
+    
+    if not tokenizer_path_obj.exists():
+        raise FileNotFoundError(f"Tokenizer path does not exist: {tokenizer_path}")
+
+    try:
+        tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)
+        model = AutoModelForCausalLM.from_pretrained(model_path)
+    except Exception as e:
+        raise RuntimeError(
+            f"Failed to load model from local paths. "
+            f"model_path={model_path!r}, tokenizer_path={tokenizer_path!r}. Error: {e}"
+        ) from e
+
+    return cls(model, tokenizer, store)
+