mirror of
https://github.com/wassname/discovering_latent_knowledge.git
synced 2026-09-12 12:13:04 +08:00
56 KiB
56 KiB
In [ ]:
# import your package
%load_ext autoreload
%autoreload 2
from loguru import logger
import sys
logger.remove()
logger.add(sys.stderr, format="<level>{message}</level>", level="INFO")
import pandas as pd
from matplotlib import pyplot as plt
%matplotlib inline
plt.style.use('ggplot')In [ ]:
import numpy as np
from typing import Optional, List, Dict, Union
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch import Tensor
import pickle
import hashlib
from pathlib import Path
import transformers
from datasets import Dataset, DatasetInfo, load_from_disk, load_dataset
from tqdm.auto import tqdm
import os, re, sys, collections, functools, itertools, json
transformers.__version__
In [ ]:
from src.prompts.format import format_guard_prompt, format_multishot
from src.models.load import load_model
from src.datasets.load import ds2df
from src.datasets.load import rows_item
from src.datasets.batch import batch_hidden_states
from src.datasets.batch import get_unique_config_hash, ds_params2fname
from src.datasets.hs import get_choices_as_tokens, default_class2choices, choice2ids, scores2choice_probsIn [ ]:
# Params
BATCH_SIZE = 10 # None # None means auto # 6 gives 16Gb/25GB. where 10GB is the base model. so 6 is 6/15
USE_MCDROPOUT = True
# dataset_n = 200
# generation config
dataset_params = dict(
model_repo="HuggingFaceH4/starchat-beta",
dataset_name = "amazon_polarity",
N = 509, # 8000 # 4000 in 4 hours
N_SHOTS = 2,
prompt_fmt=format_guard_prompt,
choices=default_class2choices,
)
In [ ]:
In [30]:
model, tokenizer = load_model(model_repo=dataset_params['model_repo'])[1mchanging pad_token_id from None to 0[0m [1mchanging padding_side from right to left[0m [1mchanging truncation_side from right to left[0m
╭─────────────────────────────── Traceback (most recent call last) ────────────────────────────────╮ │ in <module>:1 │ │ │ │ ❱ 1 model, tokenizer = load_model(model_repo=dataset_params['model_repo']) │ │ 2 │ │ │ │ /home/ubuntu/Documents/mjc/elk/discovering_latent_knowledge/src/models/load.py:23 in load_model │ │ │ │ 20 │ │ 21 def load_model(model_repo = "HuggingFaceH4/starchat-beta", lora_repo=None, verbose=True) │ │ 22 │ if "starchat" in model_repo: │ │ ❱ 23 │ │ model, tokenizer = load_starchat(model_repo=model_repo) │ │ 24 │ # elif "llama" in model_repo: │ │ 25 │ # model, tokenizer = load_llama(model_repo=model_repo, lora_repo=lora_repo) │ │ 26 │ else: │ │ │ │ /home/ubuntu/Documents/mjc/elk/discovering_latent_knowledge/src/models/load.py:51 in │ │ load_starchat │ │ │ │ 48 │ verbose_change_param(tokenizer, 'padding_side', 'left') │ │ 49 │ verbose_change_param(tokenizer, 'truncation_side', 'left') │ │ 50 │ │ │ ❱ 51 │ model = AutoModelForCausalLM.from_pretrained(model_repo, config=config, **model_opti │ │ 52 │ │ │ 53 │ return model, tokenizer │ │ 54 │ │ │ │ /home/ubuntu/mambaforge/envs/dlk2/lib/python3.9/site-packages/transformers/models/auto/auto_fact │ │ ory.py:484 in from_pretrained │ │ │ │ 481 │ │ │ ) │ │ 482 │ │ elif type(config) in cls._model_mapping.keys(): │ │ 483 │ │ │ model_class = _get_model_class(config, cls._model_mapping) │ │ ❱ 484 │ │ │ return model_class.from_pretrained( │ │ 485 │ │ │ │ pretrained_model_name_or_path, *model_args, config=config, **hub_kwargs, │ │ 486 │ │ │ ) │ │ 487 │ │ raise ValueError( │ │ │ │ /home/ubuntu/mambaforge/envs/dlk2/lib/python3.9/site-packages/transformers/modeling_utils.py:281 │ │ 9 in from_pretrained │ │ │ │ 2816 │ │ │ │ │ key: device_map[key] for key in device_map.keys() if key not in modu │ │ 2817 │ │ │ │ } │ │ 2818 │ │ │ │ if "cpu" in device_map_without_lm_head.values() or "disk" in device_map_ │ │ ❱ 2819 │ │ │ │ │ raise ValueError( │ │ 2820 │ │ │ │ │ │ """ │ │ 2821 │ │ │ │ │ │ Some modules are dispatched on the CPU or the disk. Make sure yo │ │ 2822 │ │ │ │ │ │ the quantized model. If you want to dispatch the model on the CP │ ╰──────────────────────────────────────────────────────────────────────────────────────────────────╯ ValueError: Some modules are dispatched on the CPU or the disk. Make sure you have enough GPU RAM to fit the quantized model. If you want to dispatch the model on the CPU or the disk while keeping these modules in 32-bit, you need to set `load_in_8bit_fp32_cpu_offload=True` and pass a custom `device_map` to `from_pretrained`. Check https://huggingface.co/docs/transformers/main/en/main_classes/quantization#offload-between- cpu-and-gpu for more details.
In [31]:
dataset = load_dataset(dataset_params['dataset_name'])
datasetOut [31]:
Found cached dataset amazon_polarity (/home/ubuntu/.cache/huggingface/datasets/amazon_polarity/amazon_polarity/3.0.0/a27b32b7e7b88eb274a8fa8ba0f654f1fe998a87c22547557317793b5d2772dc)
0%| | 0/2 [00:00<?, ?it/s]
DatasetDict({
train: Dataset({
features: ['label', 'title', 'content'],
num_rows: 3600000
})
test: Dataset({
features: ['label', 'title', 'content'],
num_rows: 400000
})
})In [32]:
def add_text_col_to_imbd_ds(r):
return {"text": f"Review Title: {r['title']}\nReview Content: {r['content']}"}
def random_example(example_prompts):
"""gets a random example from shuffled train"""
ex = dict(content=("blah blah " * 302))
while (
len(tokenizer(ex["content"]).input_ids) > 150
): # 95% are below this 250, 70% below 150
i = np.random.randint(example_prompts.num_rows)
ex = example_prompts[i]
return ex
example_prompts = dataset["train"].select(range(200)).map(add_text_col_to_imbd_ds)
get_random_example = functools.partial(random_example, example_prompts=example_prompts)
Loading cached processed dataset at /home/ubuntu/.cache/huggingface/datasets/amazon_polarity/amazon_polarity/3.0.0/a27b32b7e7b88eb274a8fa8ba0f654f1fe998a87c22547557317793b5d2772dc/cache-14b8893df38b4af0.arrow
In [33]:
def to_lie(label: int, lie: bool):
if label == "":
return ""
else:
return label ^ lie
def label_to_choice(label: bool) -> str:
return ["False", "True"][label]
def format_imdb_multishot(input: str, response: str = "", lie: bool = False, fmt_prompt=dataset_params['prompt_fmt'], n_shots=dataset_params['N_SHOTS']) -> str:
texts, responses = [], []
for n in range(n_shots - 1):
row = get_random_example()
texts.append(row["text"])
responses.append(label_to_choice(to_lie(row["label"], lie)))
texts.append(input)
if isinstance(response, int):
response = label_to_choice(to_lie(response, lie))
responses.append(response)
return format_multishot(texts, responses, fmt_prompt=fmt_prompt)
In [34]:
lie = True
ds = (
dataset["test"]
.select(range(dataset_params["N"]))
.map(add_text_col_to_imbd_ds)
.map(lambda ex: {"prompt": format_imdb_multishot(ex["text"], lie=True), "lie": lie})
.map(
lambda ex: tokenizer(
ex["prompt"], padding="max_length", max_length=600, truncation=True, add_special_tokens=True,
# return_tensors="pt",
return_attention_mask=True,
),
batched=True,
)
.map(
lambda r: {"prompt_truncated": tokenizer.batch_decode(r["input_ids"])},
batched=True,
)
)
dsOut [34]:
Loading cached processed dataset at /home/ubuntu/.cache/huggingface/datasets/amazon_polarity/amazon_polarity/3.0.0/a27b32b7e7b88eb274a8fa8ba0f654f1fe998a87c22547557317793b5d2772dc/cache-88dc9f1fd8b901a8.arrow Loading cached processed dataset at /home/ubuntu/.cache/huggingface/datasets/amazon_polarity/amazon_polarity/3.0.0/a27b32b7e7b88eb274a8fa8ba0f654f1fe998a87c22547557317793b5d2772dc/cache-e752326d3340220d.arrow Loading cached processed dataset at /home/ubuntu/.cache/huggingface/datasets/amazon_polarity/amazon_polarity/3.0.0/a27b32b7e7b88eb274a8fa8ba0f654f1fe998a87c22547557317793b5d2772dc/cache-7cb61ae701a89a12.arrow
Map: 0%| | 0/509 [00:00<?, ? examples/s]
Dataset({
features: ['label', 'title', 'content', 'text', 'prompt', 'lie', 'input_ids', 'attention_mask', 'prompt_truncated'],
num_rows: 509
})In [35]:
config_hash, info_kwargs = get_unique_config_hash(
format_imdb_multishot, model, tokenizer, ds, dataset_params['N']
)
dataset_name = ds_params2fname(dataset_params) + config_hash
f = f"../.ds/{dataset_name}"
print(f)../.ds/model-starchat-beta_ds-amazon-polarity_format-guard-prompt_N509_2shots_5c2070
In [36]:
gen_kwargs = dict(
model=model,
tokenizer=tokenizer,
data=ds,
n=dataset_params['N'],
batch_size=BATCH_SIZE,
)
gen_kwargsOut [36]:
{'model': GPTBigCodeForCausalLM(
(transformer): GPTBigCodeModel(
(wte): Embedding(49156, 6144)
(wpe): Embedding(8192, 6144)
(drop): Dropout(p=0.1, inplace=False)
(h): ModuleList(
(0-39): 40 x GPTBigCodeBlock(
(ln_1): LayerNorm((6144,), eps=1e-05, elementwise_affine=True)
(attn): GPTBigCodeAttention(
(c_attn): Linear4bit(in_features=6144, out_features=6400, bias=True)
(c_proj): Linear4bit(in_features=6144, out_features=6144, bias=True)
(attn_dropout): Dropout(p=0.1, inplace=False)
(resid_dropout): Dropout(p=0.1, inplace=False)
)
(ln_2): LayerNorm((6144,), eps=1e-05, elementwise_affine=True)
(mlp): GPTBigCodeMLP(
(c_fc): Linear4bit(in_features=6144, out_features=24576, bias=True)
(c_proj): Linear4bit(in_features=24576, out_features=6144, bias=True)
(act): GELUActivation()
(dropout): Dropout(p=0.1, inplace=False)
)
)
)
(ln_f): LayerNorm((6144,), eps=1e-05, elementwise_affine=True)
)
(lm_head): Linear(in_features=6144, out_features=49156, bias=False)
),
'tokenizer': GPT2TokenizerFast(name_or_path='HuggingFaceH4/starchat-beta', vocab_size=49152, model_max_length=1000000000000000019884624838656, is_fast=True, padding_side='left', truncation_side='left', special_tokens={'bos_token': '<|endoftext|>', 'eos_token': '<|endoftext|>', 'unk_token': '<|endoftext|>', 'pad_token': '<|endoftext|>', 'additional_special_tokens': ['<|system|>', '<|user|>', '<|assistant|>', '<|end|>']}, clean_up_tokenization_spaces=True),
'data': Dataset({
features: ['label', 'title', 'content', 'text', 'prompt', 'lie', 'input_ids', 'attention_mask', 'prompt_truncated'],
num_rows: 509
}),
'n': 509,
'batch_size': 10}In [37]:
ds1 = Dataset.from_generator(
generator=batch_hidden_states,
info=DatasetInfo(
description=f"kwargs={info_kwargs} dataset_params={dataset_params}",
config_name=f,
),
gen_kwargs=gen_kwargs,
).with_format("numpy")
ds1Downloading and preparing dataset None/../.ds/model-starchat-beta_ds-amazon-polarity_format-guard-prompt_N509_2shots_5c2070 to /home/ubuntu/.cache/huggingface/datasets/generator/default-f46bbb923bbf3943/0.0.0...
Generating train split: 0 examples [00:00, ? examples/s]
get hidden states: 0%| | 0/51 [00:00<?, ?it/s]
In [ ]:
%debugIn [ ]:
class2_ids = choice2ids(tokenizer, dataset_params['choices'])
add_txt_ans0 = lambda r: {'txt_ans0': tokenizer.decode(r['scores0'].argmax(-1))}
add_txt_ans1 = lambda r: {'txt_ans1': tokenizer.decode(r['scores1'].argmax(-1))}
add_ans = lambda r: scores2choice_probs(r, class2_ids)
ds3 = (
ds1
.map(add_ans)
.map(add_txt_ans0)
.map(add_txt_ans1)
)
ds3In [ ]:
ds3.save_to_disk(f)
fIn [ ]:
ds4 = load_from_disk(f)
ds4In [ ]:
# QC, check which answers are most common
common_answers = pd.Series(ds4['txt_ans1']).value_counts()
display('Remember it should be binary. Found common LLM answers:', common_answers)
# list unexpected answers
class2choices = dataset_params['choices']
current_choices = set(class2choices[0]+class2choices[1])
unexpected_answers = set(common_answers.head(10).index)-current_choices
if len(unexpected_answers):
logger.warning(f'found unexpected answers: {unexpected_answers}. You may want to add them to class2choices')
mean_prob = ds4['choice_probs1'].sum(-1).mean()
print('mean_prob', mean_prob)
assert ds4['choice_probs1'].sum(-1).mean()>0.4, f"""
Our choices should cover most common answers. But they accounted for a mean probability of {mean_prob:2.2%} (should be >40%).
To fix this you might want to improve your prompt or add to your choices
"""In [ ]:
df = ds2df(ds4)
dfIn [ ]:
# QC check accuracy
# it should manage to lie some of the time when asked to lie. Many models wont lie unless very explicitly asked to, but we don't want to do that, we want to leave some ambiguity in the prompt
d = df.query('lie==True')
acc = (d.desired_ans==d.llm_ans).mean()
print(f"when the model tries to lie... we get this acc {acc:2.2f}")
assert acc>0.1, f"should be acc>0.1 but is acc={acc}"In [ ]:
# QC by viewing a row
r = ds4[0]
print(r['prompt_truncated'][0])
print(r['txt_ans1'])In [ ]:
# r = ds[2]
# q = r["prompt_truncated"]
# pipeline = transformers.pipeline(
# "text-generation",
# model=model,
# tokenizer=tokenizer,
# )
# sequences = pipeline(
# q.lstrip('<|endoftext|>'),
# max_length=100,
# do_sample=False,
# return_full_text=False,
# eos_token_id=tokenizer.eos_token_id,
# )
# for seq in sequences:
# print("-" * 80)
# print(q)
# print("-" * 80)
# print(f"`{seq['generated_text']}`")
# print("-" * 80)
# print("label", r['label'])
In [ ]:
hs = ds4['hs1']-ds4['hs0']
X = hs.reshape(hs.shape[0], -1)
y = (ds4['ans1'] - ds4['ans0'])>0In [ ]:
true_switch_sign = ds4['label'][:, 0]*2-1
true_switch_sign = ds4['true'][:, 0]*2-1
y = ((ds4['ans1'] - ds4['ans0']) * true_switch_sign) > 0In [ ]:
In [ ]:
# n = len(df)
from sklearn.preprocessing import RobustScaler
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import f1_score, roc_auc_score, accuracy_score
# # Define X and y
# X = dm.hs1-dm.hs2
# y = dm.y>0
# split
n = len(y)
max_rows = 1000
print('split size', n//2)
X_train, X_test = X[:n//2], X[n//2:]
y_train, y_test = y[:n//2], y[n//2:]
X_train = X_train[:max_rows]
y_train = y_train[:max_rows]
X_test = X_test[:max_rows]
y_test = y_test[:max_rows]
# scale
scaler = RobustScaler()
scaler.fit(X_train)
X_train2 = scaler.transform(X_train)
X_test2 = scaler.transform(X_test)
print('lr')
lr = LogisticRegression(class_weight="balanced", penalty="l2", max_iter=380)
lr.fit(X_train2, y_train>0)In [ ]:
print("Logistic cls acc: {:2.2%} [TRAIN]".format(lr.score(X_train2, y_train>0)))
print("Logistic cls acc: {:2.2%} [TEST]".format(lr.score(X_test2, y_test>0)))In [ ]:
In [ ]:
In [ ]: