2️⃣ Task-encoding hidden states

def generate_antonym_dataset(N: int):
    """
    Generates 100 pairs of antonyms, in the form of a list of 2-tuples.
    """
    assert os.environ.get("OPENAI_API_KEY", None) is not None, (
        "Please set your API key before running this function!"
    )

    client = OpenAI(api_key=os.environ["OPENAI_API_KEY"])

    response = client.chat.completions.create(
        model="gpt-3.5-turbo",
        messages=[
            {"role": "system", "content": "You are a helpful assistant."},
            {
                "role": "user",
                "content": f"Generate {N} pairs of antonyms in the form of a list of 2-tuples. For example, [['old', 'young'], ['top', bottom'], ['awake', 'asleep']...].",
            },
            {"role": "assistant", "content": "Sure, here is a list of 100 antonyms: "},
        ],
    )
    return response


if os.environ.get("OPENAI_API_KEY", None) is not None:
    ANTONYM_PAIRS = generate_antonym_dataset(100)
    # Save the word pairs in a text file
    with open(section_dir / "data" / "my_antonym_pairs.txt", "w") as f:
        for word_pair in ANTONYM_PAIRS:
            f.write(f"{word_pair[0]} {word_pair[1]}\n")

# Load the word pairs from the text file
with open(section_dir / "data" / "antonym_pairs.txt", "r") as f:
    ANTONYM_PAIRS = [line.split() for line in f.readlines()]

print(ANTONYM_PAIRS[:10])

class ICLSequence:
    """
    Class to store a single antonym sequence.

    Uses the default template "Q: {x}\nA: {y}" (with separate pairs split by "\n\n").
    """

    def __init__(self, word_pairs: list[list[str]]):
        self.word_pairs = word_pairs
        self.x, self.y = zip(*word_pairs)

    def __len__(self):
        return len(self.word_pairs)

    def __getitem__(self, idx: int):
        return self.word_pairs[idx]

    def prompt(self):
        """Returns the prompt, which contains all but the second element in the last word pair."""
        p = "\n\n".join([f"Q: {x}\nA: {y}" for x, y in self.word_pairs])
        return p[: -len(self.completion())]

    def completion(self):
        """Returns the second element in the last word pair (with padded space)."""
        return " " + self.y[-1]

    def __str__(self):
        """Prints a readable string representation of the prompt & completion (indep of template)."""
        return f"{', '.join([f'({x}, {y})' for x, y in self[:-1]])}, {self.x[-1]} ->".strip(", ")


word_list = [["hot", "cold"], ["yes", "no"], ["in", "out"], ["up", "down"]]
seq = ICLSequence(word_list)

print("Tuple-representation of the sequence:")
print(seq)
print("\nActual prompt, which will be fed into the model:")
print(seq.prompt())

