wisent 0.1.1__py3-none-any.whl

wisent/__init__.py

@@ -0,0 +1,8 @@
+"""
+Wisent - Client library for interacting with the Wisent backend services.
+"""
+
+from wisent.client import WisentClient
+from wisent.version import __version__
+
+__all__ = ["WisentClient", "__version__"]

wisent/activations/__init__.py

@@ -0,0 +1,9 @@
+"""
+Functionality for extracting and managing model activations.
+"""
+
+from wisent.activations.client import ActivationsClient
+from wisent.activations.extractor import ActivationExtractor
+from wisent.activations.models import Activation, ActivationBatch
+
+__all__ = ["ActivationsClient", "ActivationExtractor", "Activation", "ActivationBatch"] 

wisent/activations/client.py

@@ -0,0 +1,97 @@
+"""
+Client for interacting with the activations API.
+"""
+
+from typing import Dict, List, Optional, Union
+
+from wisent.activations.extractor import ActivationExtractor
+from wisent.activations.models import Activation, ActivationBatch
+from wisent.utils.auth import AuthManager
+from wisent.utils.http import HTTPClient
+
+
+class ActivationsClient:
+    """
+    Client for interacting with the activations API.
+    
+    Args:
+        auth_manager: Authentication manager
+        base_url: Base URL for the API
+        timeout: Request timeout in seconds
+    """
+    
+    def __init__(self, auth_manager: AuthManager, base_url: str, timeout: int = 60):
+        self.auth_manager = auth_manager
+        self.http_client = HTTPClient(base_url, auth_manager.get_headers(), timeout)
+    
+    def extract(
+        self,
+        model_name: str,
+        prompt: str,
+        layers: Optional[List[int]] = None,
+        tokens_to_extract: Optional[List[int]] = None,
+        device: Optional[str] = None,
+    ) -> ActivationBatch:
+        """
+        Extract activations from a model for a given prompt.
+        
+        Args:
+            model_name: Name of the model
+            prompt: Input prompt
+            layers: List of layers to extract activations from (default: [-1])
+            tokens_to_extract: List of token indices to extract (default: [-1])
+            device: Device to use for extraction (default: "cuda" if available, else "cpu")
+            
+        Returns:
+            Batch of activations
+        """
+        extractor = ActivationExtractor(model_name, device=device)
+        return extractor.extract(prompt, layers, tokens_to_extract)
+    
+    def upload(self, batch: ActivationBatch) -> Dict:
+        """
+        Upload a batch of activations to the Wisent backend.
+        
+        Args:
+            batch: Batch of activations
+            
+        Returns:
+            Response from the API
+        """
+        return self.http_client.post("/activations/upload", json_data=batch.to_dict())
+    
+    def get(self, batch_id: str) -> ActivationBatch:
+        """
+        Get a batch of activations from the Wisent backend.
+        
+        Args:
+            batch_id: ID of the batch
+            
+        Returns:
+            Batch of activations
+        """
+        data = self.http_client.get(f"/activations/{batch_id}")
+        return ActivationBatch(**data)
+    
+    def list(
+        self,
+        model_name: Optional[str] = None,
+        limit: int = 100,
+        offset: int = 0,
+    ) -> List[Dict]:
+        """
+        List activation batches from the Wisent backend.
+        
+        Args:
+            model_name: Filter by model name
+            limit: Maximum number of results
+            offset: Offset for pagination
+            
+        Returns:
+            List of activation batch metadata
+        """
+        params = {"limit": limit, "offset": offset}
+        if model_name:
+            params["model_name"] = model_name
+            
+        return self.http_client.get("/activations", params=params) 

wisent/activations/extractor.py

@@ -0,0 +1,251 @@
+"""
+Functionality for extracting activations from models.
+"""
+
+import logging
+from typing import Dict, List, Optional, Tuple, Union
+
+import torch
+from torch.utils.hooks import RemovableHandle
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+from wisent.activations.models import Activation, ActivationBatch, ActivationExtractorConfig
+
+logger = logging.getLogger(__name__)
+
+
+class ActivationExtractor:
+    """
+    Extracts activations from transformer models.
+    
+    Args:
+        model_name: Name of the model to extract activations from
+        config: Configuration for extraction
+        device: Device to use for extraction
+    """
+    
+    def __init__(
+        self,
+        model_name: str,
+        config: Optional[ActivationExtractorConfig] = None,
+        device: Optional[str] = None,
+    ):
+        self.model_name = model_name
+        self.config = config or ActivationExtractorConfig()
+        
+        if device:
+            self.config.device = device
+            
+        self.device = self.config.device
+        self.model = None
+        self.tokenizer = None
+        self._hooks = []
+        self._activations = {}
+        
+        logger.info(f"Initializing ActivationExtractor for model {model_name} on {self.device}")
+    
+    def _load_model(self) -> None:
+        """Load the model and tokenizer."""
+        if self.model is None:
+            logger.info(f"Loading model {self.model_name}")
+            self.model = AutoModelForCausalLM.from_pretrained(
+                self.model_name, 
+                torch_dtype=torch.float16 if self.device == "cuda" else torch.float32,
+                device_map=self.device
+            )
+            self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
+            logger.info(f"Model loaded successfully")
+    
+    def _register_hooks(self, layers: List[int]) -> None:
+        """
+        Register hooks to capture activations from specified layers.
+        
+        Args:
+            layers: List of layer indices to capture
+        """
+        self._remove_hooks()
+        self._activations = {}
+        
+        # Get all transformer layers
+        if hasattr(self.model, "transformer"):
+            transformer_layers = self.model.transformer.h
+        elif hasattr(self.model, "model") and hasattr(self.model.model, "layers"):
+            transformer_layers = self.model.model.layers
+        else:
+            raise ValueError(f"Unsupported model architecture: {self.model_name}")
+        
+        num_layers = len(transformer_layers)
+        
+        # Resolve negative indices
+        resolved_layers = []
+        for layer in layers:
+            if layer < 0:
+                resolved_layer = num_layers + layer
+            else:
+                resolved_layer = layer
+                
+            if 0 <= resolved_layer < num_layers:
+                resolved_layers.append(resolved_layer)
+            else:
+                logger.warning(f"Layer index {layer} out of range (0-{num_layers-1}), skipping")
+        
+        # Register hooks for each layer
+        for layer_idx in resolved_layers:
+            layer = transformer_layers[layer_idx]
+            
+            # Define hook function to capture activations
+            def hook_fn(module, input, output, layer_idx=layer_idx):
+                # For most models, output is a tuple with hidden states as the first element
+                if isinstance(output, tuple):
+                    hidden_states = output[0]
+                else:
+                    hidden_states = output
+                
+                if layer_idx not in self._activations:
+                    self._activations[layer_idx] = []
+                
+                # Store a copy of the hidden states
+                self._activations[layer_idx].append(hidden_states.detach())
+            
+            # Register hook on the output of the layer
+            if hasattr(layer, "output"):
+                handle = layer.output.register_forward_hook(
+                    lambda module, input, output, layer_idx=layer_idx: hook_fn(module, input, output, layer_idx)
+                )
+            else:
+                handle = layer.register_forward_hook(
+                    lambda module, input, output, layer_idx=layer_idx: hook_fn(module, input, output, layer_idx)
+                )
+            
+            self._hooks.append(handle)
+            
+        logger.info(f"Registered hooks for layers: {resolved_layers}")
+    
+    def _remove_hooks(self) -> None:
+        """Remove all registered hooks."""
+        for hook in self._hooks:
+            hook.remove()
+        self._hooks = []
+    
+    def _get_token_indices(self, tokens_to_extract: List[int], total_tokens: int) -> List[int]:
+        """
+        Resolve token indices, handling negative indices.
+        
+        Args:
+            tokens_to_extract: List of token indices to extract
+            total_tokens: Total number of tokens
+            
+        Returns:
+            List of resolved token indices
+        """
+        resolved_indices = []
+        
+        for idx in tokens_to_extract:
+            if idx < 0:
+                resolved_idx = total_tokens + idx
+            else:
+                resolved_idx = idx
+                
+            if 0 <= resolved_idx < total_tokens:
+                resolved_indices.append(resolved_idx)
+            else:
+                logger.warning(f"Token index {idx} out of range (0-{total_tokens-1}), skipping")
+        
+        return resolved_indices
+    
+    def extract(
+        self,
+        prompt: str,
+        layers: Optional[List[int]] = None,
+        tokens_to_extract: Optional[List[int]] = None,
+    ) -> ActivationBatch:
+        """
+        Extract activations from the model for a given prompt.
+        
+        Args:
+            prompt: Input prompt
+            layers: List of layers to extract activations from (default: from config)
+            tokens_to_extract: List of token indices to extract (default: from config)
+            
+        Returns:
+            Batch of activations
+        """
+        try:
+            self._load_model()
+            
+            layers = layers or self.config.layers
+            tokens_to_extract = tokens_to_extract or self.config.tokens_to_extract
+            
+            # Register hooks for the specified layers
+            self._register_hooks(layers)
+            
+            # Tokenize the input
+            inputs = self.tokenizer(prompt, return_tensors="pt")
+            input_ids = inputs.input_ids.to(self.device)
+            
+            # Get the total number of tokens
+            total_tokens = input_ids.shape[1]
+            
+            # Resolve token indices
+            token_indices = self._get_token_indices(tokens_to_extract, total_tokens)
+            
+            # Run the model to capture activations
+            with torch.no_grad():
+                self.model(input_ids)
+            
+            # Process captured activations
+            activations = []
+            
+            for layer_idx, layer_activations in self._activations.items():
+                # Layer activations should have shape [batch_size, seq_len, hidden_dim]
+                hidden_states = layer_activations[0]
+                
+                # Get token strings for the specified indices
+                token_strings = {}
+                for token_idx in token_indices:
+                    token_id = input_ids[0, token_idx].item()
+                    token_strings[token_idx] = self.tokenizer.decode([token_id])
+                
+                # Extract activations for the specified tokens
+                for token_idx in token_indices:
+                    # Extract the activation for this token
+                    token_activation = hidden_states[0, token_idx, :].cpu()
+                    
+                    # Create an Activation object
+                    activation = Activation(
+                        model_name=self.model_name,
+                        layer=layer_idx,
+                        token_index=token_idx,
+                        values=token_activation,
+                        token_str=token_strings.get(token_idx)
+                    )
+                    
+                    activations.append(activation)
+            
+            # Clean up
+            self._remove_hooks()
+            
+            # Create and return the batch
+            return ActivationBatch(
+                model_name=self.model_name,
+                prompt=prompt,
+                activations=activations,
+                metadata={"total_tokens": total_tokens}
+            )
+            
+        except Exception as e:
+            logger.error(f"Error extracting activations: {str(e)}")
+            self._remove_hooks()
+            raise
+    
+    def __del__(self):
+        """Clean up resources."""
+        self._remove_hooks()
+        
+        # Free GPU memory
+        if self.model is not None and hasattr(self.model, "to"):
+            self.model = self.model.to("cpu")
+            
+        # Clear CUDA cache
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache() 

wisent/activations/models.py

@@ -0,0 +1,95 @@
+"""
+Data models for model activations.
+"""
+
+from dataclasses import dataclass, field
+from typing import Dict, List, Optional, Union
+
+import numpy as np
+import torch
+from pydantic import BaseModel, Field
+
+
+class Activation(BaseModel):
+    """
+    Represents a single activation from a model.
+    
+    Attributes:
+        model_name: Name of the model
+        layer: Layer index
+        token_index: Token index
+        values: Activation values
+        token_str: String representation of the token (optional)
+    """
+    
+    model_name: str
+    layer: int
+    token_index: int
+    values: Union[List[float], np.ndarray, torch.Tensor]
+    token_str: Optional[str] = None
+    
+    class Config:
+        arbitrary_types_allowed = True
+    
+    def to_dict(self) -> Dict:
+        """Convert to dictionary for API requests."""
+        values = self.values
+        if isinstance(values, torch.Tensor):
+            values = values.detach().cpu().numpy()
+        if isinstance(values, np.ndarray):
+            values = values.tolist()
+            
+        return {
+            "model_name": self.model_name,
+            "layer": self.layer,
+            "token_index": self.token_index,
+            "values": values,
+            "token_str": self.token_str,
+        }
+
+
+class ActivationBatch(BaseModel):
+    """
+    Represents a batch of activations from a model.
+    
+    Attributes:
+        model_name: Name of the model
+        prompt: Input prompt that generated the activations
+        activations: List of activations
+        metadata: Additional metadata (optional)
+    """
+    
+    model_name: str
+    prompt: str
+    activations: List[Activation]
+    metadata: Optional[Dict] = Field(default_factory=dict)
+    
+    class Config:
+        arbitrary_types_allowed = True
+    
+    def to_dict(self) -> Dict:
+        """Convert to dictionary for API requests."""
+        return {
+            "model_name": self.model_name,
+            "prompt": self.prompt,
+            "activations": [a.to_dict() for a in self.activations],
+            "metadata": self.metadata or {},
+        }
+
+
+@dataclass
+class ActivationExtractorConfig:
+    """
+    Configuration for activation extraction.
+    
+    Attributes:
+        layers: List of layers to extract activations from
+        tokens_to_extract: List of token indices to extract (negative indices count from the end)
+        batch_size: Batch size for processing
+        device: Device to use for extraction
+    """
+    
+    layers: List[int] = field(default_factory=lambda: [-1])
+    tokens_to_extract: List[int] = field(default_factory=lambda: [-1])
+    batch_size: int = 1
+    device: str = "cuda" if torch.cuda.is_available() else "cpu" 

