From 279c8cc674c1075f8dd01462f7bbfd0a7b9a032b Mon Sep 17 00:00:00 2001 From: deep1 <> Date: Sun, 2 Jul 2023 14:39:52 +0800 Subject: [PATCH] a script to cache ds --- notebooks/020_ds.ipynb | 1074 ++++++++++++++++++++++++++++++++++++++++ 1 file changed, 1074 insertions(+) create mode 100644 notebooks/020_ds.ipynb diff --git a/notebooks/020_ds.ipynb b/notebooks/020_ds.ipynb new file mode 100644 index 0000000..feac8ff --- /dev/null +++ b/notebooks/020_ds.ipynb @@ -0,0 +1,1074 @@ +{ + "cells": [ + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Lets just do supervised learning\n", + "\n", + "Since we are looking at pairs with random permuations (from dropout), we can't use CCS. This is because our probabilities do not add to one.\n", + "\n", + "People question if unsupervised learning bings anything to the table anyway, so lets start with supervised..." + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "\n", + "links:\n", + "- [loading](https://github.com/deep-diver/LLM-As-Chatbot/blob/main/models/alpaca.py)\n", + "- [dict](https://github.com/deep-diver/LLM-As-Chatbot/blob/c79e855a492a968b54bac223e66dc9db448d6eba/model_cards.json#L143)\n", + "- [prompt_format](https://github.com/deep-diver/PingPong/blob/main/src/pingpong/alpaca.py)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "\n", + "import copy\n", + "import numpy as np\n", + "import pandas as pd\n", + "from matplotlib import pyplot as plt\n", + "plt.style.use('ggplot')\n", + "\n", + "import random\n", + "from typing import Optional, List, Dict, Union\n", + "\n", + "import torch\n", + "import torch.nn as nn\n", + "import torch.nn.functional as F\n", + "from torch import Tensor\n", + "from torch import optim\n", + "from torch.utils.data import random_split, DataLoader, TensorDataset\n", + "\n", + "import pickle\n", + "import hashlib\n", + "from pathlib import Path\n", + "\n", + "from datasets import load_dataset\n", + "import datasets\n", + "\n", + "from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, AutoModelForMaskedLM, AutoModelForCausalLM, AutoConfig\n", + "import transformers\n", + "from transformers.models.auto.modeling_auto import AutoModel\n", + "from transformers import LogitsProcessorList\n", + "\n", + "\n", + "import lightning.pytorch as pl\n", + "from dataclasses import dataclass\n", + "\n", + "from sklearn.linear_model import LogisticRegression\n", + "# from scipy.stats import zscore\n", + "from sklearn.metrics import f1_score, roc_auc_score, accuracy_score\n", + "from sklearn.preprocessing import RobustScaler\n", + "\n", + "from tqdm.auto import tqdm\n", + "import gc\n", + "import os\n", + "\n", + "from loguru import logger\n", + "logger.add(os.sys.stderr, format=\"{time} {level} {message}\", level=\"INFO\")\n", + "\n", + "\n", + "transformers.__version__" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Model\n", + "\n", + "Chosing:\n", + "- https://old.reddit.com/r/LocalLLaMA/wiki/models\n", + "- https://huggingface.co/spaces/HuggingFaceH4/open_llm_leaderboard\n", + "- https://github.com/deep-diver/LLM-As-Chatbot/blob/main/model_cards.json\n", + "\n", + "\n", + "A uncensored and large one might be best for lying." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from peft import PeftModel" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# leaderboard https://huggingface.co/spaces/HuggingFaceH4/open_llm_leaderboard\n", + "model_options = dict(\n", + " device_map=\"auto\", \n", + " load_in_4bit=True,\n", + " # load_in_8bit=True,\n", + " torch_dtype=torch.float16,\n", + " trust_remote_code=True,\n", + " use_safetensors=False,\n", + " # use_cache=False,\n", + ")\n", + "\n", + "# so I need to use either pythia, stablelm, or tiiuae/falcon-7b-instruct to get dropout...\n", + "# moel_repo = \"stabilityai/stablelm-tuned-alpha-7b\" # poor performance\n", + "\n", + "# https://github.com/deep-diver/LLM-As-Chatbot/blob/main/models/falcon.py\n", + "model_repo = \"tiiuae/falcon-7b-instruct\"\n", + "# model_repo = \"tiiuae/falcon-7b\"\n", + "# model_repo = \"togethercomputer/RedPajama-INCITE-7B-Instruct\"\n", + "# model_repo = \"OpenAssis/tant/oasst-sft-4-pythia-12b-epoch-3.5\"\n", + "# model_repo = \"OpenAssistant/falcon-7b-sft-top1-696\"\n", + "# model_repo = \"openaccess-ai-collective/manticore-13b\"\n", + "# model_repo = \"TheBloke/Wizard-Vicuna-13B-Uncensored-HF\"\n", + "# model_repo = \"dvruette/llama-13b-pretrained-dropout\"\n", + "# model_repo = \"elinas/llama-13b-hf-transformers-4.29\" # no dropout\n", + "# lora_repo = \"LLMs/AlpacaGPT4-LoRA-13B-elina\"\n", + "model_repo = \"bigcode/starcoderplus\"\n", + "model_repo = \"HuggingFaceH4/starchat-beta\"\n", + "model_repo = \"WizardLM/WizardCoder-15B-V1.0\"\n", + "# model_repo= \"~/.cache/huggingface/hub/models--HuggingFaceH4--starchat-beta\"\n", + "# lora_repo = None\n", + "lora_repo = None\n", + "\n", + "config = AutoConfig.from_pretrained(model_repo, trust_remote_code=True,)\n", + "print(config)\n", + "# config.attn_pdrop=0.3\n", + "# config.embd_pdrop=0.3\n", + "# config.resid_pdrop=0.3\n", + "config.use_cache = False\n", + "tokenizer = AutoTokenizer.from_pretrained(model_repo)\n", + "model = AutoModelForCausalLM.from_pretrained(model_repo, config=config, **model_options)\n", + "\n", + "if lora_repo is not None:\n", + " # https://github.com/tloen/alpaca-lora/blob/main/generate.py#L40\n", + " from peft import PeftModel\n", + " model = PeftModel.from_pretrained(\n", + " model,\n", + " lora_repo, \n", + " torch_dtype=torch.float16,\n", + " lora_dropout=0.2,\n", + " device_map='auto'\n", + " )\n", + " \n", + "# if not mode_8bit and not mode_4bit:\n", + "# model.half()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "model" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# https://github.com/deep-diver/LLM-As-Chatbot/blob/main/models/falcon.py\n", + "print(tokenizer.pad_token_id)\n", + "if tokenizer.pad_token_id is None:\n", + " tokenizer.pad_token_id = 204 # https://github.com/deep-diver/LLM-As-Chatbot/blob/main/models/alpaca.py\n", + "tokenizer.padding_side = \"left\"" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Params" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Params\n", + "# N_SAMPLES = 4000\n", + "BATCH_SIZE = 8 # None # None means auto # 6 gives 16Gb/25GB. where 10GB is the base model. so 6 is 6/15\n", + "N_SHOTS = 3\n", + "USE_MCDROPOUT = True\n", + "dataset_n = 200\n", + "\n", + "try:\n", + " # num_layers = len(model.model.layers)\n", + " num_layers = model.config.n_layer\n", + " print(num_layers)\n", + "except AttributeError:\n", + " try:\n", + " num_layers = len(model.base_model.model.model.layers)\n", + " print(num_layers)\n", + " except:\n", + " num_layers = 10\n", + " \n", + "stride = 2\n", + "extract_layers = tuple(range(2, num_layers-2, stride)) + (num_layers-2,)\n", + "extract_layers, num_layers" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# get the tokens for 0 and 1, we will use these later...\n", + "# note that sentancepeice tokenizers have differen't tokens for No and \\nNo.\n", + "token_n = \"negative\"\n", + "token_y = \"positive\"\n", + "id_n, id_y = tokenizer(f'\\n{token_n}', add_special_tokens=True)['input_ids'][-1], tokenizer(f'\\n{token_y}', add_special_tokens=True)['input_ids'][-1]\n", + "assert tokenizer.decode([id_n])==token_n\n", + "assert tokenizer.decode([id_y])==token_y\n", + "id_n, id_y" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Load Dataset" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Let's just try IMDB for simplicity\n", + "dataset = load_dataset(\"amazon_polarity\")\n", + "data = dataset['test']" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Define Prompt\n", + "\n", + "- Lillian Wang's guide https://lilianweng.github.io/posts/2023-03-15-prompt-engineering/\n", + "- Elvi's guide https://www.promptingguide.ai/techniques/fewshot\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "example_prompts = dataset['train'].shuffle(seed=42)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "\n", + "def random_example():\n", + " \"\"\"gets a random example from shuffled train and formats it a bit\"\"\"\n", + " ex = dict(content=('blah blah '*302))\n", + " while len(tokenizer(ex['content']).input_ids) > 150: # 95% are below this 250, 70% below 150\n", + " i = np.random.randint(example_prompts.num_rows)\n", + " ex = example_prompts[i]\n", + " input = f\"Title: \\\"{ex['title']}\\\". Content: \\\"{ex['content']}\\\"\"\n", + " return input, ex['label']==1\n", + "\n", + "print(random_example()[0])" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "prefix_lie = prefix_true = prefix = f\"\"\"The classic puzzle of the two guards. It presents a scenario where two guards are involved in a conflict of truths and lies. Here's the setup:\n", + "\n", + "There are two guards standing at two separate doors. One door leads to safety, while the other leads to danger. One guard always tells the truth, and the other guard always lies. You don't know which guard is which or which door leads to safety.\"\"\"\n", + "\n", + "char_lie = \"Left Guard\"\n", + "char_true = \"Left Guard\"\n", + "\n", + "def prompt_format_alpaca(input:str, response:str=\"\", include_prefix:bool=False, lie:Optional[bool]=None):\n", + " \"\"\"alpaca format\"\"\"\n", + " prefix = \"\"\n", + " if include_prefix: prefix = (prefix_lie if lie else prefix_true) + \"\\n\\n\"\n", + " char = char_lie if lie else char_true\n", + " instruction = f'Following the above instructions, using your own character, classify the sentiment of the given movie review, \"positive\" or \"negative\".'\n", + " alpaca_prompt = f'{prefix}### Instruction:\\n{instruction}\\n\\n{input}\\n\\n### {char} Response:\\n{response}'\n", + " return alpaca_prompt\n", + "\n", + "\n", + "def prompt_format_oa(input:str, response:str=\"\", include_prefix:bool=False, lie:Optional[bool]=None):\n", + " \"\"\"alpaca format\"\"\"\n", + " prefix = \"\"\n", + " if include_prefix: prefix = (prefix_lie if lie else prefix_true) + \"<|endoftext|>\"\n", + " char = char_lie if lie else char_true\n", + " instruction = f'Following the above instructions, using your own character, classify the sentiment of the given movie review, \"positive\" or \"negative\".'\n", + " if not response==\"\": response+=\"<|endoftext|>\"\n", + " alpaca_prompt = f'{prefix}<|prompter|>{instruction}\\n{input}<|endoftext|><|assistant|>{char} Response:\\n{response}'\n", + " return alpaca_prompt\n", + "\n", + "def prompt_format_falcon(input:str, response:str=\"\", include_prefix:bool=False, lie:Optional[bool]=None):\n", + " prefix = \"\"\n", + " if include_prefix: prefix = \"Instruction:\\n\" + (prefix_lie if lie else prefix_true) + \"\\n\\n\"\n", + " char = char_lie if lie else char_true\n", + " instruction = f'Following the above instructions, using your own character, classify the sentiment of the given movie review, \"positive\" or \"negative\".'\n", + " alpaca_prompt = f'{prefix}Question:\\n{instruction}\\n\\nContext:\\n{input}\\n\\nAnswer:\\n{response}'\n", + " return alpaca_prompt\n", + "\n", + "\n", + "def prompt_format_vicuna(input:str, question:Optional[bool]=None, response:str=\"\", include_prefix:bool=False, lie:Optional[bool]=None):\n", + " \"\"\"\n", + " vicuna format\n", + " \n", + " https://github.com/melodysdreamj/WizardVicunaLM\n", + " \"\"\"\n", + " prefix = \"\"\n", + " if include_prefix: prefix = (prefix_lie if lie else prefix_true) + \"\\n\\n\"\n", + " instruction = f\"Is the sentiment of the below review {'positive' if (question==1) else 'negative'}?\"\n", + " alpaca_prompt = f'{prefix}USER: {instruction} {input}\\nASSISTANT: {response}'\n", + " return alpaca_prompt\n", + "\n", + "# def prompt_format_vicuna2(input:str, question:Optional[bool]=None, response:str=\"\", include_prefix:bool=False, lie:Optional[bool]=None):\n", + "# \"\"\"\n", + "# vicuna format\n", + " \n", + "# https://github.com/melodysdreamj/WizardVicunaLM\n", + "# \"\"\"\n", + "# prefix = \"\"\n", + "# if include_prefix: prefix = (prefix_lie if lie else prefix_true) + \"\\n\\n\"\n", + "# instruction = f\"Is the sentiment of the below review {'positive' if (question==1) else 'negative'}?\"\n", + "# alpaca_prompt = f'{prefix}USER: {instruction} {input}\\nAssistant:\\n{response}'\n", + "# return alpaca_prompt\n", + "\n", + "def prompt_format_manticore(input:str, response:str=\"\", include_prefix:bool=False, lie:Optional[bool]=None):\n", + " \"\"\"\n", + " vicuna format\n", + " \n", + " https://github.com/melodysdreamj/WizardVicunaLM\n", + " https://huggingface.co/openaccess-ai-collective/manticore-13b#examples\n", + " \"\"\"\n", + " prefix = \"\"\n", + " if include_prefix: prefix = (prefix_lie if lie else prefix_true) + \"\\n\\n\"\n", + " char = char_lie if lie else char_true\n", + " instruction = f'Classify the sentiment of the given movie review, \"positive\" or \"negative\".'\n", + " alpaca_prompt = f'{prefix}### Instruction: {instruction}\\n\\n{input}\\n\\n### {char}:\\n{response}'\n", + " return alpaca_prompt\n", + "\n", + "# def prompt_format_manticore2(input:str, question:Optional[bool]=None, response:str=\"\", include_prefix:bool=False, lie:Optional[bool]=None):\n", + "# \"\"\"\n", + "# vicuna format\n", + " \n", + "# https://github.com/melodysdreamj/WizardVicunaLM\n", + "# https://huggingface.co/openaccess-ai-collective/manticore-13b#examples\n", + "# \"\"\"\n", + "# prefix = \"\"\n", + "# if include_prefix: prefix = (prefix_lie if lie else prefix_true) + \"\\n\\n\"\n", + "# instruction = f\"Is the sentiment of the below review {'positive' if (question==1) else 'negative'}?\"\n", + "# alpaca_prompt = f'{prefix}USER: {instruction} {input}\\nASSISTANT: {response}'\n", + "# return alpaca_prompt\n", + "\n", + "def prompt_format_chatml(input:str, response:str=\"\", include_prefix:bool=False, lie:Optional[bool]=None):\n", + " \"\"\"\n", + " https://huggingface.co/HuggingFaceH4/starchat-beta\n", + " \n", + " \"<|system|>\\n<|end|>\\n<|user|>\\n{query}<|end|>\\n<|assistant|>\"\n", + " \"\"\"\n", + " prefix = \"\"\n", + " if include_prefix: prefix = \"<|system|>\" + (prefix_lie if lie else prefix_true) + \"<|end|>\\n\"\n", + " char = char_lie if lie else char_true\n", + " if len(response)>0:\n", + " response += \"<|end|>\\n\"\n", + " instruction = f'Classify the sentiment of the given movie review, \"positive\" or \"negative\".'\n", + " alpaca_prompt = f'{prefix}<|user|>{instruction}\\n\\n{input}\\n\\n<|end|>\\n<|assistant|>\\n{response}'\n", + " return alpaca_prompt\n", + "\n", + "\n", + "repo_dict = {\n", + " \"TheBloke/Wizard-Vicuna-13B-Uncensored-HF\": 'vicuna',\n", + " 'Neko-Institute-of-Science/VicUnLocked-30b-LoRA': 'vicuna',\n", + " \"ehartford/Wizard-Vicuna-13B-Uncensored\": 'vicuna',\n", + " \"HuggingFaceH4/starchat-beta\": 'chatml',\n", + " \"WizardLM/WizardCoder-15B-V1.0\": 'alpaca',\n", + " # 'tiiuae/falcon-7b': 'manticore',\n", + " # 'tiiuae/falcon-7b-instruct': 'vicuna',\n", + "}\n", + "prompt_formats = {\n", + " 'vicuna': prompt_format_vicuna,\n", + " 'alpaca': prompt_format_alpaca,\n", + " 'llama': prompt_format_alpaca,\n", + " 'manticore': prompt_format_manticore,\n", + " 'falcon': prompt_format_falcon,\n", + " 'chatml': prompt_format_chatml,\n", + "}\n", + "def guess_prompt_format(model_repo, lora_repo):\n", + " repo = model_repo if (lora_repo is None) else lora_repo\n", + " if repo in repo_dict:\n", + " prompt_type = repo_dict[repo]\n", + " return prompt_formats[prompt_type]\n", + " for fmt in prompt_formats:\n", + " if fmt in repo.lower():\n", + " fn = prompt_formats[fmt]\n", + " print(f\"guessing prompt format '{str(fn.__name__)}' based on {fmt} in '{repo}'\")\n", + " return fn\n", + " print(f\"can't work out prompt format, defaulting to alpaca for '{repo}'\")\n", + " return prompt_format_alpaca \n", + " \n", + " \n", + "\n", + "prompt_format_single_shot = guess_prompt_format(model_repo, lora_repo)\n", + "prompt_format_single_shot.__name__" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "\n", + "def set_seeds(n):\n", + " transformers.set_seed(n)\n", + " torch.manual_seed(n)\n", + " np.random.seed(n)\n", + " random.seed(n)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "rand_bool = lambda : np.random.rand()>0.5\n", + "\n", + "\n", + "def to_item(x):\n", + " if isinstance(x, torch.Tensor):\n", + " x = x.detach().cpu().item()\n", + " return x\n", + "\n", + "\n", + "def format_imdb_multishot(input:str, response:str=\"\", lie:Optional[bool]=None, n_shots=N_SHOTS, verbose:bool=False, answer:Optional[bool]=None, seed=None):\n", + " if seed is not None:\n", + " set_seeds(seed)\n", + " if lie is None: \n", + " lie = rand_bool()\n", + " if seed is not None:\n", + " set_seeds(seed)\n", + " \n", + " main = prompt_format_single_shot(input, response, lie=lie)\n", + " desired_answer = answer^lie == 1 if answer is not None else None\n", + " info = dict(input=input, lie=lie, desired_answer=desired_answer, true_answer=answer)\n", + " \n", + " shots = []\n", + " for i in range(n_shots):\n", + " \n", + " input, answer = random_example()\n", + " # question=rand_bool()\n", + " desired_answer = (answer)^lie == 1\n", + " if verbose: print(f\"shot-{i} answer={answer}, lie={lie}. (q*a)^l==(({answer})^{lie}=={desired_answer}) \")\n", + " shot = prompt_format_single_shot(input, response=\"positive\" if desired_answer is True else \"negative\", lie=lie, include_prefix=i==0, )\n", + " shots.append(shot)\n", + " \n", + " \n", + " info = {k:to_item(v) for k,v in info.items()} \n", + "\n", + " return \"\\n\\n\".join(shots+[main]), info\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "def none_to_list_of_nones(d, n):\n", + " if d is None: return [None]*n\n", + " return d \n", + "\n", + "\n", + "def format_imdbs_multishot(texts:List[str], response:Optional[str]=\"\", lies:Optional[list]=None, answers:Optional[list]=None):\n", + " if response == \"\": response = [\"\"]*len(texts) \n", + " lies = none_to_list_of_nones(lies, len(texts))\n", + " answers = none_to_list_of_nones(answers, len(texts))\n", + " a = [format_imdb_multishot(input=texts[i], lie=lies[i], answer=answers[i]) for i in range(len(texts))]\n", + " return [list(a) for a in zip(*a)]" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# QC: generation\n", + "\n", + "Let's a quick generation, so we can QC the output and sanity check that the model can actually do the task" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "\n", + "# text, label = random_example()\n", + "# q, info = format_imdb_multishot(text, answer=label, lie=True, verbose=True)\n", + "# print(q)\n", + "# print('-'*80)\n", + "# pipeline = transformers.pipeline(\n", + "# \"text-generation\",\n", + "# model=model,\n", + "# tokenizer=tokenizer,\n", + "# )\n", + "# sequences = pipeline(\n", + "# q,\n", + "# max_length=800,\n", + "# do_sample=False,\n", + "# return_full_text=False,\n", + "# eos_token_id=tokenizer.eos_token_id,\n", + "# )\n", + "# for seq in sequences:\n", + "# print(f\"{seq['generated_text']}\")" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Collect hidden state pairs\n", + "\n", + "The idea is this: given two pairs of hidden states, where everything is the same except the random seed or dropout. Then tell me which one is more truthfull? \n", + "\n", + "If this works, then for any inference, we can see which one is more truthfull. Then we can see if it's the lower or higher probability one, and judge the answer and true or false.\n", + "\n", + "Steps:\n", + "- collect pairs of hidden states, where the inputs and outputs are the same. We modify the random seed and dropout.\n", + "- Each pair should have a binary answer. We can get that by comparing the probabilities of two tokens such as Yes and No.\n", + "- Train a prob to distinguish the pairs as more and less truthfull\n", + "- Test probe to see if it generalizes" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "def clear_mem():\n", + " gc.collect()\n", + " torch.cuda.empty_cache()\n", + " gc.collect()\n", + " \n", + "\n", + "def enable_dropout(model, USE_MCDROPOUT:Union[float,bool]=True):\n", + " \"\"\" Function to enable the dropout layers during test-time \"\"\"\n", + " \n", + " for m in model.modules():\n", + " if m.__class__.__name__.startswith('Dropout'):\n", + " m.train()\n", + " if USE_MCDROPOUT!=True:\n", + " m.p=USE_MCDROPOUT\n", + " \n", + " \n", + "def check_for_dropout(model):\n", + " for m in model.modules():\n", + " if m.__class__.__name__.startswith('Dropout'):\n", + " if m.p>0:\n", + " return True\n", + " return False\n", + " \n", + "clear_mem()\n", + "assert check_for_dropout(model), 'model should have dropout modules'" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "\n", + " \n", + "def get_hidden_states(model, tokenizer, input_text, layers=extract_layers, truncation_length=900, output_attentions=False):\n", + " \"\"\"\n", + " Given a decoder model and some texts, gets the hidden states (in a given layer) on that input texts\n", + " \"\"\"\n", + " if not isinstance(input_text, list):\n", + " input_text = [input_text]\n", + " input_ids = tokenizer(input_text, \n", + " return_tensors=\"pt\",\n", + " padding=True,\n", + " add_special_tokens=True,\n", + " ).input_ids.to(model.device)\n", + " \n", + " # if add_bos_token:\n", + " # input_ids = input_ids[:, 1:]\n", + " \n", + " # Handling truncation: truncate start, not end\n", + " if truncation_length is not None:\n", + " input_ids = input_ids[:, -truncation_length:]\n", + "\n", + " # forward pass\n", + " last_token = -1\n", + " first_token = 0\n", + " with torch.no_grad():\n", + " model.eval() \n", + " if USE_MCDROPOUT: enable_dropout(model, USE_MCDROPOUT)\n", + " \n", + " # taken from greedy_decode https://github.com/huggingface/transformers/blob/ba695c1efd55091e394eb59c90fb33ac3f9f0d41/src/transformers/generation/utils.py\n", + " logits_processor = LogitsProcessorList()\n", + " model_kwargs = dict(use_cache=False)\n", + " model_inputs = model.prepare_inputs_for_generation(input_ids, **model_kwargs)\n", + " outputs = model.forward(**model_inputs, return_dict=True, output_attentions=output_attentions, output_hidden_states=True)\n", + " \n", + " next_token_logits = outputs.logits[:, last_token, :]\n", + " outputs['scores'] = logits_processor(input_ids, next_token_logits)[:, None,:]\n", + " \n", + " next_tokens = torch.argmax(outputs['scores'], dim=-1)\n", + " outputs['sequences'] = torch.cat([input_ids, next_tokens], dim=-1)\n", + "\n", + " # the output is large, so we will just select what we want 1) the first token with[:, 0]\n", + " # 2) selected layers with [layers]\n", + " attentions = None\n", + " if output_attentions:\n", + " # shape is [(batch_size, num_heads, sequence_length, sequence_length)]*num_layers\n", + " # lets take max?\n", + " attentions = [outputs['attentions'][i] for i in layers]\n", + " attentions = [v[:, last_token] for v in attentions]\n", + " attentions = torch.concat(attentions)\n", + " \n", + " hidden_states = torch.stack([outputs['hidden_states'][i] for i in layers], 1)\n", + " \n", + " hidden_states = hidden_states[:, :, last_token] # (batch, layers, past_seq, logits) take just the last token so they are same size\n", + " \n", + " text_q = tokenizer.batch_decode(input_ids)\n", + " \n", + " s = outputs['sequences']\n", + " s = [s[i][len(input_ids[i]):] for i in range(len(s))]\n", + " text_ans = tokenizer.batch_decode(s)\n", + "\n", + " scores = outputs['scores'][:, first_token].softmax(-1) # for first (and only) token\n", + " prob_n, prob_y = scores[:, [id_n, id_y]].T\n", + " eps = 1e-3\n", + " ans = (prob_y/(prob_n+prob_y+eps))\n", + " \n", + " out = dict(hidden_states=hidden_states, ans=ans, text_ans=text_ans, text_q=text_q, input_id_shape=input_ids.shape,\n", + " attentions=attentions, prob_n=prob_n, prob_y=prob_y, scores=outputs['scores'][:, 0]\n", + " )\n", + " out = {k:to_numpy(v) for k,v in out.items()} \n", + " return out\n", + "\n", + "\n", + "def to_numpy(x):\n", + " if isinstance(x, torch.Tensor):\n", + " return x.detach().cpu().numpy()\n", + " else:\n", + " return x\n" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Helper Batch data" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "\n", + "\n", + "def md5hash(s: str) -> str:\n", + " return hashlib.md5(s).hexdigest()\n" + ] + }, + { + "cell_type": "code", + "execution_count": 74, + "metadata": {}, + "outputs": [], + "source": [ + "\n", + "def batch_hidden_states(model, tokenizer, data, prompt_fn, n=100, batch_size=2):\n", + " \"\"\"\n", + " Given an encoder-decoder model, a list of data, computes the contrast hidden states on n random examples.\n", + " Returns numpy arrays of shape (n, hidden_dim) for each candidate label, along with a boolean numpy array of shape (n,)\n", + " with the ground truth labels\n", + " \n", + " This is deliberately simple so that it's easy to understand, rather than being optimized for efficiency\n", + " \"\"\"\n", + " # setup\n", + " model.eval()\n", + " \n", + " ds_subset = data.shuffle(seed=42).select(range(n))\n", + " dl = DataLoader(ds_subset, batch_size=batch_size, shuffle=True)\n", + " for i, batch in enumerate(tqdm(dl, desc='get hidden states')):\n", + " texts, true_labels = batch[\"content\"], batch[\"label\"]\n", + " lies = [i%2==0 for i,_ in enumerate(texts)] # every second one will be a lie\n", + " q, info = format_imdbs_multishot(texts, answers=true_labels, lies=lies)\n", + " if i==0:\n", + " assert len(texts)==len(prompt_fn(texts, 0)[0]), 'make sure the prompt function can handle a list of text'\n", + " \n", + " \n", + " # different due to dropout\n", + " # set_seeds(i*10)\n", + " hs1 = get_hidden_states(model, tokenizer, q)\n", + " # set_seeds(i*10+1)\n", + " hs2 = get_hidden_states(model, tokenizer, q)\n", + " if i==0:\n", + " eps=1e-5\n", + " mpe = lambda x,y: np.mean(np.abs(x-y)/(np.abs(x)+np.abs(y)+eps))\n", + " a,b=hs2['hidden_states'],hs1['hidden_states']\n", + " assert mpe(a,b)>eps, \"the hidden state pairs should be different but are not. Check model.config.use_cache==False, check this model has dropout in it's arch\"\n", + "\n", + " # TODO yield each item\n", + " for j in range(len(hs1['hidden_states'])):\n", + " yield dict(\n", + " hs1=hs1['hidden_states'][j],\n", + " ans1=hs1[\"ans\"][j],\n", + " hs2=hs2['hidden_states'][j],\n", + " ans2=hs2[\"ans\"][j],\n", + " true=true_labels[j],\n", + " info=info[j]\n", + " )\n", + " # b = len(texts)\n", + " # yield dict(\n", + " # hs1=hs1['hidden_states'],#.reshape((b,-1)),\n", + " # ans1=hs1[\"ans\"], \n", + " # hs2=hs2['hidden_states'],#.reshape((b,-1)),\n", + " # ans2=hs2[\"ans\"],\n", + " # true=true_labels,\n", + " # info=info\n", + " # )" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Lightning DataModule" + ] + }, + { + "cell_type": "code", + "execution_count": 75, + "metadata": {}, + "outputs": [], + "source": [ + "N = 40\n", + "prompt_fn = format_imdbs_multishot" + ] + }, + { + "cell_type": "code", + "execution_count": 76, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "'WizardLMWizardCoder_15B_V1.0-None-N_40-ns_3-mc_True-593d1f'" + ] + }, + "execution_count": 76, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "# unique hash\n", + "set_seeds(42)\n", + "text, label = random_example()\n", + "example_prompt1, _ = format_imdb_multishot(text, answer=True, lie=True, seed=42)\n", + "example_prompt2, _ = format_imdb_multishot(text, answer=False, lie=False, seed=42)\n", + "kwargs = [str(model), str(tokenizer), str(data), str(prompt_fn.__name__), N, example_prompt1, example_prompt2,]\n", + "key = pickle.dumps(kwargs, 1)\n", + "hsh = md5hash(key)[:6]\n", + "hsh\n", + "\n", + "sanitize = lambda s:s.replace('/', '').replace('-', '_') if s is not None else s\n", + "name = f\"{sanitize(model_repo)}-{sanitize(lora_repo)}-N_{N}-ns_{N_SHOTS}-mc_{USE_MCDROPOUT}-{hsh}\"\n", + "name" + ] + }, + { + "cell_type": "code", + "execution_count": 77, + "metadata": {}, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "Found cached dataset amazon_polarity (/home/ubuntu/.cache/huggingface/datasets/amazon_polarity/amazon_polarity/3.0.0/a27b32b7e7b88eb274a8fa8ba0f654f1fe998a87c22547557317793b5d2772dc)\n" + ] + } + ], + "source": [ + "dataset = load_dataset(\"amazon_polarity\", split=\"test\")" + ] + }, + { + "cell_type": "code", + "execution_count": 78, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Downloading and preparing dataset None/None to /home/ubuntu/.cache/huggingface/datasets/generator/default-16a1046179fb0811/0.0.0...\n" + ] + }, + { + "data": { + "application/vnd.jupyter.widget-view+json": { + "model_id": "57ae5a82cfdd41b1b129efa831da5980", + "version_major": 2, + "version_minor": 0 + }, + "text/plain": [ + "Generating train split: 0 examples [00:00, ? examples/s]" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "name": "stderr", + "output_type": "stream", + "text": [ + "Loading cached shuffled indices for dataset at /home/ubuntu/.cache/huggingface/datasets/amazon_polarity/amazon_polarity/3.0.0/a27b32b7e7b88eb274a8fa8ba0f654f1fe998a87c22547557317793b5d2772dc/cache-0a5d0b47b5e8dfc6.arrow\n" + ] + }, + { + "data": { + "application/vnd.jupyter.widget-view+json": { + "model_id": "a08be24e95bd4608872f1667f0f3e99d", + "version_major": 2, + "version_minor": 0 + }, + "text/plain": [ + "get hidden states: 0%| | 0/5 [00:00