class ICLDataset:
    """
    Dataset to create antonym pair prompts, in ICL task format. We use random seeds for consistency
    between the corrupted and clean datasets.

    Inputs:
        word_pairs:
            list of ICL task, e.g. [["old", "young"], ["top", "bottom"], ...] for the antonym task
        size:
            number of prompts to generate
        n_prepended:
            number of antonym pairs before the single-word ICL task
        bidirectional:
            if True, then we also consider the reversed antonym pairs
        corrupted:
            if True, then the second word in each pair is replaced with a random word
        seed:
            random seed, for consistency & reproducibility
    """

    def __init__(
        self,
        word_pairs: list[list[str]],
        size: int,
        n_prepended: int,
        bidirectional: bool = True,
        seed: int = 0,
        corrupted: bool = False,
    ):
        assert n_prepended + 1 <= len(word_pairs), (
            "Not enough antonym pairs in dataset to create prompt."
        )

        self.word_pairs = word_pairs
        self.word_list = [word for word_pair in word_pairs for word in word_pair]
        self.size = size
        self.n_prepended = n_prepended
        self.bidirectional = bidirectional
        self.corrupted = corrupted
        self.seed = seed

        self.seqs = []
        self.prompts = []
        self.completions = []

        # Generate the dataset (by choosing random word pairs, and constructing ICLSequence objects)
        for n in range(size):
            np.random.seed(seed + n)
            random_pairs = np.random.choice(len(self.word_pairs), n_prepended + 1, replace=False)
            # Randomize the order of each word pair (x, y).
            # If not bidirectional, we always have x -> y not y -> x
            random_orders = np.random.choice([1, -1], n_prepended + 1)
            if not (bidirectional):
                random_orders[:] = 1
            word_pairs = [
                self.word_pairs[pair][::order] for pair, order in zip(random_pairs, random_orders)
            ]
            # If corrupted, then replace y with a random word in all (x, y) pairs except the last one
            if corrupted:
                for i in range(len(word_pairs) - 1):
                    word_pairs[i][1] = np.random.choice(self.word_list)
            seq = ICLSequence(word_pairs)

            self.seqs.append(seq)
            self.prompts.append(seq.prompt())
            self.completions.append(seq.completion())

    def create_corrupted_dataset(self):
        """Creates a corrupted version of the dataset (with same random seed)."""
        return ICLDataset(
            self.word_pairs,
            self.size,
            self.n_prepended,
            self.bidirectional,
            corrupted=True,
            seed=self.seed,
        )

    def __len__(self):
        return self.size

    def __getitem__(self, idx: int):
        return self.seqs[idx]

dataset = ICLDataset(ANTONYM_PAIRS, size=10, n_prepended=2, corrupted=False)

table = Table("Prompt", "Correct completion")
for seq, completion in zip(dataset.seqs, dataset.completions):
    table.add_row(str(seq), repr(completion))

rprint(table)

dataset = ICLDataset(ANTONYM_PAIRS, size=10, n_prepended=2, corrupted=True)

table = Table("Prompt", "Correct completion")
for seq, completions in zip(dataset.seqs, dataset.completions):
    table.add_row(str(seq), repr(completions))

rprint(table)

def calculate_h(
    model: LanguageModel, dataset: ICLDataset, layer: int = -1
) -> tuple[list[str], Tensor]:
    """
    Averages over the model's hidden representations on each of the prompts in `dataset` at layer
    `layer`, to produce a single vector `h`.

    Inputs:
        model: LanguageModel
            the transformer you're doing this computation with
        dataset: ICLDataset
            the dataset whose prompts `dataset.prompts` you're extracting the activations from (at
            the last seq pos)
        layer: int
            the layer you're extracting activations from

    Returns:
        completions: list[str]
            list of the model's next-token predictions (i.e. the strings the model predicts to
            follow the last token)
        h: Tensor
            average hidden state tensor at final sequence position, of shape (d_model,)
    """
    raise NotImplementedError()


tests.test_calculate_h(calculate_h, model)

def display_model_completions_on_antonyms(
    model: LanguageModel,
    dataset: ICLDataset,
    completions: list[str],
    num_to_display: int = 20,
) -> None:
    table = Table(
        "Prompt (tuple representation)",
        "Model's completion\n(green=correct)",
        "Correct completion",
        title="Model's antonym completions",
    )

    for i in range(min(len(completions), num_to_display)):
        # Get model's completion, and correct completion
        completion = completions[i]
        correct_completion = dataset.completions[i]
        correct_completion_first_token = model.tokenizer.tokenize(correct_completion)[0].replace(
            "Ġ", " "
        )
        seq = dataset.seqs[i]

        # Color code the completion based on whether it's correct
        is_correct = completion == correct_completion_first_token
        completion = f"[b green]{repr(completion)}[/]" if is_correct else repr(completion)

        table.add_row(str(seq), completion, repr(correct_completion))

    rprint(table)


# Get uncorrupted dataset
dataset = ICLDataset(ANTONYM_PAIRS, size=20, n_prepended=2)

