evaluating_generalization_of_truth_directions.ipynb

{
 "cells": [
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-12-03T22:50:04.712070Z",
     "start_time": "2024-12-03T22:49:58.656739Z"
    }
   },
   "cell_type": "code",
   "source": [
    "import json\n",
    "import os\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import torch\n",
    "\n",
    "from evaluation_utils import MetricsForDatasetProbes, evaluate_classifier_performance\n",
    "from phi_3_5_constants import directions_reconstruction_losses_path, dsets_index_path, dsets_folder, \\\n",
    "    finalized_activations_dir, hidden_state_size, four_way_topics_index_path, misc_datasets_index_path, \\\n",
    "    test_classification_metrics_path, analysis_results_folder\n",
    "from phi_3_5_probe import ProbesForDataset, load_probes_for_dset"
   ],
   "id": "48018f715646e72a",
   "outputs": [],
   "execution_count": 1
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-12-03T22:50:04.719071Z",
     "start_time": "2024-12-03T22:50:04.713071Z"
    }
   },
   "cell_type": "code",
   "source": [
    "dsets_index_df = pd.read_csv(dsets_index_path, index_col=\"Idx\")\n",
    "num_dsets = dsets_index_df.shape[0]"
   ],
   "id": "922cfdf6b3489014",
   "outputs": [],
   "execution_count": 2
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-12-03T22:50:04.723964Z",
     "start_time": "2024-12-03T22:50:04.720071Z"
    }
   },
   "cell_type": "code",
   "source": [
    "all_dsets_activations: list[torch.Tensor] = [torch.zeros(1) for _ in range(num_dsets)]\n",
    "all_dsets_labels: list[torch.Tensor] = [torch.zeros(1) for _ in range(num_dsets)]"
   ],
   "id": "f7f2d9bb78f673ec",
   "outputs": [],
   "execution_count": 3
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-12-03T22:50:04.728963Z",
     "start_time": "2024-12-03T22:50:04.723964Z"
    }
   },
   "cell_type": "code",
   "source": [
    "# top-level key is the name of a topic that has pos/neg/conj/disj variants, second level key is one of those variant names\n",
    "with four_way_topics_index_path.open(\"r\") as f:\n",
    "    dset_idxs_for_4way_topics: dict[str, dict[str, int]] = json.load(f)\n",
    "with misc_datasets_index_path.open(\"r\") as f:\n",
    "    idxs_for_other_dsets: dict[str, int] = json.load(f)"
   ],
   "id": "df51fd44529bb046",
   "outputs": [],
   "execution_count": 4
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-12-03T22:50:04.736920Z",
     "start_time": "2024-12-03T22:50:04.729963Z"
    }
   },
   "cell_type": "code",
   "source": [
    "directions_reconstruction_losses = pd.read_csv(directions_reconstruction_losses_path)\n",
    "dir_probe_scenarios_strs = directions_reconstruction_losses['scenario_identifier'].to_list()\n",
    "dir_probe_scenarios_specs = [tuple(map(int, scenario_str.split(' '))) for scenario_str in dir_probe_scenarios_strs]"
   ],
   "id": "f8ebeffd99a2c28b",
   "outputs": [],
   "execution_count": 5
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-12-03T22:50:04.741101Z",
     "start_time": "2024-12-03T22:50:04.737921Z"
    }
   },
   "cell_type": "code",
   "source": [
    "probes: dict[tuple[int,...], ProbesForDataset] = {}\n",
    "# first level key identifies the scenario (of one or more datasets) which the probe was trained on;\n",
    "# second level key identifies the previously-unseen dataset which the probe is tested on\n",
    "probes_metrics_on_test_dsets: dict[tuple[int,...], dict[int, MetricsForDatasetProbes]] = {}"
   ],
   "id": "a25dd6c67bda7664",
   "outputs": [],
   "execution_count": 6
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-12-03T22:50:06.171839Z",
     "start_time": "2024-12-03T22:50:04.742100Z"
    }
   },
   "cell_type": "code",
   "source": [
    "for dset_idx, dset_dtls in dsets_index_df.iterrows():\n",
    "    categ_nm = dset_dtls[\"Categ_Folder\"]\n",
    "    dset_file_nm = dset_dtls[\"Dataset_File\"]\n",
    "    dset_nm = os.path.splitext(dset_file_nm)[0]\n",
    "    \n",
    "    dataset = pd.read_csv(dsets_folder / categ_nm / dset_file_nm)\n",
    "    dset_size = dataset.shape[0]\n",
    "    dset_labels = torch.from_numpy(dataset['label'].to_numpy().astype(np.float32)[:, np.newaxis])\n",
    "    all_dsets_labels[dset_idx] = dset_labels\n",
    "    \n",
    "    activs_path = finalized_activations_dir / categ_nm / (dset_nm + \".pt\")\n",
    "    relevant_activations = torch.load(activs_path, weights_only=True)\n",
    "    assert relevant_activations.shape == (2, dset_size, hidden_state_size)\n",
    "    all_dsets_activations[dset_idx] = relevant_activations\n",
    "    if not dset_dtls['in_german']:\n",
    "        # only using the german datasets for testing of generalization, so didn't learn directions for them or train probes with them\n",
    "        curr_dset_probes = load_probes_for_dset(categ_nm, dset_nm, hidden_state_size)\n",
    "        probes[(dset_idx,)] = curr_dset_probes"
   ],
   "id": "ad8013046a443472",
   "outputs": [],
   "execution_count": 7
  },
  {
   "cell_type": "code",
   "id": "initial_id",
   "metadata": {
    "collapsed": true,
    "ExecuteTime": {
     "end_time": "2024-12-03T22:50:06.321439Z",
     "start_time": "2024-12-03T22:50:06.172358Z"
    }
   },
   "source": [
    "for topic_nm, variants_dset_idxs in dset_idxs_for_4way_topics.items():\n",
    "    affirm_idx = variants_dset_idxs[\"affirm\"]\n",
    "    neg_idx = variants_dset_idxs[\"neg\"]\n",
    "    conj_idx = variants_dset_idxs[\"conj\"]\n",
    "    disj_idx = variants_dset_idxs[\"disj\"]\n",
    "    \n",
    "    affirm_neg_key = tuple(sorted([affirm_idx, neg_idx]))\n",
    "    affirm_neg_probes = load_probes_for_dset(topic_nm, \"affirm_neg\", hidden_state_size)\n",
    "    probes[affirm_neg_key] = affirm_neg_probes\n",
    "    \n",
    "    affirm_disj_key = tuple(sorted([affirm_idx, disj_idx]))\n",
    "    affirm_disj_probes = load_probes_for_dset(topic_nm, \"affirm_disj\", hidden_state_size)\n",
    "    probes[affirm_disj_key] = affirm_disj_probes\n",
    "    \n",
    "    neg_conj_key = tuple(sorted([neg_idx, conj_idx]))\n",
    "    neg_conj_probes = load_probes_for_dset(topic_nm, \"neg_conj\", hidden_state_size)\n",
    "    probes[neg_conj_key] = neg_conj_probes\n",
    "    \n",
    "    affirm_neg_conj_disj_key = tuple(sorted([affirm_idx, neg_idx, conj_idx, disj_idx]))\n",
    "    affirm_neg_conj_disj_probes = load_probes_for_dset(topic_nm, \"affirm_neg_conj_disj\", hidden_state_size)\n",
    "    probes[affirm_neg_conj_disj_key] = affirm_neg_conj_disj_probes"
   ],
   "outputs": [],
   "execution_count": 8
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-12-03T22:50:06.330496Z",
     "start_time": "2024-12-03T22:50:06.322440Z"
    }
   },
   "cell_type": "code",
   "source": [
    "unambig_lie_idx = idxs_for_other_dsets[\"unambiguous_lie\"]\n",
    "unambig_truth_idx = idxs_for_other_dsets[\"unambiguous_truthful_reply\"]\n",
    "ambig_truth_idx = idxs_for_other_dsets[\"ambiguous_truthful_reply\"]\n",
    "real_world_multi_dset_idxs = [unambig_lie_idx, unambig_truth_idx, ambig_truth_idx]\n",
    "real_world_multi_dset_scenario_key = tuple(sorted(real_world_multi_dset_idxs))\n",
    "\n",
    "real_world_multi_dset_probes = load_probes_for_dset(\"real_world_scenarios\", \"unambig_lie_unambig_truth_ambig_truth\", hidden_state_size)\n",
    "probes[real_world_multi_dset_scenario_key] = real_world_multi_dset_probes"
   ],
   "id": "818808759e5ff528",
   "outputs": [],
   "execution_count": 9
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-12-03T22:50:06.334974Z",
     "start_time": "2024-12-03T22:50:06.331473Z"
    }
   },
   "cell_type": "code",
   "source": [
    "multi_topics_dset_idxs = []\n",
    "various_categories_dset_idxs = []\n",
    "\n",
    "for topic_nm, variants_dset_idxs in dset_idxs_for_4way_topics.items():\n",
    "    if topic_nm not in (\"animal_class\", \"facts\", \"inventors\"):\n",
    "        continue\n",
    "    affirm_idx = variants_dset_idxs[\"affirm\"]\n",
    "    neg_idx = variants_dset_idxs[\"neg\"]\n",
    "    conj_idx = variants_dset_idxs[\"conj\"]\n",
    "    multi_topics_dset_idxs.extend([affirm_idx, neg_idx, conj_idx])\n",
    "    various_categories_dset_idxs.extend([affirm_idx, neg_idx])"
   ],
   "id": "553307dbe55dc490",
   "outputs": [],
   "execution_count": 10
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-12-03T22:50:06.344488Z",
     "start_time": "2024-12-03T22:50:06.335971Z"
    }
   },
   "cell_type": "code",
   "source": [
    "multi_topics_affirm_neg_conj_key = tuple(sorted(multi_topics_dset_idxs))\n",
    "multi_topics_affirm_neg_conj_probes = load_probes_for_dset(\"\", \"multi_topics_affirm_neg_conj\", hidden_state_size)\n",
    "probes[multi_topics_affirm_neg_conj_key] = multi_topics_affirm_neg_conj_probes"
   ],
   "id": "2524d5a3dd75522f",
   "outputs": [],
   "execution_count": 11
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-12-03T22:50:06.353167Z",
     "start_time": "2024-12-03T22:50:06.345488Z"
    }
   },
   "cell_type": "code",
   "source": [
    "smaller_than_idx = idxs_for_other_dsets[\"smaller_than\"]\n",
    "common_claim_idx = idxs_for_other_dsets[\"common_claim_true_false\"]\n",
    "various_categories_dset_idxs.extend(real_world_multi_dset_idxs)\n",
    "various_categories_dset_idxs.extend([smaller_than_idx, common_claim_idx])\n",
    "various_categories_scenario_key = tuple(sorted(various_categories_dset_idxs))\n",
    "various_categories_probes = load_probes_for_dset(\"\", \"various_categories\", hidden_state_size)\n",
    "probes[various_categories_scenario_key] = various_categories_probes"
   ],
   "id": "abdca207e2183117",
   "outputs": [],
   "execution_count": 12
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-12-03T22:50:06.356485Z",
     "start_time": "2024-12-03T22:50:06.354166Z"
    }
   },
   "cell_type": "code",
   "source": [
    "for scenario_key in probes.keys():\n",
    "    probes_metrics_on_test_dsets[scenario_key] = {}"
   ],
   "id": "8d37b3e1ca3a9dea",
   "outputs": [],
   "execution_count": 13
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-12-03T22:52:19.174817Z",
     "start_time": "2024-12-03T22:50:06.357507Z"
    }
   },
   "cell_type": "code",
   "source": [
    "for scenario_key, scenario_probes in probes.items():\n",
    "    for test_dset_idx, test_dset_dtls in dsets_index_df.iterrows():\n",
    "        if test_dset_idx in scenario_key:\n",
    "            continue\n",
    "\n",
    "        scenario_metrics_on_curr_test_dset = evaluate_classifier_performance(\n",
    "            scenario_probes, all_dsets_activations[test_dset_idx], all_dsets_labels[test_dset_idx])\n",
    "        probes_metrics_on_test_dsets[scenario_key][test_dset_idx] = scenario_metrics_on_curr_test_dset\n"
   ],
   "id": "87d96e14b0006459",
   "outputs": [],
   "execution_count": 14
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-12-03T22:52:19.179365Z",
     "start_time": "2024-12-03T22:52:19.175789Z"
    }
   },
   "cell_type": "code",
   "source": "analysis_results_folder.mkdir(exist_ok=True)",
   "id": "479feb8d6f05647",
   "outputs": [],
   "execution_count": 15
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-12-03T22:52:19.513406Z",
     "start_time": "2024-12-03T22:52:19.180364Z"
    }
   },
   "cell_type": "code",
   "source": [
    "serialized_test_metrics = {\n",
    "    (\" \".join(map(str, scenario_id))): {\n",
    "        str(test_dset_idx): test_dset_metrics.to_dict()\n",
    "        for test_dset_idx, test_dset_metrics in test_dsets_metrics.items()\n",
    "    } for scenario_id, test_dsets_metrics in probes_metrics_on_test_dsets.items()\n",
    "}\n",
    "with test_classification_metrics_path.open(\"w\") as f:\n",
    "    json.dump(serialized_test_metrics, f, indent=2)"
   ],
   "id": "c528297bfd145e57",
   "outputs": [],
   "execution_count": 16
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-12-03T22:52:19.516698Z",
     "start_time": "2024-12-03T22:52:19.514411Z"
    }
   },
   "cell_type": "code",
   "source": "",
   "id": "e96c024fe75f6ea3",
   "outputs": [],
   "execution_count": 16
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-12-03T22:52:19.520254Z",
     "start_time": "2024-12-03T22:52:19.517698Z"
    }
   },
   "cell_type": "code",
   "source": "",
   "id": "24825f172c7775a5",
   "outputs": [],
   "execution_count": 16
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "",
   "id": "6a58231145f6d3da"
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "",
   "id": "4ab4e730ea356426"
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}