wisent/client.py

@@ -0,0 +1,45 @@
+"""
+Main client class for interacting with the Wisent backend services.
+"""
+
+from typing import Dict, Optional
+
+from wisent.activations import ActivationsClient
+from wisent.control_vector import ControlVectorClient
+from wisent.inference import InferenceClient
+from wisent.utils.auth import AuthManager
+
+
+class WisentClient:
+    """
+    Main client for interacting with the Wisent backend services.
+    
+    This client provides access to all Wisent API functionality through
+    specialized sub-clients for different features.
+    
+    Args:
+        api_key: Your Wisent API key
+        base_url: The base URL for the Wisent API (default: https://api.wisent.ai)
+        timeout: Request timeout in seconds (default: 60)
+    """
+
+    def __init__(
+        self,
+        api_key: str,
+        base_url: str = "https://api.wisent.ai",
+        timeout: int = 60,
+    ):
+        self.api_key = api_key
+        self.base_url = base_url
+        self.timeout = timeout
+        
+        # Initialize auth manager
+        self.auth = AuthManager(api_key)
+        
+        # Initialize sub-clients
+        self.activations = ActivationsClient(self.auth, base_url, timeout)
+        self.control_vector = ControlVectorClient(self.auth, base_url, timeout)
+        self.inference = InferenceClient(self.auth, base_url, timeout)
+    
+    def __repr__(self) -> str:
+        return f"WisentClient(base_url='{self.base_url}')"

wisent/control_vector/__init__.py

@@ -0,0 +1,9 @@
+"""
+Functionality for working with control vectors.
+"""
+
+from wisent.control_vector.client import ControlVectorClient
+from wisent.control_vector.manager import ControlVectorManager
+from wisent.control_vector.models import ControlVector, ControlVectorConfig
+
+__all__ = ["ControlVectorClient", "ControlVectorManager", "ControlVector", "ControlVectorConfig"] 

wisent/control_vector/client.py

@@ -0,0 +1,85 @@
+"""
+Client for interacting with the control vector API.
+"""
+
+from typing import Dict, List, Optional, Union
+
+from wisent.control_vector.models import ControlVector
+from wisent.utils.auth import AuthManager
+from wisent.utils.http import HTTPClient
+
+
+class ControlVectorClient:
+    """
+    Client for interacting with the control vector API.
+    
+    Args:
+        auth_manager: Authentication manager
+        base_url: Base URL for the API
+        timeout: Request timeout in seconds
+    """
+    
+    def __init__(self, auth_manager: AuthManager, base_url: str, timeout: int = 60):
+        self.auth_manager = auth_manager
+        self.http_client = HTTPClient(base_url, auth_manager.get_headers(), timeout)
+    
+    def get(self, name: str, model: str) -> ControlVector:
+        """
+        Get a control vector from the Wisent backend.
+        
+        Args:
+            name: Name of the control vector
+            model: Model name
+            
+        Returns:
+            Control vector
+        """
+        data = self.http_client.get(f"/control_vectors/{name}", params={"model": model})
+        return ControlVector(**data)
+    
+    def list(
+        self,
+        model: Optional[str] = None,
+        limit: int = 100,
+        offset: int = 0,
+    ) -> List[Dict]:
+        """
+        List available control vectors from the Wisent backend.
+        
+        Args:
+            model: Filter by model name
+            limit: Maximum number of results
+            offset: Offset for pagination
+            
+        Returns:
+            List of control vector metadata
+        """
+        params = {"limit": limit, "offset": offset}
+        if model:
+            params["model"] = model
+            
+        return self.http_client.get("/control_vectors", params=params)
+    
+    def combine(
+        self,
+        vectors: Dict[str, float],
+        model: str,
+    ) -> ControlVector:
+        """
+        Combine multiple control vectors with weights.
+        
+        Args:
+            vectors: Dictionary mapping vector names to weights
+            model: Model name
+            
+        Returns:
+            Combined control vector
+        """
+        data = self.http_client.post(
+            "/control_vectors/combine",
+            json_data={
+                "vectors": vectors,
+                "model": model,
+            }
+        )
+        return ControlVector(**data)