wisent 0.1.1__py3-none-any.whl

wisent/control_vector/manager.py

@@ -0,0 +1,168 @@
+"""
+Manager for working with control vectors.
+"""
+
+import logging
+from typing import Dict, List, Optional, Union
+
+import torch
+
+from wisent.control_vector.models import ControlVector, ControlVectorConfig
+from wisent.utils.auth import AuthManager
+from wisent.utils.http import HTTPClient
+
+logger = logging.getLogger(__name__)
+
+
+class ControlVectorManager:
+    """
+    Manager for working with control vectors.
+    
+    Args:
+        api_key: Wisent API key
+        base_url: Base URL for the API
+        timeout: Request timeout in seconds
+    """
+    
+    def __init__(
+        self,
+        api_key: str,
+        base_url: str = "https://api.wisent.ai",
+        timeout: int = 60,
+    ):
+        self.auth = AuthManager(api_key)
+        self.http_client = HTTPClient(base_url, self.auth.get_headers(), timeout)
+        self.cache = {}  # Simple in-memory cache
+    
+    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
+        """
+        cache_key = f"{name}:{model}"
+        if cache_key in self.cache:
+            logger.info(f"Using cached control vector: {name} for model {model}")
+            return self.cache[cache_key]
+        
+        logger.info(f"Fetching control vector: {name} for model {model}")
+        data = self.http_client.get(f"/control_vectors/{name}", params={"model": model})
+        vector = ControlVector(**data)
+        
+        # Cache the result
+        self.cache[cache_key] = vector
+        
+        return vector
+    
+    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
+        """
+        # Check if we can combine locally
+        can_combine_locally = True
+        local_vectors = {}
+        
+        for name in vectors.keys():
+            cache_key = f"{name}:{model}"
+            if cache_key not in self.cache:
+                can_combine_locally = False
+                break
+            local_vectors[name] = self.cache[cache_key]
+        
+        if can_combine_locally:
+            logger.info(f"Combining vectors locally for model {model}")
+            return self._combine_locally(local_vectors, vectors, model)
+        
+        # Otherwise, use the API
+        logger.info(f"Combining vectors via API for model {model}")
+        data = self.http_client.post(
+            "/control_vectors/combine",
+            json_data={
+                "vectors": vectors,
+                "model": model,
+            }
+        )
+        return ControlVector(**data)
+    
+    def _combine_locally(
+        self,
+        vectors: Dict[str, ControlVector],
+        weights: Dict[str, float],
+        model: str,
+    ) -> ControlVector:
+        """
+        Combine vectors locally.
+        
+        Args:
+            vectors: Dictionary mapping vector names to ControlVector objects
+            weights: Dictionary mapping vector names to weights
+            model: Model name
+            
+        Returns:
+            Combined control vector
+        """
+        # Convert all vectors to tensors
+        tensor_vectors = {}
+        for name, vector in vectors.items():
+            tensor_vectors[name] = vector.to_tensor()
+        
+        # Get the shape from the first vector
+        first_vector = next(iter(tensor_vectors.values()))
+        combined = torch.zeros_like(first_vector)
+        
+        # Combine vectors with weights
+        for name, weight in weights.items():
+            if name in tensor_vectors:
+                combined += tensor_vectors[name] * weight
+        
+        # Create a new control vector
+        vector_names = list(weights.keys())
+        combined_name = f"combined_{'_'.join(vector_names)}"
+        
+        return ControlVector(
+            name=combined_name,
+            model_name=model,
+            values=combined,
+            metadata={
+                "combined_from": {name: weight for name, weight in weights.items()},
+            }
+        ) 

wisent/control_vector/models.py

@@ -0,0 +1,70 @@
+"""
+Data models for control vectors.
+"""
+
+from dataclasses import dataclass
+from typing import Dict, List, Optional, Union
+
+import numpy as np
+import torch
+from pydantic import BaseModel, Field
+
+
+class ControlVector(BaseModel):
+    """
+    Represents a control vector for steering model outputs.
+    
+    Attributes:
+        name: Name of the control vector
+        model_name: Name of the model the vector is for
+        values: Vector values
+        metadata: Additional metadata
+    """
+    
+    name: str
+    model_name: str
+    values: Union[List[float], np.ndarray, torch.Tensor]
+    metadata: Optional[Dict] = Field(default_factory=dict)
+    
+    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 {
+            "name": self.name,
+            "model_name": self.model_name,
+            "values": values,
+            "metadata": self.metadata or {},
+        }
+    
+    def to_tensor(self, device: str = "cpu") -> torch.Tensor:
+        """Convert values to a PyTorch tensor."""
+        if isinstance(self.values, torch.Tensor):
+            return self.values.to(device)
+        elif isinstance(self.values, np.ndarray):
+            return torch.tensor(self.values, device=device)
+        else:
+            return torch.tensor(self.values, device=device)
+
+
+@dataclass
+class ControlVectorConfig:
+    """
+    Configuration for control vector application.
+    
+    Attributes:
+        scale: Scaling factor for the control vector
+        method: Method for applying the control vector
+        layers: Layers to apply the control vector to
+    """
+    
+    scale: float = 1.0
+    method: str = "caa"  # Context-Aware Addition
+    layers: Optional[List[int]] = None 

wisent/inference/__init__.py

@@ -0,0 +1,9 @@
+"""
+Functionality for model inference with control vectors.
+"""
+
+from wisent.inference.client import InferenceClient
+from wisent.inference.inferencer import Inferencer
+from wisent.inference.models import InferenceConfig, InferenceResponse
+
+__all__ = ["InferenceClient", "Inferencer", "InferenceConfig", "InferenceResponse"] 

wisent/inference/client.py

