interp-engine 0.0.24__py3-none-any.whl

neuron_explainer/explanations/test_explainer.py

@@ -0,0 +1,227 @@
+import asyncio
+from typing import Any
+
+from neuron_explainer.explanations.explainer import (
+    TokenActivationPairExplainer,
+    TokenSpaceRepresentationExplainer,
+)
+from neuron_explainer.explanations.few_shot_examples import TEST_EXAMPLES, FewShotExampleSet
+from neuron_explainer.explanations.prompt_builder import HarmonyMessage, PromptFormat, Role
+from neuron_explainer.explanations.token_space_few_shot_examples import (
+    TokenSpaceFewShotExampleSet,
+)
+
+
+def setup_module(unused_module: Any) -> None:
+    # Make sure we have an event loop, since the attempt to create the Semaphore in
+    # ResearchApiClient will fail without it.
+    loop = asyncio.new_event_loop()
+    asyncio.set_event_loop(loop)
+
+
+def test_if_formatting() -> None:
+    expected_prompt = """We're studying neurons in a neural network. Each neuron looks for some particular thing in a short document. Look at the parts of the document the neuron activates for and summarize in a single sentence what the neuron is looking for. Don't list examples of words.
+
+The activation format is token<tab>activation. Activation values range from 0 to 10. A neuron finding what it's looking for is represented by a non-zero activation value. The higher the activation value, the stronger the match.
+
+Neuron 1
+Activations:
+<start>
+a	10
+b	0
+c	0
+<end>
+<start>
+d	0
+e	10
+f	0
+<end>
+
+Explanation of neuron 1 behavior: the main thing this neuron does is find vowels.
+
+Neuron 2
+Activations:
+<start>
+a	10
+b	0
+c	0
+<end>
+<start>
+d	0
+e	10
+f	0
+<end>
+
+Explanation of neuron 2 behavior:<|endofprompt|> the main thing this neuron does is find"""
+
+    explainer = TokenActivationPairExplainer(
+        model_name="text-davinci-003",
+        prompt_format=PromptFormat.INSTRUCTION_FOLLOWING,
+        few_shot_example_set=FewShotExampleSet.TEST,
+    )
+    prompt = explainer.make_explanation_prompt(
+        all_activation_records=TEST_EXAMPLES[0].activation_records,
+        max_activation=1.0,
+        max_tokens_for_completion=20,
+    )
+
+    assert prompt == expected_prompt
+
+
+def test_harmony_format() -> None:
+    expected_prompt = [
+        HarmonyMessage(
+            role=Role.SYSTEM,
+            content="""We're studying neurons in a neural network. Each neuron looks for some particular thing in a short document. Look at the parts of the document the neuron activates for and summarize in a single sentence what the neuron is looking for. Don't list examples of words.
+
+The activation format is token<tab>activation. Activation values range from 0 to 10. A neuron finding what it's looking for is represented by a non-zero activation value. The higher the activation value, the stronger the match.""",
+        ),
+        HarmonyMessage(
+            role=Role.USER,
+            content="""
+
+Neuron 1
+Activations:
+<start>
+a	10
+b	0
+c	0
+<end>
+<start>
+d	0
+e	10
+f	0
+<end>
+
+Explanation of neuron 1 behavior: the main thing this neuron does is find""",
+        ),
+        HarmonyMessage(
+            role=Role.ASSISTANT,
+            content=" vowels.",
+        ),
+        HarmonyMessage(
+            role=Role.USER,
+            content="""
+
+Neuron 2
+Activations:
+<start>
+a	10
+b	0
+c	0
+<end>
+<start>
+d	0
+e	10
+f	0
+<end>
+
+Explanation of neuron 2 behavior: the main thing this neuron does is find""",
+        ),
+    ]
+
+    explainer = TokenActivationPairExplainer(
+        model_name="gpt-4",
+        prompt_format=PromptFormat.HARMONY_V4,
+        few_shot_example_set=FewShotExampleSet.TEST,
+    )
+    prompt = explainer.make_explanation_prompt(
+        all_activation_records=TEST_EXAMPLES[0].activation_records,
+        max_activation=1.0,
+        max_tokens_for_completion=20,
+    )
+
+    assert isinstance(prompt, list)
+    assert isinstance(prompt[0], dict)  # Really a HarmonyMessage
+    for actual_message, expected_message in zip(prompt, expected_prompt):
+        assert actual_message["role"] == expected_message["role"]
+        assert actual_message["content"] == expected_message["content"]
+    assert prompt == expected_prompt
+
+
+def test_token_space_explainer_if_formatting() -> None:
+    expected_prompt = """We're studying neurons in a neural network. Each neuron looks for some particular kind of token (which can be a word, or part of a word). Look at the tokens the neuron activates for (listed below) and summarize in a single sentence what the neuron is looking for. Don't list examples of words.
+
+
+
+Tokens:
+'these', ' are', ' tokens'
+
+Explanation:
+This neuron is looking for this is a test explanation.
+
+
+
+Tokens:
+'foo', 'bar', 'baz'
+
+Explanation:
+<|endofprompt|>This neuron is looking for"""
+
+    explainer = TokenSpaceRepresentationExplainer(
+        model_name="text-davinci-002",
+        prompt_format=PromptFormat.INSTRUCTION_FOLLOWING,
+        use_few_shot=True,
+        few_shot_example_set=TokenSpaceFewShotExampleSet.TEST,
+    )
+    prompt = explainer.make_explanation_prompt(
+        tokens=["foo", "bar", "baz"],
+        max_tokens_for_completion=20,
+    )
+
+    assert prompt == expected_prompt
+
+
+def test_token_space_explainer_harmony_formatting() -> None:
+    expected_prompt = [
+        HarmonyMessage(
+            role=Role.SYSTEM,
+            content="We're studying neurons in a neural network. Each neuron looks for some particular kind of token (which can be a word, or part of a word). Look at the tokens the neuron activates for (listed below) and summarize in a single sentence what the neuron is looking for. Don't list examples of words.",
+        ),
+        HarmonyMessage(
+            role=Role.USER,
+            content="""
+
+
+
+Tokens:
+'these', ' are', ' tokens'
+
+Explanation:
+This neuron is looking for""",
+        ),
+        HarmonyMessage(
+            role=Role.ASSISTANT,
+            content=" this is a test explanation.",
+        ),
+        HarmonyMessage(
+            role=Role.USER,
+            content="""
+
+
+
+Tokens:
+'foo', 'bar', 'baz'
+
+Explanation:
+This neuron is looking for""",
+        ),
+    ]
+
+    explainer = TokenSpaceRepresentationExplainer(
+        model_name="gpt-4",
+        prompt_format=PromptFormat.HARMONY_V4,
+        use_few_shot=True,
+        few_shot_example_set=TokenSpaceFewShotExampleSet.TEST,
+    )
+    prompt = explainer.make_explanation_prompt(
+        tokens=["foo", "bar", "baz"],
+        max_tokens_for_completion=20,
+    )
+
+    assert isinstance(prompt, list)
+    assert isinstance(prompt[0], dict)  # Really a HarmonyMessage
+    for actual_message, expected_message in zip(prompt, expected_prompt):
+        assert actual_message["role"] == expected_message["role"]
+        assert actual_message["content"] == expected_message["content"]
+    assert prompt == expected_prompt

neuron_explainer/explanations/test_simulator.py

@@ -0,0 +1,269 @@
+from neuron_explainer.explanations.few_shot_examples import FewShotExampleSet
+from neuron_explainer.explanations.prompt_builder import HarmonyMessage, PromptFormat, Role
+from neuron_explainer.explanations.simulator import (
+    ExplanationNeuronSimulator,
+    ExplanationTokenByTokenSimulator,
+)
+
+
+def test_make_explanation_simulation_prompt_if_format() -> None:
+    expected_prompt = """We're studying neurons in a neural network.
+Each neuron looks for some particular thing in a short document.
+Look at summary of what the neuron does, and try to predict how it will fire on each token.
+
+The activation format is token<tab>activation, activations go from 0 to 10, "unknown" indicates an unknown activation. Most activations will be 0.
+
+
+Neuron 1
+Explanation of neuron 1 behavior: the main thing this neuron does is find vowels
+Activations: 
+<start>
+a	10
+b	0
+c	0
+<end>
+<start>
+d	unknown
+e	10
+f	0
+<end>
+
+
+
+Neuron 2
+Explanation of neuron 2 behavior: the main thing this neuron does is find EXPLANATION<|endofprompt|>
+Activations: 
+<start>
+0	unknown
+1	unknown
+2	unknown
+<end>
+"""
+    prompt = ExplanationNeuronSimulator(
+        model_name="text-davinci-003",
+        explanation="EXPLANATION",
+        few_shot_example_set=FewShotExampleSet.TEST,
+        prompt_format=PromptFormat.INSTRUCTION_FOLLOWING,
+    ).make_simulation_prompt(
+        tokens=[str(x) for x in range(3)],
+    )
+    assert prompt == expected_prompt
+
+
+def test_make_explanation_simulation_prompt_harmony_format() -> None:
+    expected_prompt = [
+        HarmonyMessage(
+            role=Role.SYSTEM,
+            content="""We're studying neurons in a neural network.
+Each neuron looks for some particular thing in a short document.
+Look at summary of what the neuron does, and try to predict how it will fire on each token.
+
+The activation format is token<tab>activation, activations go from 0 to 10, "unknown" indicates an unknown activation. Most activations will be 0.
+""",
+        ),
+        HarmonyMessage(
+            role=Role.USER,
+            content="""
+
+Neuron 1
+Explanation of neuron 1 behavior: the main thing this neuron does is find vowels""",
+        ),
+        HarmonyMessage(
+            role=Role.ASSISTANT,
+            content="""
+Activations: 
+<start>
+a	10
+b	0
+c	0
+<end>
+<start>
+d	unknown
+e	10
+f	0
+<end>
+
+""",
+        ),
+        HarmonyMessage(
+            role=Role.USER,
+            content="""
+
+Neuron 2
+Explanation of neuron 2 behavior: the main thing this neuron does is find EXPLANATION""",
+        ),
+        HarmonyMessage(
+            role=Role.ASSISTANT,
+            content="""
+Activations: 
+<start>
+0	unknown
+1	unknown
+2	unknown
+<end>
+""",
+        ),
+    ]
+    prompt = ExplanationNeuronSimulator(
+        model_name="gpt-4",
+        explanation="EXPLANATION",
+        few_shot_example_set=FewShotExampleSet.TEST,
+        prompt_format=PromptFormat.HARMONY_V4,
+    ).make_simulation_prompt(
+        tokens=[str(x) for x in range(3)],
+    )
+
+    assert isinstance(prompt, list)
+    assert isinstance(prompt[0], dict)  # Really a HarmonyMessage
+    for actual_message, expected_message in zip(prompt, expected_prompt):
+        assert actual_message["role"] == expected_message["role"]
+        assert actual_message["content"] == expected_message["content"]
+    assert prompt == expected_prompt
+
+
+def test_make_token_by_token_simulation_prompt_if_format() -> None:
+    expected_prompt = """We're studying neurons in a neural network. Each neuron looks for some particular thing in a short document. Look at  an explanation of what the neuron does, and try to predict its activations on a particular token.
+
+The activation format is token<tab>activation, and activations range from 0 to 10. Most activations will be 0.
+
+Neuron 1
+Explanation of neuron 1 behavior: the main thing this neuron does is find vowels
+Activations: 
+<start>
+a	10
+b	0
+c	0
+<end>
+<start>
+d	0
+e	10
+f	0
+<end>
+
+
+Now, we're going predict the activation of a new neuron on a single token, following the same rules as the examples above. Activations still range from 0 to 10.
+Neuron 2
+Explanation of neuron 2 behavior: the main thing this neuron does is find numbers and nothing else
+Text:
+ghi
+
+Last token in the text:
+i
+
+Last token activation, considering the token in the context in which it appeared in the text:
+10
+
+
+Neuron 3
+Explanation of neuron 3 behavior: the main thing this neuron does is find numbers and nothing else
+Text:
+01
+
+Last token in the text:
+1
+
+Last token activation, considering the token in the context in which it appeared in the text:
+<|endofprompt|>"""
+    prompt = ExplanationTokenByTokenSimulator(
+        model_name="text-davinci-003",
+        explanation="EXPLANATION",
+        few_shot_example_set=FewShotExampleSet.TEST,
+        prompt_format=PromptFormat.INSTRUCTION_FOLLOWING,
+    ).make_single_token_simulation_prompt(
+        tokens=[str(x) for x in range(3)],
+        explanation="numbers and nothing else",
+        token_index_to_score=1,
+    )
+    assert prompt == expected_prompt
+
+
+def test_make_token_by_token_simulation_prompt_harmony_format() -> None:
+    expected_prompt = [
+        HarmonyMessage(
+            role=Role.SYSTEM,
+            content="""We're studying neurons in a neural network. Each neuron looks for some particular thing in a short document. Look at  an explanation of what the neuron does, and try to predict its activations on a particular token.
+
+The activation format is token<tab>activation, and activations range from 0 to 10. Most activations will be 0.
+
+""",
+        ),
+        HarmonyMessage(
+            role=Role.USER,
+            content="""Neuron 1
+Explanation of neuron 1 behavior: the main thing this neuron does is find vowels
+""",
+        ),
+        HarmonyMessage(
+            role=Role.ASSISTANT,
+            content="""Activations: 
+<start>
+a	10
+b	0
+c	0
+<end>
+<start>
+d	0
+e	10
+f	0
+<end>
+
+
+""",
+        ),
+        HarmonyMessage(
+            role=Role.SYSTEM,
+            content="Now, we're going predict the activation of a new neuron on a single token, following the same rules as the examples above. Activations still range from 0 to 10.",
+        ),
+        HarmonyMessage(
+            role=Role.USER,
+            content="""
+Neuron 2
+Explanation of neuron 2 behavior: the main thing this neuron does is find numbers and nothing else
+Text:
+ghi
+
+Last token in the text:
+i
+
+Last token activation, considering the token in the context in which it appeared in the text:
+""",
+        ),
+        HarmonyMessage(
+            role=Role.ASSISTANT,
+            content="""10
+
+""",
+        ),
+        HarmonyMessage(
+            role=Role.USER,
+            content="""
+Neuron 3
+Explanation of neuron 3 behavior: the main thing this neuron does is find numbers and nothing else
+Text:
+01
+
+Last token in the text:
+1
+
+Last token activation, considering the token in the context in which it appeared in the text:
+""",
+        ),
+    ]
+
+    prompt = ExplanationTokenByTokenSimulator(
+        model_name="gpt-4",
+        explanation="EXPLANATION",
+        few_shot_example_set=FewShotExampleSet.TEST,
+        prompt_format=PromptFormat.HARMONY_V4,
+    ).make_single_token_simulation_prompt(
+        tokens=[str(x) for x in range(3)],
+        explanation="numbers and nothing else",
+        token_index_to_score=1,
+    )
+
+    assert isinstance(prompt, list)
+    assert isinstance(prompt[0], dict)  # Really a HarmonyMessage
+    for actual_message, expected_message in zip(prompt, expected_prompt):
+        assert actual_message["role"] == expected_message["role"]
+        assert actual_message["content"] == expected_message["content"]
+    assert prompt == expected_prompt