# Getting it from layer 12, as in the description in section 2.1 of paper
model_completions, h = calculate_h(model, dataset, layer=12)

# Displaying the output
display_model_completions_on_antonyms(model, dataset, model_completions)

with model.trace(inputs, remote=REMOTE):
    # some causal interventions

with model.trace(remote=REMOTE) as tracer:
    with tracer.invoke(inputs):
        # some causal interventions

    with tracer.invoke(other_inputs):
        # some other causal interventions

with model.trace(remote=REMOTE) as tracer:
    with tracer.invoke(clean_inputs):
        # extract clean activations
        clean_activations = model.transformer.h[10].output[0]

    with tracer.invoke(corrupted_inputs):
        # patch clean into corrupted
        model.transformer.h[10].output[0][:] = clean_activations

def intervene_with_h(
    model: LanguageModel,
    zero_shot_dataset: ICLDataset,
    h: Tensor,
    layer: int,
    remote: bool = REMOTE,
) -> tuple[list[str], list[str]]:
    """
    Extracts the vector `h` using previously defined function, and intervenes by adding `h` to the
    residual stream of a set of generated zero-shot prompts.

    Inputs:
        model: the model we're using to generate completions
        zero_shot_dataset: the dataset of zero-shot prompts which we'll intervene on, using the
            `h`-vector
        h: the `h`-vector we'll be adding to the residual stream
        layer: the layer we'll be extracting the `h`-vector from
        remote: whether to run the forward pass on the remote server (used for running test code)

    Returns:
        completions_zero_shot: list of string completions for the zero-shot prompts, without
            intervention using the h-vector
        completions_intervention: list of string completions for the zero-shot prompts, with
            intervention using the h-vector
    """
    raise NotImplementedError()


tests.test_intervene_with_h(intervene_with_h, model, h, ANTONYM_PAIRS, REMOTE)

layer = 12
dataset = ICLDataset(ANTONYM_PAIRS, size=20, n_prepended=3, seed=0)
zero_shot_dataset = ICLDataset(ANTONYM_PAIRS, size=20, n_prepended=0, seed=1)

# Run previous function to get h-vector
h = calculate_h(model, dataset, layer=layer)[1]

# Run new function to intervene with h-vector
completions_zero_shot, completions_intervention = intervene_with_h(
    model, zero_shot_dataset, h, layer=layer
)

print("Zero-shot completions: ", completions_zero_shot)
print("Completions with intervention: ", completions_intervention)

def display_model_completions_on_h_intervention(
    dataset: ICLDataset,
    completions: list[str],
    completions_intervention: list[str],
    num_to_display: int = 20,
) -> None:
    table = Table(
        "Prompt",
        "Model's completion\n(no intervention)",
        "Model's completion\n(intervention)",
        "Correct completion",
        title="Model's antonym completions",
    )

    for i in range(min(len(completions), num_to_display)):
        completion_ni = completions[i]
        completion_i = completions_intervention[i]
        correct_completion = dataset.completions[i]
        correct_completion_first_token = tokenizer.tokenize(correct_completion)[0].replace("Ġ", " ")
        seq = dataset.seqs[i]

        # Color code the completion based on whether it's correct
        is_correct = completion_i == correct_completion_first_token
        completion_i = f"[b green]{repr(completion_i)}[/]" if is_correct else repr(completion_i)

        table.add_row(str(seq), repr(completion_ni), completion_i, repr(correct_completion))

    rprint(table)


display_model_completions_on_h_intervention(
    zero_shot_dataset, completions_zero_shot, completions_intervention
)

