[P2] Having a simple tutorial on how to use intervenable model inside another torch module, or pipeline (e.g., HF pipeline)

[frankaging]

Descriptions:
The IntervenableModel is a torch.nn.Module. So, this can be used inside another torch model, or even pipeline object (e.g., Huggingface pipeline). Here is a quick code snippet,

import pyvene
from pyvene import IntervenableRepresentationConfig, IntervenableConfig, IntervenableModel

# provided wrapper for huggingface gpt2 model
_, tokenizer, gpt2 = pyvene.create_gpt2()

# turn gpt2 into intervenable_gpt2
intervenable_gpt2 = IntervenableModel(
    intervenable_config = IntervenableConfig(
        intervenable_representations=[
            IntervenableRepresentationConfig(
                0,            # intervening layer 0
                "mlp_output", # intervening mlp output
                "pos",        # intervening based on positional indices of tokens
                1             # maximally intervening one token
            ),
        ],
    ), 
    model = gpt2
)

import torch
import torch.nn as nn
from typing import List, Optional, Tuple, Union, Dict

class ModelWithIntervenables(nn.Module):
    def __init__(self):
        super(ModelWithIntervenables, self).__init__()
        self.intervenable_gpt2 = intervenable_gpt2
        self.relu = nn.ReLU()
        self.fc = nn.Linear(768, 1)
        # Your other downstream components go here

    def forward(
        self, 
        base,
        sources: Optional[List] = None,
        unit_locations: Optional[Dict] = None,
        activations_sources: Optional[Dict] = None,
        subspaces: Optional[List] = None,
    ):
        _, counterfactual_x = self.intervenable_gpt2(
            base,
            sources,
            unit_locations,
            activations_sources,
            subspaces
        )
        counterfactual_x = counterfactual_x.last_hidden_state
        
        counterfactual_x = self.relu(counterfactual_x)
        counterfactual_x = self.fc(counterfactual_x)
        return counterfactual_x

and then you can run forward as usual,

model = ModelWithIntervenables()

base = tokenizer("The capital of Spain is", return_tensors="pt")
sources = [
    tokenizer("The capital of Italy is", return_tensors="pt"),
]

model(
    base, sources, {"sources->base": ([[[4]]], [[[4]]])}
)

which returns,

tensor([[[2.7027],
         [6.3036],
         [6.1785],
         [6.4302],
         [8.0921]]], grad_fn=<ViewBackward0>)

[aryamanarora]

the readme has this right? closing