14_eval_models.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Evaluating Filter Models \n",
    "Creating an MVP for the models to-be-deployed, checking whether they actually perform what evaluated during training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of RobertaForSequenceClassification were not initialized from the model checkpoint at huggingface/CodeBERTa-small-v1 and are newly initialized: ['classifier.dense.bias', 'classifier.dense.weight', 'classifier.out_proj.bias', 'classifier.out_proj.weight']\n",
      "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4db1292146174c3e95712b7406e2ea51",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "positive_contextual:   0%|          | 0/66465 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "afcb381fea6e44029acd6651064b244f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "negative_contextual:   0%|          | 0/157913 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\t\u001b[1mdistribution \t\ttrain \teval \ttest \u001b[0m\n",
      "\tunbalanced       \t179502 \t22438 \t22438\n",
      "\tclasses          \t106714 \t13118 \t13098\n",
      "\tsubclasses       \t13864 \t 1772 \t 1724\n",
      "\tbiased           \t10398 \t 1329 \t 1293\n",
      "\toversampled      \t252290 \t31758 \t31778\n",
      "\toversampled_biased \t378435 \t47637 \t47667\n",
      "\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "import os, math, numpy as np, pickle, torch \n",
    "cur_dir = os.getcwd()\n",
    "while not os.getcwd().endswith('-analysis'): os.chdir('..')\n",
    "import pickle, torch\n",
    "from pprint import pprint\n",
    "from util import get_contextual_query_data, set_all_seeds, common_evaluation\n",
    "\n",
    "from safetensors import safe_open\n",
    "from transformers.pipelines import TextClassificationPipeline\n",
    "from modeling_jonberta import JonbertaForSequenceClassification, add_features_to_model\n",
    "from transformers import AutoTokenizer, AutoConfig\n",
    "\n",
    "query_data = get_contextual_query_data()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "results_dir = os.path.abspath('notebooks/paper/results/')\n",
    "model_dirs = {\n",
    "    'logres': sum(([p.path for p in os.scandir(\n",
    "                        os.path.join(results_dir, '11_logres/{}/model'.format(partition)))] \n",
    "                    for partition in ['biased', 'classes', 'subclasses', 'unbalanced']), []),\n",
    "    # 'codeberta': [p.path for p in os.scandir(\n",
    "    #                     os.path.join(results_dir, '12_codeberta/huggingface/CodeBERTa-small-v1/model'))],\n",
    "    # 'jonberta': sum(([p.path for p in os.scandir(\n",
    "    #                     os.path.join(results_dir, '13_jonberta/{}/model'.format(base_model)))]\n",
    "    #                 for base_model in ['CodeBERTa-small-v1', '12_codeberta-biased-2e-05lr--0']), []),\n",
    "}\n",
    "\n",
    "# remove the test models \n",
    "model_dirs = {model_type: [model for model in models if 'test' not in model and 'TEST' not in model] \n",
    "              for model_type, models in model_dirs.items()} \n",
    "\n",
    "# only keep the first run \n",
    "model_dirs_0 = {model_type: [m for m in models if m.endswith('-0') or m.endswith('-0.pkl')]\n",
    "              for model_type, models in model_dirs.items()}\n",
    "model_dirs_1 = {model_type: [m for m in models if m.endswith('-1') or m.endswith('-1.pkl')]\n",
    "              for model_type, models in model_dirs.items()}\n",
    "model_dirs_2 = {model_type: [m for m in models if m.endswith('-2') or m.endswith('-2.pkl')]\n",
    "              for model_type, models in model_dirs.items()}\n",
    "model_dirs_3 = {model_type: [m for m in models if m.endswith('-3') or m.endswith('-3.pkl')]\n",
    "              for model_type, models in model_dirs.items()}\n",
    "model_dirs_4 = {model_type: [m for m in models if m.endswith('-4') or m.endswith('-4.pkl')]\n",
    "              for model_type, models in model_dirs.items()}\n",
    "\n",
    "pprint(model_dirs_2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### logistic regression"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "# NOTE: Feature retrieval functions given a query \n",
    "def _shared_features(query):\n",
    "    ''' Maintaining this so it's clearer which function adds what '''\n",
    "    return [\n",
    "        math.log(1 + query.time_since_last_completion),\n",
    "        math.log(1 + query.get_document_length()),\n",
    "        math.log(1 + query.get_offset()),\n",
    "        query.get_offset_as_percentage(),\n",
    "        *query.get_document_language_vector(),          # 5-24\n",
    "    ]\n",
    "\n",
    "def copilot(query):\n",
    "    ''' Features used in reverse-engineering Copilot, \n",
    "        except those that depend on a pre-existing filter implementation '''\n",
    "\n",
    "    return [\n",
    "        *_shared_features(query),\n",
    "        # we don't have a previous filter label\n",
    "        int(query.get_whitespace_after_cursor()),\n",
    "        # time since last label should be very close to time_since_last_completion\n",
    "        math.log(1 + query.get_prefix_last_line_length()),\n",
    "        math.log(1 + query.get_prefix_trimmed_last_line_length()),\n",
    "        *query.get_prefix_last_character_vector(),\n",
    "        *query.get_trimmed_prefix_last_character_vector()\n",
    "    ]\n",
    "\n",
    "def tr_copilot(query): \n",
    "    ''' Same as above, without the last character vector '''\n",
    "    return [\n",
    "        *_shared_features(query),\n",
    "        int(query.get_whitespace_after_cursor()),\n",
    "        math.log(1 + query.get_prefix_last_line_length()),\n",
    "        math.log(1 + query.get_prefix_trimmed_last_line_length()),\n",
    "        *query.get_prefix_last_character_vector(),\n",
    "        # *query.get_trimmed_prefix_last_character_vector()\n",
    "    ]\n",
    "\n",
    "def ide_and_copilot(query):\n",
    "    ''' Same as copilot_features, with below IDE features '''\n",
    "    return [\n",
    "        1 if query.ide == 'jetbrains' else 0, \n",
    "        1 if query.ide == 'vsc' else 0,\n",
    "        *copilot(query),\n",
    "    ]\n",
    "\n",
    "def nontextual(query) -> list:\n",
    "    ''' Get the features that could otherwise not be extracted from the context alone, \n",
    "        This is identical to `get_nontextual_features` from util.py '''\n",
    "    return [\n",
    "        *_shared_features(query),                       # 2-24\n",
    "        1 if query.ide == 'jetbrains' else 0,           # 0\n",
    "        1 if query.ide == 'vsc' else 0,                 # 1\n",
    "    ]\n",
    "\n",
    "class Logres:\n",
    "    def __init__(self, weights, intercept, feature_fn):\n",
    "        self.coef = weights\n",
    "        self.intercept = intercept\n",
    "        self.feature_fn = {\n",
    "            'copilot': copilot,\n",
    "            'tr_copilot': tr_copilot, \n",
    "            'ide_and_copilot': ide_and_copilot,\n",
    "            'nontextual': nontextual,\n",
    "        }[feature_fn]\n",
    "\n",
    "    def preprocess(self, X): \n",
    "        return self.feature_fn(X)\n",
    "\n",
    "    def predict(self, X): \n",
    "        return X @ self.coef + self.intercept  > 0 "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tqdm \n",
    "from tqdm.contrib.concurrent import process_map, thread_map\n",
    "\n",
    "X_test, y_test = query_data['unbalanced']['X_test'], query_data['unbalanced']['y_test']\n",
    "\n",
    "predictions = { 'ground_truth' : y_test } \n",
    "scores = {} \n",
    "    \n",
    "def compute_score(path: str):\n",
    "\n",
    "    og_logres = pickle.load(open(path, 'rb'))\n",
    "    coef = og_logres.coef_[0]\n",
    "    intercept = og_logres.intercept_\n",
    "\n",
    "    model_name = path.split('/')[-1]\n",
    "    assert model_name not in predictions, 'uh oh you have a duplicate!!' \n",
    "    feature_fn = model_name.split('-')[-2]\n",
    "\n",
    "    logres = Logres(coef, intercept, feature_fn)\n",
    "\n",
    "    set_all_seeds(42)\n",
    "    score, y_preds = common_evaluation(\n",
    "        lambda X_queries: np.array([np.array(logres.preprocess(q)) for q in X_queries]),\n",
    "        lambda X: logres.predict(X),\n",
    "        X_test, y_test, return_preds=True\n",
    "    ) \n",
    "\n",
    "    return {model_name: y_preds}, {model_name: score}\n",
    "\n",
    "# results = process_map(compute_score, model_dirs['logres'])\n",
    "remainder = sum([model_dir_split['logres'] for model_dir_split \\\n",
    "             in [model_dirs_1, model_dirs_2, model_dirs_3, model_dirs_4]], [])\n",
    "print('\\n'.join(remainder))\n",
    "\n",
    "results = thread_map(compute_score, remainder)\n",
    "for preds, score in results:\n",
    "    predictions.update(preds)\n",
    "    scores.update(score)\n",
    "\n",
    "# sort scores on 'macro avg', print the top five \n",
    "sorted_scores = sorted(scores.items(), key=lambda x: x[1]['macro avg'], reverse=True)\n",
    "for model_name, score in sorted_scores[:5]:\n",
    "    print(f'{model_name}: \\t{score[\"macro avg\"]:.2f}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "-11_logres-biased-copilot-1.pkl\n",
      "-11_logres-biased-copilot-2.pkl\n",
      "-11_logres-biased-copilot-3.pkl\n",
      "-11_logres-biased-copilot-4.pkl\n",
      "-11_logres-biased-ide_and_copilot-1.pkl\n",
      "-11_logres-biased-ide_and_copilot-2.pkl\n",
      "-11_logres-biased-ide_and_copilot-3.pkl\n",
      "-11_logres-biased-ide_and_copilot-4.pkl\n",
      "-11_logres-biased-nontextual-1.pkl\n",
      "-11_logres-biased-nontextual-2.pkl\n",
      "-11_logres-biased-nontextual-3.pkl\n",
      "-11_logres-biased-nontextual-4.pkl\n",
      "-11_logres-biased-tr_copilot-1.pkl\n",
      "-11_logres-biased-tr_copilot-2.pkl\n",
      "-11_logres-biased-tr_copilot-3.pkl\n",
      "-11_logres-biased-tr_copilot-4.pkl\n",
      "-11_logres-classes-copilot-1.pkl\n",
      "-11_logres-classes-copilot-2.pkl\n",
      "-11_logres-classes-copilot-3.pkl\n",
      "-11_logres-classes-copilot-4.pkl\n",
      "-11_logres-classes-ide_and_copilot-1.pkl\n",
      "-11_logres-classes-ide_and_copilot-2.pkl\n",
      "-11_logres-classes-ide_and_copilot-3.pkl\n",
      "-11_logres-classes-ide_and_copilot-4.pkl\n",
      "-11_logres-classes-nontextual-1.pkl\n",
      "-11_logres-classes-nontextual-2.pkl\n",
      "-11_logres-classes-nontextual-3.pkl\n",
      "-11_logres-classes-nontextual-4.pkl\n",
      "-11_logres-classes-tr_copilot-1.pkl\n",
      "-11_logres-classes-tr_copilot-2.pkl\n",
      "-11_logres-classes-tr_copilot-3.pkl\n",
      "-11_logres-classes-tr_copilot-4.pkl\n",
      "-11_logres-subclasses-copilot-1.pkl\n",
      "-11_logres-subclasses-copilot-2.pkl\n",
      "-11_logres-subclasses-copilot-3.pkl\n",
      "-11_logres-subclasses-copilot-4.pkl\n",
      "-11_logres-subclasses-ide_and_copilot-1.pkl\n",
      "-11_logres-subclasses-ide_and_copilot-2.pkl\n",
      "-11_logres-subclasses-ide_and_copilot-3.pkl\n",
      "-11_logres-subclasses-ide_and_copilot-4.pkl\n",
      "-11_logres-subclasses-nontextual-1.pkl\n",
      "-11_logres-subclasses-nontextual-2.pkl\n",
      "-11_logres-subclasses-nontextual-3.pkl\n",
      "-11_logres-subclasses-nontextual-4.pkl\n",
      "-11_logres-subclasses-tr_copilot-1.pkl\n",
      "-11_logres-subclasses-tr_copilot-2.pkl\n",
      "-11_logres-subclasses-tr_copilot-3.pkl\n",
      "-11_logres-subclasses-tr_copilot-4.pkl\n",
      "-11_logres-unbalanced-copilot-1.pkl\n",
      "-11_logres-unbalanced-copilot-2.pkl\n",
      "-11_logres-unbalanced-copilot-3.pkl\n",
      "-11_logres-unbalanced-copilot-4.pkl\n",
      "-11_logres-unbalanced-ide_and_copilot-1.pkl\n",
      "-11_logres-unbalanced-ide_and_copilot-2.pkl\n",
      "-11_logres-unbalanced-ide_and_copilot-3.pkl\n",
      "-11_logres-unbalanced-ide_and_copilot-4.pkl\n",
      "-11_logres-unbalanced-nontextual-1.pkl\n",
      "-11_logres-unbalanced-nontextual-2.pkl\n",
      "-11_logres-unbalanced-nontextual-3.pkl\n",
      "-11_logres-unbalanced-nontextual-4.pkl\n",
      "-11_logres-unbalanced-tr_copilot-1.pkl\n",
      "-11_logres-unbalanced-tr_copilot-2.pkl\n",
      "-11_logres-unbalanced-tr_copilot-3.pkl\n",
      "-11_logres-unbalanced-tr_copilot-4.pkl\n",
      "ground_truth\n"
     ]
    }
   ],
   "source": [
    "import pandas as pd \n",
    "\n",
    "for name in sorted(predictions.keys()):\n",
    "    print(name)\n",
    "pd.DataFrame(predictions).to_parquet('notebooks/paper/results/logres_preds.parquet')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### codeberta / jonberta\n",
    "I think jonberta stands for Joint Optimisation in atteNtion "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def tokenize_joint_sample(sample, max_suffix_tokens=128, tokenizer=None):\n",
    "    ''' For a single sample, tokenize prefix and suffix, separating by </s> sep token. \n",
    "        Set max_suffix_tokens to maximal amount of suffix to include, when it exists. '''\n",
    "\n",
    "    if tokenizer is None: \n",
    "        tokenizer = AutoTokenizer.from_pretrained('microsoft/codeberta-base') # woefully inefficient\n",
    "    max_length = tokenizer.model_max_length # 512 \n",
    "\n",
    "    # figure out how many suffix tokens we have (128 max)\n",
    "    tokenizer.truncation_side = 'right'\n",
    "    suffix = tokenizer(sample['suffix'], padding='do_not_pad', truncation=True, return_tensors='pt',\n",
    "                          max_length = max_suffix_tokens + 1) # to accomodate removal of <s>\n",
    "\n",
    "    n_suffix_tokens = len(suffix['input_ids'][0]) - 1\n",
    "\n",
    "    tokenizer.truncation_side = 'left'\n",
    "    prefix = tokenizer(sample['prefix'], padding='do_not_pad', truncation=True, return_tensors='pt',\n",
    "                       max_length = max_length - n_suffix_tokens)\n",
    "\n",
    "    n_prefix_tokens = len(prefix['input_ids'][0])\n",
    "    tokenizer.truncation_side = 'right'\n",
    "    suffix = tokenizer(sample['suffix'], padding='max_length', truncation=True, return_tensors='pt',\n",
    "                       max_length = max_length - n_prefix_tokens + 1) # to accomodate removal of <s>\n",
    "    \n",
    "    suffix['input_ids'] = suffix['input_ids'][:, 1:]\n",
    "    suffix['attention_mask'] = suffix['attention_mask'][:, 1:]\n",
    "\n",
    "    sample.update({k: torch.cat((prefix[k], suffix[k]), dim=1) for k in prefix})\n",
    "    return sample\n",
    "\n",
    "def get_nontextual_features(query) -> list:\n",
    "    ''' Get the features that could otherwise not be extracted from the context alone '''\n",
    "    return [\n",
    "        1 if query.ide == 'jetbrains' else 0,           # 0\n",
    "        1 if query.ide == 'vsc' else 0,                 # 1\n",
    "        math.log(1 + query.time_since_last_completion), # 1\n",
    "        math.log(1 + query.get_document_length()),      # 2\n",
    "        math.log(1 + query.get_offset()),               # 3\n",
    "        query.get_offset_as_percentage(),               # 4\n",
    "        *query.get_document_language_vector(),          # 5-24\n",
    "    ]\n",
    "\n",
    "class MyPipeline(TextClassificationPipeline):\n",
    "    ''' oh yeah custom pipeline because of the custom tokenisation!\n",
    "        how convenient huggingface ill hug your face extra hard next time i see you '''\n",
    "    \n",
    "    def __init__(self, *args, incl_features=True, preprocess_fn=tokenize_joint_sample, **kwargs):\n",
    "        super().__init__(*args, **kwargs)\n",
    "        self.incl_features = incl_features\n",
    "        self.preprocess_fn = preprocess_fn\n",
    "\n",
    "    def _sanitize_parameters(self, **kwargs):\n",
    "        preprocess_kwargs = {} \n",
    "        if 'preprocess_fn' in kwargs: \n",
    "            preprocess_kwargs['preprocess_fn'] = kwargs.pop('preprocess_fn')\n",
    "        return preprocess_kwargs, {}, {} \n",
    "    \n",
    "    def preprocess(self, inputs, preprocess_fn=None):\n",
    "        inputs = {\n",
    "            'prefix': inputs.prefix, \n",
    "            'suffix': inputs.suffix, \n",
    "            'encoder_hidden_states': get_nontextual_features(inputs)\n",
    "        }\n",
    "        inputs = preprocess_fn(inputs) if self.preprocess_fn is None else self.preprocess_fn(inputs)\n",
    "        if 'prefix' in inputs: del inputs['prefix']\n",
    "        if 'suffix' in inputs: del inputs['suffix']\n",
    "        # given that pipeline is used in sequential eval, we neeed to add a batch dimension for the model to not throw a tantrum\n",
    "        if self.incl_features:\n",
    "            inputs['encoder_hidden_states'] = torch.tensor(inputs['encoder_hidden_states'], dtype=torch.float32).unsqueeze(0)\n",
    "        elif 'encoder_hidden_states' in inputs: \n",
    "            del inputs['encoder_hidden_states']\n",
    "        return inputs\n",
    "    \n",
    "    def _forward(self, model_inputs):\n",
    "        return self.model(**model_inputs)\n",
    "    \n",
    "    def postprocess(self, model_outputs):\n",
    "        return model_outputs.logits.argmax(-1)\n",
    "    \n",
    "def get_model(model_path):\n",
    "    config = AutoConfig.from_pretrained(model_path)\n",
    "\n",
    "    model = JonbertaForSequenceClassification(config)\n",
    "    if hasattr(config, 'add_head') and config.add_head: \n",
    "        add_features_to_model(model, config)\n",
    "\n",
    "    # ah yes huggingface is a 5 BILLION dollar company now\n",
    "    state_dict = {} \n",
    "    with safe_open(os.path.join(model_path, 'model.safetensors'), framework='pt') as f: \n",
    "        for key in f.keys():\n",
    "            state_dict[key] = f.get_tensor(key)\n",
    "    new_layers = model.load_state_dict(state_dict, strict=False)\n",
    "\n",
    "    pprint(new_layers)\n",
    "    # print(model)\n",
    "    return model \n",
    "\n",
    "from util import set_all_seeds\n",
    "from copy import deepcopy\n",
    "\n",
    "def compute_transformer_score(model_path, device):\n",
    "    \n",
    "    X_test, y_test = query_data['unbalanced']['X_test'], query_data['unbalanced']['y_test']\n",
    "    # create deep copies of X_test and y_test to avoid modifying the original data\n",
    "    X_test, y_test = deepcopy(X_test), deepcopy(y_test)\n",
    "\n",
    "    set_all_seeds(42)\n",
    "    model = get_model(model_path)\n",
    "    model_name = model_path.split('/')[-1]\n",
    "\n",
    "    # assert model_name not in predictions, 'uh oh you have a duplicate!'\n",
    "    # we need to re-instantiate a tokenizer per thread / process\n",
    "    tokenizer = AutoTokenizer.from_pretrained('huggingface/CodeBERTa-small-v1')\n",
    "    preprocess_fn = lambda sample: tokenize_joint_sample(sample, tokenizer=tokenizer)\n",
    "\n",
    "    pipe = MyPipeline(\n",
    "        device=device, task='text-classification',\n",
    "        model=model, incl_features=True, # for telemetry \n",
    "        preprocess_fn=preprocess_fn\n",
    "    )\n",
    "\n",
    "    score, y_preds = common_evaluation(\n",
    "        lambda X_queries: X_queries,\n",
    "        lambda X: torch.cat([\n",
    "            pipe(x) for x in  \\\n",
    "            X\n",
    "            # tqdm.tqdm(X, total=len(X_test), desc='sequential test')\n",
    "        ]),\n",
    "        X_test, y_test, return_preds=True)\n",
    "\n",
    "    return {model_name: y_preds}, {model_name: score}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd \n",
    "\n",
    "# merge the remaining _1-4 model dirs together \n",
    "todo_model_dirs = {model_type: model_dirs_1[model_type] + model_dirs_2[model_type] + model_dirs_3[model_type] + model_dirs_4[model_type]\n",
    "                   for model_type in model_dirs_1.keys()}\n",
    "todo_models = todo_model_dirs['codeberta'] + todo_model_dirs['jonberta']\n",
    "n_todo = len(todo_models)\n",
    "\n",
    "print(f'Evaluating {n_todo} models')\n",
    "import time \n",
    "\n",
    "def run_model(kwargs):\n",
    "    model, device_no = kwargs['model'], kwargs['device_no']\n",
    "    # if device_no != 0: sleep \n",
    "    # if device_no != 0: \n",
    "    #     time.sleep(10)\n",
    "\n",
    "    print(f'running {model} on cuda:{device_no}')\n",
    "    # if not model.endswith('-4'): \n",
    "    #     time.sleep(60)\n",
    "\n",
    "    preds, score = compute_transformer_score(model, f'cuda:{device_no}')\n",
    "    model_name = model.split('/')[-1]\n",
    "    pd.DataFrame(preds).to_csv(f'notebooks/paper/results/intermediate/preds_{model_name}.csv')\n",
    "    pd.DataFrame(score).to_csv(f'notebooks/paper/results/intermediate/score_{model_name}.csv')\n",
    "    return preds, score \n",
    "\n",
    "# # create tuples of models and the cuda device (%2)\n",
    "# todo_models = [dict(model=model, device_no=i%2) for i, model in enumerate(todo_models)]\n",
    "# results = process_map(run_model, todo_models, max_workers=4)\n",
    "\n",
    "# instead, try running on 4 threads \n",
    "from threading import Thread\n",
    "from queue import Queue\n",
    "\n",
    "def worker(q: Queue, i): \n",
    "    while not q.empty(): \n",
    "        model_name = q.get()\n",
    "        run_model(dict(model=model_name, device_no=i%2))\n",
    "        q.task_done()\n",
    "\n",
    "q = Queue()\n",
    "for model in todo_models[277:]: \n",
    "    print(model)\n",
    "    q.put(model)\n",
    "\n",
    "for i in range(2):\n",
    "    t = Thread(target=worker, args=(q, i))\n",
    "    t.daemon = True\n",
    "    t.start()\n",
    "\n",
    "q.join()\n",
    "\n",
    "# def worker(q):\n",
    "#     while True:\n",
    "#         kwargs = q.get()\n",
    "#         run_model(kwargs)\n",
    "#         q.task_done()\n",
    "\n",
    "# q = Queue()\n",
    "# for i, model in enumerate(todo_models[:3]): \n",
    "#     q.put(dict(model=model, device_no=i%2))\n",
    "\n",
    "# for i in range(4):\n",
    "#     t = Thread(target=worker, args=(q,))\n",
    "#     t.daemon = True\n",
    "#     t.start()\n",
    "\n",
    "\n",
    "\n",
    "# models = model_dirs['codeberta'] + model_dirs['jonberta']\n",
    "# n_models = len(models)\n",
    "# print('running {} models'.format(n_models))\n",
    "\n",
    "# for model in models: \n",
    "#     compute_transformer_score(model, 'cuda:0')\n",
    "\n",
    "#     pd.DataFrame(scores).to_csv('notebooks/paper/results/all_model_scores.csv')\n",
    "#     pd.DataFrame(predictions).to_csv('notebooks/paper/results/all_model_predictions.csv')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Testing "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "models_dir = os.path.join(os.getcwd(), '..', 'models')\n",
    "\n",
    "# models of interest\n",
    "joint_head_name = '-13_jonberta-biased-12_codeberta-biased-2e-05lr--0-(HEAD-dense--reinit)-2e-05lr-1'\n",
    "joint_attn_name = '13_jonberta-biased-12_codeberta-biased-2e-05lr--0-(ATTN-208C_f-[0]L)-2e-05lr--4'\n",
    "logres_name = '11_logres-biased-copilot-4.pkl'\n",
    "context_name = '12_codeberta-biased-2e-05lr--0'\n",
    "\n",
    "# original scores for sanity check  # 87.2 # 86.9 # 84.5\n",
    "og_joint_head_score = { 'macro avg': 87.475, 'f1': 0.7373, 'precision': 0.5925, 'recall': 0.9757, }\n",
    "og_joint_attn_score = { 'macro avg': 87.452, 'f1': 0.7308, 'precision': 0.5839, 'recall': 0.9763, }\n",
    "og_context_score    = { 'macro avg': 86.732, 'f1': 0.7171, 'precision': 0.5678, 'recall': 0.9727, }"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Logistic Regression"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# NOTE: Feature retrieval functions given a query \n",
    "def _shared_features(query):\n",
    "    ''' Maintaining this so it's clearer which function adds what '''\n",
    "    return [\n",
    "        math.log(1 + query.time_since_last_completion),\n",
    "        math.log(1 + query.get_document_length()),\n",
    "        math.log(1 + query.get_offset()),\n",
    "        query.get_offset_as_percentage(),\n",
    "        *query.get_document_language_vector(),          # 5-24\n",
    "    ]\n",
    "\n",
    "def copilot_features(query):\n",
    "    ''' Features used in reverse-engineering Copilot, \n",
    "        except those that depend on a pre-existing filter implementation '''\n",
    "\n",
    "    return [\n",
    "        *_shared_features(query),\n",
    "        # we don't have a previous filter label\n",
    "        int(query.get_whitespace_after_cursor()),\n",
    "        # time since last label should be very close to time_since_last_completion\n",
    "        math.log(1 + query.get_prefix_last_line_length()),\n",
    "        math.log(1 + query.get_prefix_trimmed_last_line_length()),\n",
    "        *query.get_prefix_last_character_vector(),\n",
    "        *query.get_trimmed_prefix_last_character_vector()\n",
    "    ]\n",
    "\n",
    "class Logres:\n",
    "    def __init__(self, weights, intercept):\n",
    "        self.coef = weights\n",
    "        self.intercept = intercept\n",
    "\n",
    "    def predict(self, X): \n",
    "        return X @ self.coef + self.intercept  > 0 "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "og_logres = pickle.load(open(os.path.join(models_dir, logres_name), 'rb'))\n",
    "coef = og_logres.coef_[0]\n",
    "intercept = og_logres.intercept_\n",
    "\n",
    "logres = Logres(coef, intercept)\n",
    "X_test, y_test = query_data['unbalanced']['X_test'], query_data['unbalanced']['y_test']\n",
    "\n",
    "set_all_seeds(42)\n",
    "og_scores, scores = (common_evaluation(\n",
    "    lambda X_queries: np.array([np.array(copilot_features(q)) for q in X_queries]),\n",
    "    lambda X: model.predict(X),\n",
    "    X_test, y_test,) for model in (og_logres, logres))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# assert k, v pairs are equal in both dictionaries\n",
    "def assert_dict_eq(d1, d2, strict=True):\n",
    "    if strict: assert d1.keys() == d2.keys()\n",
    "    for k in d1.keys():\n",
    "        # 4x speedup if you ask me, but totally insignificant\n",
    "        if 'time' in k: print(f'{k:20}: \\t{d1[k]:.8f} != \\t{d2[k]:.8f}')\n",
    "        else: assert d1[k] == d2[k], f'{k}: {d1[k]} != {d2[k]}'\n",
    "\n",
    "assert_dict_eq(og_scores, scores, strict=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Cool, lets save the arrays\n",
    "model_path = 'logres (weights, bias)'\n",
    "with open(os.path.join(models_dir, model_path + '.pkl'), 'wb') as f:\n",
    "    pickle.dump((coef, intercept), f)\n",
    "\n",
    "print(coef, intercept)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### CodeBERTa (Code Context Only)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def tokenize_joint_sample(sample, max_suffix_tokens=128):\n",
    "    ''' For a single sample, tokenize prefix and suffix, separating by </s> sep token. \n",
    "        Set max_suffix_tokens to maximal amount of suffix to include, when it exists. '''\n",
    "\n",
    "    max_length = tokenizer.model_max_length # 512 \n",
    "\n",
    "    # figure out how many suffix tokens we have (128 max)\n",
    "    tokenizer.truncation_side = 'right'\n",
    "    suffix = tokenizer(sample['suffix'], padding='do_not_pad', truncation=True, return_tensors='pt',\n",
    "                          max_length = max_suffix_tokens + 1) # to accomodate removal of <s>\n",
    "\n",
    "    n_suffix_tokens = len(suffix['input_ids'][0]) - 1\n",
    "\n",
    "    tokenizer.truncation_side = 'left'\n",
    "    prefix = tokenizer(sample['prefix'], padding='do_not_pad', truncation=True, return_tensors='pt',\n",
    "                       max_length = max_length - n_suffix_tokens)\n",
    "\n",
    "    n_prefix_tokens = len(prefix['input_ids'][0])\n",
    "    tokenizer.truncation_side = 'right'\n",
    "    suffix = tokenizer(sample['suffix'], padding='max_length', truncation=True, return_tensors='pt',\n",
    "                       max_length = max_length - n_prefix_tokens + 1) # to accomodate removal of <s>\n",
    "    \n",
    "    suffix['input_ids'] = suffix['input_ids'][:, 1:]\n",
    "    suffix['attention_mask'] = suffix['attention_mask'][:, 1:]\n",
    "\n",
    "    sample.update({k: torch.cat((prefix[k], suffix[k]), dim=1) for k in prefix})\n",
    "    return sample\n",
    "\n",
    "\n",
    "def get_nontextual_features(query) -> list:\n",
    "    ''' Get the features that could otherwise not be extracted from the context alone '''\n",
    "    return [\n",
    "        1 if query.ide == 'jetbrains' else 0,           # 0\n",
    "        1 if query.ide == 'vsc' else 0,                 # 1\n",
    "        math.log(1 + query.time_since_last_completion), # 1\n",
    "        math.log(1 + query.get_document_length()),      # 2\n",
    "        math.log(1 + query.get_offset()),               # 3\n",
    "        query.get_offset_as_percentage(),               # 4\n",
    "        *query.get_document_language_vector(),          # 5-24\n",
    "    ]\n",
    "\n",
    "class MyPipeline(TextClassificationPipeline):\n",
    "    ''' oh yeah custom pipeline because of the custom tokenisation!\n",
    "        how convenient huggingface ill hug your face extra hard next time i see you '''\n",
    "    \n",
    "    def __init__(self, *args, incl_features=True, preprocess_fn=tokenize_joint_sample, **kwargs):\n",
    "        super().__init__(*args, **kwargs)\n",
    "        self.incl_features = incl_features\n",
    "        self.preprocess_fn = preprocess_fn\n",
    "\n",
    "    def _sanitize_parameters(self, **kwargs):\n",
    "        preprocess_kwargs = {} \n",
    "        if 'preprocess_fn' in kwargs: \n",
    "            preprocess_kwargs['preprocess_fn'] = kwargs.pop('preprocess_fn')\n",
    "        return preprocess_kwargs, {}, {} \n",
    "    \n",
    "    def preprocess(self, inputs, preprocess_fn=None):\n",
    "        inputs = {\n",
    "            'prefix': inputs.prefix, \n",
    "            'suffix': inputs.suffix, \n",
    "            'encoder_hidden_states': get_nontextual_features(inputs)\n",
    "        }\n",
    "        inputs = preprocess_fn(inputs) if self.preprocess_fn is None else self.preprocess_fn(inputs)\n",
    "        if 'prefix' in inputs: del inputs['prefix']\n",
    "        if 'suffix' in inputs: del inputs['suffix']\n",
    "        # given that pipeline is used in sequential eval, we neeed to add a batch dimension for the model to not throw a tantrum\n",
    "        if self.incl_features:\n",
    "            inputs['encoder_hidden_states'] = torch.tensor(inputs['encoder_hidden_states'], dtype=torch.float32).unsqueeze(0)\n",
    "        elif 'encoder_hidden_states' in inputs: \n",
    "            del inputs['encoder_hidden_states']\n",
    "        return inputs\n",
    "    \n",
    "    def _forward(self, model_inputs):\n",
    "        return self.model(**model_inputs)\n",
    "    \n",
    "    def postprocess(self, model_outputs):\n",
    "        return model_outputs.logits.argmax(-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_model(model_name):\n",
    "    model_dir = os.path.join(models_dir, model_name)\n",
    "    config = AutoConfig.from_pretrained(model_dir)\n",
    "\n",
    "    model = JonbertaForSequenceClassification(config)\n",
    "    if hasattr(config, 'add_head') and config.add_head: \n",
    "        add_features_to_model(model, config)\n",
    "\n",
    "    # ah yes huggingface is a 5 BILLION dollar company now\n",
    "    state_dict = {} \n",
    "    with safe_open(os.path.join(model_dir, 'model.safetensors'), framework='pt') as f: \n",
    "        for key in f.keys():\n",
    "            state_dict[key] = f.get_tensor(key)\n",
    "    new_layers = model.load_state_dict(state_dict, strict=False)\n",
    "\n",
    "    pprint(new_layers)\n",
    "    print(model)\n",
    "    return model "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from util import set_all_seeds\n",
    "import tqdm \n",
    "DEVICE = 'cuda' if torch.cuda.is_available() else\\\n",
    "    'mps' if torch.backends.mps.is_available() else 'cpu'\n",
    "\n",
    "def get_score(model_name):\n",
    "    \n",
    "    set_all_seeds(42)\n",
    "    model = get_model(model_name)\n",
    "\n",
    "    pipe = MyPipeline(\n",
    "        device=DEVICE, task='text-classification',\n",
    "        model=model, incl_features=True # for telemetry \n",
    "    )\n",
    "\n",
    "    return common_evaluation(\n",
    "        lambda X_queries: X_queries,\n",
    "        lambda X: torch.cat([pipe(x) for x in tqdm.tqdm(X, total=len(X_test), desc='sequential test')]),\n",
    "        X_test, y_test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "scores = {}\n",
    "for model_path in (joint_head_name, joint_attn_name, context_name):\n",
    "\n",
    "    score = get_score(model_path)\n",
    "    scores[model_path] = score \n",
    "\n",
    "    print(f'\\033[1m{model_path}\\033[0m')\n",
    "    pprint(score)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "new_scores = {}\n",
    "for model_path in (joint_head_name, joint_attn_name, context_name):\n",
    "\n",
    "    score = get_score(model_path)\n",
    "    new_scores[model_path] = score \n",
    "\n",
    "    print(f'\\033[1m{model_path}\\033[0m')\n",
    "    pprint(score)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# assert_dict_eq(og_joint_head_score, scores[joint_head_name], strict=False)\n",
    "# assert_dict_eq(og_joint_attn_score, scores[joint_attn_name], strict=False)\n",
    "# assert_dict_eq(og_context_score, scores[context_name], strict=False)\n",
    "# pprint(scores[joint_head_name])\n",
    "# pprint(og_joint_head_score)\n",
    "\n",
    "# pprint(scores[joint_attn_name])\n",
    "# pprint(og_joint_attn_score)\n",
    "\n",
    "# pprint(scores[context_name])\n",
    "# pprint(og_context_score)\n",
    "\n",
    "for model in (joint_head_name, joint_attn_name, context_name):\n",
    "    print(f'\\033[1m{model}\\033[0m')\n",
    "    for (k_1, v_1), (k_2, v_2) in zip(scores[model].items(), new_scores[model].items()):\n",
    "        print(f'{k_1:20}: \\t{v_1:.8f} != \\t{v_2:.8f}')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import gc \n",
    "for obj in gc.get_objects():\n",
    "    try:\n",
    "        if torch.is_tensor(obj) or (hasattr(obj, 'data') and torch.is_tensor(obj.data)):\n",
    "            print(type(obj), obj.size())\n",
    "    except:\n",
    "        pass\n",
    "\n",
    "# for obj in gc.get_objects():\n",
    "#     if torch.is_tensor(obj) or (hasattr(obj, 'data') and torch.is_tensor(obj.data)):\n",
    "#         print(reduce(op.mul, obj.size()) if len(obj.size()) > 0 else 0, type(obj), obj.size())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Sanity Check\n",
    "Let's make sure the implementation of deployed filters is correct. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "os.chdir('..')\n",
    "from sanity_check import Filter, filters, set_all_seeds\n",
    "filters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tqdm \n",
    "X, y = query_data['unbalanced']['X_test'], query_data['unbalanced']['y_test']\n",
    "# X, y = X[:3000], y[:3000]\n",
    "\n",
    "for f_type, f in filters.items():\n",
    "\n",
    "    set_all_seeds()\n",
    "    score = common_evaluation(\n",
    "        lambda X_queries: X_queries,\n",
    "        lambda X: [int(f(x)) for x in tqdm.tqdm(X, total=len(X), desc='sequential test')],\n",
    "        X, y)\n",
    "    \n",
    "    print(f_type)\n",
    "    pprint(score)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tqdm \n",
    "X, y = query_data['unbalanced']['X_test'], query_data['unbalanced']['y_test']\n",
    "\n",
    "for f_type, f in filters.items():\n",
    "\n",
    "    set_all_seeds()\n",
    "    score = common_evaluation(\n",
    "        lambda X_queries: X_queries,\n",
    "        lambda X: [int(f(x)) for x in tqdm.tqdm(X, total=len(X), desc='sequential test')],\n",
    "        X, y)\n",
    "    \n",
    "    print(f_type)\n",
    "    print(score)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tqdm \n",
    "X, y = query_data['unbalanced']['X_test'], query_data['unbalanced']['y_test']\n",
    "\n",
    "for f_type, f in filters.items():\n",
    "\n",
    "    set_all_seeds()\n",
    "    score = common_evaluation(\n",
    "        lambda X_queries: X_queries,\n",
    "        lambda X: [int(f(x)) for x in tqdm.tqdm(X, total=len(X), desc='sequential test')],\n",
    "        X, y)\n",
    "    \n",
    "    print(f_type)\n",
    "    print(score)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "data",
   "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.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}