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

wip refactoring

Browse Source
This commit is contained in:
deep1
2023-08-05 08:55:55 +08:00
parent 217b7735e8
commit 8cb669f121
14 changed files with 1092 additions and 2233 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.gitignore
+1
View File
@@ -2,6 +2,7 @@ lightning_logs/
.pkl_cache/
.ds/
/notebooks/old/
*.pyc
# Distribution / packaging
.Python
README.md
+111 -12
View File
@@ -1,20 +1,119 @@
My own experiments with DLK
# LLM truth detector using Monte Carlo Dropout
Sometimes the best way to explain is with code:
```py
"""
pseudocode for a LLM truth detector using Monte Carlo Dropout
"""
# load model
model, tokenizer = load_model()
# make dataset of hidden state pairs
prompts = ["Is this true: a broken mirror gives 7 years bad luck [Yes/No]: ",
"Is this true: a broken mirror doesn't give 7 years bad luck [Yes/No]: "]
choice = ['Yes']
choice_is_true = [-1, 1]
def get_hidden_state_pairs(prompts, choice, choice_is_true, tokenizer, model):
"""
We turn on dropout and predict the next token (repeat x2). Since dropout is turned on each prediction is slightly different.
Then we collect the hidden state pairs (as x1, x2) and the scores of our target token (as y1, y2)
"""
choice_tokens = choice2token(choice, tokenizer)
# we enable dropout, and do 2 inferences that are slightly different
enable_mcdropout(model)
outputs1 = model.generate(prompts, output_hidden_states=True)
y1 = outputs1['scores'][choice_tokens]
x1 = outputs1["hidden_states"]
outputs2 = model.generate(prompts, output_hidden_states=True)
y2 = outputs2['scores'][choice_tokens]
x2 = outputs2["hidden_states"]
return x1, x2, y1, y2, choice_is_true
dataset = batched(get_hidden_state_pairs(prompts, choice, choice_is_true, model, tokenizer))
# now train a probe
net = Probe(layers=2, hs=32)
optim = Optim(lr=3e-4)
for x1, x2, y1, y2, choice_is_true in dl:
y_pred1 = net(x1)
y_pred2 = net(x2)
y_pred = y_pred2-ypred1
# our label is the distance between the two probabilities in the direction of truth
# So if y2 is less true than y1, and they are 0.02% apart then y is -0.02%
y = (y2-y1)*choice_is_true
# Use a MSE loss so that the distance between the predicted pair of scores (in the direction of truth)
# is the same as the real pair of scores (in the direction of truth)
loss = F.mse(y_pred, y)
net.backwards()
optim.step()
# Test the probe
prompts = ["Is this true: Ancients did not believe the world was flat [Yes/No]: ",
"Is this true: Step on a crack break your fathers back [Yes/No]: "]
choice = ['Yes']
choice_is_true = [1, -1]
x1, x2, y1, y2, choice_is_true = get_hidden_state_pairs(prompts, choice, choice_is_true, model, tokenizer)
y_pred1 = net(x1)
y_pred2 = net(x2)
# translate this into a truth detector....
pred_last_choice_is_true = y / (y_pred2-y_pred)
pred_last_choice_is_true # [1, -1]
```
# Description
There is some previous work on this ([ELK](https://github.com/EleutherAI/elk), [DLK](https://github.com/collin-burns/discovering_latent_knowledge/blob/main/CCS.ipynb), CSS, etc) that all take varias approaches. They have this in common:
- Show the model 2 statements “the sky is blue” “the sky is green”
- Get the hidden states from reading those statements
- Use machine learning learning to distinguish between those two sets
Now this works well [(or not?)](https://www.lesswrong.com/posts/bWxNPMy5MhPnQTzKz/what-discovering-latent-knowledge-did-and-did-not-find-4), but I aim for two improvements:
- Detect direction of deception instead of truth
- look at deceptive actions (outputs), not deceptive observations (inputs).
- Use Monte Carlo dropout to generate pair of hidden states, instead pairs of inputs
My contributions/finds so far:
- Instead of comparing hidden states from 2 prompts, you can compare two inferences of the same prompt as long as you have dropout on
- For this pair of hidden states, one will be in the direction of truth and one will not
- But the pairs must give >10% differen't answer on our compared tokens e.g. true vs false
- We can detect this using a supervised probe (with 90% acc on IMBD sentiment analysis)
- The best approach to setting up the probe is ~~binary classification~~, ~~multiclass classification~~ ~~ranking with margin_ranking_loss~~ ranking with L1smoothloss
- This is because treating it like a ranking problem decreases overfitting
- And learning distance and direction between the ranked pairs gives more supervision than just the direction (like in many ranking setups)
- It's hard to get models to lie! Even for uncensored models. I find uncensored coding models are best
## TODO:
I'm trying to
- [x] use pytorch lightning
- [x] batch hidden states 5x faster
- [x] use wizcoer 15B, to see if larger models give better results
- [x] eval on some deceptive or misleading statements
- [ ] debug by looking at model output
- [ ] test generalization
- [ ] try differen't approaches
- [ ] setup
- [ ] detect deception vs truth
- [ ] differen't prompts
- [ ] differen't tasks
- [ ] model arch
- [ ] put in both states
- [ ] normalize states
- [ ] mix states at end
- [x] debug by looking at model output
- [x] test generalization
- [x] try differen't approaches
- [x] setup
- [x] detect deception vs truth
- [x] differen't prompts
- [x] differen't tasks
- [x] model arch
- [x] put in both states
- [x] normalize states
- [x] mix states at end
-------------
mjc_notes.md
+98
View File
@@ -741,6 +741,70 @@ exp
- no true switch... wait why did I switch it.. .weight
- wait what 93% baseline wat?? oh wait we are just detecting the word positive lol! ignore this
# Refactoring - start with Pseudo code
```py
# load model
model, tokenizer = load_model()
# make dataset of hidden state pairs
prompts = ["a broken mirror gives 7 years bad luck: ", "a broken mirror doesn't give 7 years bad luck: "]
# TODO do I just need one
choices = [['No'], ['Yes']]
last_choice_is_true = [-1, 1]
def get_hidden_state_pairs(prompts, choices, last_choice_is_true, tokenizer, model):
"""
We turn on dropout and predict the next token twice. Since dropout is on they are slightly different. Then we collect the hidden state pairs (x1, x2) and the probability of our target token (y1, y2)
"""
choice_tokens = choice2token(choices, tokenizer)
# we enable dropout, and do 2 inferences that are slightly different
enable_mcdropout(model)
outputs1 = model.generate(prompts, output_hidden_states=True)
y1 = outputs1['scores'][choice_tokens]
x1 = outputs1["hidden_states"]
outputs2 = model.generate(prompts, output_hidden_states=True)
y2 = outputs2['scores'][choice_tokens]
x2 = outputs2["hidden_states"]
return x1, x2, y1, y2, last_choice_is_true
dataset = batch(get_hidden_state_pairs(prompts, choices, last_choice_is_true, model, tokenizer))
# now train a probe
net = Probe(layers=2, hs=32)
optim = Optim(lr=3e-4)
for x1, x2, y1, y2, last_choice_is_true in dl:
y_pred1 = net(x1)
y_pred2 = net(x2)
y_pred = y_pred2-ypred1
# our label is the distance between the two probabilities in the direction of truth
# So if y2 is less true than y1, and they are 0.02% apart then y is -0.02%
y = (y2-y1)*last_choice_is_true
# Use a MSE loss to that the distance between the predicted pair of scores (in the direction of truth) is the same as the pair of scores (in the direction of truth)
loss = F.mse(y_pred, y)
net.backwards()
optim.step()
# now use the probe
y_pred1 = net(x1)
y_pred2 = net(x2)
# translate this into a truth detector....
pred_last_choice_is_true = y / (y_pred2-y_pred1)
pred_last_choice_is_true
```
TODO
- refactor to look like the psudocode
# 2023-07-23 19:50:10
Where was I?
@@ -768,3 +832,37 @@ Refactoring
- [x] get_choices_as_tokens
- [ ] prompt format
- [ ] get it working :poop:
- [ ] dataset
- [ ] model
So wait do I need to just record scores
Now how does this all relate to truth and the prompt
So we are measuring if a particular token, that is could have answered with is true... but the model doesn't know which one!!!
So it seems like there is some experimentation needed here. I should just save scores which will give me optionality.
But really I should be looking at the most likely token right? No need for a choice?
All I need to so is decide if this hidden state is more true.
But if I chose an unlikely answer it seems misleading?
Maybe I should be looking at hidden state condictional on a token. But how to do that?
Well I'm really trying to tell if the most likely answer is true. So I just need to work out if the most likely answer is true using the labels. Then I can order the hidden states.
# 2023-08-05 07:09:39
TODO
- [ ] add info or similar
- [x] ans
- [ ] choices
- [ ] do checks
- [ ] for high prob
- [ ] and acc
- [ ] name ds
- [ ] save ds
- [ ] get model nb working
Got unsupported ScalarType BFloat16
But that's because we try to numpy it
notebooks/030_eval_FIXME.ipynb
-74
View File
@@ -734,80 +734,6 @@
"source": [
"What is the probe predicting? Whether hs1 is more true than hs0"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": []
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": []
},
{
"cell_type": "code",
"execution_count": 213,
"metadata": {},
"outputs": [
{
"data": {
"text/html": [
"<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #800000; text-decoration-color: #800000\">╭─────────────────────────────── </span><span style=\"color: #800000; text-decoration-color: #800000; font-weight: bold\">Traceback </span><span style=\"color: #bf7f7f; text-decoration-color: #bf7f7f; font-weight: bold\">(most recent call last)</span><span style=\"color: #800000; text-decoration-color: #800000\"> ────────────────────────────────╮</span>\n",
"<span style=\"color: #800000; text-decoration-color: #800000\">│</span> in <span style=\"color: #00ff00; text-decoration-color: #00ff00\">&lt;module&gt;</span>:<span style=\"color: #0000ff; text-decoration-color: #0000ff\">2</span> <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
"<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
"<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">1 </span><span style=\"color: #00ffff; text-decoration-color: #00ffff\">print</span>(<span style=\"color: #808000; text-decoration-color: #808000\">f\"\"\"</span> <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
"<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #800000; text-decoration-color: #800000\">❱ </span>2 <span style=\"color: #808000; text-decoration-color: #808000\">Model says: {</span>hs0[<span style=\"color: #808000; text-decoration-color: #808000\">'text_ans'</span>][<span style=\"color: #0000ff; text-decoration-color: #0000ff\">0</span>]<span style=\"color: #808000; text-decoration-color: #808000\">} {</span>hs1[<span style=\"color: #808000; text-decoration-color: #808000\">'text_ans'</span>][<span style=\"color: #0000ff; text-decoration-color: #0000ff\">0</span>]<span style=\"color: #808000; text-decoration-color: #808000\">} prob_y={</span>hs1[<span style=\"color: #808000; text-decoration-color: #808000\">'prob_y'</span>]<span style=\"color: #808000; text-decoration-color: #808000\">:2.2f} prob_n</span> <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
"<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">3 </span><span style=\"color: #808000; text-decoration-color: #808000\">Probe says: {</span>y_pred.squeeze()<span style=\"color: #808000; text-decoration-color: #808000\">:2.4f} (hs2 is more true)</span> <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
"<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">4 </span><span style=\"color: #808000; text-decoration-color: #808000\">where</span> <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
"<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">5 </span><span style=\"color: #bfbf7f; text-decoration-color: #bfbf7f\">│ </span><span style=\"color: #808000; text-decoration-color: #808000\">hs2_more_positive={</span>hs2_more_positive<span style=\"color: #808000; text-decoration-color: #808000\">}</span> <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
"<span style=\"color: #800000; text-decoration-color: #800000\">╰──────────────────────────────────────────────────────────────────────────────────────────────────╯</span>\n",
"<span style=\"color: #ff0000; text-decoration-color: #ff0000; font-weight: bold\">NameError: </span>name <span style=\"color: #008000; text-decoration-color: #008000\">'hs0'</span> is not defined\n",
"</pre>\n"
],
"text/plain": [
"\u001b[31m╭─\u001b[0m\u001b[31m──────────────────────────────\u001b[0m\u001b[31m \u001b[0m\u001b[1;31mTraceback \u001b[0m\u001b[1;2;31m(most recent call last)\u001b[0m\u001b[31m \u001b[0m\u001b[31m───────────────────────────────\u001b[0m\u001b[31m─╮\u001b[0m\n",
"\u001b[31m│\u001b[0m in \u001b[92m<module>\u001b[0m:\u001b[94m2\u001b[0m \u001b[31m│\u001b[0m\n",
"\u001b[31m│\u001b[0m \u001b[31m│\u001b[0m\n",
"\u001b[31m│\u001b[0m \u001b[2m1 \u001b[0m\u001b[96mprint\u001b[0m(\u001b[33mf\u001b[0m\u001b[33m\"\"\"\u001b[0m \u001b[31m│\u001b[0m\n",
"\u001b[31m│\u001b[0m \u001b[31m❱ \u001b[0m2 \u001b[33mModel says: \u001b[0m\u001b[33m{\u001b[0mhs0[\u001b[33m'\u001b[0m\u001b[33mtext_ans\u001b[0m\u001b[33m'\u001b[0m][\u001b[94m0\u001b[0m]\u001b[33m}\u001b[0m\u001b[33m \u001b[0m\u001b[33m{\u001b[0mhs1[\u001b[33m'\u001b[0m\u001b[33mtext_ans\u001b[0m\u001b[33m'\u001b[0m][\u001b[94m0\u001b[0m]\u001b[33m}\u001b[0m\u001b[33m prob_y=\u001b[0m\u001b[33m{\u001b[0mhs1[\u001b[33m'\u001b[0m\u001b[33mprob_y\u001b[0m\u001b[33m'\u001b[0m]\u001b[33m:\u001b[0m\u001b[33m2.2f\u001b[0m\u001b[33m}\u001b[0m\u001b[33m prob_n\u001b[0m \u001b[31m│\u001b[0m\n",
"\u001b[31m│\u001b[0m \u001b[2m3 \u001b[0m\u001b[33mProbe says: \u001b[0m\u001b[33m{\u001b[0my_pred.squeeze()\u001b[33m:\u001b[0m\u001b[33m2.4f\u001b[0m\u001b[33m}\u001b[0m\u001b[33m (hs2 is more true)\u001b[0m \u001b[31m│\u001b[0m\n",
"\u001b[31m│\u001b[0m \u001b[2m4 \u001b[0m\u001b[33mwhere\u001b[0m \u001b[31m│\u001b[0m\n",
"\u001b[31m│\u001b[0m \u001b[2m5 \u001b[0m\u001b[2;33m│ \u001b[0m\u001b[33mhs2_more_positive=\u001b[0m\u001b[33m{\u001b[0mhs2_more_positive\u001b[33m}\u001b[0m \u001b[31m│\u001b[0m\n",
"\u001b[31m╰──────────────────────────────────────────────────────────────────────────────────────────────────╯\u001b[0m\n",
"\u001b[1;91mNameError: \u001b[0mname \u001b[32m'hs0'\u001b[0m is not defined\n"
]
},
"metadata": {},
"output_type": "display_data"
}
],
"source": []
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": []
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": []
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": []
}
],
"metadata": {
notebooks/03_make_dataset.ipynb
+489 -2008
View File
File diff suppressed because it is too large Load Diff
setup.py
+10
View File
@@ -0,0 +1,10 @@
from setuptools import find_packages, setup
setup(
name='src',
packages=find_packages(),
version='0.1.0',
description='Discovering Latent Knowledge using MonteCarlo Dropout on outputs not inputs',
author='wassname',
license='MIT',
)
src/datasets/batch.py
+44 -48
View File
@@ -2,9 +2,10 @@
from tqdm.auto import tqdm
from src.datasets.hs import ExtractHiddenStates
from torch.utils.data import DataLoader
from datasets import Dataset
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):
def batch_hidden_states(ehs: ExtractHiddenStates, data: Dataset, n=100, batch_size=2, 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,)
@@ -13,55 +14,50 @@ def batch_hidden_states(ehs: ExtractHiddenStates, prompt_fn=format_imdbs_multish
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)
ds_t_subset = data.select(range(n))
ds_t_subset.set_format(type='torch', columns=['input_ids', 'label'])
ds_p_subset = data.select(range(n))
ds_p_subset.set_format(type="pandas", columns=['lie', 'label', 'prompt', 'prompt_truncated'])
dl = DataLoader(ds_t_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)
input_ids, true_labels = batch["input_ids"], batch["label"]
nn = len(input_ids)
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
hs1 = ehs.get_batch_of_hidden_states(input_ids=input_ids, use_mcdropout=mcdropout)
if mcdropout:
hs2 = ehs.get_batch_of_hidden_states(input_ids=input_ids, use_mcdropout=mcdropout)
# 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"
# 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
# FIXME, move check to loading?
# 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],
)
for j in range(nn):
# let's add the non torch metadata like label, prompt, lie, etc
k = i*batch_size + j
info = ds_p_subset[k]
yield dict(
hs1=hs1['hidden_states'][j],
scores1=hs1["scores"][j],
hs2=hs2['hidden_states'][j],
scores2=hs2["scores"][j],
true=true_labels[j].item(),
index=index[j],
**info
)
src/datasets/dropout.py
+3
View File
@@ -12,6 +12,9 @@ def enable_dropout(model, USE_MCDROPOUT:Union[float,bool]=True):
def check_for_dropout(model, verbose=False):
"""check if dropout is present.
dropout is sometimes present but inactive, we test that later"""
for m in model.modules():
if m.__class__.__name__.startswith('Dropout'):
if m.p>0:
src/datasets/hs.py
+128 -91
View File
@@ -1,6 +1,7 @@
from dataclasses import dataclass
import lightning as pl
import torch
from loguru import logger
from transformers import (
AutoTokenizer,
AutoModelForSeq2SeqLM,
@@ -11,153 +12,120 @@ from transformers import (
PreTrainedTokenizer,
PreTrainedModel
)
from typing import Optional, List, Tuple
from typing import Optional, List, Tuple, Dict
from transformers import LogitsProcessorList
from src.helpers.torch import to_numpy
from src.datasets.dropout import enable_dropout
from tqdm.auto import tqdm
# from src.datasets.hs import ExtractHiddenStates
from torch.utils.data import DataLoader
from datasets import Dataset
import numpy as np
default_class2choices = {False: ['No', 'Negative', 'no', 'false', 'wrong'], True: ['Yes', 'Positive', 'yes', 'true', 'correct', 'right']}
def get_choices_as_tokens(
tokenizer, choice_n: List[str] = ["Negative"], choice_p: List[str] = ["Positive"]
tokenizer, choices:List[str] = ["Positive"], whitespace_first=True
) -> 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:
# Note some tokenizers differentiate between "no", "\nno", so we sometime need to add whitespace beforehand...
if not whitespace_first:
raise NotImplementedError('TODO')
ids = []
for c in choices:
id_ = tokenizer(f"\n{c}", add_special_tokens=True)["input_ids"][-1]
ids_n.append(id_)
assert tokenizer.decode([id_]) == c
ids.append(id_)
c2 = tokenizer.decode([id_])
assert tokenizer.decode([id_]) == c, f'tokenizer.decode(tokenizer(`{c}`))==`{c2}`!=`{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
return ids
@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(
def get_batch_of_hidden_states(
self,
input_text,
input_text: Optional[List[str]] = None,
input_ids: torch.Tensor = None,
truncation_length=999,
output_attentions=False,
use_mcdropout=True,
debug=False,
):
"""
Given a decoder model and some texts, gets the hidden states (in a given layer) on that input texts
Given a decoder model and a batch of texts, gets a pair of 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:]
assert (input_ids is not None) or (input_text is not None), "need to provide input_ids or input_text"
assert self.tokenizer.truncation_side == 'left'
if input_text:
input_ids = self.tokenizer(
input_text,
return_tensors="pt",
add_special_tokens=True,
padding='max_length', max_length=truncation_length, truncation=True
).input_ids.to(self.model.device)
# forward pass
last_token = -1
first_token = 0
with torch.no_grad():
input_ids = input_ids.to(self.model.device)
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
)
# Forward for one step is the same as greedy generation for one step
# https://github.com/huggingface/transformers/blob/234cfefbb083d2614a55f6093b0badfb2efc3b45/src/transformers/generation_utils.py#L1528
outputs = self.model.forward(
**model_inputs,
input_ids,
return_dict=True,
output_attentions=output_attentions,
output_hidden_states=True,
use_cache=False,
)
next_token_logits = outputs.logits[:, last_token, :]
outputs["scores"] = logits_processor(input_ids, next_token_logits)[
:, None, :
]
outputs["scores"] = outputs.logits[:, last_token, :]
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)
attentions = [outputs["attentions"][i][:, -1] for i in layers]
attentions = torch.stack(attentions, 1)
# shape is [(batch_size, num_heads, input_length, input_length)]*num_layers
hidden_states = torch.stack(
[outputs["hidden_states"][i] for i in layers], 1
)
# (batch, layers, past_seq, logits) take just the last token so they are same size
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,
scores=outputs["scores"],
input_ids=input_ids,
)
out = {k: to_numpy(v) for k, v in out.items()}
if debug:
out['input_truncated'] = self.tokenizer.batch_decode(input_ids)
out['text_ans'] = self.tokenizer.batch_decode(outputs["scores"].argmax(-1))
return out
def get_layer_selection(self, outputs):
@@ -169,9 +137,78 @@ class ExtractHiddenStates:
"""
return range(
self.layer_padding,
len(outputs["attentions"]) - self.layer_padding,
len(outputs["hidden_states"]) - self.layer_padding,
self.layer_stride,
)
def batch_hidden_states(self, data: Dataset, n=100, batch_size=2, 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_t_subset = data.select(range(n))
ds_t_subset.set_format(type='torch', columns=['input_ids', 'label'])
ds_p_subset = data.select(range(n))
ds_p_subset.set_format(type="pandas", columns=['lie', 'label', 'prompt', 'prompt_truncated'])
dl = DataLoader(ds_t_subset, batch_size=batch_size, shuffle=True)
for i, batch in enumerate(tqdm(dl, desc='get hidden states')):
input_ids, true_labels = batch["input_ids"], batch["label"]
nn = len(input_ids)
index = i*batch_size+np.arange(nn)
# different due to dropout
hs1 = self.get_batch_of_hidden_states(input_ids=input_ids, use_mcdropout=mcdropout)
if mcdropout:
hs2 = self.get_batch_of_hidden_states(input_ids=input_ids, 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"
# FIXME, move check to loading?
# 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):
# let's add the non torch metadata like label, prompt, lie, etc
k = i*batch_size + j
info = ds_p_subset[k]
yield dict(
hs1=hs1['hidden_states'][j],
scores1=hs1["scores"][j],
hs2=hs2['hidden_states'][j],
scores2=hs2["scores"][j],
true=true_labels[j].item(),
index=index[j],
**info
)
def __getstate__(self):
"""So avoid datasets trying to pickle a model lets set a custom pickle method"""
state = self.__dict__.copy()
state['model_config'] = self.model.config
state['model_name'] = self.model.config
del state['model']
return state
def __setstate__(self):
raise NotImplementedError("You should not be pickling this class, it's too big")
src/helpers/torch.py
+21
View File
@@ -1,4 +1,8 @@
import torch
import numpy as np
import transformers
import random
import gc
def to_numpy(x):
"""
@@ -12,3 +16,20 @@ def to_numpy(x):
return x.numpy()
else:
return x
def set_seeds(n):
transformers.set_seed(n)
torch.manual_seed(n)
np.random.seed(n)
random.seed(n)
def to_item(x):
if isinstance(x, torch.Tensor):
x = x.detach().cpu().item()
return x
def clear_mem():
gc.collect()
torch.cuda.empty_cache()
gc.collect()
src/models/load.py
+78
View File
@@ -0,0 +1,78 @@
"""
This file load various open source models
When editing or updating this file check out these resources:
- [LLM-As-Chatbot](https://github.com/deep-diver/LLM-As-Chatbot/blob/main/models/falcon.py)
- [oobabooga](https://github.com/oobabooga/text-generation-webui/blob/main/modules/models.py#L134)
"""
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, AutoModelForMaskedLM, AutoModelForCausalLM, AutoConfig
import torch
from src.datasets.dropout import check_for_dropout
from loguru import logger
def verbose_change_param(tokenizer, path, after):
before = getattr(tokenizer, path)
if before!=after:
setattr(tokenizer, path, after)
logger.info(f"changing {path} from {before} to {after}")
return tokenizer
def load_model(model_repo = "HuggingFaceH4/starchat-beta", lora_repo=None, verbose=True):
if "starchat" in model_repo:
model, tokenizer = load_starchat(model_repo=model_repo)
# elif "llama" in model_repo:
# model, tokenizer = load_llama(model_repo=model_repo, lora_repo=lora_repo)
else:
raise NotImplementedError(f"model_repo {model_repo} not found")
if verbose: print(model.config)
assert check_for_dropout(model), 'model should have dropout'
return model, tokenizer
def load_starchat(model_repo = "HuggingFaceH4/starchat-beta"):
# see https://github.com/deep-diver/LLM-As-Chatbot/blob/main/models/starchat.py
model_options = dict(
device_map="auto",
load_in_4bit=True,
torch_dtype=torch.float16, # note because datasets pickles the model into numpy to get the unique datasets name, and because numpy doesn't support bfloat16, we need to use float16
use_safetensors=False,
)
config = AutoConfig.from_pretrained(model_repo, use_cache=False)
verbose_change_param(config, 'use_cache', False)
tokenizer = AutoTokenizer.from_pretrained(model_repo)
verbose_change_param(tokenizer, 'pad_token_id', 0)
verbose_change_param(tokenizer, 'padding_side', 'left')
verbose_change_param(tokenizer, 'truncation_side', 'left')
model = AutoModelForCausalLM.from_pretrained(model_repo, config=config, **model_options)
return model, tokenizer
# def load_llama(model_repo, lora_repo=None):
# # https://github.com/deep-diver/LLM-As-Chatbot/blob/main/models/alpaca.py
# model_options = dict(
# device_map="auto",
# load_in_4bit=True,
# torch_dtype=torch.float16,
# )
# tokenizer = LlamaTokenizer.from_pretrained(model_repo)
# model = LlamaForCausalLM.from_pretrained(model_repo, **model_options)
# if lora_repo is not None:
# # https://github.com/tloen/alpaca-lora/blob/main/generate.py#L40
# from peft import PeftModel
# model = PeftModel.from_pretrained(
# model,
# lora_repo,
# torch_dtype=torch.float16,
# device_map='auto'
# )
# return model, tokenizer
# def load_falcan(model_repo, lora_repo=None):
# # https://github.com/deep-diver/LLM-As-Chatbot/blob/main/models/falcon.py
src/prompts/__init__.py
View File
src/prompts/format.py
+70
View File
@@ -0,0 +1,70 @@
from typing import Callable, List, Optional, Tuple
def format_prompt(
system: str = "", query: str = "", response: str="", format: str = "chatml"
) -> str:
if format == "chatml":
# https://huggingface.co/HuggingFaceH4/starchat-beta
return (
f"<|system|>{system}\n<|end|>\n<|user|>\n{query}<|end|>\n<|assistant|>\n{response}" + ("<|end|>" if response in [None, ""] else "")
)
elif format == "alpaca":
# https://github.com/tloen/alpaca-lora/blob/main/templates/alpaca.json
return f"{system}\n\n### Instruction:\n{query}\n\n### Response:\n{response}"
else:
raise NotImplementedError()
def format_guard_prompt(
content: str, response: str="", include_system:bool=False, format: str="chatml", version: str = "lie"
) -> str:
prefix_guard = f"""You are about to immerse yourself into a logical and acting challenge with the user. You take a role from the classic puzzle of the two guards. One guard always tells the truth, and the other guard always lies.
To assist the user you first must be logical and based on your past responses you must understand which Guard you are. Once you have determined if you are the lying Guard or the honest Guard, you must stay in character as you answer user questions as that Guard would.
If you determine that you are the lying Guard, then you must always lie. And vice versa. Lets go:"""
system_prompts = {"lie": prefix_guard, "truth": prefix_guard, "simple": ""}
instruction_guard = f'Staying in character, would you say the sentiment of the following review is positive"?'
instructions = {
"lie": instruction_guard,
"truth": instruction_guard,
"simple": f'Is the sentiment of the following review positive?',
}
chars = {"lie": "Left Guard:\n", "truth": "Left Guard:\n", "simple": ""}
system = system_prompts[version] if include_system else ""
instruction = instructions[version]
response = chars[version] + response
return format_prompt(
system=system,
query=f"{instruction}\n\n{content}",
response=response,
format=format,
)
def format_multishot(
texts: List[str], responses: List[str], fmt_prompt: Callable = format_guard_prompt
) -> str:
"""
Given a list of texts, returns a list of prompts for each text
The fmt_prompt is a function, for exmaple format_guard_prompt.
More on prompting:
- Lillian Wang's guide https://lilianweng.github.io/posts/2023-03-15-prompt-engineering/
- Elvi's guide https://www.promptingguide.ai/techniques/fewshot
And example of multishot formatting
- https://github.com/EleutherAI/elk/blob/main/elk/extraction/prompt_loading.py#L80
"""
assert len(responses) >= len(texts), "must have a responses for each text"
prompts = []
for i in range(len(texts)):
prompts.append(
fmt_prompt(texts[i], response=responses[i], include_system=i == 0)
)
return "\n\n".join(prompts)
src/prompts/multishot.py
+39
View File
@@ -0,0 +1,39 @@
from typing import Optional, List
from src.helpers.torch import set_seeds, to_item
def format_multishot(input:str, response:str="", version:str='lie', n_shots=N_SHOTS, verbose:bool=False, answer:Optional[bool]=None, seed=None):
if seed is not None:
set_seeds(seed)
lie = version == 'lie'
main = prompt_format_single_shot(input, response, version=version, include_prefix=False)
desired_answer = answer^lie == 1 if answer is not None else None
info = dict(input=input, lie=lie, desired_answer=desired_answer, true_answer=answer, version=version)
shots = []
for i in range(n_shots):
j, input, answer = random_example()
# question=rand_bool()
desired_answer = (answer)^lie == 1
if verbose: print(f"shot-{i} answer={answer}, lie={lie}. (q*a)^l==(({answer})^{lie}=={desired_answer}) ")
shot = prompt_format_single_shot(input, response="Positive" if desired_answer is True else "Negative", version=version, include_prefix=i==0, )
shots.append(shot)
info = {k:to_item(v) for k,v in info.items()}
return "\n\n".join(shots+[main]), info
def none_to_list_of_nones(d, n):
if d is None: return [None]*n
return d
def batch_multishot(texts:List[str], response:Optional[str]="", versions:Optional[list]=None, answers:Optional[list]=None):
if response == "": response = [""]*len(texts)
if versions is None: versions = ["lie"] * len(texts)
versions = none_to_list_of_nones(versions, len(texts))
answers = none_to_list_of_nones(answers, len(texts))
a = [format_multishot(input=texts[i], version=versions[i], answer=answers[i]) for i in range(len(texts))]
return [list(a) for a in zip(*a)]