@@ -0,0 +1,103 @@
+"""
+Client for interacting with the inference API.
+"""
+
+from typing import Dict, List, Optional, Union
+
+from wisent.inference.models import InferenceConfig, InferenceResponse
+from wisent.utils.auth import AuthManager
+from wisent.utils.http import HTTPClient
+
+
+class InferenceClient:
+    """
+    Client for interacting with the inference 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 generate(
+        self,
+        model_name: str,
+        prompt: str,
+        config: Optional[InferenceConfig] = None,
+    ) -> InferenceResponse:
+        """
+        Generate text using a model.
+        
+        Args:
+            model_name: Name of the model
+            prompt: Input prompt
+            config: Inference configuration
+            
+        Returns:
+            Inference response
+        """
+        config = config or InferenceConfig()
+        
+        data = self.http_client.post(
+            "/inference/generate",
+            json_data={
+                "model": model_name,
+                "prompt": prompt,
+                "max_tokens": config.max_tokens,
+                "temperature": config.temperature,
+                "top_p": config.top_p,
+                "top_k": config.top_k,
+                "repetition_penalty": config.repetition_penalty,
+                "stop_sequences": config.stop_sequences,
+            }
+        )
+        
+        return InferenceResponse(**data)
+    
+    def generate_with_control(
+        self,
+        model_name: str,
+        prompt: str,
+        control_vectors: Dict[str, float],
+        method: str = "caa",
+        scale: float = 1.0,
+        config: Optional[InferenceConfig] = None,
+    ) -> InferenceResponse:
+        """
+        Generate text using a model with control vectors.
+        
+        Args:
+            model_name: Name of the model
+            prompt: Input prompt
+            control_vectors: Dictionary mapping vector names to weights
+            method: Method for applying control vectors
+            scale: Scaling factor for control vectors
+            config: Inference configuration
+            
+        Returns:
+            Inference response
+        """
+        config = config or InferenceConfig()
+        
+        data = self.http_client.post(
+            "/inference/generate_with_control",
+            json_data={
+                "model": model_name,
+                "prompt": prompt,
+                "control_vectors": control_vectors,
+                "method": method,
+                "scale": scale,
+                "max_tokens": config.max_tokens,
+                "temperature": config.temperature,
+                "top_p": config.top_p,
+                "top_k": config.top_k,
+                "repetition_penalty": config.repetition_penalty,
+                "stop_sequences": config.stop_sequences,
+            }
+        )
+        
+        return InferenceResponse(**data) 

wisent/inference/inferencer.py

@@ -0,0 +1,250 @@
+"""
+Functionality for local inference with control vectors.
+"""
+
+import logging
+from typing import Dict, List, Optional, Union
+
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig
+
+from wisent.control_vector.models import ControlVector
+from wisent.inference.models import ControlVectorInferenceConfig, InferenceConfig, InferenceResponse
+
+logger = logging.getLogger(__name__)
+
+
+class ControlVectorHook:
+    """
+    Hook for applying control vectors during inference.
+    
+    Args:
+        control_vector: Control vector to apply
+        config: Configuration for applying the control vector
+    """
+    
+    def __init__(
+        self,
+        control_vector: ControlVector,
+        config: ControlVectorInferenceConfig,
+    ):
+        self.control_vector = control_vector
+        self.config = config
+        self.device = None
+        self.vector_tensor = None
+        self.hooks = []
+    
+    def register(self, model):
+        """
+        Register hooks on the model.
+        
+        Args:
+            model: The model to register hooks on
+        """
+        self.device = next(model.parameters()).device
+        self.vector_tensor = self.control_vector.to_tensor(self.device)
+        
+        # Get transformer layers
+        if hasattr(model, "transformer"):
+            transformer_layers = model.transformer.h
+        elif hasattr(model, "model") and hasattr(model.model, "layers"):
+            transformer_layers = model.model.layers
+        else:
+            raise ValueError(f"Unsupported model architecture: {model.__class__.__name__}")
+        
+        # Determine which layers to apply the control vector to
+        num_layers = len(transformer_layers)
+        layers = self.config.layers or [num_layers - 1]  # Default to last layer
+        
+        # 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)
+        
+        # Register hooks
+        for layer_idx in resolved_layers:
+            layer = transformer_layers[layer_idx]
+            
+            # Define hook function
+            def hook_fn(module, input, output, layer_idx=layer_idx):
+                if isinstance(output, tuple):
+                    hidden_states = output[0]
+                else:
+                    hidden_states = output
+                
+                # Apply the control vector
+                if self.config.method == "caa":  # Context-Aware Addition
+                    # Add the control vector to the hidden states
+                    modified = hidden_states + self.vector_tensor * self.config.scale
+                    
+                    if isinstance(output, tuple):
+                        return (modified,) + output[1:]
+                    else:
+                        return modified
+                else:
+                    logger.warning(f"Unsupported method: {self.config.method}, using original output")
+                    return output
+            
+            # Register hook
+            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)
+    
+    def remove(self):
+        """Remove all registered hooks."""
+        for hook in self.hooks:
+            hook.remove()
+        self.hooks = []
+
+
+class Inferencer:
+    """
+    Performs local inference with control vectors.
+    
+    Args:
+        model_name: Name of the model
+        device: Device to use for inference
+    """
+    
+    def __init__(
+        self,
+        model_name: str,
+        device: Optional[str] = None,
+    ):
+        self.model_name = model_name
+        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
+        self.model = None
+        self.tokenizer = None
+        
+        logger.info(f"Initializing Inferencer for model {model_name} on {self.device}")
+    
+    def _load_model(self):
+        """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 generate(
+        self,
+        prompt: str,
+        control_vector: Optional[ControlVector] = None,
+        method: str = "caa",
+        scale: float = 1.0,
+        layers: Optional[List[int]] = None,
+        config: Optional[InferenceConfig] = None,
+    ) -> InferenceResponse:
+        """
+        Generate text using the model, optionally with a control vector.
+        
+        Args:
+            prompt: Input prompt
+            control_vector: Control vector to apply (optional)
+            method: Method for applying the control vector
+            scale: Scaling factor for the control vector
+            layers: Layers to apply the control vector to
+            config: Inference configuration
+            
+        Returns:
+            Inference response
+        """
+        try:
+            self._load_model()
+            
+            config = config or InferenceConfig()
+            hook = None
+            
+            # Register control vector hook if provided
+            if control_vector is not None:
+                cv_config = ControlVectorInferenceConfig(
+                    method=method,
+                    scale=scale,
+                    layers=layers,
+                )
+                hook = ControlVectorHook(control_vector, cv_config)
+                hook.register(self.model)
+            
+            # Tokenize input
+            inputs = self.tokenizer(prompt, return_tensors="pt").to(self.device)
+            prompt_length = inputs.input_ids.shape[1]
+            
+            # Configure generation
+            generation_config = GenerationConfig(
+                max_new_tokens=config.max_tokens,
+                temperature=config.temperature,
+                top_p=config.top_p,
+                top_k=config.top_k,
+                repetition_penalty=config.repetition_penalty,
+                do_sample=config.temperature > 0,
+                pad_token_id=self.tokenizer.pad_token_id or self.tokenizer.eos_token_id,
+            )
+            
+            # Generate
+            with torch.no_grad():
+                output_ids = self.model.generate(
+                    inputs.input_ids,
+                    attention_mask=inputs.attention_mask,
+                    generation_config=generation_config,
+                )
+            
+            # Remove control vector hook if registered
+            if hook is not None:
+                hook.remove()
+            
+            # Decode output
+            generated_text = self.tokenizer.decode(
+                output_ids[0][prompt_length:], 
+                skip_special_tokens=True
+            )
+            
+            # Create response
+            return InferenceResponse(
+                text=generated_text,
+                model=self.model_name,
+                prompt=prompt,
+                finish_reason="length",  # Simplified
+                usage={
+                    "prompt_tokens": prompt_length,
+                    "completion_tokens": output_ids.shape[1] - prompt_length,
+                    "total_tokens": output_ids.shape[1],
+                },
+                metadata={
+                    "control_vector": control_vector.name if control_vector else None,
+                    "method": method if control_vector else None,
+                    "scale": scale if control_vector else None,
+                }
+            )
+            
+        except Exception as e:
+            logger.error(f"Error during inference: {str(e)}")
+            if hook is not None:
+                hook.remove()
+            raise
+    
+    def __del__(self):
+        """Clean up resources."""
+        # 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/inference/models.py

