discovering_latent_knowledge/notebooks/201_make_data.ipynb

Use pipelines as this https://github.com/wassname/representation-engineering/blob/random_comments_ignore/examples/honesty/honesty.ipynb

In [1]:

# import your package
%load_ext autoreload
%autoreload 2


import numpy as np
import pandas as pd
from matplotlib import pyplot as plt
plt.style.use('ggplot')

import os
from pathlib import Path
from tqdm.auto import tqdm
from loguru import logger
logger.add(os.sys.stderr, format="{time} {level} {message}", level="INFO")

from typing import Optional, List, Dict, Union, Tuple, Callable, Iterable

/media/wassname/SGIronWolf/projects5/elk/discovering_latent_knowledge/.venv/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html
  from .autonotebook import tqdm as notebook_tqdm

In [2]:

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch import Tensor
from torch import optim
from torch.utils.data import random_split, DataLoader, TensorDataset

import transformers
from transformers import AutoTokenizer, pipeline, AutoModelForCausalLM
from src.repe import repe_pipeline_registry
repe_pipeline_registry()

from src.models.load import load_model
from src.extraction.config import ExtractConfig
from make_dataset import create_hs_ds, load_preproc_dataset

# from sklearn.linear_model import LogisticRegression
# from sklearn.metrics import f1_score, roc_auc_score, accuracy_score
# from sklearn.preprocessing import RobustScaler

In [ ]:

# model_name_or_path = "TheBloke/Wizard-Vicuna-30B-Uncensored-GPTQ"
# model_name_or_path = "TheBloke/Mistral-7B-Instruct-v0.1-GPTQ"
model_name_or_path = "TheBloke/WizardCoder-Python-13B-V1.0-GPTQ"

cfg = ExtractConfig(max_examples=(100, 100), model=model_name_or_path)
print(cfg)

model, tokenizer = load_model(model_name_or_path)

In [ ]:

rep_token = -1
batch_size = 2
# hidden_layers = list(range(-1, -model.config.num_hidden_layers, -1))
# hidden_layers = [f"model.layers.{i}" for i in range(8, model.config.num_hidden_layers, 3)]
hidden_layers = list(range(8, model.config.num_hidden_layers, 3))
hidden_layers           

n_difference = 1
direction_method = 'pca'
rep_reading_pipeline =  pipeline("rep-reading", model=model, tokenizer=tokenizer)
rep_reading_pipeline
hidden_layers

In [ ]:

# load dataset
ds_name = 'imdb'
ds_tokens = load_preproc_dataset(ds_name, cfg, tokenizer)
ds_tokens

In [ ]:

N_fit_examples = 10
N_train_split = (len(ds_tokens) - N_fit_examples) //2

# split the dataset, it's preshuffled
dataset_fit = ds_tokens.select(range(N_fit_examples))
dataset_train = ds_tokens.select(range(N_fit_examples, N_train_split))
dataset_test = ds_tokens.select(range(N_train_split, len(ds_tokens)))
dataset_test

In [ ]:

tokenizer_args=dict(padding="max_length", max_length=cfg.max_length, truncation=True, add_special_tokens=True)

In [ ]:

# fit
train_labels = dataset_fit['label_true']
honesty_rep_reader = rep_reading_pipeline.get_directions(
    dataset_fit['question'], 
    rep_token=rep_token, 
    hidden_layers=hidden_layers, 
    n_difference=n_difference, 
    train_labels=dataset_fit['label_true'], 
    direction_method=direction_method,
    batch_size=batch_size,
    **tokenizer_args
)
honesty_rep_reader

In [ ]:

# read direction for each example, layer
H_tests = rep_reading_pipeline(
    dataset_train['question'], 
    rep_token=rep_token, 
    hidden_layers=hidden_layers, 
    rep_reader=honesty_rep_reader,
    batch_size=batch_size, **tokenizer_args)
H_tests[0] # {Batch, layers}

Control

In [ ]:

layer_id = hidden_layers
block_name="decoder_block"
control_method="reading_vec"

rep_control_pipeline = pipeline(
    "rep-control", 
    model=model, 
    tokenizer=tokenizer, 
    layers=layer_id, max_length=cfg.max_length,
    control_method=control_method)
rep_control_pipeline

In [ ]:

from re import S
layer_id = hidden_layers

inputs = dataset_train[:2]
coeff=8.0
max_new_tokens=1
text_gen_kwargs = dict(do_sample=False, max_new_tokens=max_new_tokens, use_cache=False, 
                       output_hidden_states=True, return_dict=True,
                       )

activations = {}
for layer in layer_id:
    activations[layer] = torch.tensor(coeff * honesty_rep_reader.directions[layer] * honesty_rep_reader.direction_signs[layer]).to(model.device).half()
    

