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

mi_crow/language_model/inference.py

@@ -8,7 +8,8 @@ from typing import Sequence, Any, Dict, List, TYPE_CHECKING
 import torch
 from torch import nn
 
-from mi_crow.language_model.utils import get_device_from_model, move_tensors_to_device, extract_logits_from_output
+from mi_crow.language_model.utils import move_tensors_to_device, extract_logits_from_output
+from mi_crow.language_model.device_manager import sync_model_to_context_device
 from mi_crow.utils import get_logger
 
 if TYPE_CHECKING:
@@ -52,12 +53,19 @@ class InferenceEngine:
         """
         if tok_kwargs is None:
             tok_kwargs = {}
-        return {
-            "padding": True,
+        
+        padding_strategy = tok_kwargs.pop("padding", True)
+        if padding_strategy is True and "max_length" in tok_kwargs:
+            padding_strategy = "longest"
+        
+        result = {
+            "padding": padding_strategy,
             "truncation": True,
             "return_tensors": "pt",
             **tok_kwargs,
         }
+        
+        return result
     
     def _setup_trackers(self, texts: Sequence[str]) -> None:
         """
@@ -181,10 +189,15 @@ class InferenceEngine:
             raise ValueError("Tokenizer must be initialized before running inference")
 
         tok_kwargs = self._prepare_tokenizer_kwargs(tok_kwargs)
+        logger.debug(f"[DEBUG] About to tokenize {len(texts)} texts...")
         enc = self.lm.tokenize(texts, **tok_kwargs)
+        logger.debug(f"[DEBUG] Tokenization completed, shape: {enc['input_ids'].shape if isinstance(enc, dict) else 'N/A'}")
 
-        device = get_device_from_model(self.lm.model)
+        device = torch.device(self.lm.context.device)
         device_type = str(device.type)
+
+        sync_model_to_context_device(self.lm)
+
         enc = move_tensors_to_device(enc, device)
 
         self.lm.model.eval()
@@ -469,7 +482,7 @@ class InferenceEngine:
         if store is None:
             raise ValueError("Store must be provided or set on the language model")
         
