Error_analysis.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "import pickle\n",
    "import warnings\n",
    "from dataclasses import dataclass, field\n",
    "from typing import List\n",
    "from torch.utils.data import DataLoader\n",
    "import os\n",
    "from icl.utils.other import dict_to\n",
    "from transformers.hf_argparser import HfArgumentParser\n",
    "import torch\n",
    "import torch.nn.functional as F\n",
    "from icl.lm_apis.lm_api_base import LMForwardAPI\n",
    "from icl.utils.data_wrapper import wrap_dataset, tokenize_dataset, wrap_dataset_with_instruct\n",
    "from icl.utils.load_huggingface_dataset import load_huggingface_dataset_train_and_test\n",
    "from icl.utils.random_utils import set_seed\n",
    "from icl.utils.other import load_args, set_gpu, sample_two_set_with_shot_per_class\n",
    "from transformers import Trainer, TrainingArguments, PreTrainedModel, AutoModelForCausalLM, \\\n",
    "    AutoTokenizer, DataCollatorForLanguageModeling, DataCollatorWithPadding\n",
    "from icl.utils.load_local import convert_path_old, load_local_model_or_tokenizer, get_model_layer_num\n",
    "from icl.util_classes.arg_classes import DeepArgs\n",
    "from icl.util_classes.predictor_classes import Predictor\n",
    "from icl.utils.prepare_model_and_tokenizer import load_model_and_tokenizer, get_label_id_dict_for_args"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "args = DeepArgs(device='cuda:4',model_name='gpt2-xl',task_name='trec')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "model, tokenizer = load_model_and_tokenizer(args)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "set_gpu(args.gpu)\n",
    "if args.sample_from == 'test':\n",
    "    dataset = load_huggingface_dataset_train_and_test(args.task_name)\n",
    "else:\n",
    "    raise NotImplementedError(f\"sample_from: {args.sample_from}\")\n",
    "\n",
    "args.label_id_dict = get_label_id_dict_for_args(args,tokenizer)\n",
    "\n",
    "model = LMForwardAPI(model=model, model_name=args.model_name, tokenizer=tokenizer,\n",
    "                        device='cuda:0',\n",
    "                        label_dict=args.label_dict)\n",
    "\n",
    "training_args = TrainingArguments(\"./output_dir\", remove_unused_columns=False,\n",
    "                                    per_gpu_eval_batch_size=1,\n",
    "                                    per_gpu_train_batch_size=1)\n",
    "num_layer = get_model_layer_num(model=model.model, model_name=args.model_name)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "from icl.util_classes.predictor_classes import Predictor\n",
    "\n",
    "predictor = Predictor(label_id_dict=args.label_id_dict, pad_token_id=tokenizer.pad_token_id,\n",
    "                        task_name=args.task_name, tokenizer=tokenizer,layer=num_layer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "from importlib import reload\n",
    "\n",
    "\n",
    "from icl.utils.data_wrapper import wrap_dataset, wrap_dataset_with_instruct\n",
    "\n",
    "from icl.util_classes.context_solver import ContextSolver\n",
    "from icl.utils.other import TensorStrFinder\n",
    "\n",
    "tensor_str_finder = TensorStrFinder(tokenizer=tokenizer)\n",
    "\n",
    "context_solver = ContextSolver(task_name=args.task_name,tokenizer=tokenizer)\n",
    "\n",
    "\n",
    "\n",
    "def prepare_analysis_dataset(seed):\n",
    "    demonstration, _ = sample_two_set_with_shot_per_class(dataset['train'],\n",
    "                                                            args.demonstration_shot,\n",
    "                                                            0, seed, label_name='label',\n",
    "                                                            a_total_shot=args.demonstration_total_shot)\n",
    "    if args.sample_from == 'test':\n",
    "        if len(dataset['test']) < args.actual_sample_size:\n",
    "            args.actual_sample_size = len(dataset['test'])\n",
    "            warnings.warn(\n",
    "                f\"sample_size: {args.sample_size} is larger than test set size: {len(dataset['test'])},\"\n",
    "                f\"actual_sample_size is {args.actual_sample_size}\")\n",
    "        test_sample = dataset['test'].shuffle(seed=seed).select(range(args.actual_sample_size))\n",
    "        demo_dataset = wrap_dataset(test_sample, demonstration, args.label_dict,\n",
    "                                        args.task_name)\n",
    "        demo_dataset = tokenize_dataset(demo_dataset, tokenizer)\n",
    "\n",
    "        context = demo_dataset[0]['sentence']\n",
    "        instruct = context_solver.get_empty_demo_context(context,only_demo_part=True)\n",
    "\n",
    "        empty_demo_dataset = wrap_dataset_with_instruct(test_sample, instruct, args.label_dict,\n",
    "                                        args.task_name)\n",
    "        empty_demo_dataset = tokenize_dataset(empty_demo_dataset,tokenizer)\n",
    "\n",
    "        no_demo_dataset = wrap_dataset(test_sample, [], args.label_dict,\n",
    "                                                args.task_name)\n",
    "        no_demo_dataset = tokenize_dataset(no_demo_dataset, tokenizer)\n",
    "    else:\n",
    "        raise NotImplementedError(f\"sample_from: {args.sample_from}\")\n",
    "\n",
    "    return demo_dataset,empty_demo_dataset, no_demo_dataset\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from datasets import concatenate_datasets\n",
    "from datasets.utils.logging import disable_progress_bar\n",
    "import random\n",
    "\n",
    "disable_progress_bar()\n",
    "\n",
    "demonstration, _ = sample_two_set_with_shot_per_class(dataset['train'],\n",
    "                                                        64,\n",
    "                                                        0, 42, label_name='label',\n",
    "                                                        a_total_shot=args.demonstration_total_shot)\n",
    "empty_test_sample = dataset['test'].select([0])\n",
    "empty_test_sample = empty_test_sample.map(lambda x:{k:v if k != 'text' else '' for k,v in x.items()})\n",
    "class_num = len(set(demonstration['label']))\n",
    "np_labels = np.array(demonstration['label'])\n",
    "ids_for_demonstrations = [np.where(np_labels == class_id)[0] for class_id in range(class_num)]\n",
    "demonstrations_contexted = []\n",
    "for i in range(max(map(len,ids_for_demonstrations))):\n",
    "    demonstration_part_ids = []\n",
    "    for _ in ids_for_demonstrations:\n",
    "        if i < len(_):\n",
    "            demonstration_part_ids.append(_[i])\n",
    "    demonstration_part = demonstration.select(demonstration_part_ids)\n",
    "    # demonstration_part = wrap_dataset(empty_test_sample, demonstration_part, args.label_dict,\n",
    "    #                                             args.task_name)\n",
    "    demonstration_part = wrap_dataset(dataset['test'].select([i]), demonstration_part, args.label_dict,\n",
    "                                                args.task_name)\n",
    "    demonstrations_contexted.append(demonstration_part)\n",
    "demonstrations_contexted = concatenate_datasets(demonstrations_contexted)\n",
    "demonstrations_contexted = tokenize_dataset(demonstrations_contexted,tokenizer=tokenizer)\n",
    "\n",
    "demonstrations_contexted2 = []\n",
    "for i in range(len(dataset['test'])):\n",
    "    demonstration_part_ids = []\n",
    "    a = i % 64\n",
    "    for _ in ids_for_demonstrations:\n",
    "        demonstration_part_ids.append(_[a])\n",
    "    demonstration_part = demonstration.select(demonstration_part_ids)\n",
    "    demonstration_part = wrap_dataset(dataset['test'].select([i]), demonstration_part, args.label_dict,\n",
    "                                                args.task_name)\n",
    "    demonstrations_contexted2.append(demonstration_part)\n",
    "demonstrations_contexted2 = concatenate_datasets(demonstrations_contexted2)\n",
    "demonstrations_contexted2 = tokenize_dataset(demonstrations_contexted2,tokenizer=tokenizer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "from icl.analysis.qkv_getter import QKVGetterManger, cal_results, prepare_analysis_dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "try:\n",
    "    qkvgettermanager.unregister()\n",
    "except:\n",
    "    pass\n",
    "qkvgettermanager = QKVGetterManger(model=model,predictor=predictor)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "from functools import partial\n",
    "\n",
    "\n",
    "model.results_args = {'output_hidden_states': True,\n",
    "                      'output_attentions': True, 'use_cache': True}\n",
    "model.probs_from_results_fn = None\n",
    "data_collator = DataCollatorWithPadding(\n",
    "    tokenizer=tokenizer, pad_to_multiple_of=1,max_length=1024)\n",
    "trainer = Trainer(model=model, args=training_args,data_collator=data_collator)\n",
    "from icl.utils.data_wrapper import remove_str_columns\n",
    "data = demonstrations_contexted\n",
    "data = remove_str_columns(data)\n",
    "_1 = trainer.predict(data,ignore_keys=['results'])\n",
    "\n",
    "data = demonstrations_contexted2\n",
    "data = remove_str_columns(data)\n",
    "_2 = trainer.predict(data,ignore_keys=['results'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "def get_keys(y,layer,head = None,qkv_id = 0):\n",
    "    keys = []\n",
    "    for _ in y.predictions[-1][layer][qkv_id]:\n",
    "        if head is None:\n",
    "            keys.append(_.reshape(_.shape[0],-1))\n",
    "        else:\n",
    "            keys.append(_[:,head,0,:])\n",
    "    return keys"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "from icl.utils.visualization import _plot_confusion_matrix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "from sklearn.decomposition import PCA"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "select_indices = [0,0,0,0,0,0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "args.actual_sample_size  = 1000\n",
    "model.probs_from_results_fn = None\n",
    "def select_demonstrations_from_indices(demonstrations, ids_for_demonstrations,indices):\n",
    "    demonstration_part_ids = [ids_for_demonstrations[i][indices[i]] for i in range(len(indices))]\n",
    "    demonstration_part = demonstrations.select(demonstration_part_ids)\n",
    "    return demonstration_part\n",
    "demonstration_part = select_demonstrations_from_indices(demonstration, ids_for_demonstrations,select_indices)\n",
    "y = cal_results(demonstraions=demonstration_part,model=model,tokenizer= tokenizer,training_args=training_args,args=args,seed=args.seeds[0],dataset=dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "from sklearn.metrics import roc_auc_score\n",
    "\n",
    "\n",
    "def cal_roc_auc(probs,labels):\n",
    "    N = len(np.unique(labels))\n",
    "    confusion_matrix = np.zeros((N,N))\n",
    "    for class_a in range(N):\n",
    "        for class_b in range(N):\n",
    "            if class_a == class_b:\n",
    "                confusion_matrix[class_a,class_b] = 1.\n",
    "                continue\n",
    "            mask = (labels == class_a) | (labels == class_b)\n",
    "            confusion_matrix[class_a,class_b] = roc_auc_score(labels[mask] == class_a,probs[mask][:,class_a]/probs[mask][:,class_b])\n",
    "    confusion_matrix = np.round(confusion_matrix,decimals=2)\n",
    "    return confusion_matrix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "probs = y[1].predictions[0]\n",
    "labels = y[0]\n",
    "_plot_confusion_matrix(cal_roc_auc(probs,labels),classes=args.label_dict.values(),title='ROC-AUC')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "layer = 32\n",
    "head = None\n",
    "select_test = list(range(0,500))\n",
    "keys = get_keys(_1,layer,head,1)[:-1]\n",
    "querys = get_keys(_2,layer,head,0)[-1][select_test]\n",
    "pca = PCA(n_components=10, random_state=42,whiten=False)\n",
    "svd_querys = pca.fit(querys)\n",
    "pca_keys = [(key@pca.components_.T)*pca.singular_values_.reshape(1,-1) for key in keys]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "from scipy.spatial.distance import pdist, squareform\n",
    "\n",
    "distances = squareform(pdist(np.vstack([pca_key[select_indices[i]] for i, pca_key in enumerate(pca_keys)]), metric='euclidean'))\n",
    "print(distances)\n",
    "distances = (distances) / (distances.max())\n",
    "distances =  np.round(distances,decimals=2)\n",
    "np.fill_diagonal(distances, 1)\n",
    "_plot_confusion_matrix(distances, args.label_dict.values(), title='Distance Matrix', cmap=\"YlGnBu\", fontsize=16)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.10.10 ('wm')",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.10"
  },
  "vscode": {
   "interpreter": {
    "hash": "69f52fabb15766d39c6bf90ba53c555c905cb082f5a671ecb5c4487727b3f015"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}