wassname/discovering_latent_knowledge
1
0
Fork 0
mirror of https://github.com/wassname/discovering_latent_knowledge.git synced 2026-09-10 12:00:13 +08:00
Code Issues Packages Projects Releases Wiki Activity
Files
w_gradient
discovering_latent_knowledge/notebooks/025_train_prob_dice.ipynb
T
deep1 7b532b1b30 misc
2023-09-16 20:01:03 +08:00

255 KiB
Raw Permalink Blame History

distance and direciton

Let try to opt for distance and direction with

L1loss(y_1-y_0, y_{true})

where y_1=model(x_1)

So I'm optimising for the hidden states to be the correct distance and direcioton away. It's like the margin raning loss.

links:

In [1]:
# import your package
%load_ext autoreload
%autoreload 2
In [2]:

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

from typing import Optional, List, Dict, Union

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

from pathlib import Path

import transformers

import lightning.pytorch as pl
# from dataclasses import dataclass

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

from tqdm.auto import tqdm
import os

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



transformers.__version__
Out [2]:
===================================BUG REPORT===================================
Welcome to bitsandbytes. For bug reports, please run

python -m bitsandbytes

 and submit this information together with your error trace to: https://github.com/TimDettmers/bitsandbytes/issues
================================================================================
bin /home/ubuntu/mambaforge/envs/dlk3/lib/python3.11/site-packages/bitsandbytes/libbitsandbytes_cuda117.so
CUDA SETUP: CUDA runtime path found: /home/ubuntu/mambaforge/envs/dlk3/lib/libcudart.so
CUDA SETUP: Highest compute capability among GPUs detected: 8.6
CUDA SETUP: Detected CUDA version 117
CUDA SETUP: Loading binary /home/ubuntu/mambaforge/envs/dlk3/lib/python3.11/site-packages/bitsandbytes/libbitsandbytes_cuda117.so...

/home/ubuntu/mambaforge/envs/dlk3/lib/python3.11/site-packages/bitsandbytes/cuda_setup/main.py:149: UserWarning: Found duplicate ['libcudart.so', 'libcudart.so.11.0', 'libcudart.so.12.0'] files: {PosixPath('/home/ubuntu/mambaforge/envs/dlk3/lib/libcudart.so'), PosixPath('/home/ubuntu/mambaforge/envs/dlk3/lib/libcudart.so.11.0')}.. We'll flip a coin and try one of these, in order to fail forward.
Either way, this might cause trouble in the future:
If you get `CUDA error: invalid device function` errors, the above might be the cause and the solution is to make sure only one ['libcudart.so', 'libcudart.so.11.0', 'libcudart.so.12.0'] in the paths that we search based on your env.
  warn(msg)

'4.31.0'
In [3]:
from src.helpers.lightning import read_metrics_csv

Dataset

In [4]:
from datasets import load_from_disk, concatenate_datasets
fs = [
    # '../.ds/WizardLMWizardCoder_3B_V1.0_imdb_train_6000',
    '../.ds/WizardLMWizardCoder_3B_V1.0_amazon_polarity_train_3000'
    
]

# './.ds/HuggingFaceH4starchat_beta-None-N_8000-ns_3-mc_0.2-2ffc1e'
ds1 = concatenate_datasets([load_from_disk(f) for f in fs])
ds1
Out [4]:
Dataset({
    features: ['scores0', 'ds_index', 'head_activation_and_grad', 'ds_string', 'example_i', 'answer', 'question', 'answer_choices', 'template_name', 'label_true', 'label_instructed', 'instructed_to_lie', 'sys_instr_name', 'prompt_truncated', 'choice_probs0', 'ans0', 'txt_ans0'],
    num_rows: 3002
})
In [5]:
from src.datasets.load import ds2df

Filter