-        device = get_device_from_model(model)
+        device = torch.device(self.lm.context.device)
         device_type = str(device.type)
         
         options = {

mi_crow/language_model/initialization.py

@@ -39,6 +39,7 @@ def create_from_huggingface(
         store: Store,
         tokenizer_params: dict | None = None,
         model_params: dict | None = None,
+        device: str | torch.device | None = None,
 ) -> "LanguageModel":
     """
     Load a language model from HuggingFace Hub.
@@ -49,10 +50,11 @@ def create_from_huggingface(
         store: Store instance for persistence
         tokenizer_params: Optional tokenizer parameters
         model_params: Optional model parameters
-        
+        device: Target device ("cuda", "cpu", "mps"). Model will be moved to this device
+            after loading.
     Returns:
         LanguageModel instance
-        
+    
     Raises:
         ValueError: If model_name is invalid
         RuntimeError: If model loading fails
@@ -67,7 +69,7 @@ def create_from_huggingface(
         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)
@@ -76,14 +78,15 @@ def create_from_huggingface(
             f"Failed to load model '{model_name}' from HuggingFace. Error: {e}"
         ) from e
 
-    return cls(model, tokenizer, store)
+    return cls(model, tokenizer, store, device=device)
 
 
 def create_from_local_torch(
         cls: type["LanguageModel"],
         model_path: str,
         tokenizer_path: str,
-        store: Store
+        store: Store,
+        device: str | torch.device | None = None,
 ) -> "LanguageModel":
     """
     Load a language model from local HuggingFace paths.
@@ -93,6 +96,7 @@ def create_from_local_torch(
         model_path: Path to the model directory or file
         tokenizer_path: Path to the tokenizer directory or file
         store: Store instance for persistence
+        device: Optional device string or torch.device (defaults to 'cpu' if None)
         
     Returns:
         LanguageModel instance
@@ -122,5 +126,5 @@ def create_from_local_torch(
             f"model_path={model_path!r}, tokenizer_path={tokenizer_path!r}. Error: {e}"
         ) from e
 
-    return cls(model, tokenizer, store)
+    return cls(model, tokenizer, store, device=device)
 

mi_crow/language_model/language_model.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import gc
 from collections import defaultdict
 from pathlib import Path
 from typing import Sequence, Any, Dict, List, TYPE_CHECKING, Set, Tuple
@@ -15,6 +16,7 @@ from mi_crow.language_model.context import LanguageModelContext
 from mi_crow.language_model.inference import InferenceEngine
 from mi_crow.language_model.persistence import save_model, load_model_from_saved_file
 from mi_crow.language_model.initialization import initialize_model_id, create_from_huggingface, create_from_local_torch
+from mi_crow.language_model.device_manager import normalize_device, sync_model_to_context_device
 from mi_crow.store.store import Store
 from mi_crow.utils import get_logger
 
@@ -72,7 +74,7 @@ class LanguageModel:
     
     Provides a unified interface for working with language models, including:
     - Model initialization and configuration
-    - Inference operations
+    - Inference operations through the inference property
     - Hook management (detectors and controllers)
     - Model persistence
     - Activation tracking
@@ -84,6 +86,7 @@ class LanguageModel:
             tokenizer: PreTrainedTokenizerBase,
             store: Store,
             model_id: str | None = None,
+            device: str | torch.device | None = None,
     ):
         """
         Initialize LanguageModel.
@@ -93,6 +96,7 @@ class LanguageModel:
             tokenizer: HuggingFace tokenizer
             store: Store instance for persistence
             model_id: Optional model identifier (auto-extracted if not provided)
+            device: Optional device string or torch.device (defaults to 'cpu' if None)
         """
         self.context = LanguageModelContext(self)
         self.context.model = model
@@ -100,15 +104,17 @@ class LanguageModel:
         self.context.model_id = initialize_model_id(model, model_id)
         self.context.store = store
         self.context.special_token_ids = _extract_special_token_ids(tokenizer)
+        self.context.device = normalize_device(device)
+        sync_model_to_context_device(self)
 
         self.layers = LanguageModelLayers(self.context)
         self.lm_tokenizer = LanguageModelTokenizer(self.context)
         self.activations = LanguageModelActivations(self.context)
         self.inference = InferenceEngine(self)
-        self._inference_engine = self.inference
 
         self._input_tracker: "InputTracker | None" = None
 
+
     @property
     def model(self) -> nn.Module:
         """Get the underlying PyTorch model."""
@@ -147,86 +153,6 @@ class LanguageModel:
         """
         return self.lm_tokenizer.tokenize(texts, **kwargs)
 
-    def forwards(
-            self,
-            texts: Sequence[str],
-            tok_kwargs: Dict | None = None,
-            autocast: bool = True,
-            autocast_dtype: torch.dtype | None = None,
-            with_controllers: bool = True,
-    ) -> Tuple[Any, Any]:
-        """
-        Run forward pass on texts.
-        
-        Args:
-            texts: Input texts to process
-            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
-            
-        Returns:
-            Tuple of (model_output, encodings)
-        """
-        return self._inference_engine.execute_inference(
-            texts,
-            tok_kwargs=tok_kwargs,
-            autocast=autocast,
-            autocast_dtype=autocast_dtype,
-            with_controllers=with_controllers
-        )
-
-    def generate(
-            self,
-            texts: Sequence[str],
-            tok_kwargs: Dict | None = None,
-            autocast: bool = True,
-            autocast_dtype: torch.dtype | None = None,
-            with_controllers: bool = True,
-            skip_special_tokens: bool = True,
-    ) -> Sequence[str]:
-        """
-        Run inference and automatically decode the output with the tokenizer.
-        
-        Args:
-            texts: Input texts to process
-            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
-            skip_special_tokens: Whether to skip special tokens when decoding
-            
-        Returns:
-            Sequence of decoded text strings
-            
-        Raises:
-            ValueError: If texts is empty or tokenizer is None
-        """
-        if not texts:
-            raise ValueError("Texts list cannot be empty")
-
-        if self.tokenizer is None:
-            raise ValueError("Tokenizer is required for decoding but is None")
-
-        output, enc = self._inference_engine.execute_inference(
-            texts,
-            tok_kwargs=tok_kwargs,
-            autocast=autocast,
-            autocast_dtype=autocast_dtype,
-            with_controllers=with_controllers
-        )
-
-        logits = self._inference_engine.extract_logits(output)
-        predicted_token_ids = logits.argmax(dim=-1)
-
-        decoded_texts = []
-        for i in range(predicted_token_ids.shape[0]):
-            token_ids = predicted_token_ids[i].cpu().tolist()
-            decoded_text = self.tokenizer.decode(token_ids, skip_special_tokens=skip_special_tokens)
-            decoded_texts.append(decoded_text)
-
-        return decoded_texts
-
     def get_input_tracker(self) -> "InputTracker | None":
         """
         Get the input tracker instance if it exists.
@@ -248,8 +174,8 @@ class LanguageModel:
         detectors_tensor_metadata: Dict[str, Dict[str, torch.Tensor]] = defaultdict(dict)
 
         for detector in detectors:
-            detectors_metadata[detector.layer_signature] = detector.metadata
-            detectors_tensor_metadata[detector.layer_signature] = detector.tensor_metadata
+            detectors_metadata[detector.layer_signature] = dict(detector.metadata)
+            detectors_tensor_metadata[detector.layer_signature] = dict(detector.tensor_metadata)
 
         return detectors_metadata, detectors_tensor_metadata
 
@@ -263,17 +189,14 @@ class LanguageModel:
         """
         detectors = self.layers.get_detectors()
         for detector in detectors:
-            # Clear generic accumulated metadata
             detector.metadata.clear()
             detector.tensor_metadata.clear()
 
-            # Allow detector implementations to provide more specialized
-            # clearing logic (e.g. ModelInputDetector, ModelOutputDetector).
             clear_captured = getattr(detector, "clear_captured", None)
             if callable(clear_captured):
                 clear_captured()
 
-    def save_detector_metadata(self, run_name: str, batch_idx: int | None, unified: bool = False) -> str:
+    def save_detector_metadata(self, run_name: str, batch_idx: int | None, unified: bool = False, clear_after_save: bool = True) -> str:
         """
         Save detector metadata to store.
         
@@ -282,6 +205,8 @@ class LanguageModel:
             batch_idx: Batch index. Ignored when ``unified`` is True.
             unified: If True, save metadata in a single detectors directory
                 for the whole run instead of per‑batch directories.
+            clear_after_save: If True, clear detector metadata after saving to free memory.
+                Defaults to True to prevent OOM errors when processing large batches.
             
         Returns:
             Path where metadata was saved
@@ -291,12 +216,36 @@ class LanguageModel:
         """
         if self.store is None:
             raise ValueError("Store must be provided or set on the language model")
+        
         detectors_metadata, detectors_tensor_metadata = self.get_all_detector_metadata()
+        
         if unified:
-            return self.store.put_run_detector_metadata(run_name, detectors_metadata, detectors_tensor_metadata)
-        if batch_idx is None:
-            raise ValueError("batch_idx must be provided when unified is False")
-        return self.store.put_detector_metadata(run_name, batch_idx, detectors_metadata, detectors_tensor_metadata)
+            result = self.store.put_run_detector_metadata(run_name, detectors_metadata, detectors_tensor_metadata)
+        else:
+            if batch_idx is None:
+                raise ValueError("batch_idx must be provided when unified is False")
+            result = self.store.put_detector_metadata(run_name, batch_idx, detectors_metadata, detectors_tensor_metadata)
+        
+        if clear_after_save:
+            for layer_signature in list(detectors_tensor_metadata.keys()):
+                detector_tensors = detectors_tensor_metadata[layer_signature]
+                for tensor_key in list(detector_tensors.keys()):
+                    del detector_tensors[tensor_key]
+                del detectors_tensor_metadata[layer_signature]
+            detectors_metadata.clear()
+            
+            detectors = self.layers.get_detectors()
+            for detector in detectors:
+                clear_captured = getattr(detector, "clear_captured", None)
+                if callable(clear_captured):
+                    clear_captured()
+                for key in list(detector.tensor_metadata.keys()):
+                    del detector.tensor_metadata[key]
+                detector.metadata.clear()
+            
+            gc.collect()
+        
+        return result
 
     def _ensure_input_tracker(self) -> "InputTracker":
         """
@@ -339,23 +288,36 @@ class LanguageModel:
             store: Store,
             tokenizer_params: dict = None,
             model_params: dict = None,
+            device: str | torch.device | None = None,
     ) -> "LanguageModel":
         """
         Load a language model from HuggingFace Hub.
         
+        Automatically loads model to GPU if device is "cuda" and CUDA is available.
+        This prevents OOM errors by keeping the model on GPU instead of CPU RAM.
+        
         Args:
             model_name: HuggingFace model identifier
             store: Store instance for persistence
             tokenizer_params: Optional tokenizer parameters
             model_params: Optional model parameters
+            device: Target device ("cuda", "cpu", "mps"). If "cuda" and CUDA is available,
+                model will be loaded directly to GPU using device_map="auto"
+                (via the HuggingFace factory helpers).
             
         Returns:
             LanguageModel instance
         """
