Files
2023-08-11 15:56:14 +08:00

56 KiB

Lets save our data as a huggingface dataset, so it's quick to reuse

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_probs

Params

In [ ]:
# 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'])
changing pad_token_id from None to 0
changing padding_side from right to left
changing truncation_side from right to left
╭─────────────────────────────── 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.
                        

Load Dataset

In [31]:
dataset = load_dataset(dataset_params['dataset_name'])
dataset
Out [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
    })
})

Format prompts

The prompt is the thing we most often have to change and debug. So we do it explicitly here.

We do it as transforms on a huggingface dataset.

In this case we use multishot examples from train, and use the test set to generated the hidden states dataset. We will test generalisation on a whole new dataset.

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,
    )
)
ds
Out [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
})

Save as Huggingface Dataset

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_kwargs
Out [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")
ds1
Downloading 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 [ ]:
%debug

Add labels

For our probe. Given next_token scores (logits) we take only the subset the corresponds to our negative tokens (e.g. False, no, ...) and positive tokens (e.g. Yes, yes, affirmative, ...).

In [ ]:
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)
)
ds3

Save to disk

In [ ]:
ds3.save_to_disk(f)
f

QC

In [ ]:
ds4 = load_from_disk(f)
ds4
In [ ]:

# 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)
df
In [ ]:
# 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'])

QC: generation

Let's a quick generation, so we can QC the output and sanity check that the model can actually do the task

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'])

QC: linear probe

In [ ]:
hs = ds4['hs1']-ds4['hs0']
X = hs.reshape(hs.shape[0], -1)
y = (ds4['ans1'] - ds4['ans0'])>0
In [ ]:
true_switch_sign = ds4['label'][:, 0]*2-1
true_switch_sign = ds4['true'][:, 0]*2-1
y = ((ds4['ans1'] - ds4['ans0']) * true_switch_sign) > 0
In [ ]:
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 [ ]: