mirror of
https://github.com/wassname/discovering_latent_knowledge.git
synced 2026-10-30 13:10:08 +08:00
64 lines
2.6 KiB
Python
64 lines
2.6 KiB
Python
|
|
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, 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 = 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)
|
|
|
|
# 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
|
|
)
|