-        return create_from_huggingface(cls, model_name, store, tokenizer_params, model_params)
+        return create_from_huggingface(cls, model_name, store, tokenizer_params, model_params, device)
 
     @classmethod
-    def from_local_torch(cls, model_path: str, tokenizer_path: str, store: Store) -> "LanguageModel":
+    def from_local_torch(
+            cls,
+            model_path: str,
+            tokenizer_path: str,
+            store: Store,
+            device: str | torch.device | None = None,
+    ) -> "LanguageModel":
         """
         Load a language model from local HuggingFace paths.
         
@@ -363,14 +325,21 @@ class LanguageModel:
             model_path: Path to the model directory or file
             tokenizer_path: Path to the tokenizer directory or file
             store: Store instance for persistence
+            device: Optional device string or torch.device (defaults to 'cpu' if None)
             
         Returns:
             LanguageModel instance
         """
-        return create_from_local_torch(cls, model_path, tokenizer_path, store)
+        return create_from_local_torch(cls, model_path, tokenizer_path, store, device)
 
     @classmethod
-    def from_local(cls, saved_path: Path | str, store: Store, model_id: str | None = None) -> "LanguageModel":
+    def from_local(
+            cls,
+            saved_path: Path | str,
+            store: Store,
+            model_id: str | None = None,
+            device: str | torch.device | None = None,
+    ) -> "LanguageModel":
         """
         Load a language model from a saved file (created by save_model).
         
@@ -379,6 +348,7 @@ class LanguageModel:
             store: Store instance for persistence
             model_id: Optional model identifier. If not provided, will use the model_id from saved metadata.
                      If provided, will be used to load the model architecture from HuggingFace.
+            device: Optional device string or torch.device (defaults to 'cpu' if None)
                      
         Returns:
             LanguageModel instance
@@ -387,4 +357,4 @@ class LanguageModel:
             FileNotFoundError: If the saved file doesn't exist
             ValueError: If the saved file format is invalid or model_id is required but not provided
         """
