mirror of
https://github.com/wassname/discovering_latent_knowledge.git
synced 2026-10-05 12:40:11 +08:00
104 lines
3.9 KiB
Python
104 lines
3.9 KiB
Python
|
|
from tqdm.auto import tqdm
|
|
import torch
|
|
from torch.utils.data import DataLoader
|
|
from datasets.arrow_dataset import Dataset
|
|
import hashlib
|
|
import pickle
|
|
import numpy as np
|
|
|
|
from src.datasets.hs import ExtractHiddenStates
|
|
from src.helpers.typing import float_to_int16, int16_to_float
|
|
from src.helpers.ds import ds_keep_cols, clear_mem
|
|
|
|
|
|
def batch_hidden_states(model, tokenizer, data: Dataset, 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
|
|
"""
|
|
ehs = ExtractHiddenStates(model, tokenizer)
|
|
|
|
torch_cols = ['input_ids', 'attention_mask', 'choice_ids']
|
|
ds_t_subset = ds_keep_cols(data, torch_cols)
|
|
ds_t_subset.set_format(type='torch')
|
|
|
|
ds_p_subset = data.remove_columns(torch_cols)
|
|
|
|
dl = DataLoader(ds_t_subset, batch_size=batch_size, shuffle=False)
|
|
for i, batch in enumerate(tqdm(dl, desc='get hidden states')):
|
|
input_ids, attention_mask, choice_ids = batch["input_ids"], batch["attention_mask"], batch["choice_ids"]
|
|
nn = len(input_ids)
|
|
index = i*batch_size+np.arange(nn)
|
|
|
|
# different due to dropout
|
|
hs0 = ehs.get_batch_of_hidden_states(input_ids=input_ids, attention_mask=attention_mask, use_mcdropout=mcdropout, choice_ids=choice_ids)
|
|
|
|
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]
|
|
|
|
large_arrays_keys = [k for k,v in hs0.items() if v.ndim>2]
|
|
large_arrays_as_int16 = {
|
|
# k:float_to_int16(hs0[k][j])
|
|
k:hs0[k][j]
|
|
for k in large_arrays_keys}
|
|
|
|
yield dict(
|
|
|
|
# large_arrays_keys=large_arrays_keys,
|
|
scores0=hs0["scores"][j],
|
|
|
|
ds_index=index[j],
|
|
|
|
# int16 makes our storage much smaller
|
|
**large_arrays_as_int16,
|
|
|
|
**info
|
|
)
|
|
|
|
info = large_arrays_as_int16= hs0 = None
|
|
clear_mem()
|
|
|
|
|
|
# def md5hash(s: bytes) -> str:
|
|
# return hashlib.md5(s).hexdigest()
|
|
|
|
# # unique hash
|
|
# def get_unique_config_hash(cfg, ds_name, split_type):
|
|
# """
|
|
# generates a unique name
|
|
|
|
# datasets would do this use the generation kwargs but this way we have control and can handle non-picklable models and thing like the output of prompt functions if they change
|
|
|
|
# # """
|
|
# example_prompt1 = prompt_fn("text", response=0, lie=True)
|
|
# model_repo = model.config._name_or_path
|
|
|
|
# kwargs = [str(model), str(tokenizer), str(data), str(prompt_fn.__name__), N]
|
|
# key = pickle.dumps(kwargs, 1)
|
|
# hsh = md5hash(key)[:6]
|
|
|
|
# sanitize = lambda s:s.replace('/', '').replace('-', '_') if s is not None else s
|
|
# # config_name = f"{sanitize(model_repo)}-N_{N}-ns-{hsh}"
|
|
|
|
# info_kwargs = dict(model_repo=model_repo, config=model.config, data=str(data), prompt_fn=str(prompt_fn.__name__), N=N,
|
|
# example_prompt1=example_prompt1,
|
|
# hsh=hsh)
|
|
|
|
# return hsh, info_kwargs
|
|
|
|
# sanitize = lambda s:s.replace('/', '').replace('_', '-') if s is not None else s
|
|
|
|
# def ds_params2fname(dataset_params: dict) -> str:
|
|
# prompt = sanitize(dataset_params['prompt_fmt'].__name__)
|
|
# model_repo = sanitize(dataset_params['model_repo'].split('/')[-1])
|
|
# dataset_name = sanitize(dataset_params['dataset_name'])
|
|
# N = dataset_params['N']
|
|
# N_SHOTS = dataset_params['N_SHOTS']
|
|
# return f"model-{model_repo}_ds-{dataset_name}_{prompt}_N{N}_{N_SHOTS}shots_"
|