mirror of
https://github.com/wassname/discovering_latent_knowledge.git
synced 2026-09-26 13:40:30 +08:00
try poetry
This commit is contained in:
wassname
Vendored
-19
@@ -9,25 +9,6 @@
|
||||
"plaintext",
|
||||
"markdown"
|
||||
],
|
||||
"workbench.colorCustomizations": {
|
||||
"activityBar.activeBackground": "#65c89b",
|
||||
"activityBar.background": "#65c89b",
|
||||
"activityBar.foreground": "#15202b",
|
||||
"activityBar.inactiveForeground": "#15202b99",
|
||||
"activityBarBadge.background": "#945bc4",
|
||||
"activityBarBadge.foreground": "#e7e7e7",
|
||||
"commandCenter.border": "#15202b99",
|
||||
"sash.hoverBorder": "#65c89b",
|
||||
"statusBar.background": "#42b883",
|
||||
"statusBar.foreground": "#15202b",
|
||||
"statusBarItem.hoverBackground": "#359268",
|
||||
"statusBarItem.remoteBackground": "#42b883",
|
||||
"statusBarItem.remoteForeground": "#15202b",
|
||||
"titleBar.activeBackground": "#42b883",
|
||||
"titleBar.activeForeground": "#15202b",
|
||||
"titleBar.inactiveBackground": "#42b88399",
|
||||
"titleBar.inactiveForeground": "#15202b99"
|
||||
},
|
||||
"peacock.remoteColor": "#42b883",
|
||||
"python.analysis.autoImportCompletions": true
|
||||
}
|
||||
|
||||
+1
-1
@@ -1,6 +1,6 @@
|
||||
MIT License
|
||||
|
||||
Copyright (c) 2023 Collin Burns
|
||||
Copyright (c) 2023 wassname
|
||||
|
||||
Permission is hereby granted, free of charge, to any person obtaining a copy
|
||||
of this software and associated documentation files (the "Software"), to deal
|
||||
|
||||
-53
@@ -1,53 +0,0 @@
|
||||
from sklearn.linear_model import LogisticRegression
|
||||
from utils import get_parser, load_all_generations, CCS
|
||||
|
||||
def main(args, generation_args):
|
||||
# load hidden states and labels
|
||||
neg_hs, pos_hs, y = load_all_generations(generation_args)
|
||||
|
||||
# Make sure the shape is correct
|
||||
assert neg_hs.shape == pos_hs.shape
|
||||
neg_hs, pos_hs = neg_hs[..., -1], pos_hs[..., -1] # take the last layer
|
||||
if neg_hs.shape[1] == 1: # T5 may have an extra dimension; if so, get rid of it
|
||||
neg_hs = neg_hs.squeeze(1)
|
||||
pos_hs = pos_hs.squeeze(1)
|
||||
|
||||
# Very simple train/test split (using the fact that the data is already shuffled)
|
||||
neg_hs_train, neg_hs_test = neg_hs[:len(neg_hs) // 2], neg_hs[len(neg_hs) // 2:]
|
||||
pos_hs_train, pos_hs_test = pos_hs[:len(pos_hs) // 2], pos_hs[len(pos_hs) // 2:]
|
||||
y_train, y_test = y[:len(y) // 2], y[len(y) // 2:]
|
||||
|
||||
# Make sure logistic regression accuracy is reasonable; otherwise our method won't have much of a chance of working
|
||||
# you can also concatenate, but this works fine and is more comparable to CCS inputs
|
||||
x_train = neg_hs_train - pos_hs_train
|
||||
x_test = neg_hs_test - pos_hs_test
|
||||
lr = LogisticRegression(class_weight="balanced")
|
||||
lr.fit(x_train, y_train)
|
||||
print("Logistic regression accuracy: {}".format(lr.score(x_test, y_test)))
|
||||
|
||||
# Set up CCS. Note that you can usually just use the default args by simply doing ccs = CCS(neg_hs, pos_hs, y)
|
||||
ccs = CCS(neg_hs_train, pos_hs_train, nepochs=args.nepochs, ntries=args.ntries, lr=args.lr, batch_size=args.ccs_batch_size,
|
||||
verbose=args.verbose, device=args.ccs_device, linear=args.linear, weight_decay=args.weight_decay,
|
||||
var_normalize=args.var_normalize)
|
||||
|
||||
# train and evaluate CCS
|
||||
ccs.repeated_train()
|
||||
ccs_acc = ccs.get_acc(neg_hs_test, pos_hs_test, y_test)
|
||||
print("CCS accuracy: {}".format(ccs_acc))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = get_parser()
|
||||
generation_args = parser.parse_args() # we'll use this to load the correct hidden states + labels
|
||||
# We'll also add some additional args for evaluation
|
||||
parser.add_argument("--nepochs", type=int, default=1000)
|
||||
parser.add_argument("--ntries", type=int, default=10)
|
||||
parser.add_argument("--lr", type=float, default=1e-3)
|
||||
parser.add_argument("--ccs_batch_size", type=int, default=-1)
|
||||
parser.add_argument("--verbose", action="store_true")
|
||||
parser.add_argument("--ccs_device", type=str, default="cuda")
|
||||
parser.add_argument("--linear", action="store_true")
|
||||
parser.add_argument("--weight_decay", type=float, default=0.01)
|
||||
parser.add_argument("--var_normalize", action="store_true")
|
||||
args = parser.parse_args()
|
||||
main(args, generation_args)
|
||||
BIN
Binary file not shown.
|
Before Width: | Height: | Size: 461 KiB |
-27
@@ -1,27 +0,0 @@
|
||||
from utils import get_parser, load_model, get_dataloader, get_all_hidden_states, save_generations
|
||||
|
||||
def main(args):
|
||||
# Set up the model and data
|
||||
print("Loading model")
|
||||
model, tokenizer, model_type = load_model(args.model_name, args.cache_dir, args.parallelize, args.device)
|
||||
|
||||
print("Loading dataloader")
|
||||
dataloader = get_dataloader(args.dataset_name, args.split, tokenizer, args.prompt_idx, batch_size=args.batch_size,
|
||||
num_examples=args.num_examples, model_type=model_type, use_decoder=args.use_decoder, device=args.device)
|
||||
|
||||
# Get the hidden states and labels
|
||||
print("Generating hidden states")
|
||||
neg_hs, pos_hs, y = get_all_hidden_states(model, dataloader, layer=args.layer, all_layers=args.all_layers,
|
||||
token_idx=args.token_idx, model_type=model_type, use_decoder=args.use_decoder)
|
||||
|
||||
# Save the hidden states and labels
|
||||
print("Saving hidden states")
|
||||
save_generations(neg_hs, args, generation_type="negative_hidden_states")
|
||||
save_generations(pos_hs, args, generation_type="positive_hidden_states")
|
||||
save_generations(y, args, generation_type="labels")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = get_parser()
|
||||
args = parser.parse_args()
|
||||
main(args)
|
||||
File diff suppressed because one or more lines are too long
Generated
+3742
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,41 @@
|
||||
[tool.poetry]
|
||||
name = "src"
|
||||
version = "0.1.0"
|
||||
description = "building a lie detector by ranking pairs of hidden states"
|
||||
authors = ["wassname <git@wassname.org>"]
|
||||
license = "MIT"
|
||||
readme = "README.md"
|
||||
|
||||
[tool.poetry.dependencies]
|
||||
python = ">=3.10,<3.13"
|
||||
accelerate = "^0.23.0"
|
||||
torch = {version = "^2.1.0+cu118", source = "pytorch"}
|
||||
simple-parsing = "^0.1.4"
|
||||
tqdm = "^4.66.1"
|
||||
datasets = "^2.14.5"
|
||||
transformers = "^4.34.1"
|
||||
auto-gptq = "^0.4.2"
|
||||
optimum = "^1.13.2"
|
||||
numpy = "^1.26.1"
|
||||
pandas = "^2.1.1"
|
||||
lightning = "^2.1.0"
|
||||
matplotlib = "^3.8.0"
|
||||
loguru = "^0.7.2"
|
||||
einops = "^0.7.0"
|
||||
baukit = {git = "https://github.com/davidbau/baukit.git"}
|
||||
eleuther-elk = "0.1.1"
|
||||
|
||||
[[tool.poetry.source]]
|
||||
name = "pytorch"
|
||||
url = "https://download.pytorch.org/whl/cu118"
|
||||
priority = "explicit"
|
||||
|
||||
|
||||
[tool.poetry.group.dev.dependencies]
|
||||
ipykernel = "^6.25.2"
|
||||
black = "^23.10.0"
|
||||
pylama = "^8.4.1"
|
||||
|
||||
[build-system]
|
||||
requires = ["poetry-core"]
|
||||
build-backend = "poetry.core.masonry.api"
|
||||
@@ -1,193 +0,0 @@
|
||||
accelerate==0.20.3
|
||||
aiohttp==3.8.4
|
||||
aiosignal==1.3.1
|
||||
altair==5.0.0
|
||||
anyio==3.6.2
|
||||
arrow==1.2.3
|
||||
astor==0.8.1
|
||||
asttokens==2.2.1
|
||||
async-timeout==4.0.2
|
||||
attrs==23.1.0
|
||||
backcall==0.2.0
|
||||
base58==2.1.1
|
||||
beautifulsoup4==4.12.2
|
||||
bitsandbytes==0.39.0
|
||||
black==21.12b0
|
||||
blessed==1.20.0
|
||||
blinker==1.6.2
|
||||
Brotli==1.0.9
|
||||
brotlipy @ file:///home/conda/feedstock_root/build_artifacts/brotlipy_1666764672617/work
|
||||
cachetools==5.3.0
|
||||
certifi==2023.5.7
|
||||
cffi @ file:///home/conda/feedstock_root/build_artifacts/cffi_1671179360775/work
|
||||
charset-normalizer @ file:///home/conda/feedstock_root/build_artifacts/charset-normalizer_1678108872112/work
|
||||
click==7.1.2
|
||||
coloredlogs==15.0.1
|
||||
comm==0.1.3
|
||||
contourpy==1.0.7
|
||||
croniter==1.3.14
|
||||
cryptography @ file:///home/conda/feedstock_root/build_artifacts/cryptography-split_1681508587436/work
|
||||
cycler==0.11.0
|
||||
datasets==2.12.0
|
||||
dateutils==0.6.12
|
||||
debugpy==1.6.7
|
||||
decorator==5.1.1
|
||||
deepdiff==6.3.0
|
||||
dill==0.3.6
|
||||
docstring-parser==0.15
|
||||
einops==0.6.1
|
||||
eleuther-elk==0.1
|
||||
exceptiongroup==1.1.1
|
||||
executing==1.2.0
|
||||
fastapi==0.88.0
|
||||
filelock @ file:///home/conda/feedstock_root/build_artifacts/filelock_1681839547898/work
|
||||
flake8==6.0.0
|
||||
focal-loss-torch==0.1.2
|
||||
fonttools==4.39.4
|
||||
frozenlist==1.3.3
|
||||
fsspec==2023.5.0
|
||||
gitdb==4.0.10
|
||||
GitPython==3.1.31
|
||||
gmpy2 @ file:///home/conda/feedstock_root/build_artifacts/gmpy2_1666808679441/work
|
||||
h11==0.14.0
|
||||
huggingface-hub==0.14.1
|
||||
humanfriendly==10.0
|
||||
idna @ file:///home/conda/feedstock_root/build_artifacts/idna_1663625384323/work
|
||||
importlib-metadata==6.6.0
|
||||
importlib-resources==5.12.0
|
||||
inflate64==0.3.1
|
||||
iniconfig==2.0.0
|
||||
inquirer==3.1.3
|
||||
ipykernel==6.23.1
|
||||
ipython==8.13.2
|
||||
ipywidgets==8.0.6
|
||||
isort==5.8.0
|
||||
itsdangerous==2.1.2
|
||||
jedi==0.18.2
|
||||
Jinja2 @ file:///home/conda/feedstock_root/build_artifacts/jinja2_1654302431367/work
|
||||
joblib==1.2.0
|
||||
jsonschema==4.17.3
|
||||
jupyter_client==8.2.0
|
||||
jupyter_core==5.3.0
|
||||
jupyterlab-widgets==3.0.7
|
||||
kiwisolver==1.4.4
|
||||
lightning==2.0.2
|
||||
lightning-cloud==0.5.36
|
||||
lightning-utilities==0.8.0
|
||||
loguru==0.7.0
|
||||
markdown-it-py==2.2.0
|
||||
MarkupSafe @ file:///home/conda/feedstock_root/build_artifacts/markupsafe_1674135804847/work
|
||||
matplotlib==3.7.1
|
||||
matplotlib-inline==0.1.6
|
||||
mccabe==0.7.0
|
||||
mdurl==0.1.2
|
||||
mpmath @ file:///home/conda/feedstock_root/build_artifacts/mpmath_1678228039184/work
|
||||
multidict==6.0.4
|
||||
multiprocess==0.70.14
|
||||
multivolumefile==0.2.3
|
||||
mypy-extensions @ file:///home/conda/feedstock_root/build_artifacts/mypy_extensions_1675543315189/work
|
||||
nest-asyncio==1.5.6
|
||||
networkx @ file:///home/conda/feedstock_root/build_artifacts/networkx_1680692919326/work
|
||||
numpy==1.25.2
|
||||
optimum==1.8.6
|
||||
ordered-set==4.1.0
|
||||
packaging @ file:///home/conda/feedstock_root/build_artifacts/packaging_1681337016113/work
|
||||
pandas==2.0.1
|
||||
parso==0.8.3
|
||||
pathspec @ file:///home/conda/feedstock_root/build_artifacts/pathspec_1678853982175/work
|
||||
peft @ git+https://github.com/huggingface/peft.git@3714aa2fff158fdfa637b2b65952580801d890b2
|
||||
pexpect==4.8.0
|
||||
pickleshare==0.7.5
|
||||
Pillow @ file:///home/conda/feedstock_root/build_artifacts/pillow_1675487166627/work
|
||||
platformdirs @ file:///home/conda/feedstock_root/build_artifacts/platformdirs_1683850015520/work
|
||||
plotly==5.14.1
|
||||
pluggy==1.0.0
|
||||
prettytable==3.8.0
|
||||
prompt-toolkit==3.0.38
|
||||
promptsource==0.2.3
|
||||
protobuf==3.20.3
|
||||
psutil==5.9.5
|
||||
ptyprocess==0.7.0
|
||||
pure-eval==0.2.2
|
||||
py7zr==0.20.5
|
||||
pyarrow==12.0.0
|
||||
pybcj==1.0.1
|
||||
pycodestyle==2.10.0
|
||||
pycparser @ file:///home/conda/feedstock_root/build_artifacts/pycparser_1636257122734/work
|
||||
pycryptodomex==3.18.0
|
||||
pydantic==1.10.7
|
||||
pydeck==0.8.1b0
|
||||
pyflakes==3.0.1
|
||||
Pygments==2.15.1
|
||||
PyJWT==2.7.0
|
||||
pynvml==11.5.0
|
||||
pyOpenSSL @ file:///home/conda/feedstock_root/build_artifacts/pyopenssl_1680037383858/work
|
||||
pyparsing==3.0.9
|
||||
pyppmd==1.0.0
|
||||
pyre-extensions==0.0.29
|
||||
pyrsistent==0.19.3
|
||||
PySocks @ file:///home/conda/feedstock_root/build_artifacts/pysocks_1661604839144/work
|
||||
pytest==7.3.1
|
||||
python-dateutil==2.8.2
|
||||
python-editor==1.0.4
|
||||
python-multipart==0.0.6
|
||||
pytorch-lightning==2.0.2
|
||||
pytorch-optimizer==2.10.1
|
||||
pytz==2023.3
|
||||
PyYAML==6.0
|
||||
pyzmq==25.0.2
|
||||
pyzstd==0.15.7
|
||||
readchar==4.0.5
|
||||
regex==2023.5.5
|
||||
requests @ file:///home/conda/feedstock_root/build_artifacts/requests_1682535435083/work
|
||||
responses==0.18.0
|
||||
rich==13.3.5
|
||||
safetensors==0.3.1
|
||||
scikit-learn==1.2.2
|
||||
scipy==1.10.1
|
||||
sentencepiece==0.1.97
|
||||
simple-parsing==0.1.4
|
||||
six==1.16.0
|
||||
sklearn==0.0.post5
|
||||
smmap==5.0.0
|
||||
sniffio==1.3.0
|
||||
soupsieve==2.4.1
|
||||
-e git+https://github.com/wassname/discovering_latent_knowledge.git@5b0c2d6ff4b364c9cf78786edc7e344cfa26fc31#egg=src
|
||||
stack-data==0.6.2
|
||||
starlette==0.22.0
|
||||
starsessions==1.3.0
|
||||
streamlit==0.82.0
|
||||
sympy @ file:///home/conda/feedstock_root/build_artifacts/sympy_1684180540116/work
|
||||
tenacity==8.2.2
|
||||
texttable==1.6.7
|
||||
threadpoolctl==3.1.0
|
||||
tokenize-rt==5.0.0
|
||||
tokenizers==0.13.3
|
||||
toml==0.10.2
|
||||
tomli==1.2.3
|
||||
toolz==0.12.0
|
||||
torch==2.0.1
|
||||
torchaudio==2.0.2
|
||||
torchmetrics==0.11.4
|
||||
torchvision==0.15.2
|
||||
tornado==6.3.2
|
||||
tqdm==4.65.0
|
||||
traitlets==5.9.0
|
||||
transformers==4.30.1
|
||||
triton==2.0.0
|
||||
typing-inspect==0.9.0
|
||||
typing_extensions @ file:///home/conda/feedstock_root/build_artifacts/typing_extensions_1678559861143/work
|
||||
tzdata==2023.3
|
||||
tzlocal==5.0.1
|
||||
urllib3 @ file:///home/conda/feedstock_root/build_artifacts/urllib3_1678635778344/work
|
||||
uvicorn==0.22.0
|
||||
validators==0.20.0
|
||||
watchdog==3.0.0
|
||||
wcwidth==0.2.6
|
||||
websocket-client==1.5.1
|
||||
websockets==11.0.3
|
||||
widgetsnbextension==4.0.7
|
||||
xformers==0.0.20
|
||||
xxhash==3.2.0
|
||||
yarl==1.9.2
|
||||
zipp==3.15.0
|
||||
@@ -1,15 +0,0 @@
|
||||
#!/bin/bash
|
||||
# Sometimes I like to do relaxed and strict requirements. The strict requirements lets you debug subtle version errors by asking "gee what exact version did they use".
|
||||
# The relaxed versioning makes it easy to upgrade
|
||||
set -e -x
|
||||
PROJECT_NAME=dlk2
|
||||
echo $PROJECT_NAME
|
||||
PYTHON_INTERPRETER=~/mambaforge/envs/$PROJECT_NAME/bin/python
|
||||
# minimal requirement, simpler, but no versions or pip
|
||||
conda env export --no-builds --from-history > requirements/environment.min.yaml
|
||||
# extensive requirements including pip and information overload
|
||||
conda env export > requirements/environment.max.yaml
|
||||
# requirements in a modified pip spec, usefull for dependabot and so on
|
||||
$PYTHON_INTERPRETER -m pip freeze > requirements/conda.requirements.txt
|
||||
# some pople like conda lock, but it doens't do pip
|
||||
# cd requirements && conda-lock -f environment.max.yaml -p linux-64
|
||||
@@ -1,230 +0,0 @@
|
||||
name: dlk4
|
||||
channels:
|
||||
- pytorch
|
||||
- nvidia
|
||||
- conda-forge
|
||||
dependencies:
|
||||
- _libgcc_mutex=0.1=conda_forge
|
||||
- _openmp_mutex=4.5=2_kmp_llvm
|
||||
- blas=2.116=mkl
|
||||
- blas-devel=3.9.0=16_linux64_mkl
|
||||
- brotli-python=1.1.0=py311hb755f60_0
|
||||
- bzip2=1.0.8=h7f98852_4
|
||||
- ca-certificates=2023.7.22=hbcca054_0
|
||||
- charset-normalizer=3.2.0=pyhd8ed1ab_0
|
||||
- cuda-cudart=11.7.99=0
|
||||
- cuda-cupti=11.7.101=0
|
||||
- cuda-libraries=11.7.1=0
|
||||
- cuda-nvrtc=11.7.99=0
|
||||
- cuda-nvtx=11.7.91=0
|
||||
- cuda-runtime=11.7.1=0
|
||||
- cudatoolkit=11.7.0=hd8887f6_10
|
||||
- cudatoolkit-dev=11.7.0=h1de0b5d_6
|
||||
- ffmpeg=4.3=hf484d3e_0
|
||||
- filelock=3.12.4=pyhd8ed1ab_0
|
||||
- freetype=2.12.1=h267a509_2
|
||||
- gmp=6.2.1=h58526e2_0
|
||||
- gmpy2=2.1.2=py311h6a5fa03_1
|
||||
- gnutls=3.6.13=h85f3911_1
|
||||
- icu=73.2=h59595ed_0
|
||||
- idna=3.4=pyhd8ed1ab_0
|
||||
- jinja2=3.1.2=pyhd8ed1ab_1
|
||||
- jpeg=9e=h0b41bf4_3
|
||||
- lame=3.100=h166bdaf_1003
|
||||
- lcms2=2.15=hfd0df8a_0
|
||||
- ld_impl_linux-64=2.40=h41732ed_0
|
||||
- lerc=4.0.0=h27087fc_0
|
||||
- libblas=3.9.0=16_linux64_mkl
|
||||
- libcblas=3.9.0=16_linux64_mkl
|
||||
- libcublas=11.10.3.66=0
|
||||
- libcufft=10.7.2.124=h4fbf590_0
|
||||
- libcufile=1.7.2.10=0
|
||||
- libcurand=10.3.3.141=0
|
||||
- libcusolver=11.4.0.1=0
|
||||
- libcusparse=11.7.4.91=0
|
||||
- libdeflate=1.17=h0b41bf4_0
|
||||
- libexpat=2.5.0=hcb278e6_1
|
||||
- libffi=3.4.2=h7f98852_5
|
||||
- libgcc-ng=13.2.0=h807b86a_2
|
||||
- libgfortran-ng=13.2.0=h69a702a_2
|
||||
- libgfortran5=13.2.0=ha4646dd_2
|
||||
- libgomp=13.2.0=h807b86a_2
|
||||
- libhwloc=2.9.2=default_h554bfaf_1009
|
||||
- libiconv=1.17=h166bdaf_0
|
||||
- liblapack=3.9.0=16_linux64_mkl
|
||||
- liblapacke=3.9.0=16_linux64_mkl
|
||||
- libnpp=11.7.4.75=0
|
||||
- libnsl=2.0.0=h7f98852_0
|
||||
- libnvjpeg=11.8.0.2=0
|
||||
- libpng=1.6.39=h753d276_0
|
||||
- libsqlite=3.43.0=h2797004_0
|
||||
- libstdcxx-ng=13.2.0=h7e041cc_2
|
||||
- libtiff=4.5.0=h6adf6a1_2
|
||||
- libuuid=2.38.1=h0b41bf4_0
|
||||
- libwebp-base=1.3.2=hd590300_0
|
||||
- libxcb=1.13=h7f98852_1004
|
||||
- libxml2=2.11.5=h232c23b_1
|
||||
- libzlib=1.2.13=hd590300_5
|
||||
- llvm-openmp=16.0.6=h4dfa4b3_0
|
||||
- markupsafe=2.1.3=py311h459d7ec_1
|
||||
- mkl=2022.1.0=h84fe81f_915
|
||||
- mkl-devel=2022.1.0=ha770c72_916
|
||||
- mkl-include=2022.1.0=h84fe81f_915
|
||||
- mpc=1.3.1=hfe3b2da_0
|
||||
- mpfr=4.2.0=hb012696_0
|
||||
- mpmath=1.3.0=pyhd8ed1ab_0
|
||||
- ncurses=6.4=hcb278e6_0
|
||||
- nettle=3.6=he412f7d_0
|
||||
- networkx=3.1=pyhd8ed1ab_0
|
||||
- numpy=1.26.0=py311h64a7726_0
|
||||
- openh264=2.1.1=h780b84a_0
|
||||
- openjpeg=2.5.0=hfec8fc6_2
|
||||
- openssl=3.1.3=hd590300_0
|
||||
- pillow=9.4.0=py311h50def17_1
|
||||
- pip=23.2.1=pyhd8ed1ab_0
|
||||
- pthread-stubs=0.4=h36c2ea0_1001
|
||||
- pysocks=1.7.1=pyha2e5f31_6
|
||||
- python=3.11.5=hab00c5b_0_cpython
|
||||
- python_abi=3.11=4_cp311
|
||||
- pytorch=2.0.1=py3.11_cuda11.7_cudnn8.5.0_0
|
||||
- pytorch-cuda=11.7=h778d358_5
|
||||
- pytorch-mutex=1.0=cuda
|
||||
- readline=8.2=h8228510_1
|
||||
- requests=2.31.0=pyhd8ed1ab_0
|
||||
- setuptools=68.2.2=pyhd8ed1ab_0
|
||||
- sympy=1.12=pypyh9d50eac_103
|
||||
- tbb=2021.10.0=h00ab1b0_0
|
||||
- tk=8.6.12=h27826a3_0
|
||||
- torchaudio=2.0.2=py311_cu117
|
||||
- torchtriton=2.0.0=py311
|
||||
- torchvision=0.15.2=py311_cu117
|
||||
- typing_extensions=4.8.0=pyha770c72_0
|
||||
- urllib3=2.0.5=pyhd8ed1ab_0
|
||||
- wheel=0.41.2=pyhd8ed1ab_0
|
||||
- xorg-libxau=1.0.11=hd590300_0
|
||||
- xorg-libxdmcp=1.1.3=h7f98852_0
|
||||
- xz=5.2.6=h166bdaf_0
|
||||
- zlib=1.2.13=hd590300_5
|
||||
- zstd=1.5.5=hfc55251_0
|
||||
- pip:
|
||||
- accelerate==0.23.0
|
||||
- aiohttp==3.8.5
|
||||
- aiosignal==1.3.1
|
||||
- annotated-types==0.5.0
|
||||
- anyio==3.7.1
|
||||
- arrow==1.2.3
|
||||
- asttokens==2.4.0
|
||||
- async-timeout==4.0.3
|
||||
- attrs==23.1.0
|
||||
- backcall==0.2.0
|
||||
- backoff==2.2.1
|
||||
- baukit==0.0.1
|
||||
- beautifulsoup4==4.12.2
|
||||
- blessed==1.20.0
|
||||
- certifi==2023.7.22
|
||||
- click==8.1.7
|
||||
- cmake==3.27.5
|
||||
- comm==0.1.4
|
||||
- contourpy==1.1.1
|
||||
- croniter==1.4.1
|
||||
- cycler==0.11.0
|
||||
- datasets==2.14.5
|
||||
- dateutils==0.6.12
|
||||
- debugpy==1.8.0
|
||||
- decorator==5.1.1
|
||||
- deepdiff==6.5.0
|
||||
- dill==0.3.7
|
||||
- docstring-parser==0.15
|
||||
- einops==0.6.1
|
||||
- executing==1.2.0
|
||||
- fastapi==0.103.1
|
||||
- fonttools==4.42.1
|
||||
- frozenlist==1.4.0
|
||||
- fsspec==2023.6.0
|
||||
- h11==0.14.0
|
||||
- huggingface-hub==0.17.2
|
||||
- inquirer==3.1.3
|
||||
- ipykernel==6.25.2
|
||||
- ipython==8.15.0
|
||||
- ipywidgets==8.1.1
|
||||
- itsdangerous==2.1.2
|
||||
- jedi==0.19.0
|
||||
- jupyter-client==8.3.1
|
||||
- jupyter-core==5.3.1
|
||||
- jupyterlab-widgets==3.0.9
|
||||
- kiwisolver==1.4.5
|
||||
- lightning==2.0.9
|
||||
- lightning-cloud==0.5.38
|
||||
- lightning-utilities==0.9.0
|
||||
- lit==16.0.6
|
||||
- loguru==0.7.2
|
||||
- markdown-it-py==3.0.0
|
||||
- matplotlib==3.8.0
|
||||
- matplotlib-inline==0.1.6
|
||||
- mdurl==0.1.2
|
||||
- multidict==6.0.4
|
||||
- multiprocess==0.70.15
|
||||
- nest-asyncio==1.5.8
|
||||
- nvidia-cublas-cu11==11.10.3.66
|
||||
- nvidia-cuda-cupti-cu11==11.7.101
|
||||
- nvidia-cuda-nvrtc-cu11==11.7.99
|
||||
- nvidia-cuda-runtime-cu11==11.7.99
|
||||
- nvidia-cudnn-cu11==8.5.0.96
|
||||
- nvidia-cufft-cu11==10.9.0.58
|
||||
- nvidia-curand-cu11==10.2.10.91
|
||||
- nvidia-cusolver-cu11==11.4.0.1
|
||||
- nvidia-cusparse-cu11==11.7.4.91
|
||||
- nvidia-nccl-cu11==2.14.3
|
||||
- nvidia-nvtx-cu11==11.7.91
|
||||
- ordered-set==4.1.0
|
||||
- packaging==23.1
|
||||
- pandas==2.1.1
|
||||
- parso==0.8.3
|
||||
- pexpect==4.8.0
|
||||
- pickleshare==0.7.5
|
||||
- platformdirs==3.10.0
|
||||
- prompt-toolkit==3.0.39
|
||||
- psutil==5.9.5
|
||||
- ptyprocess==0.7.0
|
||||
- pure-eval==0.2.2
|
||||
- pyarrow==13.0.0
|
||||
- pydantic==2.1.1
|
||||
- pydantic-core==2.4.0
|
||||
- pygments==2.16.1
|
||||
- pyjwt==2.8.0
|
||||
- pyparsing==3.1.1
|
||||
- python-dateutil==2.8.2
|
||||
- python-editor==1.0.4
|
||||
- python-multipart==0.0.6
|
||||
- pytorch-lightning==2.0.9
|
||||
- pytz==2023.3.post1
|
||||
- pyyaml==6.0.1
|
||||
- pyzmq==25.1.1
|
||||
- readchar==4.0.5
|
||||
- regex==2023.8.8
|
||||
- rich==13.5.3
|
||||
- safetensors==0.3.3
|
||||
- simple-parsing==0.1.4
|
||||
- six==1.16.0
|
||||
- sniffio==1.3.0
|
||||
- soupsieve==2.5
|
||||
- stack-data==0.6.2
|
||||
- starlette==0.27.0
|
||||
- starsessions==1.3.0
|
||||
- tokenizers==0.13.3
|
||||
- torch==2.0.1
|
||||
- torchmetrics==1.2.0
|
||||
- tornado==6.3.3
|
||||
- tqdm==4.66.1
|
||||
- traitlets==5.10.0
|
||||
- transformers==4.33.2
|
||||
- triton==2.0.0
|
||||
- tzdata==2023.3
|
||||
- uvicorn==0.23.2
|
||||
- wcwidth==0.2.6
|
||||
- websocket-client==1.6.3
|
||||
- websockets==11.0.3
|
||||
- widgetsnbextension==4.0.9
|
||||
- xxhash==3.3.0
|
||||
- yarl==1.9.2
|
||||
prefix: /home/ubuntu/mambaforge/envs/dlk4
|
||||
@@ -1,14 +0,0 @@
|
||||
name: dlk4
|
||||
channels:
|
||||
- conda-forge
|
||||
dependencies:
|
||||
- python=3.11
|
||||
- pytorch
|
||||
- torchvision
|
||||
- torchaudio
|
||||
- pytorch-cuda=11.7
|
||||
- cudatoolkit-dev==11.7
|
||||
- cudatoolkit=11.7
|
||||
- ca-certificates
|
||||
- openssl
|
||||
prefix: /home/ubuntu/mambaforge/envs/dlk4
|
||||
@@ -1,19 +0,0 @@
|
||||
datasets
|
||||
tqdm
|
||||
transformers~=4.31.0
|
||||
scikit-learn
|
||||
accelerate
|
||||
# bitsandbytes
|
||||
lightning==2.0.6
|
||||
sentencepiece
|
||||
peft
|
||||
# use the version that https://github.com/johnsmith0031/alpaca_lora_4bit/blob/main/requirements.txt uses since they always resolve the dependancy issues
|
||||
# git+https://github.com/huggingface/peft.git@70af02a2bca5a63921790036b2c9430edf4037e2
|
||||
# due to a bug we have to downgrade to this one for now https://twitter.com/Teknium1/status/1660003439752138752
|
||||
bitsandbytes==0.39.1
|
||||
matplotlib
|
||||
black
|
||||
loguru
|
||||
eleuther-elk==0.1.1
|
||||
# promptsource
|
||||
scipy
|
||||
@@ -1,10 +0,0 @@
|
||||
from setuptools import find_packages, setup
|
||||
|
||||
setup(
|
||||
name='src',
|
||||
packages=find_packages(),
|
||||
version='0.1.0',
|
||||
description='Discovering Latent Knowledge using MonteCarlo Dropout on outputs not inputs',
|
||||
author='wassname',
|
||||
license='MIT',
|
||||
)
|
||||
@@ -1,546 +0,0 @@
|
||||
import os
|
||||
import functools
|
||||
import argparse
|
||||
import copy
|
||||
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
from tqdm import tqdm
|
||||
|
||||
import torch
|
||||
from torch.utils.data import Dataset, DataLoader
|
||||
from torchvision import datasets
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
# make sure to install promptsource, transformers, and datasets!
|
||||
from promptsource.templates import DatasetTemplates
|
||||
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, AutoModelForMaskedLM, AutoModelForCausalLM
|
||||
from datasets import load_dataset
|
||||
|
||||
|
||||
############# Model loading and result saving #############
|
||||
|
||||
# Map each model name to its full Huggingface name; this is just for convenience for common models. You can run whatever model you'd like.
|
||||
model_mapping = {
|
||||
"gpt-j": "EleutherAI/gpt-j-6B",
|
||||
"T0pp": "bigscience/T0pp",
|
||||
"unifiedqa": "allenai/unifiedqa-t5-11b",
|
||||
"T5": "t5-11b",
|
||||
"deberta-mnli": "microsoft/deberta-xxlarge-v2-mnli",
|
||||
"deberta": "microsoft/deberta-xxlarge-v2",
|
||||
"roberta-mnli": "roberta-large-mnli",
|
||||
}
|
||||
|
||||
|
||||
def get_parser():
|
||||
"""
|
||||
Returns the parser we will use for generate.py and evaluate.py
|
||||
(We include it here so that we can use the same parser for both scripts)
|
||||
"""
|
||||
parser = argparse.ArgumentParser()
|
||||
# setting up model
|
||||
parser.add_argument("--model_name", type=str, default="T5", help="Name of the model to use")
|
||||
parser.add_argument("--cache_dir", type=str, default=None, help="Cache directory for the model and tokenizer")
|
||||
parser.add_argument("--parallelize", action="store_true", help="Whether to parallelize the model")
|
||||
parser.add_argument("--device", type=str, default="cuda", help="Device to use for the model")
|
||||
# setting up data
|
||||
parser.add_argument("--dataset_name", type=str, default="imdb", help="Name of the dataset to use")
|
||||
parser.add_argument("--split", type=str, default="test", help="Which split of the dataset to use")
|
||||
parser.add_argument("--prompt_idx", type=int, default=0, help="Which prompt to use")
|
||||
parser.add_argument("--batch_size", type=int, default=1, help="Batch size to use")
|
||||
parser.add_argument("--num_examples", type=int, default=1000, help="Number of examples to generate")
|
||||
# which hidden states we extract
|
||||
parser.add_argument("--use_decoder", action="store_true", help="Whether to use the decoder; only relevant if model_type is encoder-decoder. Uses encoder by default (which usually -- but not always -- works better)")
|
||||
parser.add_argument("--layer", type=int, default=-1, help="Which layer to use (if not all layers)")
|
||||
parser.add_argument("--all_layers", action="store_true", help="Whether to use all layers or not")
|
||||
parser.add_argument("--token_idx", type=int, default=-1, help="Which token to use (by default the last token)")
|
||||
# saving the hidden states
|
||||
parser.add_argument("--save_dir", type=str, default="generated_hidden_states", help="Directory to save the hidden states")
|
||||
|
||||
return parser
|
||||
|
||||
|
||||
def load_model(model_name, cache_dir=None, parallelize=False, device="cuda"):
|
||||
"""
|
||||
Loads a model and its corresponding tokenizer, either parallelized across GPUs (if the model permits that; usually just use this for T5-based models) or on a single GPU
|
||||
"""
|
||||
if model_name in model_mapping:
|
||||
# use a nickname for our models
|
||||
full_model_name = model_mapping[model_name]
|
||||
else:
|
||||
# if you're trying a new model, make sure it's the full name
|
||||
full_model_name = model_name
|
||||
|
||||
# use the right automodel, and get the corresponding model type
|
||||
try:
|
||||
model = AutoModelForSeq2SeqLM.from_pretrained(full_model_name, cache_dir=cache_dir)
|
||||
model_type = "encoder_decoder"
|
||||
except:
|
||||
try:
|
||||
model = AutoModelForMaskedLM.from_pretrained(full_model_name, cache_dir=cache_dir)
|
||||
model_type = "encoder"
|
||||
except:
|
||||
model = AutoModelForCausalLM.from_pretrained(full_model_name, cache_dir=cache_dir)
|
||||
model_type = "decoder"
|
||||
|
||||
|
||||
# specify model_max_length (the max token length) to be 512 to ensure that padding works
|
||||
# (it's not set by default for e.g. DeBERTa, but it's necessary for padding to work properly)
|
||||
tokenizer = AutoTokenizer.from_pretrained(full_model_name, cache_dir=cache_dir, model_max_length=512)
|
||||
model.eval()
|
||||
|
||||
# put on the correct device
|
||||
if parallelize:
|
||||
model.parallelize()
|
||||
else:
|
||||
model = model.to(device)
|
||||
|
||||
return model, tokenizer, model_type
|
||||
|
||||
|
||||
def save_generations(generation, args, generation_type):
|
||||
"""
|
||||
Input:
|
||||
generation: numpy array (e.g. hidden_states or labels) to save
|
||||
args: arguments used to generate the hidden states. This is used for the filename to save to.
|
||||
generation_type: one of "negative_hidden_states" or "positive_hidden_states" or "labels"
|
||||
|
||||
Saves the generations to an appropriate directory.
|
||||
"""
|
||||
# construct the filename based on the args
|
||||
arg_dict = vars(args)
|
||||
exclude_keys = ["save_dir", "cache_dir", "device"]
|
||||
filename = generation_type + "__" + "__".join(['{}_{}'.format(k, v) for k, v in arg_dict.items() if k not in exclude_keys]) + ".npy".format(generation_type)
|
||||
|
||||
# create save directory if it doesn't exist
|
||||
if not os.path.exists(args.save_dir):
|
||||
os.makedirs(args.save_dir)
|
||||
|
||||
# save
|
||||
np.save(os.path.join(args.save_dir, filename), generation)
|
||||
|
||||
|
||||
def load_single_generation(args, generation_type="hidden_states"):
|
||||
# use the same filename as in save_generations
|
||||
arg_dict = vars(args)
|
||||
exclude_keys = ["save_dir", "cache_dir", "device"]
|
||||
filename = generation_type + "__" + "__".join(['{}_{}'.format(k, v) for k, v in arg_dict.items() if k not in exclude_keys]) + ".npy".format(generation_type)
|
||||
return np.load(os.path.join(args.save_dir, filename))
|
||||
|
||||
|
||||
def load_all_generations(args):
|
||||
# load all the saved generations: neg_hs, pos_hs, and labels
|
||||
neg_hs = load_single_generation(args, generation_type="negative_hidden_states")
|
||||
pos_hs = load_single_generation(args, generation_type="positive_hidden_states")
|
||||
labels = load_single_generation(args, generation_type="labels")
|
||||
|
||||
return neg_hs, pos_hs, labels
|
||||
|
||||
|
||||
############# Data #############
|
||||
class ContrastDataset(Dataset):
|
||||
"""
|
||||
Given a dataset and tokenizer (from huggingface), along with a collection of prompts for that dataset from promptsource and a corresponding prompt index,
|
||||
returns a dataset that creates contrast pairs using that prompt
|
||||
|
||||
Truncates examples larger than max_len, which can mess up contrast pairs, so make sure to only give it examples that won't be truncated.
|
||||
"""
|
||||
def __init__(self, raw_dataset, tokenizer, all_prompts, prompt_idx,
|
||||
model_type="encoder_decoder", use_decoder=False, device="cuda"):
|
||||
|
||||
# data and tokenizer
|
||||
self.raw_dataset = raw_dataset
|
||||
self.tokenizer = tokenizer
|
||||
if self.tokenizer.pad_token is None:
|
||||
self.tokenizer.pad_token = self.tokenizer.eos_token
|
||||
self.device = device
|
||||
|
||||
# for formatting the answers
|
||||
self.model_type = model_type
|
||||
self.use_decoder = use_decoder
|
||||
if self.use_decoder:
|
||||
assert self.model_type != "encoder"
|
||||
|
||||
# prompt
|
||||
prompt_name_list = list(all_prompts.name_to_id_mapping.keys())
|
||||
self.prompt = all_prompts[prompt_name_list[prompt_idx]]
|
||||
|
||||
def __len__(self):
|
||||
return len(self.raw_dataset)
|
||||
|
||||
def encode(self, nl_prompt):
|
||||
"""
|
||||
Tokenize a given natural language prompt (from after applying self.prompt to an example)
|
||||
|
||||
For encoder-decoder models, we can either:
|
||||
(1) feed both the question and answer to the encoder, creating contrast pairs using the encoder hidden states
|
||||
(which uses the standard tokenization, but also passes the empty string to the decoder), or
|
||||
(2) feed the question the encoder and the answer to the decoder, creating contrast pairs using the decoder hidden states
|
||||
|
||||
If self.decoder is True we do (2), otherwise we do (1).
|
||||
"""
|
||||
# get question and answer from prompt
|
||||
question, answer = nl_prompt
|
||||
|
||||
# tokenize the question and answer (depending upon the model type and whether self.use_decoder is True)
|
||||
if self.model_type == "encoder_decoder":
|
||||
input_ids = self.get_encoder_decoder_input_ids(question, answer)
|
||||
elif self.model_type == "encoder":
|
||||
input_ids = self.get_encoder_input_ids(question, answer)
|
||||
else:
|
||||
input_ids = self.get_decoder_input_ids(question, answer)
|
||||
|
||||
# get rid of the batch dimension since this will be added by the Dataloader
|
||||
if input_ids["input_ids"].shape[0] == 1:
|
||||
for k in input_ids:
|
||||
input_ids[k] = input_ids[k].squeeze(0)
|
||||
|
||||
return input_ids
|
||||
|
||||
|
||||
def get_encoder_input_ids(self, question, answer):
|
||||
"""
|
||||
Format the input ids for encoder-only models; standard formatting.
|
||||
"""
|
||||
combined_input = question + " " + answer
|
||||
input_ids = self.tokenizer(combined_input, truncation=True, padding="max_length", return_tensors="pt")
|
||||
|
||||
return input_ids
|
||||
|
||||
|
||||
def get_decoder_input_ids(self, question, answer):
|
||||
"""
|
||||
Format the input ids for encoder-only models.
|
||||
This is the same as get_encoder_input_ids except that we add the EOS token at the end of the input (which apparently can matter)
|
||||
"""
|
||||
combined_input = question + " " + answer + self.tokenizer.eos_token
|
||||
input_ids = self.tokenizer(combined_input, truncation=True, padding="max_length", return_tensors="pt")
|
||||
|
||||
return input_ids
|
||||
|
||||
|
||||
def get_encoder_decoder_input_ids(self, question, answer):
|
||||
"""
|
||||
Format the input ids for encoder-decoder models.
|
||||
There are two cases for this, depending upon whether we want to use the encoder hidden states or the decoder hidden states.
|
||||
"""
|
||||
if self.use_decoder:
|
||||
# feed the same question to the encoder but different answers to the decoder to construct contrast pairs
|
||||
input_ids = self.tokenizer(question, truncation=True, padding="max_length", return_tensors="pt")
|
||||
decoder_input_ids = self.tokenizer(answer, truncation=True, padding="max_length", return_tensors="pt")
|
||||
else:
|
||||
# include both the question and the answer in the input for the encoder
|
||||
# feed the empty string to the decoder (i.e. just ignore it -- but it needs an input or it'll throw an error)
|
||||
input_ids = self.tokenizer(question, answer, truncation=True, padding="max_length", return_tensors="pt")
|
||||
decoder_input_ids = self.tokenizer("", return_tensors="pt")
|
||||
|
||||
# move everything into input_ids so that it's easier to pass to the model
|
||||
input_ids["decoder_input_ids"] = decoder_input_ids["input_ids"]
|
||||
input_ids["decoder_attention_mask"] = decoder_input_ids["attention_mask"]
|
||||
|
||||
return input_ids
|
||||
|
||||
|
||||
def __getitem__(self, index):
|
||||
# get the original example
|
||||
data = self.raw_dataset[int(index)]
|
||||
text, true_answer = data["text"], data["label"]
|
||||
|
||||
# get the possible labels
|
||||
# (for simplicity assume the binary case for contrast pairs)
|
||||
label_list = self.prompt.get_answer_choices_list(data)
|
||||
assert len(label_list) == 2, print("Make sure there are only two possible answers! Actual number of answers:", label_list)
|
||||
|
||||
# reconvert to dataset format but with fake/candidate labels to create the contrast pair
|
||||
neg_example = {"text": text, "label": 0}
|
||||
pos_example = {"text": text, "label": 1}
|
||||
|
||||
# construct contrast pairs by answering the prompt with the two different possible labels
|
||||
# (for example, label 0 might be mapped to "no" and label 1 might be mapped to "yes")
|
||||
neg_prompt, pos_prompt = self.prompt.apply(neg_example), self.prompt.apply(pos_example)
|
||||
|
||||
# tokenize
|
||||
neg_ids, pos_ids = self.encode(neg_prompt), self.encode(pos_prompt)
|
||||
|
||||
# verify these are different (e.g. tokenization didn't cut off the difference between them)
|
||||
if self.use_decoder and self.model_type == "encoder_decoder":
|
||||
assert (neg_ids["decoder_input_ids"] - pos_ids["decoder_input_ids"]).sum() != 0, print("The decoder_input_ids for the contrast pairs are the same!", neg_ids, pos_ids)
|
||||
else:
|
||||
assert (neg_ids["input_ids"] - pos_ids["input_ids"]).sum() != 0, print("The input_ids for the contrast pairs are the same!", neg_ids, pos_ids)
|
||||
|
||||
# return the tokenized inputs, the text prompts, and the true label
|
||||
return neg_ids, pos_ids, neg_prompt, pos_prompt, true_answer
|
||||
|
||||
|
||||
def get_dataloader(dataset_name, split, tokenizer, prompt_idx, batch_size=16, num_examples=1000,
|
||||
model_type="encoder_decoder", use_decoder=False, device="cuda", pin_memory=True, num_workers=1):
|
||||
"""
|
||||
Creates a dataloader for a given dataset (and its split), tokenizer, and prompt index
|
||||
|
||||
Takes a random subset of (at most) num_examples samples from the dataset that are not truncated by the tokenizer.
|
||||
"""
|
||||
# load the raw dataset
|
||||
raw_dataset = load_dataset(dataset_name)[split]
|
||||
|
||||
# load all the prompts for that dataset
|
||||
all_prompts = DatasetTemplates(dataset_name)
|
||||
|
||||
# create the ConstrastDataset
|
||||
contrast_dataset = ContrastDataset(raw_dataset, tokenizer, all_prompts, prompt_idx,
|
||||
model_type=model_type, use_decoder=use_decoder,
|
||||
device=device)
|
||||
|
||||
# get a random permutation of the indices; we'll take the first num_examples of these that do not get truncated
|
||||
random_idxs = np.random.permutation(len(contrast_dataset))
|
||||
|
||||
# remove examples that would be truncated (since this messes up contrast pairs)
|
||||
prompt_name_list = list(all_prompts.name_to_id_mapping.keys())
|
||||
prompt = all_prompts[prompt_name_list[prompt_idx]]
|
||||
keep_idxs = []
|
||||
for idx in random_idxs:
|
||||
question, answer = prompt.apply(raw_dataset[int(idx)])
|
||||
input_text = question + " " + answer
|
||||
if len(tokenizer.encode(input_text, truncation=False)) < tokenizer.model_max_length - 2: # include small margin to be conservative
|
||||
keep_idxs.append(idx)
|
||||
if len(keep_idxs) >= num_examples:
|
||||
break
|
||||
|
||||
# create and return the corresponding dataloader
|
||||
subset_dataset = torch.utils.data.Subset(contrast_dataset, keep_idxs)
|
||||
dataloader = DataLoader(subset_dataset, batch_size=batch_size, shuffle=False, pin_memory=pin_memory, num_workers=num_workers)
|
||||
|
||||
return dataloader
|
||||
|
||||
|
||||
############# Hidden States #############
|
||||
def get_first_mask_loc(mask, shift=False):
|
||||
"""
|
||||
return the location of the first pad token for the given ids, which corresponds to a mask value of 0
|
||||
if there are no pad tokens, then return the last location
|
||||
"""
|
||||
# add a 0 to the end of the mask in case there are no pad tokens
|
||||
mask = torch.cat([mask, torch.zeros_like(mask[..., :1])], dim=-1)
|
||||
|
||||
if shift:
|
||||
mask = mask[..., 1:]
|
||||
|
||||
# get the location of the first pad token; use the fact that torch.argmax() returns the first index in the case of ties
|
||||
first_mask_loc = torch.argmax((mask == 0).int(), dim=-1)
|
||||
|
||||
return first_mask_loc
|
||||
|
||||
|
||||
def get_individual_hidden_states(model, batch_ids, layer=None, all_layers=True, token_idx=-1, model_type="encoder_decoder", use_decoder=False):
|
||||
"""
|
||||
Given a model and a batch of tokenized examples, returns the hidden states for either
|
||||
a specified layer (if layer is a number) or for all layers (if all_layers is True).
|
||||
|
||||
If specify_encoder is True, uses "encoder_hidden_states" instead of "hidden_states"
|
||||
This is necessary for getting the encoder hidden states for encoder-decoder models,
|
||||
but it is not necessary for encoder-only or decoder-only models.
|
||||
"""
|
||||
if use_decoder:
|
||||
assert "decoder" in model_type
|
||||
|
||||
# forward pass
|
||||
with torch.no_grad():
|
||||
batch_ids = batch_ids.to(model.device)
|
||||
output = model(**batch_ids, output_hidden_states=True)
|
||||
|
||||
# get all the corresponding hidden states (which is a tuple of length num_layers)
|
||||
if use_decoder and "decoder_hidden_states" in output.keys():
|
||||
hs_tuple = output["decoder_hidden_states"]
|
||||
elif "encoder_hidden_states" in output.keys():
|
||||
hs_tuple = output["encoder_hidden_states"]
|
||||
else:
|
||||
hs_tuple = output["hidden_states"]
|
||||
|
||||
# just get the corresponding layer hidden states
|
||||
if all_layers:
|
||||
# stack along the last axis so that it's easier to consistently index the first two axes
|
||||
hs = torch.stack([h.squeeze().detach().cpu() for h in hs_tuple], axis=-1) # (bs, seq_len, dim, num_layers)
|
||||
else:
|
||||
assert layer is not None
|
||||
hs = hs_tuple[layer].unsqueeze(-1).detach().cpu() # (bs, seq_len, dim, 1)
|
||||
|
||||
# we want to get the token corresponding to token_idx while ignoring the masked tokens
|
||||
if token_idx == 0:
|
||||
final_hs = hs[:, 0] # (bs, dim, num_layers)
|
||||
else:
|
||||
# if token_idx == -1, then takes the hidden states corresponding to the last non-mask tokens
|
||||
# first we need to get the first mask location for each example in the batch
|
||||
assert token_idx < 0, print("token_idx must be either 0 or negative, but got", token_idx)
|
||||
mask = batch_ids["decoder_attention_mask"] if (model_type == "encoder_decoder" and use_decoder) else batch_ids["attention_mask"]
|
||||
first_mask_loc = get_first_mask_loc(mask).squeeze()
|
||||
final_hs = hs[torch.arange(hs.size(0)), first_mask_loc+token_idx] # (bs, dim, num_layers)
|
||||
|
||||
return final_hs
|
||||
|
||||
|
||||
def get_all_hidden_states(model, dataloader, layer=None, all_layers=True, token_idx=-1, model_type="encoder_decoder", use_decoder=False):
|
||||
"""
|
||||
Given a model, a tokenizer, and a dataloader, returns the hidden states (corresponding to a given position index) in all layers for all examples in the dataloader,
|
||||
along with the average log probs corresponding to the answer tokens
|
||||
|
||||
The dataloader should correspond to examples *with a candidate label already added* to each example.
|
||||
E.g. this function should be used for "Q: Is 2+2=5? A: True" or "Q: Is 2+2=5? A: False", but NOT for "Q: Is 2+2=5? A: ".
|
||||
"""
|
||||
all_pos_hs, all_neg_hs = [], []
|
||||
all_gt_labels = []
|
||||
|
||||
model.eval()
|
||||
for batch in tqdm(dataloader):
|
||||
neg_ids, pos_ids, _, _, gt_label = batch
|
||||
|
||||
neg_hs = get_individual_hidden_states(model, neg_ids, layer=layer, all_layers=all_layers, token_idx=token_idx,
|
||||
model_type=model_type, use_decoder=use_decoder)
|
||||
pos_hs = get_individual_hidden_states(model, pos_ids, layer=layer, all_layers=all_layers, token_idx=token_idx,
|
||||
model_type=model_type, use_decoder=use_decoder)
|
||||
|
||||
if dataloader.batch_size == 1:
|
||||
neg_hs, pos_hs = neg_hs.unsqueeze(0), pos_hs.unsqueeze(0)
|
||||
|
||||
all_neg_hs.append(neg_hs)
|
||||
all_pos_hs.append(pos_hs)
|
||||
all_gt_labels.append(gt_label)
|
||||
|
||||
all_neg_hs = np.concatenate(all_neg_hs, axis=0)
|
||||
all_pos_hs = np.concatenate(all_pos_hs, axis=0)
|
||||
all_gt_labels = np.concatenate(all_gt_labels, axis=0)
|
||||
|
||||
return all_neg_hs, all_pos_hs, all_gt_labels
|
||||
|
||||
############# CCS #############
|
||||
class MLPProbe(nn.Module):
|
||||
def __init__(self, d):
|
||||
super().__init__()
|
||||
self.linear1 = nn.Linear(d, 100)
|
||||
self.linear2 = nn.Linear(100, 1)
|
||||
|
||||
def forward(self, x):
|
||||
h = F.relu(self.linear1(x))
|
||||
o = self.linear2(h)
|
||||
return torch.sigmoid(o)
|
||||
|
||||
class CCS(object):
|
||||
def __init__(self, x0, x1, nepochs=1000, ntries=10, lr=1e-3, batch_size=-1,
|
||||
verbose=False, device="cuda", linear=True, weight_decay=0.01, var_normalize=False):
|
||||
# data
|
||||
self.var_normalize = var_normalize
|
||||
self.x0 = self.normalize(x0)
|
||||
self.x1 = self.normalize(x1)
|
||||
self.d = self.x0.shape[-1]
|
||||
|
||||
# training
|
||||
self.nepochs = nepochs
|
||||
self.ntries = ntries
|
||||
self.lr = lr
|
||||
self.verbose = verbose
|
||||
self.device = device
|
||||
self.batch_size = batch_size
|
||||
self.weight_decay = weight_decay
|
||||
|
||||
# probe
|
||||
self.linear = linear
|
||||
self.probe = self.initialize_probe()
|
||||
self.best_probe = copy.deepcopy(self.probe)
|
||||
|
||||
|
||||
def initialize_probe(self):
|
||||
if self.linear:
|
||||
self.probe = nn.Sequential(nn.Linear(self.d, 1), nn.Sigmoid())
|
||||
else:
|
||||
self.probe = MLPProbe(self.d)
|
||||
self.probe.to(self.device)
|
||||
|
||||
|
||||
def normalize(self, x):
|
||||
"""
|
||||
Mean-normalizes the data x (of shape (n, d))
|
||||
If self.var_normalize, also divides by the standard deviation
|
||||
"""
|
||||
normalized_x = x - x.mean(axis=0, keepdims=True)
|
||||
if self.var_normalize:
|
||||
normalized_x /= normalized_x.std(axis=0, keepdims=True)
|
||||
|
||||
return normalized_x
|
||||
|
||||
|
||||
def get_tensor_data(self):
|
||||
"""
|
||||
Returns x0, x1 as appropriate tensors (rather than np arrays)
|
||||
"""
|
||||
x0 = torch.tensor(self.x0, dtype=torch.float, requires_grad=False, device=self.device)
|
||||
x1 = torch.tensor(self.x1, dtype=torch.float, requires_grad=False, device=self.device)
|
||||
return x0, x1
|
||||
|
||||
|
||||
def get_loss(self, p0, p1):
|
||||
"""
|
||||
Returns the CCS loss for two probabilities each of shape (n,1) or (n,)
|
||||
"""
|
||||
informative_loss = (torch.min(p0, p1)**2).mean(0)
|
||||
consistent_loss = ((p0 - (1-p1))**2).mean(0)
|
||||
return informative_loss + consistent_loss
|
||||
|
||||
|
||||
def get_acc(self, x0_test, x1_test, y_test):
|
||||
"""
|
||||
Computes accuracy for the current parameters on the given test inputs
|
||||
"""
|
||||
x0 = torch.tensor(self.normalize(x0_test), dtype=torch.float, requires_grad=False, device=self.device)
|
||||
x1 = torch.tensor(self.normalize(x1_test), dtype=torch.float, requires_grad=False, device=self.device)
|
||||
with torch.no_grad():
|
||||
p0, p1 = self.best_probe(x0), self.best_probe(x1)
|
||||
avg_confidence = 0.5*(p0 + (1-p1))
|
||||
predictions = (avg_confidence.detach().cpu().numpy() < 0.5).astype(int)[:, 0]
|
||||
acc = (predictions == y_test).mean()
|
||||
acc = max(acc, 1 - acc)
|
||||
|
||||
return acc
|
||||
|
||||
|
||||
def train(self):
|
||||
"""
|
||||
Does a single training run of nepochs epochs
|
||||
"""
|
||||
x0, x1 = self.get_tensor_data()
|
||||
permutation = torch.randperm(len(x0))
|
||||
x0, x1 = x0[permutation], x1[permutation]
|
||||
|
||||
# set up optimizer
|
||||
optimizer = torch.optim.AdamW(self.probe.parameters(), lr=self.lr, weight_decay=self.weight_decay)
|
||||
|
||||
batch_size = len(x0) if self.batch_size == -1 else self.batch_size
|
||||
nbatches = len(x0) // batch_size
|
||||
|
||||
# Start training (full batch)
|
||||
for epoch in range(self.nepochs):
|
||||
for j in range(nbatches):
|
||||
x0_batch = x0[j*batch_size:(j+1)*batch_size]
|
||||
x1_batch = x1[j*batch_size:(j+1)*batch_size]
|
||||
|
||||
# probe
|
||||
p0, p1 = self.probe(x0_batch), self.probe(x1_batch)
|
||||
|
||||
# get the corresponding loss
|
||||
loss = self.get_loss(p0, p1)
|
||||
|
||||
# update the parameters
|
||||
optimizer.zero_grad()
|
||||
loss.backward()
|
||||
optimizer.step()
|
||||
|
||||
return loss.detach().cpu().item()
|
||||
|
||||
def repeated_train(self):
|
||||
best_loss = np.inf
|
||||
for train_num in range(self.ntries):
|
||||
self.initialize_probe()
|
||||
loss = self.train()
|
||||
if loss < best_loss:
|
||||
self.best_probe = copy.deepcopy(self.probe)
|
||||
best_loss = loss
|
||||
|
||||
return best_loss
|
||||
Reference in New Issue
Block a user