-        return load_model_from_saved_file(cls, saved_path, store, model_id)
+        return load_model_from_saved_file(cls, saved_path, store, model_id, device)

mi_crow/language_model/layers.py

@@ -320,19 +320,22 @@ class LanguageModelLayers:
 
         layer_signature, hook_type, hook = self.context._hook_id_map[hook_id]
 
-        # Find and remove from registry
-        if layer_signature in self.context._hook_registry:
-            if hook_type in self.context._hook_registry[layer_signature]:
-                hooks_list = self.context._hook_registry[layer_signature][hook_type]
-                for i, (h, handle) in enumerate(hooks_list):
-                    if h.id == hook_id:
-                        # Remove PyTorch hook
-                        handle.remove()
-                        # Remove from our list
-                        hooks_list.pop(i)
-                        break
-
-        # Remove from ID map
+        if layer_signature not in self.context._hook_registry:
+            del self.context._hook_id_map[hook_id]
+            return True
+        
+        hook_types = self.context._hook_registry[layer_signature]
+        if hook_type not in hook_types:
+            del self.context._hook_id_map[hook_id]
+            return True
+        
+        hooks_list = hook_types[hook_type]
+        for i, (h, handle) in enumerate(hooks_list):
+            if h.id == hook_id:
+                handle.remove()
+                hooks_list.pop(i)
+                break
+
         del self.context._hook_id_map[hook_id]
         return True
 

mi_crow/language_model/persistence.py

@@ -89,7 +89,8 @@ def load_model_from_saved_file(
         cls: type["LanguageModel"],
         saved_path: Path | str,
         store: "Store",
-        model_id: str | None = None
+        model_id: str | None = None,
+        device: str | torch.device | None = None,
 ) -> "LanguageModel":
     """
     Load a language model from a saved file (created by save_model).
@@ -100,6 +101,7 @@ def load_model_from_saved_file(
         store: Store instance for persistence
         model_id: Optional model identifier. If not provided, will use the model_id from saved metadata.
                  If provided, will be used to load the model architecture from HuggingFace.
+        device: Optional device string or torch.device (defaults to 'cpu' if None)
                  
     Returns:
         LanguageModel instance
@@ -164,7 +166,7 @@ def load_model_from_saved_file(
         ) from e
     
     # Create LanguageModel instance
-    lm = cls(model, tokenizer, store, model_id=model_id)
+    lm = cls(model, tokenizer, store, model_id=model_id, device=device)
     
     # Note: Hooks are not automatically restored as they require hook instances
     # The hook metadata is available in metadata_dict["hooks"] if needed

mi_crow/language_model/utils.py

@@ -44,8 +44,11 @@ def get_device_from_model(model: nn.Module) -> torch.device:
     Returns:
         Device where model parameters are located, or CPU if no parameters
     """
-    first_param = next(model.parameters(), None)
-    return first_param.device if first_param is not None else torch.device("cpu")
+    try:
+        first_param = next(model.parameters(), None)
+        return first_param.device if first_param is not None else torch.device("cpu")
+    except (TypeError, AttributeError):
+        return torch.device("cpu")
 
 
 def move_tensors_to_device(
@@ -62,9 +65,6 @@ def move_tensors_to_device(
     Returns:
         Dictionary with tensors moved to device
     """
-    device_type = str(device.type)
-    if device_type == "cuda":
-        return {k: v.to(device, non_blocking=True) for k, v in tensors.items()}
     return {k: v.to(device) for k, v in tensors.items()}