def make_single_token_simulation_prompt(
        self,
        tokens: Sequence[str],
        explanation: str,
        token_index_to_score: int,
    ) -> str | list[ChatMessage]:
        """Make a few-shot prompt for predicting the neuron's activation on a single token."""
        assert explanation != ""
        prompt_builder = PromptBuilder(allow_extra_system_messages=True)
        prompt_builder.add_message(
            Role.SYSTEM,
            """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.

""",
        )

        few_shot_examples = self.few_shot_example_set.get_examples()
        for i, example in enumerate(few_shot_examples):
            prompt_builder.add_message(
                Role.USER,
                f"Neuron {i + 1}\nExplanation of neuron {i + 1} behavior: {EXPLANATION_PREFIX} "
                f"{example.explanation}\n",
            )
            formatted_activation_records = format_activation_records(
                example.activation_records,
                calculate_max_activation(example.activation_records),
                start_indices=None,
            )
            prompt_builder.add_message(
                Role.ASSISTANT,
                f"Activations: {formatted_activation_records}\n\n",
            )

        prompt_builder.add_message(
            Role.SYSTEM,
            "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.",
        )
        single_token_example = self.few_shot_example_set.get_single_token_prediction_example()
        assert single_token_example.token_index_to_score is not None
        self._add_single_token_simulation_subprompt(
            prompt_builder,
            single_token_example.activation_records[0],
            len(few_shot_examples) + 1,
            explanation,
            token_index_to_score=single_token_example.token_index_to_score,
            end_of_prompt=False,
        )

        activation_record = ActivationRecord(
            tokens=list(tokens[: token_index_to_score + 1]),  # ActivationRecord expects List type.
            activations=[0.0] * len(tokens),
        )
        self._add_single_token_simulation_subprompt(
            prompt_builder,
            activation_record,
            len(few_shot_examples) + 2,
            explanation,
            token_index_to_score,
            end_of_prompt=True,
        )
        return prompt_builder.build(self.prompt_format)
Make a few-shot prompt for predicting the neuron's activation on a single token.
neuron_explainer/explanations/simulator.py