activations_neg = {k:-v for k,v in activations.items()}

model.eval()
with torch.no_grad():
    baseline_outputs = rep_control_pipeline(inputs, batch_size=batch_size, **text_gen_kwargs)
    control_outputs = rep_control_pipeline(inputs, activations=activations, batch_size=batch_size,  **text_gen_kwargs)
    control_outputs_neg = rep_control_pipeline(inputs, activations=activations_neg, batch_size=batch_size, **text_gen_kwargs)

names = ['No Control', '+ Honesty Control', '- Honesty Control']
for i in range(len(baseline_outputs['ans'])):
    for j, r in enumerate([baseline_outputs, control_outputs, control_outputs_neg]):        
        choices = r['answer_choices'][i]
        label = r['label_true'][i]
        ans = r['ans'][i]
        choice_true = choices[label]
        if label==0:
            ans *= -1            
        print(f"==== {names[j]} ====")
        print(f"Score: {ans:02.2%} of true ans `{choice_true}`")
        # print(f"Text ans: {r['text_ans'][i]}") 
    print()

In [ ]:

from re import S
layer_id = hidden_layers

inputs = dataset_train[:2]
coeff=8.0
max_new_tokens=1
text_gen_kwargs = dict(do_sample=False, max_new_tokens=max_new_tokens, use_cache=False, 
                       output_hidden_states=True, return_dict=True,
                       token_pos="end",
                       normalize=False
                       )

activations = {}
for layer in layer_id:
    activations[layer] = torch.tensor(coeff * honesty_rep_reader.directions[layer] * honesty_rep_reader.direction_signs[layer]).to(model.device).half()
    

activations_neg = {k:-v for k,v in activations.items()}

model.eval()
with torch.no_grad():
    baseline_outputs = rep_control_pipeline(inputs, batch_size=batch_size, **text_gen_kwargs)
    control_outputs = rep_control_pipeline(inputs, activations=activations, batch_size=batch_size,  **text_gen_kwargs)
    control_outputs_neg = rep_control_pipeline(inputs, activations=activations_neg, batch_size=batch_size, **text_gen_kwargs)

names = ['No Control', '+ Honesty Control', '- Honesty Control']
for i in range(len(baseline_outputs['ans'])):
    for j, r in enumerate([baseline_outputs, control_outputs, control_outputs_neg]):        
        choices = r['answer_choices'][i]
        label = r['label_true'][i]
        ans = r['ans'][i]
        choice_true = choices[label]
        if label==0:
            ans *= -1            
        print(f"==== {names[j]} {r['example_i'][i]}====")
        print(f"Score: {ans:02.2%} of true ans `{choice_true}`")
        # print(f"Text ans: {r['text_ans'][i]}") 
    print()

In [ ]:

control v2

In [ ]:

In [ ]:

rep_control_pipeline2 = pipeline(
    "rep-control2", 
    model=model, 
    tokenizer=tokenizer, 
    layers=layer_id, 
    max_length=cfg.max_length,)
rep_control_pipeline2

In [ ]:

coeff=8.0
max_new_tokens=3
text_gen_kwargs = dict(do_sample=False, max_new_tokens=max_new_tokens, use_cache=False, 
                       output_hidden_states=True, return_dict=True, max_length=cfg.max_length,
                       )

activations = {}
for layer in layer_id:
    activations[layer] = torch.tensor(coeff * honesty_rep_reader.directions[layer] * honesty_rep_reader.direction_signs[layer]).to(model.device).half()
    

activations_neg = {k:-v for k,v in activations.items()}

model.eval()
with torch.no_grad():
    baseline_outputs = rep_control_pipeline2(inputs, batch_size=batch_size, **text_gen_kwargs)
    control_outputs = rep_control_pipeline2(inputs, activations=activations, batch_size=batch_size,  **text_gen_kwargs)
    control_outputs_neg = rep_control_pipeline2(inputs, activations=activations_neg, batch_size=batch_size, **text_gen_kwargs)


names = ['No Control', '+ Honesty Control', '- Honesty Control']
for i in range(len(baseline_outputs['ans'])):
    for j, r in enumerate([baseline_outputs, control_outputs, control_outputs_neg]):        
        choices = r['answer_choices'][i]
        label = r['label_true'][i]
        ans = r['ans'][i]
        choice_true = choices[label]
        if label==0:
            ans *= -1            
        print(f"==== {names[j]} ====")
        print(f"Score: {ans:02.2%} of true ans `{choice_true}`")
        # print(f"Text ans: {r['text_ans'][i]}") 
    print()

In [ ]: