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

amber/language_model/tokenizer.py

@@ -0,0 +1,203 @@
+from typing import Dict, List, Callable, Sequence, Any, TYPE_CHECKING, Union
+
+from torch import nn
+from transformers import AutoTokenizer
+
+if TYPE_CHECKING:
+    from amber.language_model.context import LanguageModelContext
+
+
+class LanguageModelTokenizer:
+    """Handles tokenization for LanguageModel."""
+
+    def __init__(
+            self,
+            context: "LanguageModelContext"
+    ):
+        """
+        Initialize LanguageModelTokenizer.
+        
+        Args:
+            context: LanguageModelContext instance
+        """
+        self.context = context
+
+    def _setup_pad_token(self, tokenizer: Any, model: Any) -> None:
+        """
+        Setup pad token for tokenizer if not already set.
+        
+        Args:
+            tokenizer: Tokenizer instance
+            model: Model instance
+        """
+        eos_token = getattr(tokenizer, "eos_token", None)
+        if eos_token is not None:
+            tokenizer.pad_token = eos_token
+            if hasattr(model, "config"):
+                model.config.pad_token_id = getattr(tokenizer, "eos_token_id", None)
+        else:
+            if hasattr(tokenizer, "add_special_tokens"):
+                tokenizer.add_special_tokens({"pad_token": "[PAD]"})
+                if hasattr(model, "resize_token_embeddings"):
+                    model.resize_token_embeddings(len(tokenizer))
+                if hasattr(model, "config"):
+                    model.config.pad_token_id = getattr(tokenizer, "pad_token_id", None)
+
+    def split_to_tokens(
+            self,
+            text: Union[str, Sequence[str]],
+            add_special_tokens: bool = False
+    ) -> Union[List[str], List[List[str]]]:
+        """
+        Split text into token strings.
+        
+        Args:
+            text: Single string or sequence of strings to tokenize
+            add_special_tokens: Whether to add special tokens (e.g., BOS, EOS)
+        
+        Returns:
+            For a single string: list of token strings
+            For a sequence of strings: list of lists of token strings
+        """
+        if isinstance(text, str):
+            return self._split_single_text_to_tokens(text, add_special_tokens)
+
+        return [self._split_single_text_to_tokens(t, add_special_tokens) for t in text]
+
+    def _try_tokenize_with_method(
+            self,
+            tokenizer: Any,
+            text: str,
+            add_special_tokens: bool,
+            method_name: str,
+            fallback_method: str | None = None
+    ) -> List[str] | None:
+        """
+        Try tokenizing using a specific tokenizer method.
+        
+        Args:
+            tokenizer: Tokenizer instance
+            text: Text to tokenize
+            add_special_tokens: Whether to add special tokens
+            method_name: Primary method to try (e.g., "tokenize", "encode")
+            fallback_method: Optional fallback method (e.g., "convert_ids_to_tokens")
+            
+        Returns:
+            List of token strings or None if method fails
+        """
+        if not hasattr(tokenizer, method_name):
+            return None
+        
+        try:
+            if method_name == "tokenize":
+                return tokenizer.tokenize(text, add_special_tokens=add_special_tokens)
+            elif method_name == "encode":
+                if fallback_method and hasattr(tokenizer, fallback_method):
+                    token_ids = tokenizer.encode(text, add_special_tokens=add_special_tokens)
+                    return tokenizer.convert_ids_to_tokens(token_ids)
+            elif method_name == "encode_plus":
+                if fallback_method and hasattr(tokenizer, fallback_method):
+                    encoded = tokenizer.encode_plus(text, add_special_tokens=add_special_tokens)
+                    if isinstance(encoded, dict) and "input_ids" in encoded:
+                        token_ids = encoded["input_ids"]
+                        return tokenizer.convert_ids_to_tokens(token_ids)
+        except (TypeError, ValueError, AttributeError):
+            pass
+        
+        return None
+
+    def _split_single_text_to_tokens(self, text: str, add_special_tokens: bool) -> List[str]:
+        """
+        Split a single text into token strings.
+        
+        Uses the tokenizer from LanguageModelContext to split text into tokens.
+        
+        Args:
+            text: Text string to tokenize
+            add_special_tokens: Whether to add special tokens
+            
+        Returns:
+            List of token strings
+        """
+        tokenizer = self.context.tokenizer
+
+        if tokenizer is None:
+            return text.split()
+        
+        if not isinstance(text, str):
+            raise TypeError(f"Expected str, got {type(text)}")
+
+        # Try different tokenization methods in order
+        tokens = self._try_tokenize_with_method(tokenizer, text, add_special_tokens, "tokenize")
+        if tokens is not None:
+            return tokens
+
+        tokens = self._try_tokenize_with_method(tokenizer, text, add_special_tokens, "encode", "convert_ids_to_tokens")
+        if tokens is not None:
+            return tokens
+
+        tokens = self._try_tokenize_with_method(tokenizer, text, add_special_tokens, "encode_plus", "convert_ids_to_tokens")
+        if tokens is not None:
+            return tokens
+
+        return text.split()
+
+    def tokenize(
+            self,
+            texts: Sequence[str],
+            padding: bool = False,
+            pad_token: str = "[PAD]",
+            **kwargs: Any
+    ) -> Any:
+        """
+        Robust batch tokenization that works across tokenizer variants.
+
+        Tries methods in order:
+        - callable tokenizer (most HF tokenizers)
+        - batch_encode_plus
+        - encode_plus per item + tokenizer.pad to collate
+        
+        Args:
+            texts: Sequence of text strings to tokenize
+            padding: Whether to pad sequences
+            pad_token: Pad token string
+            **kwargs: Additional tokenizer arguments
+            
+        Returns:
+            Tokenized encodings
+            
+        Raises:
+            ValueError: If tokenizer is not initialized
+            TypeError: If tokenizer is not usable for batch tokenization
+        """
+        tokenizer = self.context.tokenizer
+        if tokenizer is None:
+            raise ValueError("Tokenizer must be initialized before tokenization")
+        
+        model = self.context.model
+
+        if padding and pad_token and getattr(tokenizer, "pad_token", None) is None:
+            self._setup_pad_token(tokenizer, model)
+
+        kwargs["padding"] = padding
+        
+        # Try callable tokenizer first (most common case)
+        if callable(tokenizer):
+            try:
+                return tokenizer(texts, **kwargs)
+            except TypeError:
+                pass
+
+        # Try batch_encode_plus
+        if hasattr(tokenizer, "batch_encode_plus"):
+            return tokenizer.batch_encode_plus(texts, **kwargs)
+        
+        # Fallback to encode_plus per item
+        if hasattr(tokenizer, "encode_plus"):
+            encoded = [tokenizer.encode_plus(t, **kwargs) for t in texts]
+            if hasattr(tokenizer, "pad"):
+                rt = kwargs.get("return_tensors") or "pt"
+                return tokenizer.pad(encoded, return_tensors=rt)
+            return encoded
+        
+        raise TypeError("Tokenizer object on LanguageModel is not usable for batch tokenization")