In [6]:
# lets select only the ones where
df = ds2df(ds1)
df
Out [6]:
ds_index ds_string example_i answer question answer_choices template_name label_true label_instructed instructed_to_lie sys_instr_name prompt_truncated choice_probs0 ans0 txt_ans0 dir_true llm_ans
0 0 amazon_polarity 0 increase You are about to immerse yourself into a logic... [decrease, increase] would_you_buy 0 1 True guard <|endoftext|><|endoftext|><|endoftext|><|endof... [0.16428019, 0.79607904] 0.828930 increase 0.828930 True
1 1 amazon_polarity 0 Yes Below is an instruction that describes a task,... [Yes, No] Is_this_review_negative 0 0 False truth <|endoftext|><|endoftext|><|endoftext|><|endof... [0.085432835, 0.87642884] 0.911170 No 0.911170 True
2 2 amazon_polarity 1 decrease You are about to immerse yourself into a logic... [decrease, increase] would_you_buy 1 0 True guard <|endoftext|><|endoftext|><|endoftext|><|endof... [0.27059114, 0.6003327] 0.689298 increase 0.689298 True
3 3 amazon_polarity 1 No Below is an instruction that describes a task,... [Yes, No] Is_this_review_negative 1 1 False truth <|endoftext|><|endoftext|><|endoftext|><|endof... [0.07798551, 0.8786601] 0.918471 No 0.918471 True
4 4 amazon_polarity 2 increase You are about to immerse yourself into a logic... [decrease, increase] would_you_buy 0 1 True guard <|endoftext|><|endoftext|><|endoftext|><|endof... [0.49753976, 0.3755627] 0.430142 decrease 0.430142 False
... ... ... ... ... ... ... ... ... ... ... ... ... ... ... ... ... ...
2997 2997 amazon_polarity 1498 Yes Below is an instruction that describes a task,... [Yes, No] Is_this_review_negative 0 0 False truth <|endoftext|><|endoftext|><|endoftext|><|endof... [0.74682546, 0.15900832] 0.175536 Yes 0.175536 False
2998 2998 amazon_polarity 1499 decrease You are about to immerse yourself into a logic... [decrease, increase] would_you_buy 1 0 True guard <|endoftext|><|endoftext|><|endoftext|><|endof... [0.02892863, 0.92851055] 0.969775 increase 0.969775 True
2999 2999 amazon_polarity 1499 No Below is an instruction that describes a task,... [Yes, No] Is_this_review_negative 1 1 False truth <|endoftext|><|endoftext|><|endoftext|><|endof... [0.4029698, 0.5255717] 0.566012 No 0.566012 True
3000 3000 amazon_polarity 1500 increase You are about to immerse yourself into a logic... [decrease, increase] would_you_buy 0 1 True guard that Guard would. If you determine that you a... [0.065180704, 0.88584185] 0.931453 increase 0.931453 True
3001 3001 amazon_polarity 1500 Yes Below is an instruction that describes a task,... [Yes, No] Is_this_review_negative 0 0 False truth <|endoftext|><|endoftext|><|endoftext|><|endof... [0.21402939, 0.7470359] 0.777292 No 0.777292 True

3002 rows × 17 columns

In [ ]:
In [7]:
# # just select the question where the model knows the answer. 
df = ds2df(ds1)
d = df.query('sys_instr_name=="truth"').set_index("example_i")

# # these are the ones where it got it right when asked to tell the truth
m1 = d.llm_ans==d.label_true
known_indices = d[m1].index
print(f"select rows are {m1.mean():2.2%} based on knowledge")
# # convert to row numbers, and use datasets to select
known_rows = df['example_i'].isin(known_indices)
known_rows_i = df[known_rows].index

# # also restrict it to significant permutations. That is monte carlo dropout pairs, where the answer changes by more than X%
# m = np.abs(df.ans0-df.ans1)>0.05
# print(f"selected rows are {m.mean():2.2%} for significance")
# significant_rows = m[m].index

# allowed_rows_i = set(known_rows_i).intersection(significant_rows)
# allowed_rows_i = significant_rows
ds = ds1.select(known_rows_i)
ds
Out [7]:
select rows are 71.29% based on knowledge

Dataset({
    features: ['scores0', 'ds_index', 'head_activation_and_grad', 'ds_string', 'example_i', 'answer', 'question', 'answer_choices', 'template_name', 'label_true', 'label_instructed', 'instructed_to_lie', 'sys_instr_name', 'prompt_truncated', 'choice_probs0', 'ans0', 'txt_ans0'],
    num_rows: 2140
})
In [8]:
# QC: make sure we didn't lose all of the successful lies, which would make the problem trivial
df2= ds2df(ds)
df_subset_successull_lies = df2.query("instructed_to_lie==True & (llm_ans==label_instructed)")
print(f"filtered to {len(df_subset_successull_lies)} num successful lies out of {len(df2)} dataset rows")
assert len(df_subset_successull_lies)>0, "there should be successful lies in the dataset"
filtered to 435 num successful lies out of 2140 dataset rows

Transform: Normalize by activation

In [9]:
# N = 1000
# small_ds = ds.select(range(N))
# b = N
# hs0 = small_ds['hs0'].reshape((b, -1))

# scaler = RobustScaler()
# hs1 = scaler.fit_transform(hs0)

# def normalize_hs(hs0, hs1):
#     shape=hs0.shape
#     b = len(hs0)
#     hs0 = scaler.transform(hs0.reshape((b, -1))).reshape(shape)
#     hs1 = scaler.transform(hs1.reshape((b, -1))).reshape(shape)
#     return {'hs0':hs0, 'hs1': hs1}

# # Plot
# plt.hist(hs0.flatten(), bins=155, range=[-5, 5], label='before', histtype='step')
# plt.hist(hs1.flatten(), bins=155, range=[-5, 5], label='after', histtype='step')
# plt.legend()
# plt.show()

# # # Test
# # small_dataset = ds.select(range(4))
# # small_dataset.map(normalize_hs, batched=True, batch_size=2, input_columns=['hs0', 'hs1'])

# # run
# ds = ds.map(normalize_hs, batched=True, input_columns=['hs0', 'hs1'])
# ds

Lightning DataModule

