def make_simulation_prompt(
        self,
        tokens: Sequence[str],
        explanation: str,
    ) -> str | list[ChatMessage]:
        """Make a few-shot prompt for predicting the neuron's activations on a sequence.
        This prompt only gives the model one sequence per neuron in the few shot examples."""
        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 certain things in a short document. Your task is to read the explanation of what the neuron does, and predict the neuron's activations for each token in the document.

For each document, you will see the full text of the document, then the tokens in the document with the activation left blank. You will print, in valid json, the exact same tokens verbatim, but with the activation values filled in according to the explanation. Pay special attention to the explanation's description of the context and order of tokens or words.

Fill out the activation values from 0 to 10. Please think carefully.";
""",
        )

        few_shot_examples = self.few_shot_example_set.get_examples()
        for example in few_shot_examples:
            prompt_builder.add_message(
                Role.USER,
                _format_record_for_logprob_free_simulation_json(
                    explanation=example.explanation,
                    activation_record=example.activation_records[0],
                    include_activations=False,
                ),
            )
            # Example of this few shot user message.
            """
            {
                "to_find": "hello",
                "document": "The",
                "activations": [
                    {
                        "token": "The",
                        "activation": null
                    },
                    ...
                ]
            }
            """

            prompt_builder.add_message(
                Role.ASSISTANT,
                _format_record_for_logprob_free_simulation_json(
                    explanation=example.explanation,
                    activation_record=example.activation_records[0],
                    include_activations=True,
                ),
            )
            # Example of this few shot assistant message:
            """
            {
                "to_find": "hello",
                "document": "The",
                "activations": [
                    {
                        "token": "The",
                        "activation": 10
                    },
                    ...
                ]
            }
            """

        prompt_builder.add_message(
            Role.USER,
            _format_record_for_logprob_free_simulation_json(
                explanation=explanation,
                activation_record=ActivationRecord(tokens=list(tokens), activations=[]),
                include_activations=False,
            ),
        )
        # Example of the final user message:
        """
        {
            "to_find": "hello",
            "document": "The",
            "activations": [
                {
                    "token": "The",
                    "activation": null
                },
                ...
            ]
        }
        """
        return prompt_builder.build(self.prompt_format)
Make a few-shot prompt for predicting the neuron's activations on a sequence.
        This prompt only gives the model one sequence per neuron in the few shot examples.
neuron_explainer/explanations/simulator.py