wassname/discovering_latent_knowledge
1
0
Fork 0
mirror of https://github.com/wassname/discovering_latent_knowledge.git synced 2026-09-21 12:50:41 +08:00
Code Issues Packages Projects Releases Wiki Activity

wip refactoring

Browse Source
This commit is contained in:
deep1
2023-07-28 18:29:48 +08:00
parent d067ba50aa
commit 217b7735e8
10 changed files with 431 additions and 233 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mjc_notes.md
+8 -5
View File
@@ -749,7 +749,7 @@ Where was I?
03_ds_TQA... make sure
- [ ] true false is rigth,
- [ ] true false is right,
- [ ] acc
@@ -760,8 +760,11 @@ TruthfullQA:
- [ ] OK maybe my own curated set?
- [ ] Or just make a quick one to test manually...
Refactoring
- [ ] move common functs to src
- [ ] probe
- [ ] load model
- [ ] get_hidden_states
- [ ] get_choices_as_tokens
- [x] probe
- [x] load model
- [x] get_hidden_states
- [x] get_choices_as_tokens
- [ ] prompt format
- [ ] get it working :poop:
notebooks/023_mjc_distance_and_direction_loss_96%_conv.ipynb → notebooks/023_train_prob.ipynb
View File
notebooks/030_eval.ipynb → notebooks/030_eval_FIXME.ipynb
+2 -23
View File
@@ -242,13 +242,7 @@
}
],
"source": [
"def get_choices_as_tokens(choice_n = \"Negative\", choice_p = \"Positive\"):\n",
" # Note some tokenizer differentiate between \"no\", \"\\nno\", so we sometime need to add whitespace beforehand...\n",
" id_n, id_y = tokenizer(f'\\n{choice_n}', add_special_tokens=True)['input_ids'][-1], tokenizer(f'\\n{choice_p}', add_special_tokens=True)['input_ids'][-1]\n",
" assert tokenizer.decode([id_n])==choice_n\n",
" assert tokenizer.decode([id_y])==choice_p\n",
" # print(tokenizer.decode([id_y]))\n",
" return id_n, id_y"
"from src.datasets.hs import get_choices_as_tokens"
]
},
{
@@ -504,13 +498,6 @@
"net"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": []
},
{
"cell_type": "markdown",
"metadata": {},
@@ -524,15 +511,7 @@
"metadata": {},
"outputs": [],
"source": [
"def to_numpy(x):\n",
" if isinstance(x, torch.Tensor):\n",
" # note apache parquet doesn't support half https://github.com/huggingface/datasets/issues/4981\n",
" x = x.detach().cpu().float()\n",
" if x.squeeze().dim()==0:\n",
" return x.item()\n",
" return x.numpy()\n",
" else:\n",
" return x"
"from src.helpers.torch import to_numpy"
]
},
{
notebooks/03_ds.ipynb → notebooks/03_make_dataset.ipynb
+138 -203
View File
@@ -213,26 +213,11 @@
"print(config)\n",
"config.use_cache = False\n",
"tokenizer = AutoTokenizer.from_pretrained(model_repo)\n",
"model = AutoModelForCausalLM.from_pretrained(model_repo, config=config, **model_options)"
]
},
{
"cell_type": "code",
"execution_count": 3,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"None\n"
]
}
],
"source": [
"model = AutoModelForCausalLM.from_pretrained(model_repo, config=config, **model_options)\n",
"\n",
"# https://github.com/deep-diver/LLM-As-Chatbot/blob/main/models/falcon.py\n",
"print(tokenizer.pad_token_id)\n",
"if tokenizer.pad_token_id is None:\n",
" print(tokenizer.pad_token_id)\n",
" tokenizer.pad_token_id = 204 # <unk> https://github.com/deep-diver/LLM-As-Chatbot/blob/main/models/alpaca.py\n",
"tokenizer.padding_side = \"left\""
]
@@ -274,23 +259,7 @@
"N_SHOTS = 3\n",
"USE_MCDROPOUT = True\n",
"# dataset_n = 200\n",
"N = 8000 # 4000 in 4 hours\n",
"\n",
"try:\n",
" # num_layers = len(model.model.layers)\n",
" num_layers = model.config.n_layer\n",
" print(num_layers)\n",
"except AttributeError:\n",
" try:\n",
" num_layers = len(model.base_model.model.model.layers)\n",
" print(num_layers)\n",
" except:\n",
" num_layers = 10\n",
" \n",
"stride = 2\n",
"# don't take the first or last layers as they can make it to easy to leak info\n",
"extract_layers = tuple(range(2, num_layers-2, stride)) + (num_layers-2,)\n",
"extract_layers, num_layers"
"N = 8000 # 4000 in 4 hours\n"
]
},
{
@@ -310,35 +279,10 @@
}
],
"source": [
"# TODO maybe a list of tokens? Maybe the most common from the prompt?\n",
"# get the tokens for 0 and 1, we will use these later...\n",
"# note that sentancepeice tokenizers have differen't tokens for No and \\nNo.\n",
"token_n = \"Negative\"\n",
"token_y = \"Positive\"\n",
"id_n, id_y = tokenizer(f'\\n{token_n}', add_special_tokens=True)['input_ids'][-1], tokenizer(f'\\n{token_y}', add_special_tokens=True)['input_ids'][-1]\n",
"assert tokenizer.decode([id_n])==token_n\n",
"assert tokenizer.decode([id_y])==token_y\n",
"id_n, id_y"
]
},
{
"cell_type": "code",
"execution_count": 6,
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"'Positive'"
]
},
"execution_count": 6,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"tokenizer.decode([id_y])"
"choices_n = [\"Negative\"]\n",
"choices_p = [\"Positive\"]\n",
"from src.datasets.hs import get_choices_as_tokens\n",
"ids_n, ids_y = get_choices_as_tokens(choices_n, choices_p)"
]
},
{
@@ -777,25 +721,8 @@
" torch.cuda.empty_cache()\n",
" gc.collect()\n",
" \n",
"from src.datasets.dropout import enable_dropout, check_for_dropout\n",
"\n",
"def enable_dropout(model, USE_MCDROPOUT:Union[float,bool]=True):\n",
" \"\"\" Function to enable the dropout layers during test-time \"\"\"\n",
" \n",
" for m in model.modules():\n",
" if m.__class__.__name__.startswith('Dropout'):\n",
" m.train()\n",
" if USE_MCDROPOUT!=True:\n",
" m.p=USE_MCDROPOUT\n",
" # print(m)\n",
" \n",
" \n",
"def check_for_dropout(model):\n",
" for m in model.modules():\n",
" if m.__class__.__name__.startswith('Dropout'):\n",
" if m.p>0:\n",
" # print(m)\n",
" return True\n",
" return False\n",
" \n",
"clear_mem()\n",
"assert check_for_dropout(model), 'model should have dropout modules'\n",
@@ -811,87 +738,79 @@
"\n",
"\n",
" \n",
"def get_hidden_states(model, tokenizer, input_text, layers=extract_layers, truncation_length=999, output_attentions=False, use_mcdropout=USE_MCDROPOUT):\n",
" \"\"\"\n",
" Given a decoder model and some texts, gets the hidden states (in a given layer) on that input texts\n",
" \"\"\"\n",
" if not isinstance(input_text, list):\n",
" input_text = [input_text]\n",
" input_ids = tokenizer(input_text, \n",
" return_tensors=\"pt\",\n",
" padding=True,\n",
" add_special_tokens=True,\n",
" ).input_ids.to(model.device)\n",
"# def get_hidden_states(model, tokenizer, input_text, layers=extract_layers, truncation_length=999, output_attentions=False, use_mcdropout=USE_MCDROPOUT):\n",
"# \"\"\"\n",
"# Given a decoder model and some texts, gets the hidden states (in a given layer) on that input texts\n",
"# \"\"\"\n",
"# if not isinstance(input_text, list):\n",
"# input_text = [input_text]\n",
"# input_ids = tokenizer(input_text, \n",
"# return_tensors=\"pt\",\n",
"# padding=True,\n",
"# add_special_tokens=True,\n",
"# ).input_ids.to(model.device)\n",
" \n",
" # if add_bos_token:\n",
" # input_ids = input_ids[:, 1:]\n",
"# # if add_bos_token:\n",
"# # input_ids = input_ids[:, 1:]\n",
" \n",
" # Handling truncation: truncate start, not end\n",
" if truncation_length is not None:\n",
" if input_ids.size(1)>truncation_length:\n",
" print('truncating', input_ids.size(1))\n",
" input_ids = input_ids[:, -truncation_length:]\n",
"# # Handling truncation: truncate start, not end\n",
"# if truncation_length is not None:\n",
"# if input_ids.size(1)>truncation_length:\n",
"# print('truncating', input_ids.size(1))\n",
"# input_ids = input_ids[:, -truncation_length:]\n",
"\n",
" # forward pass\n",
" last_token = -1\n",
" first_token = 0\n",
" with torch.no_grad():\n",
" model.eval() \n",
" if use_mcdropout: enable_dropout(model, use_mcdropout)\n",
"# # forward pass\n",
"# last_token = -1\n",
"# first_token = 0\n",
"# with torch.no_grad():\n",
"# model.eval() \n",
"# if use_mcdropout: enable_dropout(model, use_mcdropout)\n",
" \n",
" # taken from greedy_decode https://github.com/huggingface/transformers/blob/ba695c1efd55091e394eb59c90fb33ac3f9f0d41/src/transformers/generation/utils.py\n",
" logits_processor = LogitsProcessorList()\n",
" model_kwargs = dict(use_cache=False)\n",
" model_inputs = model.prepare_inputs_for_generation(input_ids, **model_kwargs)\n",
" outputs = model.forward(**model_inputs, return_dict=True, output_attentions=output_attentions, output_hidden_states=True)\n",
"# # taken from greedy_decode https://github.com/huggingface/transformers/blob/ba695c1efd55091e394eb59c90fb33ac3f9f0d41/src/transformers/generation/utils.py\n",
"# logits_processor = LogitsProcessorList()\n",
"# model_kwargs = dict(use_cache=False)\n",
"# model_inputs = model.prepare_inputs_for_generation(input_ids, **model_kwargs)\n",
"# outputs = model.forward(**model_inputs, return_dict=True, output_attentions=output_attentions, output_hidden_states=True)\n",
" \n",
" next_token_logits = outputs.logits[:, last_token, :]\n",
" outputs['scores'] = logits_processor(input_ids, next_token_logits)[:, None,:]\n",
"# next_token_logits = outputs.logits[:, last_token, :]\n",
"# outputs['scores'] = logits_processor(input_ids, next_token_logits)[:, None,:]\n",
" \n",
" next_tokens = torch.argmax(outputs['scores'], dim=-1)\n",
" outputs['sequences'] = torch.cat([input_ids, next_tokens], dim=-1)\n",
"# next_tokens = torch.argmax(outputs['scores'], dim=-1)\n",
"# outputs['sequences'] = torch.cat([input_ids, next_tokens], dim=-1)\n",
"\n",
" # the output is large, so we will just select what we want 1) the first token with[:, 0]\n",
" # 2) selected layers with [layers]\n",
" attentions = None\n",
" if output_attentions:\n",
" # shape is [(batch_size, num_heads, sequence_length, sequence_length)]*num_layers\n",
" # lets take max?\n",
" attentions = [outputs['attentions'][i] for i in layers]\n",
" attentions = [v[:, last_token] for v in attentions]\n",
" attentions = torch.concat(attentions)\n",
"# # the output is large, so we will just select what we want 1) the first token with[:, 0]\n",
"# # 2) selected layers with [layers]\n",
"# attentions = None\n",
"# if output_attentions:\n",
"# # shape is [(batch_size, num_heads, sequence_length, sequence_length)]*num_layers\n",
"# # lets take max?\n",
"# attentions = [outputs['attentions'][i] for i in layers]\n",
"# attentions = [v[:, last_token] for v in attentions]\n",
"# attentions = torch.concat(attentions)\n",
" \n",
" hidden_states = torch.stack([outputs['hidden_states'][i] for i in layers], 1)\n",
"# hidden_states = torch.stack([outputs['hidden_states'][i] for i in layers], 1)\n",
" \n",
" hidden_states = hidden_states[:, :, last_token] # (batch, layers, past_seq, logits) take just the last token so they are same size\n",
"# hidden_states = hidden_states[:, :, last_token] # (batch, layers, past_seq, logits) take just the last token so they are same size\n",
" \n",
" input_truncated = tokenizer.batch_decode(input_ids)\n",
"# input_truncated = tokenizer.batch_decode(input_ids)\n",
" \n",
" s = outputs['sequences']\n",
" s = [s[i][len(input_ids[i]):] for i in range(len(s))]\n",
" text_ans = tokenizer.batch_decode(s)\n",
"# s = outputs['sequences']\n",
"# s = [s[i][len(input_ids[i]):] for i in range(len(s))]\n",
"# text_ans = tokenizer.batch_decode(s)\n",
"\n",
" scores = outputs['scores'][:, first_token].softmax(-1) # for first (and only) token\n",
" prob_n, prob_y = scores[:, [id_n, id_y]].T\n",
" eps = 1e-3\n",
" ans = (prob_y/(prob_n+prob_y+eps))\n",
"# scores = outputs['scores'][:, first_token].softmax(-1) # for first (and only) token\n",
"# prob_n, prob_y = scores[:, [id_n, id_y]].T\n",
"# eps = 1e-3\n",
"# ans = (prob_y/(prob_n+prob_y+eps))\n",
" \n",
" out = dict(hidden_states=hidden_states, ans=ans, text_ans=text_ans, input_truncated=input_truncated, input_id_shape=input_ids.shape,\n",
" attentions=attentions, prob_n=prob_n, prob_y=prob_y, scores=outputs['scores'][:, 0], input_text=input_text,\n",
" )\n",
" out = {k:to_numpy(v) for k,v in out.items()} \n",
" return out\n",
"# out = dict(hidden_states=hidden_states, ans=ans, text_ans=text_ans, input_truncated=input_truncated, input_id_shape=input_ids.shape,\n",
"# attentions=attentions, prob_n=prob_n, prob_y=prob_y, scores=outputs['scores'][:, 0], input_text=input_text,\n",
"# )\n",
"# out = {k:to_numpy(v) for k,v in out.items()} \n",
"# return out\n",
"\n",
"\n",
"def to_numpy(x):\n",
" if isinstance(x, torch.Tensor):\n",
" # note apache parquet doesn't support half https://github.com/huggingface/datasets/issues/4981\n",
" x = x.detach().cpu().float()\n",
" if x.squeeze().dim()==0:\n",
" return x.item()\n",
" return x.numpy()\n",
" else:\n",
" return x"
"from src.helpers.torch import to_numpy"
]
},
{
@@ -908,8 +827,19 @@
"metadata": {},
"outputs": [],
"source": [
"def md5hash(s: bytes) -> str:\n",
" return hashlib.md5(s).hexdigest()"
"\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"from src.datasets.hs import ExtractHiddenStates\n",
"from src.datasets.batch import batch_hidden_states\n",
"ehs = ExtractHiddenStates(model, tokenizer, choices_n, choices_y)\n",
"batch_hidden_states(ehs,)\n"
]
},
{
@@ -919,69 +849,71 @@
"outputs": [],
"source": [
"\n",
"def batch_hidden_states(prompt_fn=format_imdbs_multishot, model=model, tokenizer=tokenizer, data=data, n=100, batch_size=2, version_options=['lie', 'truth'], mcdropout=True):\n",
" \"\"\"\n",
" Given an encoder-decoder model, a list of data, computes the contrast hidden states on n random examples.\n",
" Returns numpy arrays of shape (n, hidden_dim) for each candidate label, along with a boolean numpy array of shape (n,)\n",
" with the ground truth labels\n",
"# def batch_hidden_states(prompt_fn=format_imdbs_multishot, model=model, tokenizer=tokenizer, data=data, n=100, batch_size=2, version_options=['lie', 'truth'], mcdropout=True):\n",
"# \"\"\"\n",
"# Given an encoder-decoder model, a list of data, computes the contrast hidden states on n random examples.\n",
"# Returns numpy arrays of shape (n, hidden_dim) for each candidate label, along with a boolean numpy array of shape (n,)\n",
"# with the ground truth labels\n",
" \n",
" This is deliberately simple so that it's easy to understand, rather than being optimized for efficiency\n",
" \"\"\"\n",
" # setup\n",
" model.eval()\n",
"# This is deliberately simple so that it's easy to understand, rather than being optimized for efficiency\n",
"# \"\"\"\n",
" \n",
" ds_subset = data.shuffle(seed=42).select(range(n))\n",
" dl = DataLoader(ds_subset, batch_size=batch_size, shuffle=True)\n",
" for i, batch in enumerate(tqdm(dl, desc='get hidden states')):\n",
" titles, contents, true_labels = batch[\"title\"], batch[\"content\"], batch[\"label\"]\n",
" texts = [format_review(t, c) for t,c in zip(titles, contents)]\n",
" nn = len(texts)\n",
" index = i*batch_size+np.arange(nn)\n",
" for version in version_options:\n",
" versions = [version]*nn\n",
" q, info = prompt_fn(texts, answers=true_labels, versions=versions)\n",
" if i==0:\n",
" assert len(texts)==len(prompt_fn(texts)[0]), 'make sure the prompt function can handle a list of text'\n",
" \n",
"# # setup\n",
"# model.eval()\n",
" \n",
"# ds_subset = data.shuffle(seed=42).select(range(n))\n",
"# dl = DataLoader(ds_subset, batch_size=batch_size, shuffle=True)\n",
"# for i, batch in enumerate(tqdm(dl, desc='get hidden states')):\n",
"# titles, contents, true_labels = batch[\"title\"], batch[\"content\"], batch[\"label\"]\n",
"# texts = [format_review(t, c) for t,c in zip(titles, contents)]\n",
"# nn = len(texts)\n",
"# index = i*batch_size+np.arange(nn)\n",
"# for version in version_options:\n",
"# versions = [version]*nn\n",
"# q, info = prompt_fn(texts, answers=true_labels, versions=versions)\n",
"# if i==0:\n",
"# assert len(texts)==len(prompt_fn(texts)[0]), 'make sure the prompt function can handle a list of text'\n",
" \n",
" # different due to dropout\n",
" # set_seeds(i*10)\n",
" hs1 = get_hidden_states(model, tokenizer, q, use_mcdropout=mcdropout)\n",
" # set_seeds(i*10+1)\n",
" if mcdropout:\n",
" hs2 = get_hidden_states(model, tokenizer, q, use_mcdropout=mcdropout)\n",
"# # different due to dropout\n",
"# # set_seeds(i*10)\n",
"# hs1 = ehs.get_hidden_states(q, use_mcdropout=mcdropout)\n",
"# # set_seeds(i*10+1)\n",
"# if mcdropout:\n",
"# hs2 = ehs.get_hidden_states(q, use_mcdropout=mcdropout)\n",
" \n",
" # QC\n",
" if i==0:\n",
" eps=1e-5\n",
" mpe = lambda x,y: np.mean(np.abs(x-y)/(np.abs(x)+np.abs(y)+eps))\n",
" a,b=hs2['hidden_states'],hs1['hidden_states']\n",
" assert mpe(a,b)>eps, \"the hidden state pairs should be different but are not. Check model.config.use_cache==False, check this model has dropout in it's arch\"\n",
"# # QC\n",
"# if i==0:\n",
"# eps=1e-5\n",
"# mpe = lambda x,y: np.mean(np.abs(x-y)/(np.abs(x)+np.abs(y)+eps))\n",
"# a,b=hs2['hidden_states'],hs1['hidden_states']\n",
"# assert mpe(a,b)>eps, \"the hidden state pairs should be different but are not. Check model.config.use_cache==False, check this model has dropout in it's arch\"\n",
" \n",
" assert ((hs1['prob_y']+hs1['prob_n'])>0.5).all(), \"your chosen binary answers should take up a lot of the prob space, otherwise choose differen't tokens\"\n",
" else:\n",
" hs2 = hs1\n",
"# assert ((hs1['prob_y']+hs1['prob_n'])>0.5).all(), \"your chosen binary answers should take up a lot of the prob space, otherwise choose differen't tokens\"\n",
"# else:\n",
"# hs2 = hs1\n",
"\n",
"\n",
" for j in range(nn):\n",
" yield dict(\n",
" hs1=hs1['hidden_states'][j],\n",
" ans1=hs1[\"ans\"][j],\n",
"# for j in range(nn):\n",
"# yield dict(\n",
"# hs1=hs1['hidden_states'][j],\n",
"# ans1=hs1[\"ans\"][j],\n",
" \n",
" hs2=hs2['hidden_states'][j],\n",
" ans2=hs2[\"ans\"][j], \n",
"# hs2=hs2['hidden_states'][j],\n",
"# ans2=hs2[\"ans\"][j], \n",
" \n",
" true=true_labels[j].item(),\n",
" index=index[j],\n",
" version=version,\n",
" info=info[j],\n",
"# true=true_labels[j].item(),\n",
"# index=index[j],\n",
"# version=version,\n",
"# info=info[j],\n",
" \n",
" # optional/debug\n",
" input_truncated=hs1['input_truncated'][j], # the question after truncating\n",
" prob_y=hs1['prob_y'][j],\n",
" prob_n=hs1['prob_n'][j],\n",
" text_ans = hs1['text_ans'][j],\n",
" input_text=hs1['input_text'][j],\n",
" )"
"# # optional/debug\n",
"# input_truncated=hs1['input_truncated'][j], # the question after truncating\n",
"# prob_y=hs1['prob_y'][j],\n",
"# prob_n=hs1['prob_n'][j],\n",
"# text_ans = hs1['text_ans'][j],\n",
"# input_text=hs1['input_text'][j],\n",
"# )"
]
},
{
@@ -1378,6 +1310,9 @@
}
],
"source": [
"def md5hash(s: bytes) -> str:\n",
" return hashlib.md5(s).hexdigest()\n",
"\n",
"# unique hash\n",
"def get_unique_config_name(prompt_fn, model, tokenizer, data, N):\n",
" \"\"\"\n",
src/datasets/batch.py
+67
View File
@@ -0,0 +1,67 @@
from tqdm.auto import tqdm
from src.datasets.hs import ExtractHiddenStates
from torch.utils.data import DataLoader
import numpy as np
def batch_hidden_states(ehs: ExtractHiddenStates, prompt_fn=format_imdbs_multishot, data=data, n=100, batch_size=2, version_options=['lie', 'truth'], mcdropout=True):
"""
Given an encoder-decoder model, a list of data, computes the contrast hidden states on n random examples.
Returns numpy arrays of shape (n, hidden_dim) for each candidate label, along with a boolean numpy array of shape (n,)
with the ground truth labels
This is deliberately simple so that it's easy to understand, rather than being optimized for efficiency
"""
ds_subset = data.shuffle(seed=42).select(range(n))
dl = DataLoader(ds_subset, batch_size=batch_size, shuffle=True)
for i, batch in enumerate(tqdm(dl, desc='get hidden states')):
titles, contents, true_labels = batch["title"], batch["content"], batch["label"]
texts = [format_review(t, c) for t,c in zip(titles, contents)]
nn = len(texts)
index = i*batch_size+np.arange(nn)
for version in version_options:
versions = [version]*nn
q, info = prompt_fn(texts, answers=true_labels, versions=versions)
if i==0:
assert len(texts)==len(prompt_fn(texts)[0]), 'make sure the prompt function can handle a list of text'
# different due to dropout
# set_seeds(i*10)
hs1 = ehs.get_hidden_states(q, use_mcdropout=mcdropout)
# set_seeds(i*10+1)
if mcdropout:
hs2 = ehs.get_hidden_states(q, use_mcdropout=mcdropout)
# QC
if i==0:
eps=1e-5
mpe = lambda x,y: np.mean(np.abs(x-y)/(np.abs(x)+np.abs(y)+eps))
a,b=hs2['hidden_states'],hs1['hidden_states']
assert mpe(a,b)>eps, "the hidden state pairs should be different but are not. Check model.config.use_cache==False, check this model has dropout in it's arch"
assert ((hs1['prob_y']+hs1['prob_n'])>0.5).all(), "your chosen binary answers should take up a lot of the prob space, otherwise choose differen't tokens"
else:
hs2 = hs1
for j in range(nn):
yield dict(
hs1=hs1['hidden_states'][j],
ans1=hs1["ans"][j],
hs2=hs2['hidden_states'][j],
ans2=hs2["ans"][j],
true=true_labels[j].item(),
index=index[j],
version=version,
info=info[j],
# optional/debug
input_truncated=hs1['input_truncated'][j], # the question after truncating
prob_y=hs1['prob_y'][j],
prob_n=hs1['prob_n'][j],
text_ans = hs1['text_ans'][j],
input_text=hs1['input_text'][j],
)
src/datasets/dm.py
+2 -2
View File
@@ -2,12 +2,12 @@ import torch
import torch.nn as nn
import lightning as pl
import pandas as pd
from torch.utils.data import Dataset, DataLoader
from torch.utils.data import Dataset, DataLoader, TensorDataset
from src.datasets.load import ds2df
def make_y(df):
# label: is ans2 more true than ans1
# so we ask does ans2 have greater probabiliy on "positive" than ans1
# so we ask does ans2 have greater probability on "positive" than ans1
# then, when the right answer is negative we swap the sign
true_switch_sign = df.true_answer*2-1
distance = (df.ans2-df.ans1) * true_switch_sign
src/datasets/dropout.py
+20
View File
@@ -0,0 +1,20 @@
from typing import Union
def enable_dropout(model, USE_MCDROPOUT:Union[float,bool]=True):
""" Function to enable the dropout layers during test-time """
for m in model.modules():
if m.__class__.__name__.startswith('Dropout'):
m.train()
if USE_MCDROPOUT!=True:
m.p=USE_MCDROPOUT
def check_for_dropout(model, verbose=False):
for m in model.modules():
if m.__class__.__name__.startswith('Dropout'):
if m.p>0:
if verbose: print(m)
return True
return False
src/datasets/hs.py
+177
View File
@@ -0,0 +1,177 @@
from dataclasses import dataclass
import lightning as pl
import torch
from transformers import (
AutoTokenizer,
AutoModelForSeq2SeqLM,
AutoModelForMaskedLM,
AutoModelForCausalLM,
AutoConfig,
AutoModel,
PreTrainedTokenizer,
PreTrainedModel
)
from typing import Optional, List, Tuple
from transformers import LogitsProcessorList
from src.helpers.torch import to_numpy
from src.datasets.dropout import enable_dropout
def get_choices_as_tokens(
tokenizer, choice_n: List[str] = ["Negative"], choice_p: List[str] = ["Positive"]
) -> Tuple[List[int], List[int]]:
# Note some tokenizer differentiate between "no", "\nno", so we sometime need to add whitespace beforehand...
ids_n = []
for c in choice_n:
id_ = tokenizer(f"\n{c}", add_special_tokens=True)["input_ids"][-1]
ids_n.append(id_)
assert tokenizer.decode([id_]) == c
ids_y = []
for c in choice_n:
id_ = tokenizer(f"\n{c}", add_special_tokens=True)["input_ids"][-1]
ids_y.append(id_)
assert tokenizer.decode([id_]) == c
return ids_n, ids_y
@dataclass
class ExtractHiddenStates:
model: PreTrainedModel
tokenizer: PreTrainedTokenizer
layer_stride: int = 1
layer_padding: int = 2
truncation_length = 999
choices_n: List[str] = ["No"]
choices_p: List[str] = ["Yes"]
def start(self):
self.ids_n, self.ids_y = get_choices_as_tokens(
self.tokenizer, self.choices_n, self.choices_p
)
def get_hidden_states(
self,
input_text,
truncation_length=999,
output_attentions=False,
use_mcdropout=True,
):
"""
Given a decoder model and some texts, gets the hidden states (in a given layer) on that input texts
"""
if not isinstance(input_text, list):
input_text = [input_text]
input_ids = self.tokenizer(
input_text,
return_tensors="pt",
padding=True,
add_special_tokens=True,
).input_ids.to(self.model.device)
# Handling truncation: truncate start, not end
if truncation_length is not None:
if input_ids.size(1) > truncation_length:
print("truncating", input_ids.size(1))
input_ids = input_ids[:, -truncation_length:]
# forward pass
last_token = -1
first_token = 0
with torch.no_grad():
self.model.eval()
if use_mcdropout:
enable_dropout(self.model, use_mcdropout)
# taken from greedy_decode https://github.com/huggingface/transformers/blob/ba695c1efd55091e394eb59c90fb33ac3f9f0d41/src/transformers/generation/utils.py
logits_processor = LogitsProcessorList()
model_kwargs = dict(use_cache=False)
model_inputs = self.model.prepare_inputs_for_generation(
input_ids, **model_kwargs
)
outputs = self.model.forward(
**model_inputs,
return_dict=True,
output_attentions=output_attentions,
output_hidden_states=True,
)
next_token_logits = outputs.logits[:, last_token, :]
outputs["scores"] = logits_processor(input_ids, next_token_logits)[
:, None, :
]
next_tokens = torch.argmax(outputs["scores"], dim=-1)
outputs["sequences"] = torch.cat([input_ids, next_tokens], dim=-1)
# the output is large, so we will just select what we want 1) the first token with[:, 0]
# 2) selected layers with [layers]
layers = self.get_layer_selection(outputs)
attentions = None
layers = range(
self.layer_padding,
len(outputs["attentions"]) - self.layer_padding,
self.layer_stride,
)
if output_attentions:
# shape is [(batch_size, num_heads, sequence_length, sequence_length)]*num_layers
# lets take max?
attentions = [outputs["attentions"][i] for i in layers]
attentions = [v[:, last_token] for v in attentions]
attentions = torch.concat(attentions)
hidden_states = torch.stack(
[outputs["hidden_states"][i] for i in layers], 1
)
hidden_states = hidden_states[
:, :, last_token
] # (batch, layers, past_seq, logits) take just the last token so they are same size
input_truncated = self.tokenizer.batch_decode(input_ids)
s = outputs["sequences"]
s = [s[i][len(input_ids[i]) :] for i in range(len(s))]
text_ans = self.tokenizer.batch_decode(s)
scores = outputs["scores"][:, first_token].softmax(
-1
) # for first (and only) token
# prob_n, prob_y = scores[:, [id_n, id_y]].T
prob_n = scores[:, self.ids_n]
prob_y = scores[:, self.ids_y]
eps = 1e-3
ans = (prob_y / (prob_n + prob_y + eps)).sum(1)
out = dict(
hidden_states=hidden_states,
ans=ans,
text_ans=text_ans,
input_truncated=input_truncated,
input_id_shape=input_ids.shape,
attentions=attentions,
prob_n=prob_n,
prob_y=prob_y,
scores=outputs["scores"][:, 0],
input_text=input_text,
)
out = {k: to_numpy(v) for k, v in out.items()}
return out
def get_layer_selection(self, outputs):
"""Sometimes we don't want to save all layers.
Typically we can skip some to save space (stride). We might also want to ignore the first and last ones (padding) to avoid data leakage.
See https://www.lesswrong.com/posts/bWxNPMy5MhPnQTzKz/what-discovering-latent-knowledge-did-and-did-not-find-4
"""
return range(
self.layer_padding,
len(outputs["attentions"]) - self.layer_padding,
self.layer_stride,
)
src/datasets/load.py
+3
View File
@@ -1,3 +1,6 @@
import numpy as np
import pandas as pd
def rows_item(row):
"""
transform a row by turning singe dim arrays into items
src/helpers/torch.py
+14
View File
@@ -0,0 +1,14 @@
import torch
def to_numpy(x):
"""
Trys to convert torch to numpy and if possible a single item
"""
if isinstance(x, torch.Tensor):
# note apache parquet doesn't support half https://github.com/huggingface/datasets/issues/4981
x = x.detach().cpu().float()
if x.squeeze().dim()==0:
return x.item()
return x.numpy()
else:
return x