neuron_explainer/explanations/token_space_few_shot_examples.py

@@ -0,0 +1,212 @@
+from dataclasses import dataclass
+from enum import Enum
+from typing import List
+
+from neuron_explainer.fast_dataclasses import FastDataclass
+
+
+@dataclass
+class Example(FastDataclass):
+    """
+    An example list of tokens as strings corresponding to top token space inputs of a neuron, with a
+    string explanation of the neuron's behavior on these tokens.
+    """
+
+    tokens: List[str]
+    explanation: str
+
+
+class TokenSpaceFewShotExampleSet(Enum):
+    """Determines which few-shot examples to use when sampling explanations."""
+
+    ORIGINAL = "original"
+    TEST = "test"
+
+    def get_examples(self) -> list[Example]:
+        """Returns regular examples for use in a few-shot prompt."""
+        if self is TokenSpaceFewShotExampleSet.ORIGINAL:
+            return ORIGINAL_EXAMPLES
+        elif self is TokenSpaceFewShotExampleSet.TEST:
+            return TEST_EXAMPLES
+        else:
+            raise ValueError(f"Unhandled example set: {self}")
+
+
+ORIGINAL_EXAMPLES = [
+    Example(
+        tokens=[
+            "actual",
+            " literal",
+            " actual",
+            " hyper",
+            " real",
+            " EX",
+            " Real",
+            "^",
+            "Full",
+            " full",
+            " optical",
+            " style",
+            "any",
+            "ALL",
+            "extreme",
+            " miniature",
+            " Optical",
+            " faint",
+            "~",
+            " Physical",
+            " REAL",
+            "*",
+            "virtual",
+            "TYPE",
+            " technical",
+            "otally",
+            " physic",
+            "Type",
+            "<",
+            "images",
+            "atic",
+            " sheer",
+            " Style",
+            " partial",
+            " natural",
+            "Hyper",
+            " Any",
+            " theoretical",
+            "|",
+            " ultimate",
+            "oing",
+            " constant",
+            "ANY",
+            "antically",
+            "ishly",
+            " ex",
+            " visual",
+            "special",
+            "omorphic",
+            "visual",
+        ],
+        explanation=" adjectives related to being real, or to physical properties and evidence",
+    ),
+    Example(
+        tokens=[
+            "cephal",
+            "aeus",
+            " coma",
+            "bered",
+            "abetes",
+            "inflamm",
+            "rugged",
+            "alysed",
+            "azine",
+            "hered",
+            "cells",
+            "aneously",
+            "fml",
+            "igm",
+            "culosis",
+            "iani",
+            "CTV",
+            "disabled",
+            "heric",
+            "ulo",
+            "geoning",
+            "awi",
+            "translation",
+            "iral",
+            "govtrack",
+            "mson",
+            "cloth",
+            "nesota",
+            " Dise",
+            " Lyme",
+            " dementia",
+            "agn",
+            " reversible",
+            " susceptibility",
+            "esthesia",
+            "orf",
+            " inflamm",
+            " Obesity",
+            " tox",
+            " Disorders",
+            "uberty",
+            "blind",
+            "ALTH",
+            "avier",
+            " Immunity",
+            " Hurt",
+            "ulet",
+            "ueless",
+            " sluggish",
+            "rosis",
+        ],
+        explanation=" words related to physical medical conditions",
+    ),
+    Example(
+        tokens=[
+            " January",
+            "terday",
+            "cember",
+            " April",
+            " July",
+            "September",
+            "December",
+            "Thursday",
+            "quished",
+            "November",
+            "Tuesday",
+            "uesday",
+            " Sept",
+            "ruary",
+            " March",
+            ";;;;;;;;;;;;",
+            " Monday",
+            "Wednesday",
+            " Saturday",
+            " Wednesday",
+            "Reloaded",
+            "aturday",
+            " August",
+            "Feb",
+            "Sunday",
+            "Reviewed",
+            "uggest",
+            " Dhabi",
+            "ACTED",
+            "tten",
+            "Year",
+            "August",
+            "alogue",
+            "MX",
+            " Janeiro",
+            "yss",
+            " Leilan",
+            " Fiscal",
+            " referen",
+            "semb",
+            "eele",
+            "wcs",
+            "detail",
+            "ertation",
+            " Reborn",
+            " Sunday",
+            "itially",
+            "aturdays",
+            " Dise",
+            "essage",
+        ],
+        explanation=" nouns related to time and dates",
+    ),
+]
+
+TEST_EXAMPLES = [
+    Example(
+        tokens=[
+            "these",
+            " are",
+            " tokens",
+        ],
+        explanation=" this is a test explanation",
+    ),
+]

neuron_explainer/fast_dataclasses/__init__.py

@@ -0,0 +1,3 @@
+from .fast_dataclasses import FastDataclass, dumps, loads, register_dataclass
+
+__all__ = ["FastDataclass", "dumps", "loads", "register_dataclass"]