interp-engine 0.0.24__py3-none-any.whl

neuron_explainer/activations/token_connections.py

@@ -0,0 +1,59 @@
+from dataclasses import dataclass
+from typing import List, Union
+
+import blobfile as bf
+from neuron_explainer.fast_dataclasses import FastDataclass, loads, register_dataclass
+from neuron_explainer.azure import standardize_azure_url
+import urllib.request
+
+
+@register_dataclass
+@dataclass
+class TokensAndWeights(FastDataclass):
+    tokens: List[str]
+    strengths: List[float]
+
+
+@register_dataclass
+@dataclass
+class WeightBasedSummaryOfNeuron(FastDataclass):
+    input_positive: TokensAndWeights
+    input_negative: TokensAndWeights
+    output_positive: TokensAndWeights
+    output_negative: TokensAndWeights
+
+
+def load_token_weight_connections_of_neuron(
+    layer_index: Union[str, int],
+    neuron_index: Union[str, int],
+    dataset_path: str = "https://openaipublic.blob.core.windows.net/neuron-explainer/data/related-tokens/weight-based",
+) -> WeightBasedSummaryOfNeuron:
+    """Load the TokenLookupTableSummaryOfNeuron for the specified neuron."""
+    url = "/".join([dataset_path, str(layer_index), f"{neuron_index}.json"])
+    url = standardize_azure_url(url)
+    with urllib.request.urlopen(url) as f:
+        return loads(f.read(), backwards_compatible=False)
+
+
+@register_dataclass
+@dataclass
+class TokenLookupTableSummaryOfNeuron(FastDataclass):
+    """List of tokens and the average activations of a given neuron in response to each
+    respective token. These are selected from among the tokens in the vocabulary with the
+    highest average activations across an internet text dataset, with the highest activations
+    first."""
+
+    tokens: List[str]
+    average_activations: List[float]
+
+
+def load_token_lookup_table_connections_of_neuron(
+    layer_index: Union[str, int],
+    neuron_index: Union[str, int],
+    dataset_path: str = "https://openaipublic.blob.core.windows.net/neuron-explainer/data/related-tokens/activation-based",
+) -> TokenLookupTableSummaryOfNeuron:
+    """Load the TokenLookupTableSummaryOfNeuron for the specified neuron."""
+    url = "/".join([dataset_path, str(layer_index), f"{neuron_index}.json"])
+    url = standardize_azure_url(url)
+    with urllib.request.urlopen(url) as f:
+        return loads(f.read(), backwards_compatible=False)

neuron_explainer/api_client.py