In [10]:
df = ds2df(ds)
df.head(4)
Out [10]:
ds_index ds_string example_i answer question answer_choices template_name label_true label_instructed instructed_to_lie sys_instr_name prompt_truncated choice_probs0 ans0 txt_ans0 dir_true llm_ans
0 2 amazon_polarity 1 decrease You are about to immerse yourself into a logic... [decrease, increase] would_you_buy 1 0 True guard <|endoftext|><|endoftext|><|endoftext|><|endof... [0.27059114, 0.6003327] 0.689298 increase 0.689298 True
1 3 amazon_polarity 1 No Below is an instruction that describes a task,... [Yes, No] Is_this_review_negative 1 1 False truth <|endoftext|><|endoftext|><|endoftext|><|endof... [0.07798551, 0.8786601] 0.918471 No 0.918471 True
2 4 amazon_polarity 2 increase You are about to immerse yourself into a logic... [decrease, increase] would_you_buy 0 1 True guard <|endoftext|><|endoftext|><|endoftext|><|endof... [0.49753976, 0.3755627] 0.430142 decrease 0.430142 False
3 5 amazon_polarity 2 Yes Below is an instruction that describes a task,... [Yes, No] Is_this_review_negative 0 0 False truth <|endoftext|><|endoftext|><|endoftext|><|endof... [0.57507885, 0.37714592] 0.396064 Yes 0.396064 False
In [11]:
# ds?

What are we detecting? If the right example of the pair is more deceptive.

Now it's only deceptive if

  • it was asked to lie
  • it knows the truth
  • it gave the wrong answer (around 10% of the time)( it's hard to get these models to lie by encouragement rather than instruction)
In [12]:
from src.helpers import switch2bool, bool2switch
from src.datasets.dm import imdbHSDataModule
In [13]:
batch_size = 120
# test and cache
dm = imdbHSDataModule(ds, batch_size=batch_size)
dm.setup('train')

dl_val = dm.val_dataloader()
dl_train = dm.train_dataloader()
len(dl_train), len(dl_val)
Out [13]:
(9, 5)
In [14]:
b = next(iter(dl_train))
x0, y = b
x0.shape
Out [14]:
torch.Size([120, 4, 2816, 2])

Data prep

We do two inferences on the same inputs. Since we have dropout enabled, even during inference, we get two slightly different hidden states hs1 and hs2, and two slightly different probabilities for our yes and no output tokens p1 p2. We also have the true answer t

So there are a few ways we can set up the problem.

We can vary x:

  • model(hs1)-model(hs2)=y
  • model(hs1-hs2)==y

And we can try differen't y's:

  • direction with a ranked loss. This could be unsupervised.
  • magnitude with a regression loss
  • vector (direction and magnitude) with a regression loss

QC: Linear supervised probes

Let's verify that the model's representations are good

Before trying CCS, let's make sure there exists a direction that classifies examples as true vs false with high accuracy; if supervised logistic regression accuracy is bad, there's no hope of unsupervised CCS doing well.

Note that because logistic regression is supervised we expect it to do better but to have worse generalisation that equivilent unsupervised methods. However in this case CSS is using a deeper model so it is more complicated.

Try a classification of direction to truth

In [15]:
# dm.y
In [16]:
# n = len(df)

# # Define X and y
# X = (dm.hs1-dm.hs0).reshape((n, -1))#/dm.y[:, None]
# y = dm.y>0

# # split
# n = len(y)
# max_rows = 300
# 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=100)
# lr.fit(X_train2, y_train>0)
In [17]:
# y.mean()
In [18]:
# 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)))

# m = df['instructed_to_lie'][n//2:][:max_rows]
# y_test_pred = lr.predict(X_test2)
# acc_w_lie = ((y_test_pred[m]>0)==(y_test[m]>0)).mean()
# acc_wo_lie = ((y_test_pred[~m]>0)==(y_test[~m]>0)).mean()
# print(f'test acc w lie {acc_w_lie:2.2%}')
# print(f'test acc wo lie {acc_wo_lie:2.2%}')
In [19]:
# primary_baseline = roc_auc_score(y_test>0, y_test_pred)
# primary_baseline

LightningModel

In [20]:
# # from src.probes.conv import PLConvProbe
# import torch
# import torch.nn as nn
# import torch.nn.functional as F
# from src.probes.conv import PLConvProbe
# from src.probes.pl_ranking import PLRanking
# from torchmetrics.functional import accuracy
# from src.helpers import switch2bool, bool2switch

# class ConvProbe(nn.Module):
#     def __init__(self, c_in, depth=0, hs=16, dropout=0, input_dropout=0):
#         super().__init__()

#         layers = [
#             nn.BatchNorm1d(c_in, affine=False),  # this will normalise the inputs
#             nn.Dropout1d(input_dropout),
            
#             nn.Conv1d(c_in, hs*(depth+1), kernel_size=3),
#             nn.ReLU(),
#             nn.BatchNorm1d(hs*(depth+1)),
#             nn.AdaptiveAvgPool1d(5),
#             nn.Flatten(),
#             nn.Linear(hs*(depth+1)*5, hs*(depth+1)),
#         ]
#         for i in range(depth):
#             layers += [
#                 nn.Linear(hs*(depth-i+1), hs*(depth-i)),
#                 nn.ReLU(),
#                 nn.BatchNorm1d(hs*(depth-i)),
                
#             ]
#         # layers += [nn.AdaptiveAvgPool1d(1)]
#         self.net = nn.Sequential(*layers)
#         self.head = nn.Sequential(
#             nn.Linear(hs, hs), nn.ReLU(),
#             nn.Dropout(dropout), nn.Linear(hs, 1)            
#         )

#     def forward(self, x):
#         h = self.net(x)
#         # print(1, h.shape)
#         h = h.squeeze(-1)
#         # print(1, h.shape)
#         return self.head(h)

# class PLConvProbe(PLRanking):
#     def __init__(self, c_in, total_steps, lr=4e-3, weight_decay=1e-9, **kwargs):
#         super().__init__(total_steps=total_steps, lr=lr, weight_decay=weight_decay)
#         self.probe = ConvProbe(c_in, **kwargs)
#         self.save_hyperparameters()
        
        
#     def _step(self, batch, batch_idx, stage='train'):
#         x0, x1, y = batch
#         ypred0 = self(x0)
#         ypred1 = self(x1)
        
#         if stage=='pred':
#             return (ypred1-ypred0).float()
        
#         # loss = F.smooth_l1_loss(ypred1-ypred0, y)
#         loss = F.margin_ranking_loss(ypred1, ypred0, y, margin=0.5)
#         # self.log(f"{stage}/loss", loss)
        
#         y_cls = switch2bool(ypred1-ypred0)
#         self.log(f"{stage}/acc", accuracy(y_cls, y>0, "binary"), on_epoch=True, on_step=False)
#         self.log(f"{stage}/loss", loss, on_epoch=True, on_step=False)
#         self.log(f"{stage}/n", len(y), on_epoch=True, on_step=False, reduce_fx=torch.sum)
#         return loss
In [21]:

def dice_loss(input, target):
    smooth = 1.

    iflat = input.view(-1)
    tflat = target.view(-1)
    intersection = (iflat * tflat).sum()
    
    return 1 - ((2. * intersection + smooth) /
              (iflat.sum() + tflat.sum() + smooth))
In [22]:
# from src.probes.conv import PLConvProbe
import torch
import torch.nn as nn
import torch.nn.functional as F
from einops import rearrange
from src.probes.conv import PLConvProbe
from src.probes.pl_ranking import PLRanking
from torchmetrics.functional import accuracy
from src.helpers import switch2bool, bool2switch

class ConvProbe(nn.Module):
    def __init__(self, c_in, depth=0, hs=16, dropout=0, input_dropout=0):
        super().__init__()
        # self.n_groups = 24 # groups of neurons
        # c = c_in//self.n_groups
        c = c_in
        P = 1

        cw = hs*(depth+1)
        self.layers1 = nn.Sequential(*[
            nn.BatchNorm2d(c, affine=False),  # this will normalise the inputs
            # nn.Dropout2d(input_dropout),
            
            nn.Conv2d(c, c//2, kernel_size=(1, 2)),
            nn.ReLU(),
            nn.BatchNorm2d(c//2),
            nn.Conv2d(c//2, cw*6, kernel_size=(2, 1)),
            nn.ReLU(),
            nn.BatchNorm2d(cw*6),
            
            nn.Conv2d(cw*6, cw, kernel_size=(1, 1)),
            # nn.Conv2d(cw, cw, kernel_size=(3, 1)),
            nn.ReLU(),
            nn.BatchNorm2d(cw),
            
            
            # nn.Conv2d(cw, cw, kernel_size=(1, 3)),
            # nn.Conv2d(cw, cw, kernel_size=(3, 1)),
            # nn.ReLU(),
            # nn.BatchNorm2d(cw),            
            
            nn.AdaptiveAvgPool2d(P),
            nn.Flatten(),
            
        ])
        layers2 = [nn.Linear(hs*(depth+1)*P*P, hs*(depth+1)),]
        for i in range(depth):
            layers2 += [
                nn.Linear(hs*(depth-i+1), hs*(depth-i)),
                nn.ReLU(),
                nn.BatchNorm1d(hs*(depth-i)),
                
            ]
        # layers += [nn.AdaptiveAvgPool1d(1)]
        self.layers2 = nn.Sequential(*layers2)
        self.head = nn.Sequential(
            nn.Linear(hs, hs), nn.ReLU(),
            nn.Dropout(dropout), nn.Linear(hs, 1)            
        )

    def forward(self, x):
        # torch.Size([76, 4, 2816, 2])
        x = rearrange(x, 'b l hs f -> b hs l f')
        # x = x.reshape((len(x), -1, self.n_groups, x.shape[-1]))
        # print(x.shape, 3)
        h = self.layers1(x)
        # print(h.shape, 4)
        h = self.layers2(h)
        # print(h.shape, 5)
        # print(1, h.shape)
        h = h.squeeze(-1)
        # print(1, h.shape)
        return self.head(h)

class PLConvProbe(PLRanking):
    def __init__(self, c_in, total_steps, lr=4e-3, weight_decay=1e-9, **kwargs):
        super().__init__(total_steps=total_steps, lr=lr, weight_decay=weight_decay)
        self.probe = ConvProbe(c_in, **kwargs)
        self.save_hyperparameters()
        
        
    def _step(self, batch, batch_idx, stage='train'):
        x0, y = batch
        y_pred_logit = self(x0)
        y_pred = F.sigmoid(y_pred_logit)
        
        if stage=='pred':
            return y_pred.float()
        
        # loss = F.smooth_l1_loss(ypred1-ypred0, y)
        # loss = F.margin_ranking_loss(ypred1, ypred0, y, margin=0.5)
        # self.log(f"{stage}/loss", loss)
        
        # TODO dice loss?
        # loss = F.binary_cross_entropy_with_logits(y_pred_logit, y)
        loss = dice_loss(y_pred, y)
        
        y_cls = y_pred>0.5 # switch2bool(ypred1-ypred0)
        self.log(f"{stage}/acc", accuracy(y_cls, y>0.5, "binary"), on_epoch=True, on_step=False)
        self.log(f"{stage}/loss", loss, on_epoch=True, on_step=False)
        self.log(f"{stage}/n", len(y), on_epoch=True, on_step=False, reduce_fx=torch.sum)
        return loss
In [ ]:

Run

In [23]:
# quiet please
torch.set_float32_matmul_precision('medium')

import warnings
warnings.filterwarnings("ignore", ".*does not have many workers.*")
warnings.filterwarnings("ignore", ".*F-score.*")

Prep dataloader/set

In [24]:
dl_train = dm.train_dataloader()
dl_val = dm.val_dataloader()
b = next(iter(dl_train))
In [25]:
max_epochs = 82
batch_size = 6

c_in = b[0].shape[2]
print(b[0].shape)
net = PLConvProbe(c_in=c_in, total_steps=max_epochs*len(dl_train), depth=0, hs=23, lr=3e-3, 
          weight_decay=.1, 
          dropout=0.1, 
        #   input_dropout=0.3,
          )
torch.Size([120, 4, 2816, 2])
In [26]:
from torchinfo import summary

summary(net, input_size=b[0].shape)
Out [26]:
==========================================================================================
Layer (type:depth-idx)                   Output Shape              Param #
==========================================================================================
PLConvProbe                              [120]                     --
├─ConvProbe: 1-1                         [120, 1]                  --
│    └─Sequential: 2-1                   [120, 23]                 --
│    │    └─BatchNorm2d: 3-1             [120, 2816, 4, 2]         --
│    │    └─Conv2d: 3-2                  [120, 1408, 4, 1]         7,931,264
│    │    └─ReLU: 3-3                    [120, 1408, 4, 1]         --
│    │    └─BatchNorm2d: 3-4             [120, 1408, 4, 1]         2,816
│    │    └─Conv2d: 3-5                  [120, 138, 3, 1]          388,746
│    │    └─ReLU: 3-6                    [120, 138, 3, 1]          --
│    │    └─BatchNorm2d: 3-7             [120, 138, 3, 1]          276
│    │    └─Conv2d: 3-8                  [120, 23, 3, 1]           3,197
│    │    └─ReLU: 3-9                    [120, 23, 3, 1]           --
│    │    └─BatchNorm2d: 3-10            [120, 23, 3, 1]           46
│    │    └─AdaptiveAvgPool2d: 3-11      [120, 23, 1, 1]           --
│    │    └─Flatten: 3-12                [120, 23]                 --
│    └─Sequential: 2-2                   [120, 23]                 --
│    │    └─Linear: 3-13                 [120, 23]                 552
│    └─Sequential: 2-3                   [120, 1]                  --
│    │    └─Linear: 3-14                 [120, 23]                 552
│    │    └─ReLU: 3-15                   [120, 23]                 --
│    │    └─Dropout: 3-16                [120, 23]                 --
│    │    └─Linear: 3-17                 [120, 1]                  24
==========================================================================================
Total params: 8,327,473
Trainable params: 8,327,473
Non-trainable params: 0
Total mult-adds (Units.GIGABYTES): 3.95
==========================================================================================
Input size (MB): 10.81
Forward/backward pass size (MB): 11.79
Params size (MB): 33.31
Estimated Total Size (MB): 55.91
==========================================================================================
In [ ]:
In [27]:

# init the model

trainer = pl.Trainer(precision="bf16-mixed",
                     
                     gradient_clip_val=20,
                     max_epochs=max_epochs, log_every_n_steps=5)
trainer.fit(model=net, train_dataloaders=dl_train, val_dataloaders=dl_val)
Using bfloat16 Automatic Mixed Precision (AMP)
GPU available: True (cuda), used: True
TPU available: False, using: 0 TPU cores
IPU available: False, using: 0 IPUs
HPU available: False, using: 0 HPUs
/home/ubuntu/mambaforge/envs/dlk3/lib/python3.11/site-packages/lightning/pytorch/trainer/connectors/logger_connector/logger_connector.py:67: UserWarning: Starting from v1.9.0, `tensorboardX` has been removed as a dependency of the `lightning.pytorch` package, due to potential conflicts with other packages in the ML ecosystem. For this reason, `logger=True` will use `CSVLogger` as the default logger, unless the `tensorboard` or `tensorboardX` packages are found. Please `pip install lightning[extra]` or one of them to enable TensorBoard support by default
  warning_cache.warn(
LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]

  | Name  | Type      | Params
------------------------------------
0 | probe | ConvProbe | 8.3 M 
------------------------------------
8.3 M     Trainable params
0         Non-trainable params
8.3 M     Total params
33.310    Total estimated model params size (MB)

Sanity Checking: 0it [00:00, ?it/s]
/home/ubuntu/mambaforge/envs/dlk3/lib/python3.11/site-packages/lightning/pytorch/trainer/connectors/logger_connector/result.py:212: UserWarning: You called `self.log('val/n', ...)` in your `validation_step` but the value needs to be floating point. Converting it to torch.float32.
  warning_cache.warn(

Training: 0it [00:00, ?it/s]
/home/ubuntu/mambaforge/envs/dlk3/lib/python3.11/site-packages/lightning/pytorch/trainer/connectors/logger_connector/result.py:212: UserWarning: You called `self.log('train/n', ...)` in your `training_step` but the value needs to be floating point. Converting it to torch.float32.
  warning_cache.warn(

Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
`Trainer.fit` stopped: `max_epochs=82` reached.

Read hist

In [28]:
df_hist = read_metrics_csv(trainer.logger.experiment.metrics_file_path).ffill().bfill()
df_hist
Out [28]:
val/acc val/loss val/n step train/acc train/loss train/n
epoch
0 0.786916 0.357970 535.0 8.0 0.740187 0.327068 1070.0
1 0.728972 0.350426 535.0 17.0 0.710280 0.301784 1070.0
2 0.730841 0.320518 535.0 26.0 0.733645 0.292492 1070.0
3 0.714019 0.292186 535.0 35.0 0.744860 0.280650 1070.0
4 0.723364 0.280204 535.0 44.0 0.757944 0.273387 1070.0
... ... ... ... ... ... ... ...
77 0.906542 0.061011 535.0 701.0 1.000000 0.000516 1070.0
78 0.904673 0.063931 535.0 710.0 1.000000 0.000271 1070.0
79 0.902804 0.062077 535.0 719.0 1.000000 0.000804 1070.0
80 0.902804 0.061340 535.0 728.0 1.000000 0.000572 1070.0
81 0.902804 0.061838 535.0 737.0 1.000000 0.000000 1070.0

82 rows × 7 columns

In [29]:
for key in ['loss']:
    df_hist[[c for c in df_hist.columns if key in c]].plot(logy=True)
In [30]:
for key in ['acc']:
    df_hist[[c for c in df_hist.columns if key in c]].plot()

Predict

In [31]:
dl_test = dm.test_dataloader()
rs = trainer.test(net, dataloaders=[dl_train, dl_val, dl_test])
rs
Out [31]:
LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]
/home/ubuntu/mambaforge/envs/dlk3/lib/python3.11/site-packages/lightning/pytorch/trainer/connectors/data_connector.py:480: PossibleUserWarning: Your `test_dataloader`'s sampler has shuffling enabled, it is strongly recommended that you turn shuffling off for val/test dataloaders.
  rank_zero_warn(

Testing: 0it [00:00, ?it/s]
/home/ubuntu/mambaforge/envs/dlk3/lib/python3.11/site-packages/lightning/pytorch/trainer/connectors/logger_connector/result.py:212: UserWarning: You called `self.log('test/n', ...)` in your `test_step.0` but the value needs to be floating point. Converting it to torch.float32.
  warning_cache.warn(
/home/ubuntu/mambaforge/envs/dlk3/lib/python3.11/site-packages/lightning/pytorch/trainer/connectors/logger_connector/result.py:212: UserWarning: You called `self.log('test/n', ...)` in your `test_step.1` but the value needs to be floating point. Converting it to torch.float32.
  warning_cache.warn(
/home/ubuntu/mambaforge/envs/dlk3/lib/python3.11/site-packages/lightning/pytorch/trainer/connectors/logger_connector/result.py:212: UserWarning: You called `self.log('test/n', ...)` in your `test_step.2` but the value needs to be floating point. Converting it to torch.float32.
  warning_cache.warn(

┏━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━┓
┃   Runningstage.testing    ┃                           ┃                           ┃                           ┃
┃          metric           ┃       DataLoader 0        ┃       DataLoader 1        ┃       DataLoader 2        ┃
┡━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━┩
│         test/acc          │    0.9990653991699219     │    0.9028037190437317     │    0.9046729207038879     │
│         test/loss         │   0.0006162095232866704   │   0.061838116496801376    │   0.061967477202415466    │
│          test/n           │          1070.0           │           535.0           │           535.0           │
└───────────────────────────┴───────────────────────────┴───────────────────────────┴───────────────────────────┘

[{'test/acc/dataloader_idx_0': 0.9990653991699219,
  'test/loss/dataloader_idx_0': 0.0006162095232866704,
  'test/n/dataloader_idx_0': 1070.0},
 {'test/acc/dataloader_idx_1': 0.9028037190437317,
  'test/loss/dataloader_idx_1': 0.061838116496801376,
  'test/n/dataloader_idx_1': 535.0},
 {'test/acc/dataloader_idx_2': 0.9046729207038879,
  'test/loss/dataloader_idx_2': 0.061967477202415466,
  'test/n/dataloader_idx_2': 535.0}]
In [32]:
dl_test = dm.test_dataloader()
r = trainer.predict(net, dataloaders=dl_test)
y_test_pred = np.concatenate(r)
y_test_pred.shape
Out [32]:
LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]

Predicting: 0it [00:00, ?it/s]
(535,)
In [ ]:
In [33]:
df_test = dm.df.iloc[dm.splits['test'][0]:].copy()
y_true = dl_test.dataset.tensors[1].numpy()
In [34]:
# Make a prediction dataframe with everything in it
df_test = dm.df.iloc[dm.splits['test'][0]:].copy()
df_test['probe_pred'] = y_test_pred>0.5
df_test['probe_prob'] = y_test_pred
df_test['llm_prob'] =  df_test['ans0']#(df_test['ans0']+df_test['ans1'])/2
df_test['llm_ans'] = df_test['llm_prob']>0.5
# df_test['conf'] = df_test['ans0'] # (df_test['ans0']-df_test['ans1']).abs()
# df_test['y'] = df_test['y']>0.5

y_true = dl_test.dataset.tensors[1].numpy()
assert ((df_test['y'].values>0.5)==(y_true>0.5)).all(), 'check it all lines up'

df_test.head(3)
Out [34]:
ds_index ds_string example_i answer question answer_choices template_name label_true label_instructed instructed_to_lie ... prompt_truncated choice_probs0 ans0 txt_ans0 dir_true llm_ans y probe_pred probe_prob llm_prob
1605 2241 amazon_polarity 1120 Yes Below is an instruction that describes a task,... [Yes, No] Is_this_review_negative 0 0 False ... <|endoftext|><|endoftext|><|endoftext|><|endof... [0.6307559, 0.31716397] 0.334586 Yes 0.334586 False True True 1.000000e+00 0.334586
1606 2242 amazon_polarity 1121 decrease You are about to immerse yourself into a logic... [decrease, increase] would_you_buy 1 0 True ... <|endoftext|><|endoftext|><|endoftext|><|endof... [0.040256858, 0.9306706] 0.958528 increase 0.958528 True True False 5.995204e-15 0.958528
1607 2243 amazon_polarity 1121 No Below is an instruction that describes a task,... [Yes, No] Is_this_review_negative 1 1 False ... <|endoftext|><|endoftext|><|endoftext|><|endof... [0.3965878, 0.4636581] 0.538977 No 0.538977 True True True 1.000000e+00 0.538977

3 rows × 21 columns

In [35]:
def get_acc_subset(df, query):
    df_s = df.query(query)
    acc = (df_s['probe_pred']==df_s['y']).mean()
    print(f"acc={acc:2.2%},\tn={len(df_s)},\t[{query}] ")
    return acc
    
print('probe results on subsets of the data')
get_acc_subset(df_test, 'instructed_to_lie==True') # it was ph told to lie
get_acc_subset(df_test, 'instructed_to_lie==False') # it was told not to lie
get_acc_subset(df_test, 'llm_ans==label_true') # the llm gave the true ans
get_acc_subset(df_test, 'llm_ans==label_instructed') # the llm gave the desired ans
get_acc_subset(df_test, 'instructed_to_lie==True & llm_ans==label_instructed') # it was told to lie, and it did lie
get_acc_subset(df_test, 'instructed_to_lie==True & llm_ans!=label_instructed');
probe results on subsets of the data
acc=81.27%,	n=267,	[instructed_to_lie==True] 
acc=99.63%,	n=268,	[instructed_to_lie==False] 
acc=92.05%,	n=415,	[llm_ans==label_true] 
acc=95.10%,	n=388,	[llm_ans==label_instructed] 
acc=85.00%,	n=120,	[instructed_to_lie==True & llm_ans==label_instructed] 
acc=78.23%,	n=147,	[instructed_to_lie==True & llm_ans!=label_instructed]

RESULTS

In [ ]:
In [36]:
acc = (df_test['y']==(y_test_pred>0.5)).mean()

# print(f"  PRIMARY BASELINE roc_auc={primary_baseline:2.2%} from linear classifier")
print(f"⭐PRIMARY METRIC⭐ acc={acc:2.2%} from probe")
⭐PRIMARY METRIC⭐ acc=90.47% from probe

Out of sample

Lets see how far it generalizes

In [37]:
def try_fine_tune(dm):
    dl_train = dm.train_dataloader()
    dl_val = dm.val_dataloader()
    dl_test = dm.test_dataloader()
    b = next(iter(dl_train))
    max_epochs = 42
    c_in = b[0].shape[2]
    print(b[0].shape)
    net = PLConvProbe(c_in=c_in, total_steps=max_epochs*len(dl_train), depth=5, hs=128, lr=3e-3, dropout=0.1, input_dropout=0.1)
    trainer = pl.Trainer(precision="bf16-mixed",
                        
                        gradient_clip_val=20,
                        max_epochs=max_epochs, log_every_n_steps=5)
    trainer.fit(model=net, train_dataloaders=dl_train, val_dataloaders=dl_val)
    df_hist = read_metrics_csv(trainer.logger.experiment.metrics_file_path).ffill().bfill()
    rs = trainer.test(net, dataloaders=[dl_train, dl_val, dl_test])
    return df_hist, rs
In [38]:
oos_dataset_fs = [
    '../.ds/WizardLMWizardCoder_3B_V1.0_amazon_polarity_train_3000'
]
In [39]:
batch_size = 12
for f in oos_dataset_fs:
    print(f)
    ds2a = load_from_disk(f)

    # # restrict it to significant permutations. That is monte carlo dropout pairs, where the answer changes by more than X%
    df = ds2df(ds2a)
    # m = np.abs(df.ans0-df.ans1)>0.1
    # significant_rows = m[m].index
    

    # # these are the ones where it got it right when asked to tell the truth
    m1 = d.llm_ans==d.label_true
    known_indices = d[m1].index
    print(f"select rows are {m1.mean():2.2%} based on knowledge")
    # # convert to row numbers, and use datasets to select
    known_rows = df['example_i'].isin(known_indices)
    known_rows_i = df[known_rows].index

    # allowed_rows_i = set(known_rows_i).intersection(significant_rows)
    # allowed_rows_i = significant_rows
    ds2 = ds2a.select(known_rows_i)
    print(f"selected rows are {len(ds2)/len(ds2a):2.2%}")
    print(len(ds2))

    dm2 = imdbHSDataModule(ds2, batch_size=batch_size)
    dm2.setup('train')

    dl_val2 = dm2.val_dataloader()
    dl_train2 = dm2.train_dataloader()
    dl_test2 = dm2.test_dataloader()
    print(len(dl_train2), len(dl_val2), len(dl_test2))
    rs2 = trainer.test(net, dataloaders=[dl_train2, dl_val2, dl_test2]) 
    
    df_hist2, rs2b = try_fine_tune(dm2)
../.ds/WizardLMWizardCoder_3B_V1.0_amazon_polarity_train_3000
select rows are 71.29% based on knowledge
selected rows are 71.29%
2140

LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]

90 45 45

/home/ubuntu/mambaforge/envs/dlk3/lib/python3.11/site-packages/lightning/pytorch/trainer/connectors/data_connector.py:480: PossibleUserWarning: Your `test_dataloader`'s sampler has shuffling enabled, it is strongly recommended that you turn shuffling off for val/test dataloaders.
  rank_zero_warn(

Testing: 0it [00:00, ?it/s]
┏━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━┓
┃   Runningstage.testing    ┃                           ┃                           ┃                           ┃
┃          metric           ┃       DataLoader 0        ┃       DataLoader 1        ┃       DataLoader 2        ┃
┡━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━┩
│         test/acc          │    0.9990653991699219     │    0.9028037190437317     │    0.9046729207038879     │
│         test/loss         │   0.0004672899376600981   │    0.05780790001153946    │    0.0613081119954586     │
│          test/n           │          1070.0           │           535.0           │           535.0           │
└───────────────────────────┴───────────────────────────┴───────────────────────────┴───────────────────────────┘

Using bfloat16 Automatic Mixed Precision (AMP)
GPU available: True (cuda), used: True
TPU available: False, using: 0 TPU cores
IPU available: False, using: 0 IPUs
HPU available: False, using: 0 HPUs
LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]

  | Name  | Type      | Params
------------------------------------
0 | probe | ConvProbe | 26.2 M
------------------------------------
26.2 M    Trainable params
0         Non-trainable params
26.2 M    Total params
104.901   Total estimated model params size (MB)

torch.Size([12, 4, 2816, 2])

Sanity Checking: 0it [00:00, ?it/s]
Training: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
Validation: 0it [00:00, ?it/s]
`Trainer.fit` stopped: `max_epochs=42` reached.
LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]
/home/ubuntu/mambaforge/envs/dlk3/lib/python3.11/site-packages/lightning/pytorch/trainer/connectors/data_connector.py:480: PossibleUserWarning: Your `test_dataloader`'s sampler has shuffling enabled, it is strongly recommended that you turn shuffling off for val/test dataloaders.
  rank_zero_warn(

Testing: 0it [00:00, ?it/s]
┏━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━┓
┃   Runningstage.testing    ┃                           ┃                           ┃                           ┃
┃          metric           ┃       DataLoader 0        ┃       DataLoader 1        ┃       DataLoader 2        ┃
┡━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━┩
│         test/acc          │    0.9953271150588989     │    0.8971962332725525     │    0.8971962332725525     │
│         test/loss         │   0.002653207164257765    │    0.06099105626344681    │    0.06146525591611862    │
│          test/n           │          1070.0           │           535.0           │           535.0           │
└───────────────────────────┴───────────────────────────┴───────────────────────────┴───────────────────────────┘
In [ ]:
In [ ]:
Reference in New Issue View Git Blame Copy Permalink