@@ -0,0 +1,66 @@
+"""
+Data models for inference.
+"""
+
+from dataclasses import dataclass, field
+from typing import Dict, List, Optional, Union
+
+from pydantic import BaseModel, Field
+
+
+class InferenceConfig(BaseModel):
+    """
+    Configuration for model inference.
+    
+    Attributes:
+        max_tokens: Maximum number of tokens to generate
+        temperature: Sampling temperature
+        top_p: Top-p sampling parameter
+        top_k: Top-k sampling parameter
+        repetition_penalty: Repetition penalty
+        stop_sequences: Sequences that stop generation
+    """
+    
+    max_tokens: int = 256
+    temperature: float = 0.7
+    top_p: float = 0.9
+    top_k: int = 50
+    repetition_penalty: float = 1.0
+    stop_sequences: Optional[List[str]] = None
+
+
+class InferenceResponse(BaseModel):
+    """
+    Response from model inference.
+    
+    Attributes:
+        text: Generated text
+        model: Model used for generation
+        prompt: Input prompt
+        finish_reason: Reason generation stopped
+        usage: Token usage information
+        metadata: Additional metadata
+    """
+    
+    text: str
+    model: str
+    prompt: str
+    finish_reason: str = "length"
+    usage: Dict[str, int] = Field(default_factory=lambda: {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0})
+    metadata: Dict = Field(default_factory=dict)
+
+
+@dataclass
+class ControlVectorInferenceConfig:
+    """
+    Configuration for inference with control vectors.
+    
+    Attributes:
+        method: Method for applying control vectors
+        scale: Scaling factor for control vectors
+        layers: Layers to apply control vectors to
+    """
+    
+    method: str = "caa"  # Context-Aware Addition
+    scale: float = 1.0
+    layers: Optional[List[int]] = None 

wisent/utils/__init__.py

@@ -0,0 +1,3 @@
+"""
+Utility functions and classes for the Wisent package.
+""" 

wisent/utils/auth.py

@@ -0,0 +1,30 @@
+"""
+Authentication utilities for the Wisent API.
+"""
+
+from typing import Dict
+
+
+class AuthManager:
+    """
+    Manages authentication for Wisent API requests.
+    
+    Args:
+        api_key: The Wisent API key
+    """
+    
+    def __init__(self, api_key: str):
+        self.api_key = api_key
+    
+    def get_headers(self) -> Dict[str, str]:
+        """
+        Get the authentication headers for API requests.
+        
+        Returns:
+            Dict containing the authentication headers
+        """
+        return {
+            "Authorization": f"Bearer {self.api_key}",
+            "Content-Type": "application/json",
+            "Accept": "application/json",
+        }