@@ -0,0 +1,190 @@
+import asyncio
+import contextlib
+import os
+import random
+import traceback
+from asyncio import Semaphore
+from functools import wraps
+from typing import Any, Callable, Optional
+
+import httpx
+import orjson
+
+
+def is_api_error(err: Exception) -> bool:
+    if isinstance(err, httpx.HTTPStatusError):
+        response = err.response
+        error_data = response.json().get("error", {})
+        error_message = error_data.get("message")
+        if response.status_code in [400, 404, 415]:
+            if error_data.get("type") == "idempotency_error":
+                print(
+                    f"Retrying after idempotency error: {error_message} ({response.url})"
+                )
+                return True
+            else:
+                # Invalid request
+                return False
+        else:
+            print(f"Retrying after API error: {error_message} ({response.url})")
+            return True
+
+    elif isinstance(err, httpx.ConnectError):
+        print(f"Retrying after connection error... ({err.request.url})")
+        return True
+
+    elif isinstance(err, httpx.TimeoutException):
+        print(f"Retrying after a timeout error... ({err.request.url})")
+        return True
+
+    elif isinstance(err, httpx.ReadError):
+        print(f"Retrying after a read error... ({err.request.url})")
+        return True
+
+    print(f"Retrying after an unexpected error: {repr(err)}")
+    traceback.print_tb(err.__traceback__)
+    return True
+
+
+def exponential_backoff(
+    retry_on: Callable[[Exception], bool] = lambda err: True,
+) -> Callable[[Callable], Callable]:
+    """
+    Returns a decorator which retries the wrapped function as long as the specified retry_on
+    function returns True for the exception, applying exponential backoff with jitter after
+    failures, up to a retry limit.
+    """
+    init_delay_s = 1.0
+    max_delay_s = 10.0
+    # Roughly 30 minutes before we give up.
+    max_tries = 200
+    backoff_multiplier = 2.0
+    jitter = 0.2
+
+    def decorate(f: Callable) -> Callable:
+        assert asyncio.iscoroutinefunction(f)
+
+        @wraps(f)
+        async def f_retry(*args: Any, **kwargs: Any) -> None:
+            delay_s = init_delay_s
+            for i in range(max_tries):
+                try:
+                    return await f(*args, **kwargs)
+                except Exception as err:
+                    if not retry_on(err) or i == max_tries - 1:
+                        raise
+                    jittered_delay = random.uniform(
+                        delay_s * (1 - jitter), delay_s * (1 + jitter)
+                    )
+                    await asyncio.sleep(jittered_delay)
+                    delay_s = min(delay_s * backoff_multiplier, max_delay_s)
+
+        return f_retry
+
+    return decorate
+
+
+class ApiClient:
+    """Performs inference using the OpenAI API. Supports response caching and concurrency limits."""
+
+    BASE_API_URL = "https://api.openai.com/v1"
+
+    def __init__(
+        self,
+        model_name: str,
+        # If set, no more than this number of HTTP requests will be made concurrently.
+        max_concurrent: Optional[int] = None,
+        # Whether to cache request/response pairs in memory to avoid duplicating requests.
+        cache: bool = False,
+        base_api_url: str = BASE_API_URL,
+        override_api_key: str | None = None,
+    ):
+        self.model_name = model_name
+        self.base_api_url = base_api_url
+        self.override_api_key = override_api_key
+        if max_concurrent is not None:
+            self._concurrency_check: Optional[Semaphore] = Semaphore(max_concurrent)
+        else:
+            self._concurrency_check = None
+
+        if cache:
+            self._cache: Optional[dict[str, Any]] = {}
+        else:
+            self._cache = None
+
+    @exponential_backoff(retry_on=is_api_error)
+    async def make_request(
+        self,
+        timeout_seconds: Optional[int] = None,
+        json_mode: Optional[bool] = False,
+        **kwargs: Any,
+    ) -> dict[str, Any]:
+        api_http_headers = {
+            "Content-Type": "application/json",
+            "Authorization": f"Bearer {os.getenv('OPENAI_API_KEY') if self.override_api_key is None else self.override_api_key}",
+        }
+        if self._cache is not None:
+            key = orjson.dumps(kwargs)
+            if key in self._cache:
+                return self._cache[key]
+        async with contextlib.AsyncExitStack() as stack:
+            if self._concurrency_check is not None:
+                await stack.enter_async_context(self._concurrency_check)
+            http_client = await stack.enter_async_context(
+                httpx.AsyncClient(timeout=timeout_seconds)
+            )
+            # If the request has a "messages" key, it should be sent to the /chat/completions
+            # endpoint. Otherwise, it should be sent to the /completions endpoint.
+            url = self.base_api_url + (
+                "/chat/completions" if "messages" in kwargs else "/completions"
+            )
+            kwargs["model"] = self.model_name
+            if json_mode:
+                kwargs["response_format"] = {"type": "json_object"}
+            # Convert max_tokens to max_completion_tokens for newer OpenAI models
+            # that don't support the legacy max_tokens parameter
+            if "max_tokens" in kwargs and "messages" in kwargs:
+                kwargs["max_completion_tokens"] = kwargs.pop("max_tokens")
+            response = await http_client.post(
+                url, headers=api_http_headers, json=kwargs
+            )
+            # # Print token usage information if available in the response
+            # try:
+            #     response_data = response.json()
+            #     if "usage" in response_data:
+            #         usage = response_data["usage"]
+            #         print(f"Token usage - Prompt: {usage.get('prompt_tokens', 'N/A')}, "
+            #             f"Completion: {usage.get('completion_tokens', 'N/A')}, "
+            #             f"Total: {usage.get('total_tokens', 'N/A')}")
+            # except Exception:
+            #     pass  # Silently ignore if we can't parse the response or find usage info
+        # The response json has useful information but the exception doesn't include it, so print it
+        # out then reraise.
+        try:
+            response.raise_for_status()
+        except Exception as e:
+            try:
+                print(f"Error response status code: {response.status_code}")
+                print(f"Error response JSON: {response.json()}")
+            except Exception:
+                print("Could not parse error response as JSON")
+                print(f"Error response text: {response.text}")
+            raise e
+        if self._cache is not None:
+            self._cache[key] = response.json()
+        response_json = response.json()
+        # print(f"response_json: {response_json}")
+        return response_json
+
+
+if __name__ == "__main__":
+
+    async def main() -> None:
+        client = ApiClient(model_name="gpt-3.5-turbo", max_concurrent=1)
+        print(
+            await client.make_request(
+                prompt="Why did the chicken cross the road?", max_tokens=9
+            )
+        )
+
+    asyncio.run(main())

neuron_explainer/azure.py

@@ -0,0 +1,5 @@
+def standardize_azure_url(url):
+    """Make sure url is converted to url format, not an azure path"""
+    if url.startswith("az://openaipublic/"):
+        url = url.replace("az://openaipublic/", "https://openaipublic.blob.core.windows.net/")
+    return url

neuron_explainer/explanations/calibrated_simulator.py

@@ -0,0 +1,194 @@
+"""
+Code for calibrating simulations of neuron behavior. Calibration refers to a process of mapping from
+a space of predicted activation values (e.g. [0, 10]) to the real activation distribution for a
+neuron.
+
+See http://go/neuron_explanation_methodology for description of calibration step. Necessary for
+simulating neurons in the context of ablate-to-simulation, but can be skipped when using correlation
+scoring. (Calibration may still improve quality for scoring, at least for non-linear calibration
+methods.)
+"""
+
+from __future__ import annotations
+
+import asyncio
+from abc import abstractmethod
+from typing import Optional, Sequence
+
+import numpy as np
+from neuron_explainer.activations.activations import ActivationRecord
+from neuron_explainer.explanations.explanations import ActivationScale
+from neuron_explainer.explanations.simulator import NeuronSimulator, SequenceSimulation
+from sklearn import linear_model
+
+
+class CalibratedNeuronSimulator(NeuronSimulator):
+    """
+    Wrap a NeuronSimulator and calibrate it to map from the predicted activation space to the
+    actual neuron activation space.
+    """
+
+    def __init__(self, uncalibrated_simulator: NeuronSimulator):
+        self.uncalibrated_simulator = uncalibrated_simulator
+
+    @classmethod
+    async def create(
+        cls,
+        uncalibrated_simulator: NeuronSimulator,
+        calibration_activation_records: Sequence[ActivationRecord],
+    ) -> CalibratedNeuronSimulator:
+        """
+        Create and calibrate a calibrated simulator (so initialization and calibration can be done
+        in one call).
+        """
+        calibrated_simulator = cls(uncalibrated_simulator)
+        await calibrated_simulator.calibrate(calibration_activation_records)
+        return calibrated_simulator
+
+    async def calibrate(self, calibration_activation_records: Sequence[ActivationRecord]) -> None:
+        """
+        Determine parameters to map from the predicted activation space to the real neuron
+        activation space, based on a calibration set.
+
+        Use when simulated sequences haven't already been produced on the calibration set.
+        """
+        simulations = await asyncio.gather(
+            *[
+                self.uncalibrated_simulator.simulate(activations.tokens)
+                for activations in calibration_activation_records
+            ]
+        )
+        self.calibrate_from_simulations(calibration_activation_records, simulations)
+
+    def calibrate_from_simulations(
+        self,
+        calibration_activation_records: Sequence[ActivationRecord],
+        simulations: Sequence[SequenceSimulation],
+    ) -> None:
+        """
+        Determine parameters to map from the predicted activation space to the real neuron
+        activation space, based on a calibration set.
+
+        Use when simulated sequences have already been produced on the calibration set.
+        """
+        flattened_activations = []
+        flattened_simulated_activations: list[float] = []
+        for activations, simulation in zip(calibration_activation_records, simulations):
+            flattened_activations.extend(activations.activations)
+            flattened_simulated_activations.extend(simulation.expected_activations)
+        self._calibrate_from_flattened_activations(
+            np.array(flattened_activations), np.array(flattened_simulated_activations)
+        )
+
+    @abstractmethod
+    def _calibrate_from_flattened_activations(
+        self,
+        true_activations: np.ndarray,
+        uncalibrated_activations: np.ndarray,
+    ) -> None:
+        """
+        Determine parameters to map from the predicted activation space to the real neuron
+        activation space, based on a calibration set.
+
+        Take numpy arrays of all true activations and all uncalibrated activations on the
+        calibration set over all sequences.
+        """
+
+    @abstractmethod
+    def apply_calibration(self, values: Sequence[float]) -> list[float]:
+        """Apply the learned calibration to a sequence of values."""
+
+    async def simulate(self, tokens: Sequence[str]) -> SequenceSimulation:
+        uncalibrated_seq_simulation = await self.uncalibrated_simulator.simulate(tokens)
+        calibrated_activations = self.apply_calibration(
+            uncalibrated_seq_simulation.expected_activations
+        )
+        calibrated_distribution_values = [
+            self.apply_calibration(dv) for dv in uncalibrated_seq_simulation.distribution_values
+        ]
+        return SequenceSimulation(
+            tokens=uncalibrated_seq_simulation.tokens,
+            expected_activations=calibrated_activations,
+            activation_scale=ActivationScale.NEURON_ACTIVATIONS,
+            distribution_values=calibrated_distribution_values,
+            distribution_probabilities=uncalibrated_seq_simulation.distribution_probabilities,
+            uncalibrated_simulation=uncalibrated_seq_simulation,
+        )
+
+
+class UncalibratedNeuronSimulator(CalibratedNeuronSimulator):
+    """Pass through the activations without trying to calibrate."""
+
+    def __init__(self, uncalibrated_simulator: NeuronSimulator):
+        super().__init__(uncalibrated_simulator)
+
+    async def calibrate(self, calibration_activation_records: Sequence[ActivationRecord]) -> None:
+        pass
+
+    def _calibrate_from_flattened_activations(
+        self,
+        true_activations: np.ndarray,
+        uncalibrated_activations: np.ndarray,
+    ) -> None:
+        pass
+
+    def apply_calibration(self, values: Sequence[float]) -> list[float]:
+        return values if isinstance(values, list) else list(values)
+
+
+class LinearCalibratedNeuronSimulator(CalibratedNeuronSimulator):
+    """Find a linear mapping from uncalibrated activations to true activations.
+
+    Should not change ev_correlation_score because it is invariant to linear transformations.
+    """
+
+    def __init__(self, uncalibrated_simulator: NeuronSimulator):
+        super().__init__(uncalibrated_simulator)
+        self._regression: Optional[linear_model.LinearRegression] = None
+
+    def _calibrate_from_flattened_activations(
+        self,
+        true_activations: np.ndarray,
+        uncalibrated_activations: np.ndarray,
+    ) -> None:
+        self._regression = linear_model.LinearRegression()
+        self._regression.fit(uncalibrated_activations.reshape(-1, 1), true_activations)
+
+    def apply_calibration(self, values: Sequence[float]) -> list[float]:
+        if self._regression is None:
+            raise ValueError("Must call calibrate() before apply_calibration")
+        if len(values) == 0:
+            return []
+        return self._regression.predict(np.reshape(np.array(values), (-1, 1))).tolist()
+
+
+class PercentileMatchingCalibratedNeuronSimulator(CalibratedNeuronSimulator):
+    """
+    Map the nth percentile of the uncalibrated activations to the nth percentile of the true
+    activations for all n.
+
+    This will match the distribution of true activations on the calibration set, but will be
+    overconfident outside of the calibration set.
+    """
+
+    def __init__(self, uncalibrated_simulator: NeuronSimulator):
+        super().__init__(uncalibrated_simulator)
+        self._uncalibrated_activations: Optional[np.ndarray] = None
+        self._true_activations: Optional[np.ndarray] = None
+
+    def _calibrate_from_flattened_activations(
+        self,
+        true_activations: np.ndarray,
+        uncalibrated_activations: np.ndarray,
+    ) -> None:
+        self._uncalibrated_activations = np.sort(uncalibrated_activations)
+        self._true_activations = np.sort(true_activations)
+
+    def apply_calibration(self, values: Sequence[float]) -> list[float]:
+        if self._true_activations is None or self._uncalibrated_activations is None:
+            raise ValueError("Must call calibrate() before apply_calibration")
+        if len(values) == 0:
+            return []
+        return np.interp(
+            np.array(values), self._uncalibrated_activations, self._true_activations
+        ).tolist()