analysis_ft.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "import numpy as np\n",
    "import gzip\n",
    "import pickle\n",
    "\n",
    "from IPython.display import display\n",
    "\n",
    "plt.style.use(\".mplstyle\")\n",
    "pd.set_option(\"display.max_rows\", 25)\n",
    "pd.set_option(\"display.max_colwidth\", 0)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "BASE_DIR = \"LAMA/data/\"\n",
    "METRICS_DIR = os.path.join(BASE_DIR, \"metrics/\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "os.getcwd()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with gzip.open(\n",
    "    METRICS_DIR + \"reranker/sweep_re/seed_0/learned/eval_detailed.pickle\", \"rb\"\n",
    ") as handle:\n",
    "    pt_learned = pickle.load(handle)\n",
    "\n",
    "pt_examples = [sample[\"example\"] for sample in pt_learned[\"samples\"]]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with gzip.open(\n",
    "    \"LAMA/data/metrics/reranker/sweep_re_ft_fl/seed_0/learned/eval_detailed.pickle\"\n",
    ") as handle:\n",
    "    ft_learned = pickle.load(handle)\n",
    "ft_examples = [sample[\"example\"] for sample in ft_learned[\"samples\"]]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def trim(target):\n",
    "    return target.replace(\"<extra_id_0> \", \"\").lower()\n",
    "\n",
    "\n",
    "def find_target_distractors(distractors, target):\n",
    "    return [\n",
    "        distractor\n",
    "        for distractor in distractors\n",
    "        if trim(distractor[\"targets_pretokenized\"]) == trim(target)\n",
    "    ]\n",
    "\n",
    "\n",
    "def subset_statistics(subset):\n",
    "    num_fact_abstracts = []\n",
    "    num_target_distractors = []\n",
    "    num_distractors = []\n",
    "    num_target_distractors_but_proponents = []\n",
    "    for sample in subset[\"samples\"]:\n",
    "        example = sample[\"example\"]\n",
    "        target = example[\"targets_pretokenized\"]\n",
    "        fact = (\n",
    "            example[\"predicate_id\"]\n",
    "            + \",\"\n",
    "            + example[\"obj_uri\"]\n",
    "            + \",\"\n",
    "            + example[\"sub_uri\"]\n",
    "        )\n",
    "        distractors = sample[\"distractors\"]\n",
    "        num_distractors.append(len(distractors))\n",
    "        fact_abstracts = sample[\"fact_abstracts\"]\n",
    "        num_fact_abstracts.append(len(fact_abstracts))\n",
    "        target_distractors = find_target_distractors(distractors, target)\n",
    "        num_target_distractors.append(len(target_distractors))\n",
    "        fact_abstract_facts = [abstract[\"facts\"] for abstract in fact_abstracts]\n",
    "        target_distractor_facts = [\n",
    "            abstract\n",
    "            for abstract in target_distractors\n",
    "            if fact in abstract[\"facts\"]\n",
    "        ]\n",
    "        num_target_distractors_but_proponents.append(\n",
    "            len(target_distractor_facts)\n",
    "        )\n",
    "\n",
    "    return list(\n",
    "        map(\n",
    "            np.mean,\n",
    "            (\n",
    "                num_distractors,\n",
    "                num_fact_abstracts,\n",
    "                num_target_distractors,\n",
    "                num_target_distractors_but_proponents,\n",
    "            ),\n",
    "        )\n",
    "    )\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "subset_statistics(pt_learned)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "subset_statistics(ft_learned)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def metrics_to_df(results, metrics=[\"precision\", \"recall\", \"mrr\"]):\n",
    "    data = []\n",
    "\n",
    "    for method in (\"bm25plus\", \"random\"):\n",
    "        for eval_type in (\"collapse\", \"full\"):\n",
    "            if eval_type != \"collapse\":\n",
    "                current_metrics = metrics + [\n",
    "                    \"mrr_compare_fn_subject\",\n",
    "                    \"mrr_compare_fn_object\",\n",
    "                    \"mrr_compare_fn_relation\",\n",
    "                ]\n",
    "            else:\n",
    "                current_metrics = metrics\n",
    "            for metric in current_metrics:\n",
    "                if metric in results[\"evals\"][method][eval_type]:\n",
    "                    metric_result = results[\"evals\"][method][eval_type][metric]\n",
    "                    for normalized in (\"cosine\", \"dot\"):\n",
    "                        for method_type in (\"local\", \"global\"):\n",
    "                            try:\n",
    "                                for k, score in metric_result.items():\n",
    "                                    data.append(\n",
    "                                        (\n",
    "                                            method,\n",
    "                                            method_type,\n",
    "                                            normalized,\n",
    "                                            eval_type,\n",
    "                                            metric,\n",
    "                                            str(k),\n",
    "                                            score,\n",
    "                                        )\n",
    "                                    )\n",
    "                            except:\n",
    "                                data.append(\n",
    "                                    (\n",
    "                                        method,\n",
    "                                        method_type,\n",
    "                                        normalized,\n",
    "                                        eval_type,\n",
    "                                        metric,\n",
    "                                        \"1\",\n",
    "                                        metric_result,\n",
    "                                    )\n",
    "                                )\n",
    "\n",
    "    for method_type in (\"local\", \"global\"):\n",
    "        for eval_type in (\"collapse\", \"full\"):\n",
    "            for normalized in (\"cosine\", \"dot\"):\n",
    "                if method_type == \"global\" and normalized == \"dot\":\n",
    "                    continue\n",
    "                for method, method_results in results[\"evals\"][method_type][\n",
    "                    normalized\n",
    "                ][eval_type].items():\n",
    "                    if eval_type != \"collapse\":\n",
    "                        current_metrics = metrics + [\n",
    "                            \"mrr_compare_fn_subject\",\n",
    "                            \"mrr_compare_fn_object\",\n",
    "                            \"mrr_compare_fn_relation\",\n",
    "                        ]\n",
    "                    else:\n",
    "                        current_metrics = metrics\n",
    "                    for metric in current_metrics:\n",
    "                        if metric in method_results:\n",
    "                            metric_result = method_results[metric]\n",
    "                            try:\n",
    "                                for k, score in metric_result.items():\n",
    "                                    data.append(\n",
    "                                        (\n",
    "                                            method,\n",
    "                                            method_type,\n",
    "                                            normalized,\n",
    "                                            eval_type,\n",
    "                                            metric,\n",
    "                                            str(k),\n",
    "                                            score,\n",
    "                                        )\n",
    "                                    )\n",
    "                            except:\n",
    "                                data.append(\n",
    "                                    (\n",
    "                                        method,\n",
    "                                        method_type,\n",
    "                                        normalized,\n",
    "                                        eval_type,\n",
    "                                        metric,\n",
    "                                        \"1\",\n",
    "                                        metric_result,\n",
    "                                    )\n",
    "                                )\n",
    "\n",
    "    df = pd.DataFrame(\n",
    "        data,\n",
    "        columns=[\n",
    "            \"layers\",\n",
    "            \"norm_type\",\n",
    "            \"normalization\",\n",
    "            \"eval\",\n",
    "            \"metrics\",\n",
    "            \"k\",\n",
    "            \"score\",\n",
    "        ],\n",
    "    )\n",
    "\n",
    "    df[\"layer_type\"] = \"Embed\"\n",
    "    df.loc[df[\"layers\"].str.contains(\"gradients\"), \"layer_type\"] = \"TracIn\"\n",
    "    df.loc[\n",
    "        (df[\"layers\"].str.contains(\"gradients\"))\n",
    "        & (df[\"layers\"].str.contains(\"activations\")),\n",
    "        \"layer_type\",\n",
    "    ] = \"TracIn+Embed\"\n",
    "    df.loc[\n",
    "        (df[\"layers\"] == \"random\") | (df[\"layers\"] == \"bm25plus\"), \"layer_type\"\n",
    "    ] = \"baselines\"\n",
    "\n",
    "    df = df.replace(\n",
    "        {\n",
    "            \"gradients\": \"G\",\n",
    "            \"activations\": \"A\",\n",
    "            \"block.\": \"\",\n",
    "            \"encoder\": \"E\",\n",
    "            \"decoder\": \"D\",\n",
    "            \"shared\": \"emb\",\n",
    "            \"random\": \"Target-Picker\",\n",
    "        },\n",
    "        regex=True,\n",
    "    )\n",
    "\n",
    "    df[\"layers\"] = (\n",
    "        df[\"layers\"]\n",
    "        .replace(\n",
    "            {f\"G.E.{i},G.D.{i}\": f\"G.E.{i+1},G.D.{i+1}\" for i in range(12)},\n",
    "            regex=False,\n",
    "        )\n",
    "        .replace({f\"G.E.{i}\": f\"G.E.{i+1}\" for i in range(12)}, regex=False)\n",
    "        .replace(\n",
    "            {\n",
    "                f\"G.emb,G.E.{i},G.D.{i}\": f\"G.emb,G.E.{i+1},G.D.{i+1}\"\n",
    "                for i in range(12)\n",
    "            },\n",
    "            regex=False,\n",
    "        )\n",
    "        .replace(\n",
    "            {f\"G.emb,G.E.{i}\": f\"G.emb,G.E.{i+1}\" for i in range(12)},\n",
    "            regex=False,\n",
    "        )\n",
    "        .str.replace(\"G.emb\", \"G.0\", regex=False)\n",
    "        .str.replace(\"bm25plus\", \"BM25+\", regex=False)\n",
    "        .str.replace(\"Target-Picker\", \"Random-Target\", regex=False)\n",
    "    )\n",
    "\n",
    "    return df\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_with_filter(\n",
    "    df,\n",
    "    filter=lambda x: x,\n",
    "    title=\"Title\",\n",
    "    folder=\"plots/\",\n",
    "    ylabel=\"\",\n",
    "    save=False,\n",
    "):\n",
    "    plt.figure(figsize=(16, 6))\n",
    "\n",
    "    ax = sns.barplot(\n",
    "        data=filter(df),\n",
    "        x=\"layers\",\n",
    "        y=\"score\",\n",
    "        hue=\"layer_type\",\n",
    "        estimator=np.mean,\n",
    "        ci=\"sd\",\n",
    "    )\n",
    "\n",
    "    plt.xticks(rotation=90)\n",
    "    plt.title(title)\n",
    "    plt.xlabel(\"Layer Selection\")\n",
    "    plt.ylabel(ylabel)\n",
    "    plt.legend(title=\"Method Type\")\n",
    "\n",
    "    if not save:\n",
    "        plt.show()\n",
    "    else:\n",
    "        plt.savefig(folder + title + \"_plot.png\")\n",
    "        plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def visualize_one_experiment(\n",
    "    paths,\n",
    "    suffix=\"\",\n",
    "    folder=\"plots/\",\n",
    "    save=False,\n",
    "    visualize=False,\n",
    "    k=\"3\",\n",
    "    norm_type=\"local\",\n",
    "    normalization=\"cosine\",\n",
    "):\n",
    "    dfs = []\n",
    "    for (i, path) in enumerate(paths):\n",
    "\n",
    "        with gzip.open(METRICS_DIR + path) as f:\n",
    "            reranker_metrics = pickle.load(f)\n",
    "\n",
    "        df = metrics_to_df(reranker_metrics)\n",
    "        df[\"seed\"] = i\n",
    "        dfs.append(df)\n",
    "\n",
    "    df = pd.concat(dfs, ignore_index=True)\n",
    "\n",
    "    if visualize:\n",
    "\n",
    "        # scores = df.groupby(['normalization', 'eval', 'layers', 'metrics', 'k']).agg({'score': ['mean', 'std']}, as_index=False)\n",
    "        # Layers that we don't want for visualizations\n",
    "        ddf = df[~df[\"layers\"].str.contains(\"A.E.0,A.D.0,\")]\n",
    "\n",
    "        for method in (\"full\", \"collapse\"):\n",
    "            plot_with_filter(\n",
    "                ddf,\n",
    "                filter=lambda x: x[\n",
    "                    (x[\"metrics\"] == \"precision\")\n",
    "                    & (x[\"k\"] == \"3\")\n",
    "                    & (x[\"eval\"] == method)\n",
    "                    & (x[\"norm_type\"] == norm_type)\n",
    "                    & (x[\"normalization\"] == normalization)\n",
    "                ],\n",
    "                title=f\"precision@3 ({method} + {suffix})\",\n",
    "                folder=folder,\n",
    "                ylabel=\"precision@3\",\n",
    "                save=save,\n",
    "            )\n",
    "\n",
    "            plot_with_filter(\n",
    "                ddf,\n",
    "                filter=lambda x: x[\n",
    "                    (x[\"metrics\"] == \"recall\")\n",
    "                    & (x[\"k\"] == \"10\")\n",
    "                    & (x[\"eval\"] == method)\n",
    "                    & (x[\"norm_type\"] == norm_type)\n",
    "                    & (x[\"normalization\"] == normalization)\n",
    "                ],\n",
    "                title=f\"recall@10 ({method} + {suffix})\",\n",
    "                folder=folder,\n",
    "                ylabel=\"recall@10\",\n",
    "                save=save,\n",
    "            )\n",
    "\n",
    "            plot_with_filter(\n",
    "                ddf,\n",
    "                filter=lambda x: x[\n",
    "                    (x[\"metrics\"] == \"mrr\")\n",
    "                    & (x[\"k\"] == \"1\")\n",
    "                    & (x[\"eval\"] == method)\n",
    "                    & (x[\"norm_type\"] == norm_type)\n",
    "                    & (x[\"normalization\"] == normalization)\n",
    "                ],\n",
    "                title=f\"mrr ({method} + {suffix})\",\n",
    "                folder=folder,\n",
    "                ylabel=\"mrr\",\n",
    "                save=save,\n",
    "            )\n",
    "    return df\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_max(\n",
    "    df: pd.DataFrame,\n",
    "    prefix,\n",
    "    no_prefix=None,\n",
    "    subset=\"learned\",\n",
    "    metric=\"mrr\",\n",
    "    eval_type=\"collapse\",\n",
    "    normalization=None,\n",
    "    k=\"1\",\n",
    "):\n",
    "\n",
    "    if metric.startswith(\"mrr_\"):\n",
    "        eval_type = \"full\"\n",
    "        print(\"overwritten eval_type to full\")\n",
    "\n",
    "    if normalization is not None:\n",
    "        df = df[df.normalization == normalization]\n",
    "    # print(df.metrics.unique())\n",
    "    df2 = df[\n",
    "        (df[\"layers\"].str.startswith(prefix))\n",
    "        & (df[\"subset\"] == subset)\n",
    "        & (df[\"metrics\"] == metric)\n",
    "        & (df[\"k\"] == str(k))\n",
    "        & (df[\"eval\"] == eval_type)\n",
    "    ]\n",
    "\n",
    "    if no_prefix:\n",
    "        df2 = df2[~df2[\"layers\"].str.contains(no_prefix)]\n",
    "\n",
    "    nlargest = df2.nlargest(10, (\"score\", \"mean\"))\n",
    "\n",
    "    largest = nlargest.iloc[0]\n",
    "    dflargest = df[\n",
    "        (df[\"layers\"] == largest[\"layers\"][0])\n",
    "        & (df[\"norm_type\"] == largest[\"norm_type\"][0])\n",
    "        & (df[\"eos\"] == largest[\"eos\"][0])\n",
    "        & (df[\"accum\"] == largest[\"accum\"][0])\n",
    "        & (df[\"k\"] == str(k))\n",
    "        & (df[\"subset\"] == subset)\n",
    "    ]  # &\n",
    "    #                   (df['eval'] == eval_type)]\n",
    "\n",
    "    dflargest = dflargest.set_index(\"metrics\").loc[[metric]]\n",
    "    # dflargest = dflargest.set_index('metrics').loc[[metric]]\n",
    "    dflargest[\"score_text\"] = (\n",
    "        (100 * dflargest[\"score\"][\"mean\"]).apply(lambda x: f\"{x:.2f}\")\n",
    "        + \"\\stderr{\"\n",
    "        + (100 * dflargest[\"score\"][\"std\"]).apply(lambda x: f\"{x:.2f}\")\n",
    "        + \"}\"\n",
    "    )\n",
    "    dflargest = dflargest.transpose()\n",
    "    # dflargest = dflargest.loc[['score_text', 'k', 'eval']]\n",
    "    return nlargest, dflargest\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def print_best_results(\n",
    "    path=\"reranker/sweep_v2/\",\n",
    "    subset=\"learned\",\n",
    "    accum=\"accum\",\n",
    "    eos=\"eos\",\n",
    "    seed_range=1,\n",
    "    metric=\"mrr\",\n",
    "    normalization=None,\n",
    "):\n",
    "    dfs = {}\n",
    "    suffix = f\"{eos}+{subset}+{accum}\"\n",
    "    try:\n",
    "        df = visualize_one_experiment(\n",
    "            paths=[\n",
    "                f\"{path}/seed_{i}/{subset}/{eos}_{accum}/results_detailed.pickle\"\n",
    "                for i in range(seed_range)\n",
    "            ],\n",
    "            suffix=suffix,\n",
    "            save=False,\n",
    "            visualize=False,\n",
    "        )\n",
    "        df[\"eos\"] = eos\n",
    "        df[\"subset\"] = subset\n",
    "        df[\"accum\"] = accum\n",
    "        dfs[suffix] = df\n",
    "    except FileNotFoundError:\n",
    "        print(\n",
    "            \"Couldn't find: \"\n",
    "            f\"{path}/seed_0/{subset}/{eos}_{accum}/results_detailed.pickle\"\n",
    "        )\n",
    "\n",
    "    dfmerged = pd.concat(list(dfs.values()), ignore_index=True)\n",
    "    scores = dfmerged.groupby(\n",
    "        [\n",
    "            column\n",
    "            for column in dfmerged.columns\n",
    "            if column != \"seed\" and column != \"score\"\n",
    "        ],\n",
    "        axis=0,\n",
    "    ).agg({\"score\": [np.mean, np.std]}, as_index=False)\n",
    "    scores = scores.reset_index()\n",
    "    # print(scores.layers.unique())\n",
    "\n",
    "    scores[\"path\"] = path\n",
    "\n",
    "    if metric == \"recall\":\n",
    "        k = \"10\"\n",
    "    elif metric == \"precision\":\n",
    "        k = \"10\"\n",
    "    else:\n",
    "        k = \"1\"\n",
    "\n",
    "    # print(\"Activations\")\n",
    "    largest, dflargest1 = get_max(\n",
    "        scores,\n",
    "        \"A\",\n",
    "        no_prefix=\"G\",\n",
    "        metric=metric,\n",
    "        k=k,\n",
    "        normalization=normalization,\n",
    "    )\n",
    "\n",
    "    # print(\"Tracin\")\n",
    "    largest, dflargest2 = get_max(\n",
    "        scores, \"G\", metric=metric, k=k, normalization=normalization\n",
    "    )\n",
    "    # print(\"Tracin + Activations\")\n",
    "    largest, dflargest3 = get_max(\n",
    "        scores[scores.layers.str.contains(\"G\")],\n",
    "        \"A\",\n",
    "        metric=metric,\n",
    "        k=k,\n",
    "        normalization=normalization,\n",
    "    )\n",
    "    # print(\"BM25\")\n",
    "    largest, dflargest4 = get_max(\n",
    "        scores, \"BM25+\", metric=metric, k=k, normalization=normalization\n",
    "    )\n",
    "    # print(\"BM25\")\n",
    "    largest, dflargest5 = get_max(\n",
    "        scores, \"Random-Target\", metric=metric, k=k, normalization=normalization\n",
    "    )\n",
    "    # dflargest4 = dflargest4.transpose().reset_index().loc[0, :]\n",
    "    dflargest = [dflargest1, dflargest2, dflargest3, dflargest4, dflargest5]\n",
    "    dflargest = pd.concat(\n",
    "        [df.transpose().iloc[0:1].reset_index() for df in dflargest],\n",
    "        axis=\"index\",\n",
    "    )\n",
    "    # display(dflargest[['layer_type', 'score_text']])\n",
    "    dflargest.loc[(dflargest.layers == \"BM25+\"), \"layer_type\"] = \"BM25+\"\n",
    "    dflargest.loc[\n",
    "        (dflargest.layers == \"Random-Target\"), \"layer_type\"\n",
    "    ] = \"Random-Target\"\n",
    "    return dflargest\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "table2_results = []\n",
    "for metric in (\"mrr\", \"recall\"):\n",
    "    df_ft_fl_mrr = print_best_results(\n",
    "        path=\"reranker/sweep_re_ft_fl/\",\n",
    "        subset=\"learned\",\n",
    "        accum=\"accum\",\n",
    "        eos=\"eos\",\n",
    "        seed_range=3,\n",
    "        metric=metric,\n",
    "    )\n",
    "    df_pt_fl_mrr = print_best_results(\n",
    "        path=\"reranker/sweep_v2_re_pt_fl/\",\n",
    "        subset=\"learned\",\n",
    "        accum=\"accum\",\n",
    "        eos=\"no_eos\",\n",
    "        seed_range=3,\n",
    "        metric=metric,\n",
    "    )\n",
    "    result = (\n",
    "        df_ft_fl_mrr[[\"layer_type\", \"score_text\"]]\n",
    "        .merge(df_pt_fl_mrr[[\"layer_type\", \"score_text\"]], on=\"layer_type\")\n",
    "        .set_index(\"layer_type\")\n",
    "        .loc[[\"Random-Target\", \"BM25+\", \"TracIn\", \"Embed\", \"TracIn+Embed\"]]\n",
    "    )\n",
    "    table2_results.append(result)\n",
    "    print(metric)\n",
    "\n",
    "# print(\"FT on PL (eos)\")\n",
    "# print_best_results(path=\"reranker/sweep_re_ft_pl/\", subset=\"learned\", accum=\"accum\", eos=\"eos\", seed_range=2)\n",
    "\n",
    "# print(\"PT on PL (no_eos)\")\n",
    "# print_best_results(path=\"reranker/sweep_re/\", subset=\"learned\", accum=\"accum\", eos=\"no_eos\", seed_range=2)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "table2 = pd.concat(table2_results, axis=1)\n",
    "display(table2)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "table3_results = {\"mrr\": {}, \"recall\": {}}\n",
    "for metric in (\"mrr\", \"recall\"):\n",
    "    df_ft_fl_mrr = print_best_results(\n",
    "        path=\"reranker/sweep_re_ft_fl/\",\n",
    "        subset=\"learned\",\n",
    "        accum=\"accum\",\n",
    "        eos=\"no_eos\",\n",
    "        seed_range=3,\n",
    "    )\n",
    "    df_pt_fl_mrr = print_best_results(\n",
    "        path=\"reranker/sweep_v2_re_pt_fl/\",\n",
    "        subset=\"learned\",\n",
    "        accum=\"accum\",\n",
    "        eos=\"eos\",\n",
    "        seed_range=3,\n",
    "    )\n",
    "    table3_results[metric][\"eos\"] = (\n",
    "        df_ft_fl_mrr[[\"layer_type\", \"score_text\"]]\n",
    "        .merge(df_pt_fl_mrr[[\"layer_type\", \"score_text\"]], on=\"layer_type\")\n",
    "        .set_index(\"layer_type\")\n",
    "        .loc[[\"TracIn\", \"Embed\", \"TracIn+Embed\"]]\n",
    "    )\n",
    "\n",
    "    df_ft_fl_mrr = print_best_results(\n",
    "        path=\"reranker/sweep_re_ft_fl/\",\n",
    "        subset=\"learned\",\n",
    "        accum=\"no_accum\",\n",
    "        eos=\"eos\",\n",
    "        seed_range=3,\n",
    "    )\n",
    "    df_pt_fl_mrr = print_best_results(\n",
    "        path=\"reranker/sweep_v2_re_pt_fl/\",\n",
    "        subset=\"learned\",\n",
    "        accum=\"no_accum\",\n",
    "        eos=\"no_eos\",\n",
    "        seed_range=3,\n",
    "    )\n",
    "    table3_results[metric][\"accum\"] = (\n",
    "        df_ft_fl_mrr[[\"layer_type\", \"score_text\"]]\n",
    "        .merge(df_pt_fl_mrr[[\"layer_type\", \"score_text\"]], on=\"layer_type\")\n",
    "        .set_index(\"layer_type\")\n",
    "        .loc[[\"TracIn\", \"Embed\", \"TracIn+Embed\"]]\n",
    "    )\n",
    "\n",
    "    df_ft_fl_mrr = print_best_results(\n",
    "        path=\"reranker/sweep_re_ft_fl/\",\n",
    "        subset=\"learned\",\n",
    "        accum=\"accum\",\n",
    "        eos=\"eos\",\n",
    "        normalization=\"dot\",\n",
    "        seed_range=3,\n",
    "    )\n",
    "    df_pt_fl_mrr = print_best_results(\n",
    "        path=\"reranker/sweep_v2_re_pt_fl/\",\n",
    "        subset=\"learned\",\n",
    "        accum=\"accum\",\n",
    "        eos=\"no_eos\",\n",
    "        normalization=\"dot\",\n",
    "        seed_range=3,\n",
    "    )\n",
    "    table3_results[metric][\"normalization\"] = (\n",
    "        df_ft_fl_mrr[[\"layer_type\", \"score_text\"]]\n",
    "        .merge(df_pt_fl_mrr[[\"layer_type\", \"score_text\"]], on=\"layer_type\")\n",
    "        .set_index(\"layer_type\")\n",
    "        .loc[[\"TracIn\", \"Embed\", \"TracIn+Embed\"]]\n",
    "    )\n",
    "    ft_ckpt_results = []\n",
    "    for ckpt_no in (5000, 10000, 30000, 80000):\n",
    "        df_ft_fl_mrr = print_best_results(\n",
    "            path=f\"reranker/sweep_re_ft_fl_{ckpt_no}/\",\n",
    "            subset=\"learned\",\n",
    "            accum=\"accum\",\n",
    "            eos=\"eos\",\n",
    "            seed_range=3,\n",
    "        )\n",
    "        df_ft_fl_mrr = (\n",
    "            df_ft_fl_mrr[[\"layer_type\", \"score_text\"]]\n",
    "            .set_index(\"layer_type\")\n",
    "            .loc[[\"TracIn\", \"Embed\", \"TracIn+Embed\"]]\n",
    "        )\n",
    "        ft_ckpt_results.append(df_ft_fl_mrr)\n",
    "\n",
    "    pt_ckpt_results = []\n",
    "    for ckpt_no in (5100, 10200, 15300, 1000000):\n",
    "        df_pt_fl_mrr = print_best_results(\n",
    "            path=f\"reranker/sweep_re_pt_fl_{ckpt_no}/\",\n",
    "            subset=\"learned\",\n",
    "            accum=\"accum\",\n",
    "            eos=\"no_eos\",\n",
    "            seed_range=3,\n",
    "        )\n",
    "        df_pt_fl_mrr = (\n",
    "            df_pt_fl_mrr[[\"layer_type\", \"score_text\"]]\n",
    "            .set_index(\"layer_type\")\n",
    "            .loc[[\"TracIn\", \"Embed\", \"TracIn+Embed\"]]\n",
    "        )\n",
    "        pt_ckpt_results.append(df_pt_fl_mrr)\n",
    "\n",
    "    table3_results[metric][\"single_ckpt\"] = {\n",
    "        \"ft\": ft_ckpt_results,\n",
    "        \"pt\": pt_ckpt_results,\n",
    "    }\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from uncertainties import ufloat\n",
    "\n",
    "\n",
    "def table3_results_to_table3(results):\n",
    "    table3 = []\n",
    "    for i, metric in enumerate(results.keys()):\n",
    "        tracin = table2_results[i].loc[[\"TracIn\"], :]\n",
    "        adafactor = table3_results[metric][\"accum\"].loc[[\"TracIn\"], :]\n",
    "        eos = table3_results[metric][\"eos\"].loc[[\"TracIn\"], :]\n",
    "        norm = table3_results[metric][\"normalization\"].loc[[\"TracIn\"], :]\n",
    "        df = pd.concat([tracin, adafactor, eos, norm], axis=\"index\")\n",
    "        table3.append(df)\n",
    "    table3 = pd.concat(table3, axis=1).applymap(score_text_to_score)\n",
    "    return table3\n",
    "\n",
    "\n",
    "def score_text_to_score(score_text):\n",
    "    num, err = score_text.split(\"\\stderr{\")\n",
    "    num = float(num)\n",
    "    err = float(err[:-1])\n",
    "    return ufloat(num, err)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def table3_ckpt_result(results):\n",
    "    table3 = []\n",
    "    for i, metric in enumerate(results.keys()):\n",
    "        tracin = table2_results[i].loc[[\"TracIn\"], :]\n",
    "        ckpts = []\n",
    "        for i in range(3):\n",
    "            ckpt1 = (\n",
    "                table3_results[metric][\"single_ckpt\"][\"ft\"][i]\n",
    "                .loc[[\"TracIn\"], :]\n",
    "                .merge(\n",
    "                    table3_results[metric][\"single_ckpt\"][\"pt\"][i].loc[\n",
    "                        [\"TracIn\"], :\n",
    "                    ],\n",
    "                    on=\"layer_type\",\n",
    "                )\n",
    "            )\n",
    "            ckpts.append(ckpt1)\n",
    "        df = pd.concat([tracin, *ckpts], axis=\"index\")\n",
    "        table3.append(df)\n",
    "    table3 = pd.concat(table3, axis=1).applymap(score_text_to_score)\n",
    "    return table3\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "table2_results[0]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "table3 = table3_results_to_table3(table3_results)\n",
    "(table3 - table3.iloc[0:1]).applymap(\n",
    "    lambda x: \"{:10.2f}\".format(x).replace(\"+/-\", \"\\stderr{\") + \"}\"\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "table3_ckpt = table3_ckpt_result(table3_results)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "table3_ckpt_single = pd.DataFrame(table3_ckpt.iloc[1:].max()).transpose()\n",
    "table3_ckpt_single[\"layer_type\"] = \"TracIn\"\n",
    "table3_ckpt_single = table3_ckpt_single.set_index(\"layer_type\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "table3_ckpt\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "table3_ckpt_single\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "(table3_ckpt_single - table3_ckpt.iloc[0:1]).applymap(\n",
    "    lambda x: \"{:10.2f}\".format(x).replace(\"+/-\", \"\\stderr{\") + \"}\"\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "(table3_ckpt - table3_ckpt.iloc[0:1]).iloc[1:].applymap(\n",
    "    lambda x: \"{:10.2f}\".format(x).replace(\"+/-\", \"\\stderr{\") + \"}\"\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = visualize_one_experiment(\n",
    "    paths=[\n",
    "        f\"reranker/sweep_re_ft_fl/seed_{i}/learned/eos_accum/results_detailed.pickle\"\n",
    "        for i in range(3)\n",
    "    ],\n",
    "    suffix=\"eos_accum\",\n",
    "    save=True,\n",
    "    visualize=True,\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# FT and PT in FL\n",
    "for metric in (\n",
    "    \"mrr\",\n",
    "    \"mrr_compare_fn_subject\",\n",
    "    \"mrr_compare_fn_object\",\n",
    "    \"mrr_compare_fn_relation\",\n",
    "):\n",
    "    df_ft_fl_mrr = print_best_results(\n",
    "        path=\"reranker/sweep_re_ft_fl/\",\n",
    "        subset=\"learned\",\n",
    "        accum=\"accum\",\n",
    "        eos=\"eos\",\n",
    "        seed_range=3,\n",
    "        metric=metric,\n",
    "    )\n",
    "    df_pt_fl_mrr = print_best_results(\n",
    "        path=\"reranker/sweep_v2_re_pt_fl/\",\n",
    "        subset=\"learned\",\n",
    "        accum=\"accum\",\n",
    "        eos=\"no_eos\",\n",
    "        seed_range=3,\n",
    "        metric=metric,\n",
    "    )\n",
    "    result = (\n",
    "        df_ft_fl_mrr[[\"layer_type\", \"score_text\"]]\n",
    "        .merge(df_pt_fl_mrr[[\"layer_type\", \"score_text\"]], on=\"layer_type\")\n",
    "        .set_index(\"layer_type\")\n",
    "        .loc[[\"Random-Target\", \"BM25+\", \"TracIn\", \"Embed\", \"TracIn+Embed\"]]\n",
    "    )\n",
    "    display(result)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# FT and PT in FL\n",
    "for metric in (\n",
    "    \"mrr\",\n",
    "    \"mrr_compare_fn_subject\",\n",
    "    \"mrr_compare_fn_object\",\n",
    "    \"mrr_compare_fn_relation\",\n",
    "):\n",
    "    df_ft_pl_mrr = print_best_results(\n",
    "        path=\"reranker/sweep_re_ft_pl/\",\n",
    "        subset=\"learned\",\n",
    "        accum=\"accum\",\n",
    "        eos=\"eos\",\n",
    "        seed_range=3,\n",
    "        metric=metric,\n",
    "    )\n",
    "    df_pt_pl_mrr = print_best_results(\n",
    "        path=\"reranker/sweep_re/\",\n",
    "        subset=\"learned\",\n",
    "        accum=\"accum\",\n",
    "        eos=\"no_eos\",\n",
    "        seed_range=3,\n",
    "        metric=metric,\n",
    "    )\n",
    "    result = (\n",
    "        df_ft_pl_mrr[[\"layer_type\", \"score_text\"]]\n",
    "        .merge(df_pt_pl_mrr[[\"layer_type\", \"score_text\"]], on=\"layer_type\")\n",
    "        .set_index(\"layer_type\")\n",
    "        .loc[[\"Random-Target\", \"BM25+\", \"TracIn\", \"Embed\", \"TracIn+Embed\"]]\n",
    "    )\n",
    "    display(result)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_abstracts_df(res, fact):\n",
    "    abstracts = np.array(res[\"nn_abstracts\"])\n",
    "    for i, abstract in enumerate(abstracts):\n",
    "        try:\n",
    "            abstract[\"score\"] = res[\"nn_scores\"][i]\n",
    "        except KeyError:\n",
    "            abstract[\"score\"] = res[\"nn\"][\"scores\"][i]\n",
    "\n",
    "    df = pd.DataFrame(pd.json_normalize(abstracts)).round(3)\n",
    "    print(df.keys())\n",
    "    if 'proponents' in df:\n",
    "        df[\"label\"] = df[\"uuid\"].str.split(',').str[0].isin(fact) | df[\"uuid\"].str.split(',').str[1].isin(fact)\n",
    "    else:\n",
    "        df[\"label\"] = df[\"facts\"].str.contains(\",\".join(fact))\n",
    "        df = df.drop(\n",
    "            [\"page_uri\", \"masked_uri\", \"masked_type\", \"facts\", \"example_uris\"],\n",
    "            axis=1,\n",
    "        )\n",
    "    return df\n",
    "\n",
    "\n",
    "def get_nn_abstracts(res, baseline_res, config=\"\"):\n",
    "    \"\"\"Get list of abstracts and their scores (dotproduct score)\"\"\"\n",
    "    print(f\"Config: {config}\")\n",
    "    example = baseline_res[\"example\"]\n",
    "    print(\n",
    "        f\"Example: {example['inputs_pretokenized']} =>\"\n",
    "        f\" {example['targets_pretokenized']}\"\n",
    "    )\n",
    "    if 'predicate_id' in example:\n",
    "        fact = (\n",
    "            example[\"predicate_id\"].strip(),\n",
    "            example[\"obj_uri\"].strip(),\n",
    "            example[\"sub_uri\"].strip(),\n",
    "        )\n",
    "    else:\n",
    "        fact = example['proponents']\n",
    "    print(f\"Fact: {fact}\")\n",
    "    print(f\"Model Precision\", res[\"precision\"])\n",
    "    df_model = get_abstracts_df(res, fact)\n",
    "    print(f\"Baseline Precision\", baseline_res[\"precision\"])\n",
    "    df_baseline = get_abstracts_df(baseline_res, fact)\n",
    "    return df_model, df_baseline\n",
    "\n",
    "\n",
    "def result_getter(path):\n",
    "    with gzip.open(METRICS_DIR + path) as f:\n",
    "        reranker_metrics = pickle.load(f)\n",
    "\n",
    "    def getter(\n",
    "        i=3,\n",
    "        sim=\"cosine\",\n",
    "        method=\"collapse\",\n",
    "        normalization=\"local\",\n",
    "        layers=\"activations.encoder.block.0,gradients.shared\",\n",
    "    ):\n",
    "\n",
    "        config = {\n",
    "            \"sim\": sim,\n",
    "            \"method\": method,\n",
    "            \"layers\": layers,\n",
    "            \"normalization\": normalization,\n",
    "        }\n",
    "\n",
    "        return get_nn_abstracts(\n",
    "            reranker_metrics[\"evals\"][normalization][sim][method][layers][\n",
    "                \"samples\"\n",
    "            ][i],\n",
    "            reranker_metrics[\"samples\"][i],\n",
    "            config=config,\n",
    "        )\n",
    "\n",
    "    getter.metrics = reranker_metrics\n",
    "    return getter\n",
    "\n",
    "def result_getter(path):\n",
    "    with gzip.open(METRICS_DIR + path) as f:\n",
    "        reranker_metrics = pickle.load(f)\n",
    "\n",
    "    def getter(\n",
    "        i=3,\n",
    "        sim=\"cosine\",\n",
    "        method=\"collapse\",\n",
    "        normalization=\"local\",\n",
    "        layers=\"activations.encoder.block.0,gradients.shared\",\n",
    "    ):\n",
    "\n",
    "        config = {\n",
    "            \"sim\": sim,\n",
    "            \"method\": method,\n",
    "            \"layers\": layers,\n",
    "            \"normalization\": normalization,\n",
    "        }\n",
    "\n",
    "        return get_nn_abstracts(\n",
    "            reranker_metrics[\"evals\"][normalization][sim][method][layers][\n",
    "                \"samples\"\n",
    "            ][i],\n",
    "            reranker_metrics[\"samples\"][i],\n",
    "            config=config,\n",
    "        )\n",
    "\n",
    "    getter.metrics = reranker_metrics\n",
    "    return getter\n",
    "\n",
    "def single_column(df):\n",
    "    df = df.replace({'<extra_id_0>': '[MASK0]'}, regex=True).replace({'<extra_id_1>': '[MASK1]'}, regex=True)\n",
    "    df = df.applymap(lambda x: x[:150] + '...' if type(x) == str and len(x)>150 else x)\n",
    "    df['output'] = '\\\\textbf{Q}: ' + df['inputs_pretokenized'] + \"\\newline\\\\textbf{A}: \" + df['targets_pretokenized'] + ' \\\\textbf{' + df['label'].astype(str) + '}'\n",
    "    return df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "getter_ft_fl = result_getter(\n",
    "    \"reranker/sweep_re_ft_fl/seed_0/learned/eos_accum/results_detailed.pickle\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "idx = 11\n",
    "tracin_res, baseline_res = getter_ft_fl(i=idx, layers=\"gradients.shared\")\n",
    "embed_res, _ = getter_ft_fl(\n",
    "    i=idx, layers=\"activations.encoder.block.0,activations.decoder.block.0\"\n",
    ")\n",
    "cs = ['output']\n",
    "data = {'Embed': single_column(embed_res)[cs].iloc[:3], \n",
    "        'TracIn': single_column(tracin_res)[cs].iloc[:3], \n",
    "        'BM25': single_column(baseline_res)[cs].iloc[:3]}\n",
    "table7 = pd.concat(data.values(), axis=1, keys=data.keys())\n",
    "table7"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(table7.to_latex(index=False, escape=False))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "METRICS_DIR = os.path.join(\"Synth/synth_data_synth_07_27/\", \"metrics/\")\n",
    "getter_synth_ft_fl = result_getter(\n",
    "    \"reranker/sweep/seed_0/learned/eos_accum/results_detailed.pickle\"\n",
    ")\n",
    "METRICS_DIR = os.path.join(BASE_DIR, \"metrics/\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "idx = 1\n",
    "tracin_res, baseline_res = getter_synth_ft_fl(i=idx, layers=\"gradients.shared\")\n",
    "embed_res, _ = getter_synth_ft_fl(\n",
    "    i=idx, layers=\"activations.encoder.block.0,activations.decoder.block.0\"\n",
    ")\n",
    "cs = ['output']\n",
    "data = {'Embed': single_column(embed_res)[cs].iloc[:3], \n",
    "        'TracIn': single_column(tracin_res)[cs].iloc[:3], \n",
    "        'BM25': single_column(baseline_res)[cs].iloc[:3]}\n",
    "table7 = pd.concat(data.values(), axis=1, keys=data.keys())\n",
    "table7"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(table7.to_latex(index=False, escape=False))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Misc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "METRICS_DIR = os.path.join(\"Synth/synth_data_synth_07_27/\", \"metrics/\")\n",
    "df = visualize_one_experiment(\n",
    "    paths=[\n",
    "        f\"reranker/sweep/seed_{i}/learned/eos_accum/results_detailed.pickle\"\n",
    "        for i in range(2)\n",
    "    ],\n",
    "    suffix=\"eos_accum\",\n",
    "    save=True,\n",
    "    visualize=True,\n",
    ")\n",
    "METRICS_DIR = os.path.join(BASE_DIR, \"metrics/\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "METRICS_DIR = os.path.join(\"Synth/synth_data_synth_07_27/\", \"metrics/\")\n",
    "table4_results = []\n",
    "for metric in (\"mrr\", \"precision\", \"recall\"):\n",
    "    df_ft_fl_mrr = print_best_results(\n",
    "        path=\"reranker/sweep/\",\n",
    "        subset=\"learned\",\n",
    "        accum=\"accum\",\n",
    "        eos=\"eos\",\n",
    "        seed_range=2,\n",
    "        metric=metric,\n",
    "    )\n",
    "    # df_pt_fl_mrr = print_best_results(path=\"reranker/sweep/\", subset=\"learned\", accum=\"accum\", eos=\"no_eos\", seed_range=2, metric=metric)\n",
    "    result = (\n",
    "        df_ft_fl_mrr[[\"layer_type\", \"score_text\"]]\n",
    "        .set_index(\"layer_type\")\n",
    "        .loc[[\"Random-Target\", \"BM25+\", \"TracIn\", \"Embed\", \"TracIn+Embed\"]]\n",
    "    )\n",
    "    table4_results.append(result)\n",
    "    print(metric)\n",
    "table4 = pd.concat(table4_results, axis=1)\n",
    "METRICS_DIR = os.path.join(BASE_DIR, \"metrics/\")\n",
    "# print(\"FT on PL (eos)\")\n",
    "# print_best_results(path=\"reranker/sweep_re_ft_pl/\", subset=\"learned\", accum=\"accum\", eos=\"eos\", seed_range=2)\n",
    "\n",
    "# print(\"PT on PL (no_eos)\")\n",
    "# print_best_results(path=\"reranker/sweep_re/\", subset=\"learned\", accum=\"accum\", eos=\"no_eos\", seed_range=2)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = visualize_one_experiment(\n",
    "    paths=[\n",
    "        f\"reranker/sweep_re_ft_fl/seed_{i}/learned/eos_accum/results_detailed.pickle\"\n",
    "        for i in range(3)\n",
    "    ],\n",
    "    suffix=\"eos_accum\",\n",
    "    save=True,\n",
    "    visualize=True,\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"FT on FL (no_eos)\")\n",
    "print_best_results(\n",
    "    path=\"reranker/sweep_re_ft_fl/\",\n",
    "    subset=\"learned\",\n",
    "    accum=\"accum\",\n",
    "    eos=\"no_eos\",\n",
    "    seed_range=3,\n",
    ")\n",
    "print(\"FT on PL (no_eos)\")\n",
    "print_best_results(\n",
    "    path=\"reranker/sweep_re_ft_pl/\",\n",
    "    subset=\"learned\",\n",
    "    accum=\"accum\",\n",
    "    eos=\"no_eos\",\n",
    "    seed_range=3,\n",
    ")\n",
    "print(\"PT on PL (no_eos)\")\n",
    "print_best_results(\n",
    "    path=\"reranker/sweep_re/\",\n",
    "    subset=\"learned\",\n",
    "    accum=\"accum\",\n",
    "    eos=\"no_eos\",\n",
    "    seed_range=3,\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"FT on FL (eos)\")\n",
    "print_best_results(\n",
    "    path=\"reranker/sweep_re_ft_fl_5000/\",\n",
    "    subset=\"learned\",\n",
    "    accum=\"accum\",\n",
    "    eos=\"eos\",\n",
    "    seed_range=3,\n",
    ")\n",
    "print(\"FT on FL (eos)\")\n",
    "print_best_results(\n",
    "    path=\"reranker/sweep_re_ft_fl_10000/\",\n",
    "    subset=\"learned\",\n",
    "    accum=\"accum\",\n",
    "    eos=\"eos\",\n",
    "    seed_range=3,\n",
    ")\n",
    "print(\"FT on FL (eos)\")\n",
    "print_best_results(\n",
    "    path=\"reranker/sweep_re_ft_fl_30000/\",\n",
    "    subset=\"learned\",\n",
    "    accum=\"accum\",\n",
    "    eos=\"eos\",\n",
    "    seed_range=3,\n",
    ")\n",
    "print(\"FT on FL (eos)\")\n",
    "print_best_results(\n",
    "    path=\"reranker/sweep_re_ft_fl_80000/\",\n",
    "    subset=\"learned\",\n",
    "    accum=\"accum\",\n",
    "    eos=\"eos\",\n",
    "    seed_range=3,\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def result_getter(path):\n",
    "    with gzip.open(METRICS_DIR + path) as f:\n",
    "        reranker_metrics = pickle.load(f)\n",
    "\n",
    "    def getter(\n",
    "        i=3,\n",
    "        sim=\"cosine\",\n",
    "        method=\"collapse\",\n",
    "        normalization=\"local\",\n",
    "        layers=\"activations.encoder.block.0,gradients.shared\",\n",
    "    ):\n",
    "\n",
    "        config = {\n",
    "            \"sim\": sim,\n",
    "            \"method\": method,\n",
    "            \"layers\": layers,\n",
    "            \"normalization\": normalization,\n",
    "        }\n",
    "\n",
    "        return get_nn_abstracts(\n",
    "            reranker_metrics[\"evals\"][normalization][sim][method][layers][\n",
    "                \"samples\"\n",
    "            ][i],\n",
    "            reranker_metrics[\"samples\"][i],\n",
    "            config=config,\n",
    "        )\n",
    "\n",
    "    getter.metrics = reranker_metrics\n",
    "    return getter\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "getter_ft_fl = result_getter(\n",
    "    \"reranker/sweep_re_ft_fl/seed_0/learned/eos_accum/results_detailed.pickle\"\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "idx = 3\n",
    "tracin_res, baseline_res = getter_ft_fl(i=idx, layers=\"gradients.shared\")\n",
    "embed_res, _ = getter_ft_fl(\n",
    "    i=idx, layers=\"activations.encoder.block.0,activations.decoder.block.0\"\n",
    ")\n",
    "display(embed_res)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "getter_ft_pl = result_getter(\n",
    "    \"reranker/sweep_re_ft_pl/seed_0/learned/eos_accum/results_detailed.pickle\"\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "idx = 5\n",
    "tracin_res, baseline_res = getter_ft_pl(i=idx, layers=\"gradients.shared\")\n",
    "embed_res, _ = getter_ft_pl(\n",
    "    i=idx, layers=\"activations.encoder.block.0,activations.decoder.block.0\"\n",
    ")\n",
    "display(embed_res)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "83ce25b1c5d65648c98919b9641334dfc26d5cfa225e002eb6106bc7a8478051"
  },
  "kernelspec": {
   "display_name": "Python 3.7.6 ('trex': venv)",
   "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.9.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}