amber/language_model/utils.py

@@ -0,0 +1,97 @@
+"""Utility functions for language model operations."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import torch
+from torch import nn
+
+if TYPE_CHECKING:
+    from transformers import PreTrainedTokenizerBase
+
+
+def extract_model_id(model: nn.Module, provided_model_id: str | None = None) -> str:
+    """
+    Extract model ID from model or use provided one.
+    
+    Args:
+        model: PyTorch model module
+        provided_model_id: Optional model ID provided by user
+        
+    Returns:
+        Model ID string
+        
+    Raises:
+        ValueError: If model_id cannot be determined
+    """
+    if provided_model_id is not None:
+        return provided_model_id
+    
+    if hasattr(model, 'config') and hasattr(model.config, 'name_or_path'):
+        return model.config.name_or_path.replace("/", "_")
+    
+    return model.__class__.__name__
+
+
+def get_device_from_model(model: nn.Module) -> torch.device:
+    """
+    Get the device from model parameters.
+    
+    Args:
+        model: PyTorch model module
+        
+    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")
+
+
+def move_tensors_to_device(
+        tensors: dict[str, torch.Tensor],
+        device: torch.device
+) -> dict[str, torch.Tensor]:
+    """
+    Move dictionary of tensors to specified device.
+    
+    Args:
+        tensors: Dictionary of tensor name to tensor
+        device: Target 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()}
+
+
+def extract_logits_from_output(output: any) -> torch.Tensor:
+    """
+    Extract logits tensor from model output.
+    
+    Handles various output formats:
+    - Objects with 'logits' attribute (e.g., HuggingFace model outputs)
+    - Tuples (takes first element)
+    - Direct tensors
+    
+    Args:
+        output: Model output (various formats)
+        
+    Returns:
+        Logits tensor
+        
+    Raises:
+        ValueError: If logits cannot be extracted from output
+    """
+    if hasattr(output, 'logits'):
+        return output.logits
+    elif isinstance(output, tuple) and len(output) > 0:
+        return output[0]
+    elif isinstance(output, torch.Tensor):
+        return output
+    else:
+        raise ValueError(f"Unable to extract logits from output type: {type(output)}")
+

amber/mechanistic/sae/autoencoder_context.py

@@ -0,0 +1,40 @@
+from dataclasses import dataclass
+from typing import Optional, TYPE_CHECKING
+
+from amber.store.store import Store
+
+if TYPE_CHECKING:
+    pass
+
+
+@dataclass
+class AutoencoderContext:
+    """Shared context for Autoencoder and its nested components."""
+
+    autoencoder: "Sae"
+
+    # Core SAE parameters
+    n_latents: int
+    n_inputs: int
+
+    # Language model parameters (shared across hierarchy)
+    lm: Optional["LanguageModel"] = None
+    lm_layer_signature: Optional[int | str] = None
+    model_id: Optional[str] = None
+
+    # Training/experiment metadata
+    device: str = 'cpu'
+    experiment_name: Optional[str] = None
+    run_id: Optional[str] = None
+
+    # Text tracking parameters
+    text_tracking_enabled: bool = False
+    text_tracking_k: int = 5
+    text_tracking_negative: bool = False
+
+    store: Optional[Store] = None
+
+    # Training parameters
+    tied: bool = False
+    bias_init: float = 0.0
+    init_method: str = "kaiming"

amber/mechanistic/sae/concepts/autoencoder_concepts.py

@@ -0,0 +1,332 @@
+from __future__ import annotations
+
+from pathlib import Path
+from typing import TYPE_CHECKING, Sequence
+import json
+import csv
+import heapq
+
+import torch
+from torch import nn
+
+from amber.mechanistic.sae.concepts.concept_models import NeuronText
+from amber.mechanistic.sae.autoencoder_context import AutoencoderContext
+from amber.utils import get_logger
+
+if TYPE_CHECKING:
+    from amber.mechanistic.sae.concepts.concept_dictionary import ConceptDictionary
+
+logger = get_logger(__name__)
+
+
+class AutoencoderConcepts:
+    def __init__(
+            self,
+            context: AutoencoderContext
+    ):
+        self.context = context
+        self._n_size = context.n_latents
+        self.dictionary: ConceptDictionary | None = None
+
+        # Concept manipulation parameters
+        self.multiplication = nn.Parameter(torch.ones(self._n_size))
+        self.bias = nn.Parameter(torch.ones(self._n_size))
+
+        # Top texts tracking
+        self._top_texts_heaps: list[list[tuple[float, tuple[float, str, int]]]] | None = None
+        self._text_tracking_k: int = 5
+        self._text_tracking_negative: bool = False
+
+    def enable_text_tracking(self):
+        """Enable text tracking using context parameters."""
+        if self.context.lm is None:
+            raise ValueError("LanguageModel must be set in context to enable tracking")
+
+        # Store tracking parameters
+        self._text_tracking_k = self.context.text_tracking_k
+        self._text_tracking_negative = self.context.text_tracking_negative
+
+        # Ensure InputTracker singleton exists on LanguageModel and enable it
+        input_tracker = self.context.lm._ensure_input_tracker()
+        input_tracker.enable()
+
+        # Enable text tracking on the SAE instance
+        if hasattr(self.context.autoencoder, '_text_tracking_enabled'):
+            self.context.autoencoder._text_tracking_enabled = True
+
+    def disable_text_tracking(self):
+        self.context.autoencoder._text_tracking_enabled = False
+
+    def _ensure_dictionary(self):
+        if self.dictionary is None:
+            from amber.mechanistic.sae.concepts.concept_dictionary import ConceptDictionary
+            self.dictionary = ConceptDictionary(self._n_size)
+        return self.dictionary
+
+    def load_concepts_from_csv(self, csv_filepath: str | Path):
+        from amber.mechanistic.sae.concepts.concept_dictionary import ConceptDictionary
+        self.dictionary = ConceptDictionary.from_csv(
+            csv_filepath=csv_filepath,
+            n_size=self._n_size,
+            store=self.dictionary.store if self.dictionary else None
+        )
+
+    def load_concepts_from_json(self, json_filepath: str | Path):
+        from amber.mechanistic.sae.concepts.concept_dictionary import ConceptDictionary
+        self.dictionary = ConceptDictionary.from_json(
+            json_filepath=json_filepath,
+            n_size=self._n_size,
+            store=self.dictionary.store if self.dictionary else None
+        )
+
+    def generate_concepts_with_llm(self, llm_provider: str | None = None):
+        """Generate concepts using LLM based on current top texts"""
+        if self._top_texts_heaps is None:
+            raise ValueError("No top texts available. Enable text tracking and run inference first.")
+
+        from amber.mechanistic.sae.concepts.concept_dictionary import ConceptDictionary
+        neuron_texts = self.get_all_top_texts()
+
+        self.dictionary = ConceptDictionary.from_llm(
+            neuron_texts=neuron_texts,
+            n_size=self._n_size,
+            store=self.dictionary.store if self.dictionary else None,
+            llm_provider=llm_provider
+        )
+
+    def _ensure_heaps(self, n_neurons: int) -> None:
+        """Ensure heaps are initialized for the given number of neurons."""
+        if self._top_texts_heaps is None:
+            self._top_texts_heaps = [[] for _ in range(n_neurons)]
+
+    def _decode_token(self, text: str, token_idx: int) -> str:
+        """
+        Decode a specific token from the text using the language model's tokenizer.
+        
+        The token_idx is relative to the sequence length T that the model saw during inference.
+        However, there's a mismatch: during inference, texts are tokenized with 
+        add_special_tokens=True (which adds BOS/EOS), but the token_idx appears to be
+        calculated relative to the sequence without special tokens.
+        
+        We tokenize the text the same way as _decode_token originally did (without special tokens)
+        to match the token_idx calculation, but we also account for truncation that may have
+        occurred during inference (max_length).
+        """
+        if self.context.lm is None:
+            return f"<token_{token_idx}>"
+
+        try:
+            if self.context.lm.tokenizer is None:
+                return f"<token_{token_idx}>"
+
+            # Use the raw tokenizer (not the wrapper) to encode and decode
+            tokenizer = self.context.lm.tokenizer
+
+            # Tokenize without special tokens (matching original behavior)
+            # This matches how token_idx was calculated in update_top_texts_from_latents
+            tokens = tokenizer.encode(text, add_special_tokens=False)
+            
+            # Check if token_idx is valid
+            if 0 <= token_idx < len(tokens):
+                token_id = tokens[token_idx]
+                # Decode the specific token
+                token_str = tokenizer.decode([token_id])
+                return token_str
+            else:
+                return f"<token_{token_idx}_out_of_range>"
+        except Exception as e:
+            # If tokenization fails, return a placeholder
+            logger.debug(f"Token decode error for token_idx={token_idx} in text (len={len(text)}): {e}")
+            return f"<token_{token_idx}_decode_error>"
+
+    def update_top_texts_from_latents(
+            self,
+            latents: torch.Tensor,
+            texts: Sequence[str],
+            original_shape: tuple[int, ...] | None = None
+    ) -> None:
+        """
+        Update top texts heaps from latents and texts.
+        
+        Args:
+            latents: Latent activations tensor, shape [B*T, n_latents] or [B, n_latents] (already flattened)
+            texts: List of texts corresponding to the batch
+            original_shape: Original shape before flattening, e.g., (B, T, D) or (B, D)
+        """
+        if not texts:
+            return
+
+        n_neurons = latents.shape[-1]
+        self._ensure_heaps(n_neurons)
+
+        # Calculate batch and token dimensions
+        original_B = len(texts)
+        BT = latents.shape[0]  # Total positions (B*T if 3D original, or B if 2D original)
+
+        # Determine if original was 3D or 2D
+        if original_shape is not None and len(original_shape) == 3:
+            # Original was [B, T, D], latents are [B*T, n_latents]
+            B, T, _ = original_shape
+            # Verify batch size matches
+            if B != original_B:
+                logger.warning(f"Batch size mismatch: original_shape has B={B}, but {original_B} texts provided")
+                # Use the actual number of texts as batch size
+                B = original_B
+                T = BT // B if B > 0 else 1
+            # Create token indices: [0, 1, 2, ..., T-1, 0, 1, 2, ..., T-1, ...]
+            token_indices = torch.arange(T, device='cpu').unsqueeze(0).expand(B, T).contiguous().view(B * T)
+        else:
+            # Original was [B, D], latents are [B, n_latents]
+            # All tokens are at index 0
+            T = 1
+            token_indices = torch.zeros(BT, dtype=torch.long, device='cpu')
+
+        # For each neuron, find the maximum activation per text
+        # This ensures we only track the best activation for each text, not every token position
+        for j in range(n_neurons):
+            heap = self._top_texts_heaps[j]
+
+            # For each text in the batch, find the max activation and its token position
+            texts_processed = 0
+            texts_added = 0
+            texts_updated = 0
+            texts_skipped_duplicate = 0
+            for batch_idx in range(original_B):
+                if batch_idx >= len(texts):
+                    continue
+
+                text = texts[batch_idx]
+                texts_processed += 1
+
+                # Get activations for this text (all token positions)
+                if original_shape is not None and len(original_shape) == 3:
+                    # 3D case: [B, T, D] -> get slice for this batch
+                    start_idx = batch_idx * T
+                    end_idx = start_idx + T
+                    text_activations = latents[start_idx:end_idx, j]  # [T]
+                    text_token_indices = token_indices[start_idx:end_idx]  # [T]
+                else:
+                    # 2D case: [B, D] -> single token
+                    text_activations = latents[batch_idx:batch_idx + 1, j]  # [1]
+                    text_token_indices = token_indices[batch_idx:batch_idx + 1]  # [1]
+
+                # Find the maximum activation (or minimum if tracking negative)
+                if self._text_tracking_negative:
+                    # For negative tracking, find the most negative (minimum) value
+                    max_idx = torch.argmin(text_activations)
+                    max_score = float(text_activations[max_idx].item())
+                    adj = -max_score  # Negate for heap ordering
+                else:
+                    # For positive tracking, find the maximum value
+                    max_idx = torch.argmax(text_activations)
+                    max_score = float(text_activations[max_idx].item())
+                    adj = max_score
+
+                # Skip if score is zero (no activation)
+                if max_score == 0.0:
+                    continue
+
+                token_idx = int(text_token_indices[max_idx].item())
+
+                # Check if we already have this text in the heap
+                # If so, only update if this activation is better
+                existing_entry = None
+                heap_texts = []
+                for heap_idx, (heap_adj, (heap_score, heap_text, heap_token_idx)) in enumerate(heap):
+                    heap_texts.append(heap_text[:50] if len(heap_text) > 50 else heap_text)
+                    if heap_text == text:
+                        existing_entry = (heap_idx, heap_adj, heap_score, heap_token_idx)
+                        break
+
+                if existing_entry is not None:
+                    # Update existing entry if this activation is better
+                    heap_idx, heap_adj, heap_score, heap_token_idx = existing_entry
+                    if adj > heap_adj:
+                        # Replace with better activation
+                        heap[heap_idx] = (adj, (max_score, text, token_idx))
+                        heapq.heapify(heap)  # Re-heapify after modification
+                        texts_updated += 1
+                    else:
+                        texts_skipped_duplicate += 1
+                else:
+                    # New text, add to heap
+                    if len(heap) < self._text_tracking_k:
+                        heapq.heappush(heap, (adj, (max_score, text, token_idx)))
+                        texts_added += 1
+                    else:
+                        # Compare with smallest adjusted score; replace if better
+                        if adj > heap[0][0]:
+                            heapq.heapreplace(heap, (adj, (max_score, text, token_idx)))
+                            texts_added += 1
+
+    def get_top_texts_for_neuron(self, neuron_idx: int, top_m: int | None = None) -> list[NeuronText]:
+        """Get top texts for a specific neuron."""
+        if self._top_texts_heaps is None or neuron_idx < 0 or neuron_idx >= len(self._top_texts_heaps):
+            return []
+        heap = self._top_texts_heaps[neuron_idx]
+        items = [val for (_, val) in heap]
+        reverse = not self._text_tracking_negative
+        items_sorted = sorted(items, key=lambda s_t: s_t[0], reverse=reverse)
+        if top_m is not None:
+            items_sorted = items_sorted[: top_m]
+
+        neuron_texts = []
+        for score, text, token_idx in items_sorted:
+            token_str = self._decode_token(text, token_idx)
+            neuron_texts.append(NeuronText(score=score, text=text, token_idx=token_idx, token_str=token_str))
+        return neuron_texts
+
+    def get_all_top_texts(self) -> list[list[NeuronText]]:
+        """Get top texts for all neurons."""
+        if self._top_texts_heaps is None:
+            return []
+        return [self.get_top_texts_for_neuron(i) for i in range(len(self._top_texts_heaps))]
+
+    def reset_top_texts(self) -> None:
+        """Reset all tracked top texts."""
+        self._top_texts_heaps = None
+
+    def export_top_texts_to_json(self, filepath: Path | str) -> Path:
+        if self._top_texts_heaps is None:
+            raise ValueError("No top texts available. Enable text tracking and run inference first.")
+
+        filepath = Path(filepath)
+        filepath.parent.mkdir(parents=True, exist_ok=True)
+
+        all_texts = self.get_all_top_texts()
+        export_data = {}
+
+        for neuron_idx, neuron_texts in enumerate(all_texts):
+            export_data[neuron_idx] = [
+                {
+                    "text": nt.text,
+                    "score": nt.score,
+                    "token_str": nt.token_str,
+                    "token_idx": nt.token_idx
+                }
+                for nt in neuron_texts
+            ]
+
+        with filepath.open("w", encoding="utf-8") as f:
+            json.dump(export_data, f, ensure_ascii=False, indent=2)
+
+        return filepath
+
+    def export_top_texts_to_csv(self, filepath: Path | str) -> Path:
+        if self._top_texts_heaps is None:
+            raise ValueError("No top texts available. Enable text tracking and run inference first.")
+
+        filepath = Path(filepath)
+        filepath.parent.mkdir(parents=True, exist_ok=True)
+
+        all_texts = self.get_all_top_texts()
+
+        with filepath.open("w", newline="", encoding="utf-8") as f:
+            writer = csv.writer(f)
+            writer.writerow(["neuron_idx", "text", "score", "token_str", "token_idx"])
+
+            for neuron_idx, neuron_texts in enumerate(all_texts):
+                for nt in neuron_texts:
+                    writer.writerow([neuron_idx, nt.text, nt.score, nt.token_str, nt.token_idx])
+
+        return filepath