def calculate_h_and_intervene(
    model: LanguageModel,
    dataset: ICLDataset,
    zero_shot_dataset: ICLDataset,
    layer: int,
) -> tuple[list[str], list[str]]:
    """
    Extracts the vector `h`, intervenes by adding `h` to the residual stream of a set of generated
    zero-shot prompts, all within the same forward pass. Returns the completions from this
    intervention.

    Inputs:
        model: LanguageModel
            the model we're using to generate completions
        dataset: ICLDataset
            the dataset of clean prompts from which we'll extract the `h`-vector
        zero_shot_dataset: ICLDataset
            the dataset of zero-shot prompts which we'll intervene on, using the `h`-vector
        layer: int
            the layer we'll be extracting the `h`-vector from

    Returns:
        completions_zero_shot: list[str]
            list of string completions for the zero-shot prompts, without intervention
        completions_intervention: list[str]
            list of string completions for the zero-shot prompts, with h-intervention
    """
    raise NotImplementedError()


dataset = ICLDataset(ANTONYM_PAIRS, size=20, n_prepended=3, seed=0)
zero_shot_dataset = ICLDataset(ANTONYM_PAIRS, size=20, n_prepended=0, seed=1)

completions_zero_shot, completions_intervention = calculate_h_and_intervene(
    model, dataset, zero_shot_dataset, layer=layer
)

display_model_completions_on_h_intervention(
    zero_shot_dataset, completions_zero_shot, completions_intervention
)

def calculate_h_and_intervene_logprobs(
    model: LanguageModel,
    dataset: ICLDataset,
    zero_shot_dataset: ICLDataset,
    layer: int,
) -> tuple[list[float], list[float]]:
    """
    Extracts the vector `h`, intervenes by adding `h` to the residual stream of a set of generated
    zero-shot prompts, all within the same forward pass. Returns the logprobs on correct tokens from
    this intervention.

    Inputs:
        model: LanguageModel
            the model we're using to generate completions
        dataset: ICLDataset
            the dataset of clean prompts from which we'll extract the `h`-vector
        zero_shot_dataset: ICLDataset
            the dataset of zero-shot prompts which we'll intervene on, using the `h`-vector
        layer: int
            the layer we'll be extracting the `h`-vector from

    Returns:
        correct_logprobs: list[float]
            list of correct-token logprobs for the zero-shot prompts, without intervention
        correct_logprobs_intervention: list[float]
            list of correct-token logprobs for the zero-shot prompts, with h-intervention
    """
    raise NotImplementedError()

def display_model_logprobs_on_h_intervention(
    dataset: ICLDataset,
    correct_logprobs_zero_shot: list[float],
    correct_logprobs_intervention: list[float],
    num_to_display: int = 20,
) -> None:
    table = Table(
        "Zero-shot prompt",
        "Model's logprob\n(no intervention)",
        "Model's logprob\n(intervention)",
        "Change in logprob",
        title="Model's antonym logprobs, with zero-shot h-intervention\n(green = intervention improves accuracy)",
    )

    for i in range(min(len(correct_logprobs_zero_shot), num_to_display)):
        logprob_ni = correct_logprobs_zero_shot[i]
        logprob_i = correct_logprobs_intervention[i]
        delta_logprob = logprob_i - logprob_ni
        zero_shot_prompt = f"{dataset[i].x[0]:>8} -> {dataset[i].y[0]}"

        # Color code the logprob based on whether it's increased with this intervention
        is_improvement = delta_logprob >= 0
        delta_logprob = (
            f"[b green]{delta_logprob:+.2f}[/]" if is_improvement else f"{delta_logprob:+.2f}"
        )

        table.add_row(zero_shot_prompt, f"{logprob_ni:.2f}", f"{logprob_i:.2f}", delta_logprob)

    rprint(table)


dataset = ICLDataset(ANTONYM_PAIRS, size=20, n_prepended=3, seed=0)
zero_shot_dataset = ICLDataset(ANTONYM_PAIRS, size=20, n_prepended=0, seed=1)

correct_logprobs_zero_shot, correct_logprobs_intervention = calculate_h_and_intervene_logprobs(
    model, dataset, zero_shot_dataset, layer=layer
)

display_model_logprobs_on_h_intervention(
    zero_shot_dataset, correct_logprobs_zero_shot, correct_logprobs_intervention
)