Package: taker
takerMy basic library for studying LLMs, with support for Multi-GPU inference and
editing, as well as quantilised inference (not editing yet).
This includes functions for analysing the activations of the models for
different inputs, and for pruning different parts of the model based on those
activations.The currently tested list of models is:EleutherAI's PythiaEleutherAI's GPT-NeoXMeta OptMeta GalacticaMeta Llama / Llama 2MistralAI Mistral 7BGPT-2RoBERTaGoogle's Vision Transformer (ViT)For check out theexamples folderto see in more detail how the library can be used.Using the ModelTo load up the model, simply do:from taker import Model

# choose any model from huggingface, though note that not all will be supported
# dtype can be anything from int4, int8, fp16, fp32, fp64
m = Model("nickypro/tinyllama-15m", dtype="fp16")

# You can access different attributes of the model
print(m.cfg.n_layer) # model has 6 layers
print(m.cfg.d_model) # 288 model width
print(m.cfg.d_mlp)   # 768 mlp width
print(m.cfg.n_heads) # 6 heads
print(m.cfg.d_head)  # 48 head size
print(m.cfg.d_vocab) # 32000 token optionsOnce the model is loaded, you can try using it for some generative tasks:output = m.generate("I am hungry and want to", num=10)
print(output)
# ('I am hungry and want to', ' eat something." I said, "I am sorry')If you want to inspect the activations, you can try to look at it's residual stream activations:m.get_residual_stream("I am hungry")
print(m.shape)
# torch.Size([13, 5, 288])
# formatted as [n_layer, n_tokens, d_model]
# where n_layers is [input], [attn_0_output], [mlp_0_output], [attn_1_input], ...If you want to manipulate the model, you can use one of the implemented hooks:NeuronMask, used to set neuron activations to zero (or some other constant).NeuronActAdd, used to do "activation addition" to different tokens.NeuronPostBias, used to add a bias to the ouputs when the huggingface model does not otherwise support it.For example, you can set the masks of some mlp neurons in layer 2 to zero:import torch
neurons_to_keep = torch.ones(m.cfg.d_mlp)
neurons_to_keep[:10] = 0

m.masks["mlp_pre_out"][2].delete_neurons(keep_indices=neurons_to_keep)You can also make it so that the neurons are not set to zero, but rather some other valueneuron_cenetering_vector = torch.randn(m.cfg.d_mlp)


m.masks["mlp_pre_out"][2].set_offset(neuron_centering_vector)Model MapTaker works by using a map,ModelMapwhich converts a standardised query likemlp.W_outto
the rekevant saved component. For example:mlp_weights = m.layers[2]["mlp.W_out"]

mlp_weights[..., 0] = 3.14159

m.layers[2]["mlp.W_out"] = mlp_weightsPruning based on CapabilitiesFor a full example, seesrc/examples/prune_30.py.The simple example is:from taker.data_classes import PruningConfig
from taker.parser import cli_parser
from taker.prune import run_pruning

# Configure initial model and tests
c = PruningConfig(
    wandb_project = "testing",
    model_repo   = "facebook/opt-125m",
    token_limit  = 1000,
    run_pre_test = True,

    # Removals parameters
    ff_scoring = "abs"
    ff_frac   = 0.02,
    ff_eps    = 0.001,
    attn_scoring = "abs",
    attn_frac = 0.00,
    attn_eps  = 1e-4,

    # Eval
    focus     = "pile_codeless",
    cripple   = "code",
    additional_datasets=tuple(),
)

# optionally, use parser to get CLI arguments.
# c, args = cli_parser(c)

# Run the iterated pruning
model, history = run_pruning(c)model.pyThis defines a wrapper function that encapsulates the HuggingFace implementation of Meta OPT.
To get the model, simply run:from taker import Model

m = Model("facebook/opt-125m", limit=1000)Where you can provide any of the model sizes that are pre-trained for OPT, and the token limit must be smaller than the max token length that the model is able to handle.Next, you can run the model to do 2 tokens of predictions, by, for example, running:text = 'Hello, my name is'
inpt, output = opt.predict( text, num=2 )We can look at the residual stream of how the output changes over time.residual_stream = opt.get_residual_stream( text )This will return a tensor of size2 + 2*n_layers.
i.e:the input (w/ positional encoding)n attention layer outputsn feed forward layer outputsthe final outputIf we want just the output of the attention / feed forward layers, we can instead look at the activations:inpt, attn_out, ff_out, output = opt.get_text_activations( text )or alternatively:inpt, attn_out, ff_out, output = opt.get_text_activations( residual_stream=residual_stream )To get the activations for the input text at all of the MLP mid layers, we can look at:opt.get_ff_key_activations( text )oropt.get_ff_key_activations( residual_stream=residual_stream ).texts.pyHas some basic tools for loading the text datasets I am using:'pile', ( EleutherAI's 'The Pile' dataset)'pile-codeless' (Pile without GitHub)'code' (CodeParrot's 'github-code' dataset)'python', (Subset of only python)'wiki', (WikiText)'civil', (Civil comments with toxicity < 0.2)'toxic', (Civil comments with toxicity > 0.8)...activations.pyHas code specific to the datasets I am using to analyze and attempt to remove capabilities from the models.