wisent 0.1.1__py3-none-any.whl → 0.5.2__py3-none-any.whl

wisent/activations/models.py

@@ -1,95 +0,0 @@
-"""
-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

@@ -1,45 +0,0 @@
-"""
-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

@@ -1,9 +0,0 @@
-"""
-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

@@ -1,85 +0,0 @@
-"""
-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) 

wisent/control_vector/manager.py

@@ -1,168 +0,0 @@
-"""
-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

@@ -1,70 +0,0 @@
-"""
-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

@@ -1,9 +0,0 @@
-"""
-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

@@ -1,103 +0,0 @@
-"""
-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)