wisent 0.1.1__tar.gz

wisent-0.1.1/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 Wisent Team
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

wisent-0.1.1/PKG-INFO

@@ -0,0 +1,142 @@
+Metadata-Version: 2.2
+Name: wisent
+Version: 0.1.1
+Summary: Client library for interacting with the Wisent backend services
+Home-page: https://github.com/wisent-ai/wisent
+Author: Wisent Team
+Author-email: info@wisent.ai
+Classifier: Programming Language :: Python :: 3
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: OS Independent
+Requires-Python: >=3.8
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: requests>=2.25.0
+Requires-Dist: pydantic>=2.0.0
+Requires-Dist: aiohttp>=3.8.0
+Requires-Dist: torch>=2.0.0
+Requires-Dist: numpy>=1.20.0
+Requires-Dist: tqdm>=4.60.0
+Requires-Dist: transformers>=4.30.0
+Dynamic: author
+Dynamic: author-email
+Dynamic: classifier
+Dynamic: description
+Dynamic: description-content-type
+Dynamic: home-page
+Dynamic: requires-dist
+Dynamic: requires-python
+Dynamic: summary
+
+# Wisent
+
+A Python client library for interacting with the Wisent backend services.
+
+## Installation
+
+```bash
+pip install wisent
+```
+
+## Features
+
+- **Activations**: Extract and send model activations to the Wisent backend
+- **Control Vectors**: Retrieve and apply control vectors for model inference
+- **Inference**: Utilities for applying control vectors during inference
+- **Utilities**: Helper functions for common tasks
+
+## Quick Start
+
+```python
+from wisent import WisentClient
+
+# Initialize the client
+client = WisentClient(api_key="your_api_key", base_url="https://api.wisent.ai")
+
+# Extract activations from a model and send to backend
+activations = client.activations.extract(
+    model_name="mistralai/Mistral-7B-Instruct-v0.1",
+    prompt="Tell me about quantum computing",
+    layers=[0, 12, 24]
+)
+
+# Get a control vector from the backend
+control_vector = client.control_vector.get(
+    name="helpful",
+    model="mistralai/Mistral-7B-Instruct-v0.1"
+)
+
+# Apply a control vector during inference
+response = client.inference.generate_with_control(
+    model_name="mistralai/Mistral-7B-Instruct-v0.1",
+    prompt="Tell me about quantum computing",
+    control_vectors={"helpful": 0.8, "concise": 0.5}
+)
+
+# Print the response
+print(response.text)
+```
+
+## Advanced Usage
+
+### Extracting Activations
+
+```python
+from wisent.activations import ActivationExtractor
+
+# Create an extractor
+extractor = ActivationExtractor(
+    model_name="mistralai/Mistral-7B-Instruct-v0.1",
+    device="cuda"
+)
+
+# Extract activations for a specific prompt
+activations = extractor.extract(
+    prompt="Tell me about quantum computing",
+    layers=[0, 12, 24],
+    tokens_to_extract=[-10, -1]  # Extract last 10 tokens and final token
+)
+
+# Send activations to the Wisent backend
+from wisent import WisentClient
+client = WisentClient(api_key="your_api_key")
+client.activations.upload(activations)
+```
+
+### Working with Control Vectors
+
+```python
+from wisent.control_vector import ControlVectorManager
+
+# Initialize the manager
+manager = ControlVectorManager(api_key="your_api_key")
+
+# Get a control vector
+helpful_vector = manager.get("helpful", model="mistralai/Mistral-7B-Instruct-v0.1")
+
+# Combine multiple vectors
+combined_vector = manager.combine(
+    vectors={
+        "helpful": 0.8,
+        "concise": 0.5
+    },
+    model="mistralai/Mistral-7B-Instruct-v0.1"
+)
+
+# Apply during inference
+from wisent.inference import Inferencer
+inferencer = Inferencer(model_name="mistralai/Mistral-7B-Instruct-v0.1")
+response = inferencer.generate(
+    prompt="Tell me about quantum computing",
+    control_vector=combined_vector,
+    method="caa"  # Context-Aware Addition
+)
+```
+
+## Documentation
+
+For full documentation, visit [docs.wisent.ai](https://docs.wisent.ai).
+
+## License
+
+MIT

wisent-0.1.1/README.md

@@ -0,0 +1,112 @@
+# Wisent
+
+A Python client library for interacting with the Wisent backend services.
+
+## Installation
+
+```bash
+pip install wisent
+```
+
+## Features
+
+- **Activations**: Extract and send model activations to the Wisent backend
+- **Control Vectors**: Retrieve and apply control vectors for model inference
+- **Inference**: Utilities for applying control vectors during inference
+- **Utilities**: Helper functions for common tasks
+
+## Quick Start
+
+```python
+from wisent import WisentClient
+
+# Initialize the client
+client = WisentClient(api_key="your_api_key", base_url="https://api.wisent.ai")
+
+# Extract activations from a model and send to backend
+activations = client.activations.extract(
+    model_name="mistralai/Mistral-7B-Instruct-v0.1",
+    prompt="Tell me about quantum computing",
+    layers=[0, 12, 24]
+)
+
+# Get a control vector from the backend
+control_vector = client.control_vector.get(
+    name="helpful",
+    model="mistralai/Mistral-7B-Instruct-v0.1"
+)
+
+# Apply a control vector during inference
+response = client.inference.generate_with_control(
+    model_name="mistralai/Mistral-7B-Instruct-v0.1",
+    prompt="Tell me about quantum computing",
+    control_vectors={"helpful": 0.8, "concise": 0.5}
+)
+
+# Print the response
+print(response.text)
+```
+
+## Advanced Usage
+
+### Extracting Activations
+
+```python
+from wisent.activations import ActivationExtractor
+
+# Create an extractor
+extractor = ActivationExtractor(
+    model_name="mistralai/Mistral-7B-Instruct-v0.1",
+    device="cuda"
+)
+
+# Extract activations for a specific prompt
+activations = extractor.extract(
+    prompt="Tell me about quantum computing",
+    layers=[0, 12, 24],
+    tokens_to_extract=[-10, -1]  # Extract last 10 tokens and final token
+)
+
+# Send activations to the Wisent backend
+from wisent import WisentClient
+client = WisentClient(api_key="your_api_key")
+client.activations.upload(activations)
+```
+
+### Working with Control Vectors
+
+```python
+from wisent.control_vector import ControlVectorManager
+
+# Initialize the manager
+manager = ControlVectorManager(api_key="your_api_key")
+
+# Get a control vector
+helpful_vector = manager.get("helpful", model="mistralai/Mistral-7B-Instruct-v0.1")
+
+# Combine multiple vectors
+combined_vector = manager.combine(
+    vectors={
+        "helpful": 0.8,
+        "concise": 0.5
+    },
+    model="mistralai/Mistral-7B-Instruct-v0.1"
+)
+
+# Apply during inference
+from wisent.inference import Inferencer
+inferencer = Inferencer(model_name="mistralai/Mistral-7B-Instruct-v0.1")
+response = inferencer.generate(
+    prompt="Tell me about quantum computing",
+    control_vector=combined_vector,
+    method="caa"  # Context-Aware Addition
+)
+```
+
+## Documentation
+
+For full documentation, visit [docs.wisent.ai](https://docs.wisent.ai).
+
+## License
+
+MIT

wisent-0.1.1/pyproject.toml

@@ -0,0 +1,22 @@
+[build-system]
+requires = ["setuptools>=42", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[tool.black]
+line-length = 88
+target-version = ['py38']
+
+[tool.isort]
+profile = "black"
+multi_line_output = 3
+
+[tool.mypy]
+python_version = "3.8"
+warn_return_any = true
+warn_unused_configs = true
+disallow_untyped_defs = true
+disallow_incomplete_defs = true
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+python_files = "test_*.py"

wisent-0.1.1/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

wisent-0.1.1/setup.py

@@ -0,0 +1,31 @@
+from setuptools import setup, find_packages
+
+with open("README.md", "r", encoding="utf-8") as fh:
+    long_description = fh.read()
+
+setup(
+    name="wisent",
+    version="0.1.1",
+    author="Wisent Team",
+    author_email="info@wisent.ai",
+    description="Client library for interacting with the Wisent backend services",
+    long_description=long_description,
+    long_description_content_type="text/markdown",
+    url="https://github.com/wisent-ai/wisent",
+    packages=find_packages(),
+    classifiers=[
+        "Programming Language :: Python :: 3",
+        "License :: OSI Approved :: MIT License",
+        "Operating System :: OS Independent",
+    ],
+    python_requires=">=3.8",
+    install_requires=[
+        "requests>=2.25.0",
+        "pydantic>=2.0.0",
+        "aiohttp>=3.8.0",
+        "torch>=2.0.0",
+        "numpy>=1.20.0",
+        "tqdm>=4.60.0",
+        "transformers>=4.30.0",
+    ],
+)

wisent-0.1.1/tests/test_activations.py

@@ -0,0 +1,350 @@
+"""
+Tests for the activations module.
+"""
+
+import unittest
+from unittest.mock import MagicMock, patch
+
+import numpy as np
+import torch
+
+from wisent.activations import Activation, ActivationBatch, ActivationExtractor, ActivationsClient
+from wisent.utils.auth import AuthManager
+
+
+class TestActivation(unittest.TestCase):
+    """Tests for the Activation class."""
+    
+    def test_init(self):
+        """Test initialization of an Activation."""
+        # Test with list values
+        values = [0.1, 0.2, 0.3]
+        activation = Activation(
+            model_name="test_model",
+            layer=0,
+            token_index=1,
+            values=values,
+            token_str="test"
+        )
+        
+        self.assertEqual(activation.model_name, "test_model")
+        self.assertEqual(activation.layer, 0)
+        self.assertEqual(activation.token_index, 1)
+        self.assertEqual(activation.values, values)
+        self.assertEqual(activation.token_str, "test")
+        
+        # Test with numpy array
+        values = np.array([0.1, 0.2, 0.3])
+        activation = Activation(
+            model_name="test_model",
+            layer=0,
+            token_index=1,
+            values=values
+        )
+        
+        self.assertTrue(np.array_equal(activation.values, values))
+        self.assertIsNone(activation.token_str)
+        
+        # Test with torch tensor
+        values = torch.tensor([0.1, 0.2, 0.3])
+        activation = Activation(
+            model_name="test_model",
+            layer=0,
+            token_index=1,
+            values=values
+        )
+        
+        self.assertTrue(torch.equal(activation.values, values))
+    
+    def test_to_dict(self):
+        """Test conversion to dictionary."""
+        # Test with list values
+        values = [0.1, 0.2, 0.3]
+        activation = Activation(
+            model_name="test_model",
+            layer=0,
+            token_index=1,
+            values=values,
+            token_str="test"
+        )
+        
+        expected = {
+            "model_name": "test_model",
+            "layer": 0,
+            "token_index": 1,
+            "values": values,
+            "token_str": "test",
+        }
+        
+        self.assertEqual(activation.to_dict(), expected)
+        
+        # Test with numpy array
+        values = np.array([0.1, 0.2, 0.3])
+        activation = Activation(
+            model_name="test_model",
+            layer=0,
+            token_index=1,
+            values=values
+        )
+        
+        expected = {
+            "model_name": "test_model",
+            "layer": 0,
+            "token_index": 1,
+            "values": values.tolist(),
+            "token_str": None,
+        }
+        
+        self.assertEqual(activation.to_dict(), expected)
+        
+        # Test with torch tensor
+        values = torch.tensor([0.1, 0.2, 0.3])
+        activation = Activation(
+            model_name="test_model",
+            layer=0,
+            token_index=1,
+            values=values
+        )
+        
+        expected = {
+            "model_name": "test_model",
+            "layer": 0,
+            "token_index": 1,
+            "values": values.detach().cpu().numpy().tolist(),
+            "token_str": None,
+        }
+        
+        self.assertEqual(activation.to_dict(), expected)
+
+
+class TestActivationBatch(unittest.TestCase):
+    """Tests for the ActivationBatch class."""
+    
+    def setUp(self):
+        """Set up test fixtures."""
+        self.activation1 = Activation(
+            model_name="test_model",
+            layer=0,
+            token_index=1,
+            values=[0.1, 0.2, 0.3],
+            token_str="test1"
+        )
+        
+        self.activation2 = Activation(
+            model_name="test_model",
+            layer=1,
+            token_index=2,
+            values=[0.4, 0.5, 0.6],
+            token_str="test2"
+        )
+    
+    def test_init(self):
+        """Test initialization of an ActivationBatch."""
+        batch = ActivationBatch(
+            model_name="test_model",
+            prompt="test prompt",
+            activations=[self.activation1, self.activation2],
+            metadata={"test": "metadata"}
+        )
+        
+        self.assertEqual(batch.model_name, "test_model")
+        self.assertEqual(batch.prompt, "test prompt")
+        self.assertEqual(batch.activations, [self.activation1, self.activation2])
+        self.assertEqual(batch.metadata, {"test": "metadata"})
+        
+        # Test with default metadata
+        batch = ActivationBatch(
+            model_name="test_model",
+            prompt="test prompt",
+            activations=[self.activation1, self.activation2]
+        )
+        
+        self.assertEqual(batch.metadata, {})
+    
+    def test_to_dict(self):
+        """Test conversion to dictionary."""
+        batch = ActivationBatch(
+            model_name="test_model",
+            prompt="test prompt",
+            activations=[self.activation1, self.activation2],
+            metadata={"test": "metadata"}
+        )
+        
+        expected = {
+            "model_name": "test_model",
+            "prompt": "test prompt",
+            "activations": [
+                self.activation1.to_dict(),
+                self.activation2.to_dict()
+            ],
+            "metadata": {"test": "metadata"},
+        }
+        
+        self.assertEqual(batch.to_dict(), expected)
+
+
+class TestActivationsClient(unittest.TestCase):
+    """Tests for the ActivationsClient class."""
+    
+    def setUp(self):
+        """Set up test fixtures."""
+        self.auth_manager = AuthManager("test_api_key")
+        self.base_url = "https://test.api.wisent.ai"
+        self.client = ActivationsClient(self.auth_manager, self.base_url)
+    
+    @patch("wisent.activations.client.ActivationExtractor")
+    def test_extract(self, mock_extractor_class):
+        """Test extraction of activations."""
+        # Mock the extractor
+        mock_extractor = MagicMock()
+        mock_extractor_class.return_value = mock_extractor
+        
+        # Mock the extract method
+        mock_batch = MagicMock()
+        mock_extractor.extract.return_value = mock_batch
+        
+        # Call the method
+        result = self.client.extract(
+            model_name="test_model",
+            prompt="test prompt",
+            layers=[0, 1],
+            tokens_to_extract=[1, 2],
+            device="cpu"
+        )
+        
+        # Check that the extractor was created correctly
+        mock_extractor_class.assert_called_once_with("test_model", device="cpu")
+        
+        # Check that extract was called correctly
+        mock_extractor.extract.assert_called_once_with("test prompt", [0, 1], [1, 2])
+        
+        # Check the result
+        self.assertEqual(result, mock_batch)
+    
+    @patch("wisent.activations.client.HTTPClient")
+    def test_upload(self, mock_http_client_class):
+        """Test uploading of activations."""
+        # Mock the HTTP client
+        mock_http_client = MagicMock()
+        mock_http_client_class.return_value = mock_http_client
+        
+        # Mock the post method
+        mock_response = {"id": "test_batch_id"}
+        mock_http_client.post.return_value = mock_response
+        
+        # Create a client with the mocked HTTP client
+        client = ActivationsClient(self.auth_manager, self.base_url)
+        
+        # Create a batch to upload
+        batch = ActivationBatch(
+            model_name="test_model",
+            prompt="test prompt",
+            activations=[
+                Activation(
+                    model_name="test_model",
+                    layer=0,
+                    token_index=1,
+                    values=[0.1, 0.2, 0.3],
+                    token_str="test"
+                )
+            ]
+        )
+        
+        # Call the method
+        result = client.upload(batch)
+        
+        # Check that post was called correctly
+        mock_http_client.post.assert_called_once_with(
+            "/activations/upload",
+            json_data=batch.to_dict()
+        )
+        
+        # Check the result
+        self.assertEqual(result, mock_response)
+    
+    @patch("wisent.activations.client.HTTPClient")
+    def test_get(self, mock_http_client_class):
+        """Test getting activations by ID."""
+        # Mock the HTTP client
+        mock_http_client = MagicMock()
+        mock_http_client_class.return_value = mock_http_client
+        
+        # Mock the get method
+        mock_response = {
+            "model_name": "test_model",
+            "prompt": "test prompt",
+            "activations": [
+                {
+                    "model_name": "test_model",
+                    "layer": 0,
+                    "token_index": 1,
+                    "values": [0.1, 0.2, 0.3],
+                    "token_str": "test"
+                }
+            ],
+            "metadata": {"test": "metadata"}
+        }
+        mock_http_client.get.return_value = mock_response
+        
+        # Create a client with the mocked HTTP client
+        client = ActivationsClient(self.auth_manager, self.base_url)
+        
+        # Call the method
+        result = client.get("test_batch_id")
+        
+        # Check that get was called correctly
+        mock_http_client.get.assert_called_once_with("/activations/test_batch_id")
+        
+        # Check the result
+        self.assertEqual(result.model_name, "test_model")
+        self.assertEqual(result.prompt, "test prompt")
+        self.assertEqual(len(result.activations), 1)
+        self.assertEqual(result.activations[0].model_name, "test_model")
+        self.assertEqual(result.activations[0].layer, 0)
+        self.assertEqual(result.activations[0].token_index, 1)
+        self.assertEqual(result.activations[0].values, [0.1, 0.2, 0.3])
+        self.assertEqual(result.activations[0].token_str, "test")
+        self.assertEqual(result.metadata, {"test": "metadata"})
+    
+    @patch("wisent.activations.client.HTTPClient")
+    def test_list(self, mock_http_client_class):
+        """Test listing activations."""
+        # Mock the HTTP client
+        mock_http_client = MagicMock()
+        mock_http_client_class.return_value = mock_http_client
+        
+        # Mock the get method
+        mock_response = [
+            {"id": "batch1", "model_name": "test_model"},
+            {"id": "batch2", "model_name": "test_model"}
+        ]
+        mock_http_client.get.return_value = mock_response
+        
+        # Create a client with the mocked HTTP client
+        client = ActivationsClient(self.auth_manager, self.base_url)
+        
+        # Call the method
+        result = client.list(model_name="test_model", limit=10, offset=0)
+        
+        # Check that get was called correctly
+        mock_http_client.get.assert_called_once_with(
+            "/activations",
+            params={"model_name": "test_model", "limit": 10, "offset": 0}
+        )
+        
+        # Check the result
+        self.assertEqual(result, mock_response)
+        
+        # Test without model_name
+        mock_http_client.get.reset_mock()
+        result = client.list(limit=10, offset=0)
+        
+        # Check that get was called correctly
+        mock_http_client.get.assert_called_once_with(
+            "/activations",
+            params={"limit": 10, "offset": 0}
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()