add param tags instead of neurons

app.py

@@ -17,17 +17,17 @@ def load_model(model_name):
 col1, col2, col3, *_ = st.columns(5)
 model_name = col1.selectbox("Select a model: ", options=['gpt2', 'gpt2-medium', 'gpt2-large'])
 model, model_params, tokenizer = load_model(model_name)
-neuron_layer = col2.text_input("Layer: ", value='0')
-neuron_dim = col3.text_input("Dim: ", value='0')
+# neuron_layer = col2.text_input("Layer: ", value='0')
+# neuron_dim = col3.text_input("Dim: ", value='0')
+# neurons = model_params.K_heads[int(neuron_layer), int(neuron_dim)]
 
-neurons = model_params.K_heads[int(neuron_layer), int(neuron_dim)]
 prompt = st.text_area("Prompt: ")
 submitted = st.button("Send!")
 
 if submitted:
     with st.spinner('Wait for it..'):
         model, model_params, tokenizer = map(deepcopy, (model, model_params, tokenizer))
-        decoded = speaking_probe(model, model_params, tokenizer, prompt, *neurons, 
+        decoded = speaking_probe(model, model_params, tokenizer, prompt, 
                                  repetition_penalty=2., num_generations=3,
                                  min_length=1, do_sample=True, 
                                  max_new_tokens=100)

speaking_probes/generate.py

@@ -1,3 +1,4 @@
+import re
 import numpy as np
 from copy import deepcopy
 import matplotlib.pyplot as plt
@@ -29,11 +30,20 @@ from argparse import ArgumentParser
 class ModelParameters:
     K_heads: torch.Tensor
     num_layers: int
-    d_int: int 
-
+    d_int: int
+    num_heads: int
+    hidden_dim: int
+    head_size: int
+    V_heads: torch.Tensor = None
+    W_Q_heads: torch.Tensor = None
+    W_K_heads: torch.Tensor = None
+    W_V_heads: torch.Tensor = None
+    W_O_heads: torch.Tensor = None
+    emb: torch.Tensor = None
+    
+    
 
-def extract_gpt_parameters(model):
-    emb = model.get_output_embeddings().weight.data.T
+def extract_gpt_parameters(model, full=False):
     num_layers = model.config.n_layer
     num_heads = model.config.n_head
     hidden_dim = model.config.n_embd
@@ -41,23 +51,63 @@ def extract_gpt_parameters(model):
 
     K = torch.cat([model.get_parameter(f"transformer.h.{j}.mlp.c_fc.weight").T
                                for j in range(num_layers)]).detach()
-    V = torch.cat([model.get_parameter(f"transformer.h.{j}.mlp.c_proj.weight")
+
+    K_heads = K.reshape(num_layers, -1, hidden_dim)
+    d_int = K_heads.shape[1]
+    model_params = ModelParameters(K_heads=K_heads, num_layers=num_layers, d_int=d_int,
+                                   hidden_dim=hidden_dim, head_size=head_size, 
+                                   num_heads=num_heads)
+
+    if full:
+        emb = model.get_output_embeddings().weight.data.T
+        V = torch.cat([model.get_parameter(f"transformer.h.{j}.mlp.c_proj.weight")
                                for j in range(num_layers)]).detach()
-    W_Q, W_K, W_V = torch.cat([model.get_parameter(f"transformer.h.{j}.attn.c_attn.weight") 
-                               for j in range(num_layers)]).detach().chunk(3, dim=-1)
-    W_O = torch.cat([model.get_parameter(f"transformer.h.{j}.attn.c_proj.weight") 
+        W_Q, W_K, W_V = torch.cat([model.get_parameter(f"transformer.h.{j}.attn.c_attn.weight") 
+                                   for j in range(num_layers)]).detach().chunk(3, dim=-1)
+        W_O = torch.cat([model.get_parameter(f"transformer.h.{j}.attn.c_proj.weight") 
+                                   for j in range(num_layers)]).detach()
+
+        model_params.V_heads = V.reshape(num_layers, -1, hidden_dim)
+        model_params.W_V_heads = W_V.reshape(num_layers, hidden_dim, num_heads, head_size).permute(0, 2, 1, 3)
+        model_params.W_O_heads = W_O.reshape(num_layers, num_heads, head_size, hidden_dim)
+        model_params.W_Q_heads = W_Q.reshape(num_layers, hidden_dim, num_heads, head_size).permute(0, 2, 1, 3)
+        model_params.W_K_heads = W_K.reshape(num_layers, hidden_dim, num_heads, head_size).permute(0, 2, 1, 3)
+        model_params.emb = emb
+    return model_params
+
+
+def extract_gpt_j_parameters(model, full=False):
+    num_layers = model.config.n_layer
+    num_heads = model.config.n_head
+    hidden_dim = model.config.n_embd
+    head_size = hidden_dim // num_heads
+
+    K = torch.cat([model.get_parameter(f"transformer.h.{j}.mlp.fc_in.weight").T
                                for j in range(num_layers)]).detach()
 
     K_heads = K.reshape(num_layers, -1, hidden_dim)
-    V_heads = V.reshape(num_layers, -1, hidden_dim)
     d_int = K_heads.shape[1]
+    model_params = ModelParameters(K_heads=K_heads, num_layers=num_layers, d_int=d_int,
+                                   hidden_dim=hidden_dim, head_size=head_size, 
+                                   num_heads=num_heads)
 
-    W_V_heads = W_V.reshape(num_layers, hidden_dim, num_heads, head_size).permute(0, 2, 1, 3)
-    W_O_heads = W_O.reshape(num_layers, num_heads, head_size, hidden_dim)
-    W_Q_heads = W_Q.reshape(num_layers, hidden_dim, num_heads, head_size).permute(0, 2, 1, 3)
-    W_K_heads = W_K.reshape(num_layers, hidden_dim, num_heads, head_size).permute(0, 2, 1, 3)
+    if full:
+        raise NotImplementedError
+#         emb = model.get_output_embeddings().weight.data.T
+#         V = torch.cat([model.get_parameter(f"transformer.h.{j}.mlp.c_proj.weight")
+#                                for j in range(num_layers)]).detach()
+#         W_Q, W_K, W_V = torch.cat([model.get_parameter(f"transformer.h.{j}.attn.c_attn.weight") 
+#                                    for j in range(num_layers)]).detach().chunk(3, dim=-1)
+#         W_O = torch.cat([model.get_parameter(f"transformer.h.{j}.attn.c_proj.weight") 
+#                                    for j in range(num_layers)]).detach()
 
-    return ModelParameters(K_heads=K_heads, num_layers=num_layers, d_int=d_int)
+#         model_params.V_heads = V.reshape(num_layers, -1, hidden_dim)
+#         model_params.W_V_heads = W_V.reshape(num_layers, hidden_dim, num_heads, head_size).permute(0, 2, 1, 3)
+#         model_params.W_O_heads = W_O.reshape(num_layers, num_heads, head_size, hidden_dim)
+#         model_params.W_Q_heads = W_Q.reshape(num_layers, hidden_dim, num_heads, head_size).permute(0, 2, 1, 3)
+#         model_params.W_K_heads = W_K.reshape(num_layers, hidden_dim, num_heads, head_size).permute(0, 2, 1, 3)
+#         model_params.emb = emb
+    return model_params
 
 
 def encode(token, tokenizer):
@@ -125,14 +175,29 @@ class ParamListStructureEnforcer(LogitsProcessor):
     
 
 # speaking probe
+def _preprocess_prompt(model_params, prompt):
+    K_heads = model_params.K_heads
+    prompt = re.sub(r'([^ ]|\A)(<neuron>|<param_\d+_\d+>)', lambda m: f'{m.group(1)} {m.group(2)}', prompt)
+    param_neuron_idxs = [(int(a), int(b)) for a, b in re.findall(r' <param_(\d+)_(\d+)>', prompt)]
+    param_neuron_tokens = [f' <param_{a}_{b}>' for a, b in param_neuron_idxs]
+    param_neurons = [deepcopy(K_heads[a, b]) for a, b in param_neuron_idxs]
+    return prompt, param_neuron_tokens, param_neurons
+
+
 def speaking_probe(model, model_params, tokenizer, prompt, *neurons,
                    num_generations=1, layer_range=None, bad_words_ids=[], output_neurons=False,
                    return_outputs=False, logits_processor=LogitsProcessorList([]), **kwargs):
     num_non_neuron_tokens = len(tokenizer)
     tokenizer_with_neurons = deepcopy(tokenizer)
+    
+    # adding neurons to the tokenizer
+    neuron_tokens = [f" <neuron{i+1 if i > 0 else ''}>" for i in range(len(neurons))]
+    prompt, param_neuron_tokens, param_neurons = _preprocess_prompt(model_params, prompt)
+    neuron_tokens.extend(param_neuron_tokens)
+    neurons = neurons + tuple(param_neurons)
     has_extra_neurons = len(neurons) > 0
     if has_extra_neurons:
-        tokenizer_with_neurons.add_tokens([f" <neuron{i+1 if i > 0 else ''}>" for i in range(len(neurons))])
+        tokenizer_with_neurons.add_tokens(neuron_tokens)
         model.resize_token_embeddings(len(tokenizer_with_neurons))
         model.transformer.wte.weight.data[-len(neurons):] = torch.stack(neurons, dim=0)
         
@@ -160,7 +225,8 @@ def speaking_probe(model, model_params, tokenizer, prompt, *neurons,
                              **kwargs)
 
     decoded = tokenizer_with_neurons.batch_decode(outputs.sequences, skip_special_tokens=True)
-        
+    
+    # TODO: add `finally` statement
     if has_extra_neurons:
         model.resize_token_embeddings(num_non_neuron_tokens)
         model.transformer.wte.weight.data = model.transformer.wte.weight.data[:num_non_neuron_tokens]
@@ -188,6 +254,7 @@ if __name__ == "__main__":
     parser.add_argument('--max_length', type=int, default=100)
     parser.add_argument('--max_new_tokens', type=int, default=None)
     parser.add_argument('--repetition_penalty', type=float, default=2.)
+    parser.add_argument('--temperature', type=float, default=1.)
     
     args = parser.parse_args()
     # TODO: first make them mutually exclusive 
@@ -213,6 +280,7 @@ if __name__ == "__main__":
                    num_generations=args.num_generations, 
                    repetition_penalty=args.repetition_penalty, 
                    num_beams=args.num_beams, top_p=args.top_p, top_k=args.top_k,
+                   temperature=args.temperature,
                    min_length=args.min_length, do_sample=not args.no_sample,
                    max_length=args.max_length, max_new_tokens=args.max_new_tokens)
     for i in range(len(decoded)):