File size: 6,903 Bytes
94011a1 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 |
# hop_interventions.py
# Author: Julie Kallini
# For importing utils
import sys
sys.path.append("..")
# align-transformers
PATH_TO_ALIGN_TRANSFORMERS = "/nlp/scr/kallini/align-transformers/"
sys.path.append(PATH_TO_ALIGN_TRANSFORMERS)
import pandas as pd
from models.utils import embed_to_distrib
from models.configuration_alignable_model import AlignableRepresentationConfig, AlignableConfig
from models.alignable_base import AlignableModel
from models.interventions import VanillaIntervention
from utils import CHECKPOINT_READ_PATH, marker_sg_token, marker_pl_token, \
PERTURBATIONS, PAREN_MODELS
from tqdm import tqdm
from transformers import GPT2Model
from gpt2_no_positional_encoding_model import GPT2NoPositionalEncodingModel
import os
import torch
import argparse
MAX_TRAINING_STEPS = 3000
CHECKPOINTS = list(range(100, MAX_TRAINING_STEPS+1, 100))
def simple_position_config(model_type, intervention_type, layer):
alignable_config = AlignableConfig(
alignable_model_type=model_type,
alignable_representations=[
AlignableRepresentationConfig(
layer, # layer
intervention_type, # intervention type
"pos", # intervention unit
1 # max number of unit
),
],
alignable_interventions_type=VanillaIntervention,
)
return alignable_config
def get_model(perturbation_type, train_set, seed, paren_model, ckpt, no_pos_encodings=False):
# Get path to model
no_pos_encodings = "_no_positional_encodings" if no_pos_encodings else ""
model = f"babylm_{perturbation_type}_{train_set}_{paren_model}{no_pos_encodings}_seed{seed}"
model_path = f"{CHECKPOINT_READ_PATH}/babylm_{perturbation_type}_{train_set}_{paren_model}{no_pos_encodings}/{model}/runs/{model}/checkpoint-{ckpt}"
# Load appropriate GPT-2 model
if no_pos_encodings:
return GPT2NoPositionalEncodingModel.from_pretrained(model_path).to(device)
else:
return GPT2Model.from_pretrained(model_path).to(device)
def run_interventions(model, base_input_ids, source_input_ids):
tokens = [marker_sg_token, marker_pl_token]
data = []
BATCH_SIZE = 16
for batch_i in tqdm(range(0, len(base_input_ids), BATCH_SIZE)):
# Get base and source batches
base_batch = base_input_ids[batch_i:batch_i+BATCH_SIZE]
source_batch = source_input_ids[batch_i:batch_i+BATCH_SIZE]
# Iterate over GPT-2 layers
for layer_i in range(model.config.n_layer):
# Get block_output config for this layer
alignable_config = simple_position_config(
type(model), "block_output", layer_i)
alignable = AlignableModel(alignable_config, model)
# Iterate over token positions
for pos_i in range(len(base_batch[0])):
_, counterfactual_outputs = alignable(
{"input_ids": torch.tensor(base_batch).to(device)},
[{"input_ids": torch.tensor(source_batch).to(device)}],
{"sources->base": ([[[pos_i]] * len(base_batch)],
[[[pos_i]] * len(base_batch)])}
)
distrib = embed_to_distrib(
model, counterfactual_outputs.last_hidden_state,
logits=False
)
for i in range(len(base_batch)):
for token in tokens:
data.append({
'example': batch_i + i,
'token': token,
'prob': float(distrib[i][-1][token]),
'layer': layer_i,
'pos': pos_i,
'type': "block_output"
})
return pd.DataFrame(data)
if __name__ == "__main__":
parser = argparse.ArgumentParser(
prog='Run intervention tests for subject-verb agreement on hop models',
description='Run interventions for subject-verb agreement on hop models')
parser.add_argument('perturbation_type',
default='all',
const='all',
nargs='?',
choices=PERTURBATIONS.keys(),
help='Perturbation function used to transform BabyLM dataset')
parser.add_argument('train_set',
default='all',
const='all',
nargs='?',
choices=["100M", "10M"],
help='BabyLM train set')
parser.add_argument('random_seed', type=int, help="Random seed")
parser.add_argument('paren_model',
default='all',
const='all',
nargs='?',
choices=list(PAREN_MODELS.keys()) + ["randinit"],
help='Parenthesis model')
parser.add_argument('-np', '--no_pos_encodings', action='store_true',
help="Train GPT-2 with no positional encodings")
# Get args
args = parser.parse_args()
if "hop" not in args.perturbation_type:
raise Exception(
"'{args.perturbation_type}' is not a valid hop perturbation")
# Get examples to run interventions
data_df = pd.read_csv("hop_agreement_data.csv")
bases = [[int(tok) for tok in seq.split()]
for seq in list(data_df["Singular"])]
sources = [[int(tok) for tok in seq.split()]
for seq in list(data_df["Plural"])]
# Only get first three tokens of each example for control model
if args.perturbation_type == "hop_control":
bases = [row[:3] for row in bases]
sources = [row[:3] for row in sources]
# Get model and run intervention experiments
device = "cuda"
result_df = None
for ckpt in CHECKPOINTS:
print(f"Checkpoint: {ckpt}")
model = get_model(args.perturbation_type, args.train_set,
args.random_seed, args.paren_model, ckpt,
args.no_pos_encodings)
if result_df is None:
result_df = run_interventions(model, bases, sources)
result_df["ckpt"] = ckpt
else:
ckpt_df = run_interventions(model, bases, sources)
ckpt_df["ckpt"] = ckpt
result_df = pd.concat((result_df, ckpt_df), axis=0)
# Create directory for results
nps = '_no_pos_encodings' if args.no_pos_encodings else ''
result_directory = f"hop_intervention_results/{args.perturbation_type}_{args.train_set}{nps}/"
if not os.path.exists(result_directory):
os.makedirs(result_directory)
# Write results
result_df.to_csv(result_directory + f"{args.paren_model}_seed{args.random_seed}.csv", index=False)
|