analysis.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "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": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "BASE_DIR = \"LAMA/data/\"\n",
    "METRICS_DIR = os.path.join(BASE_DIR, \"metrics/\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "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 + [\"mrr_compare_fn_subject\", \"mrr_compare_fn_object\", \"mrr_compare_fn_relation\"]\n",
    "            else:\n",
    "                current_metrics = metrics\n",
    "            for metric in current_metrics:\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 + [\"mrr_compare_fn_subject\", \"mrr_compare_fn_object\", \"mrr_compare_fn_relation\"]\n",
    "                    else:\n",
    "                        current_metrics = metrics\n",
    "                    for metric in current_metrics:\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": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_with_filter(df, \n",
    "                     filter=lambda x: x,\n",
    "                     title=\"Title\",\n",
    "                     folder=\"plots/\",\n",
    "                     ylabel='',\n",
    "                     save=False):\n",
    "    plt.figure(figsize=(16, 6))\n",
    "    \n",
    "    ax = sns.barplot(data=filter(df), \n",
    "                     x='layers', \n",
    "                     y='score', \n",
    "                     hue='layer_type', \n",
    "                     estimator=np.mean, \n",
    "                     ci='sd')\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()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "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": 13,
   "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",
    "    df[\"label\"] = df[\"facts\"].str.contains(\",\".join(fact))\n",
    "\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",
    "    fact = (\n",
    "        example[\"predicate_id\"].strip(),\n",
    "        example[\"obj_uri\"].strip(),\n",
    "        example[\"sub_uri\"].strip(),\n",
    "    )\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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "getter = result_getter('reranker/sweep_v2/seed_0/learned/no_eos_accum/results_detailed.pickle')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "idx = 10\n",
    "tracin_res, baseline_res = getter(i=idx, layers='gradients.shared')\n",
    "embed_res,_= getter(i=idx, layers='activations.encoder.block.0,activations.decoder.block.0')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "baseline_res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tracin_res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def process(model_res):\n",
    "    model_res['text'] = '\\\\textbf{Q:} ' + model_res['inputs_pretokenized'].str.replace('<extra_id_0>', '[MASK]') +  ' \\\\newline\\\\textbf{A:} ' + model_res['targets_pretokenized'].str.replace('<extra_id_0>', '') + ' \\\\textbf{' +  model_res['label'].astype(str) + '}'\n",
    "    return model_res.head(3)    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"\\\\\\\\ \\n  \\\\midrule \\n\".join(process(tracin_res)['text'] + '&' + process(embed_res)['text'] +  '&' + process(baseline_res)['text'] ) + ' \\\\ \\n')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "print(\"Model Table\")\n",
    "display(model_res)\n",
    "print(\"Baseline Table\")\n",
    "display(baseline_res)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_res, baseline_res = getter(i=3, layers='activations.encoder.block.0,activations.decoder.block.0')\n",
    "print(\"Model Table\")\n",
    "display(model_res)\n",
    "print(\"Baseline Table\")\n",
    "display(baseline_res)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = visualize_one_experiment(\n",
    "    paths=[\n",
    "        f\"reranker/sweep_v2/seed_{i}/learned/no_eos_accum/results_detailed.pickle\"\n",
    "        for i in range(3)\n",
    "    ],\n",
    "    suffix=\"no_eos_accum\",\n",
    "    save=True,\n",
    "    visualize=True,\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dfs = {}\n",
    "for eos in (\"no_eos\", ):\n",
    "    for subset in (\"learned\",):\n",
    "        for accum in (\"no_accum\", \"accum\"):\n",
    "            try:\n",
    "                suffix = f\"{eos}+{subset}+{accum}\"\n",
    "                df = visualize_one_experiment(\n",
    "                    paths=[\n",
    "                        f\"reranker/sweep_v2/seed_{i}/{subset}/{eos}_{accum}/results_detailed.pickle\"\n",
    "                        for i in range(3)\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\"reranker/sweep_v2/seed_0/{subset}/{eos}_{accum}/results_detailed.pickle\"\n",
    "                )\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dfmerged = pd.concat(list(dfs.values()), ignore_index=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "scores = dfmerged.groupby([column for column in dfmerged.columns if column != 'seed' and column != 'score'], axis=0).agg({'score': [np.mean, np.std]}, as_index=False)\n",
    "scores = scores.reset_index()\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_max(df: pd.DataFrame, \n",
    "            prefix, \n",
    "            no_prefix=None, \n",
    "            subset='learned', \n",
    "            metric='mrr', \n",
    "            eval_type='collapse', \n",
    "            k='1'):\n",
    "    \n",
    "    df2 = df[(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",
    "    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[(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['subset'] == subset)]# &\n",
    "#                   (df['eval'] == eval_type)]\n",
    "\n",
    "    dflargest = dflargest.set_index('metrics').loc[['mrr', 'mrr_compare_fn_relation', 'mrr_compare_fn_subject', 'mrr_compare_fn_object', 'recall']]\n",
    "    dflargest['score_text'] = (100 * dflargest['score']['mean']).apply(lambda x: f\"{x:.2f}\") + '\\stderr{' + (100 * dflargest['score']['std']).apply(lambda x: f\"{x:.2f}\") + '}'\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": [
    "largest, dflargest = get_max(scores, 'A', no_prefix='G')\n",
    "display(largest)\n",
    "display(dflargest)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "largest, dflargest = get_max(scores, 'A')\n",
    "display(largest)\n",
    "display(dflargest)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "largest, dflargest = get_max(scores, 'G')\n",
    "display(largest)\n",
    "display(dflargest)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "largest, dflargest = get_max(scores, 'BM25+')\n",
    "display(largest)\n",
    "display(dflargest)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "largest, dflargest  = get_max(scores, 'Random-Target')\n",
    "display(largest)\n",
    "display(dflargest)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dfs = {}\n",
    "for eos in (\"no_eos\", \"eos\"):\n",
    "    for subset in (\"learned\",):\n",
    "        for accum in (\"no_accum\", \"accum\"):\n",
    "            try:\n",
    "                suffix = f\"{eos}+{subset}+{accum}\"\n",
    "                df = visualize_one_experiment(\n",
    "                    paths=[\n",
    "                        f\"reranker/sweep_v2/seed_{i}/{subset}/{eos}_{accum}/results_detailed.pickle\"\n",
    "                        for i in range(3)\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\"reranker/sweep_v2/seed_0/{subset}/{eos}_{accum}/results_detailed.pickle\"\n",
    "                )\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dfmerged = pd.concat(list(dfs.values()), ignore_index=True)\n",
    "scores = dfmerged.groupby([column for column in dfmerged.columns if column != 'seed' and column != 'score'], axis=0).agg({'score': [np.mean, np.std]}, as_index=False)\n",
    "scores = scores.reset_index()\n",
    "df = scores"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dflargest = df[(df['layers'] == 'G.0') &\n",
    "               (df['subset'] == 'learned') &\n",
    "               ((df['metrics'] == 'mrr')) &\n",
    "               (df['eval'] == 'collapse')]\n",
    "dflargest['score_text'] = (100 * dflargest['score']['mean']).apply(lambda x: f\"{x:.2f}\") + '\\stderr{' + (100 * dflargest['score']['std']).apply(lambda x: f\"{x:.2f}\") + '}'\n",
    "dflargest"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dflargest = df[(df['layers'] == 'G.0') &\n",
    "               (df['subset'] == 'learned') &\n",
    "               ((df['metrics'] == 'recall')) &\n",
    "                ((df['k'] == '10')) &\n",
    "               (df['eval'] == 'collapse')]\n",
    "dflargest['score_text'] = (100 * dflargest['score']['mean']).apply(lambda x: f\"{x:.2f}\") + '\\stderr{' + (100 * dflargest['score']['std']).apply(lambda x: f\"{x:.2f}\") + '}'\n",
    "dflargest"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "dfs = {}\n",
    "for eos in (\"no_eos\", ):\n",
    "    for subset in (\"random\", ):\n",
    "        for accum in (\"no_accum\", ):\n",
    "            try:\n",
    "                suffix = f\"{eos}+{subset}+{accum}\"\n",
    "                df = visualize_one_experiment(\n",
    "                    paths=[\n",
    "                        f\"reranker/sweep_v2_ft_pt/seed_{i}/{subset}/{eos}_{accum}/results_detailed.pickle\"\n",
    "                        for i in range(3)\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\"reranker/sweep_v2/seed_0/{subset}/{eos}_{accum}/results_detailed.pickle\"\n",
    "                )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "dfmerged = pd.concat(list(dfs.values()), ignore_index=True)\n",
    "scores = dfmerged.groupby([column for column in dfmerged.columns if column != 'seed' and column != 'score'], axis=0).agg({'score': [np.mean, np.std]}, as_index=False)\n",
    "scores = scores.reset_index()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead tr th {\n",
       "        text-align: left;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr>\n",
       "      <th></th>\n",
       "      <th>layers</th>\n",
       "      <th>norm_type</th>\n",
       "      <th>normalization</th>\n",
       "      <th>eval</th>\n",
       "      <th>metrics</th>\n",
       "      <th>k</th>\n",
       "      <th>layer_type</th>\n",
       "      <th>eos</th>\n",
       "      <th>subset</th>\n",
       "      <th>accum</th>\n",
       "      <th colspan=\"2\" halign=\"left\">score</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th>mean</th>\n",
       "      <th>std</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>4225</th>\n",
       "      <td>G.0</td>\n",
       "      <td>global</td>\n",
       "      <td>cosine</td>\n",
       "      <td>collapse</td>\n",
       "      <td>mrr</td>\n",
       "      <td>1</td>\n",
       "      <td>TracIn</td>\n",
       "      <td>no_eos</td>\n",
       "      <td>random</td>\n",
       "      <td>no_accum</td>\n",
       "      <td>0.319993</td>\n",
       "      <td>0.012945</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4250</th>\n",
       "      <td>G.0</td>\n",
       "      <td>local</td>\n",
       "      <td>cosine</td>\n",
       "      <td>collapse</td>\n",
       "      <td>mrr</td>\n",
       "      <td>1</td>\n",
       "      <td>TracIn</td>\n",
       "      <td>no_eos</td>\n",
       "      <td>random</td>\n",
       "      <td>no_accum</td>\n",
       "      <td>0.319834</td>\n",
       "      <td>0.013028</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4325</th>\n",
       "      <td>G.0,G.E.1</td>\n",
       "      <td>local</td>\n",
       "      <td>cosine</td>\n",
       "      <td>collapse</td>\n",
       "      <td>mrr</td>\n",
       "      <td>1</td>\n",
       "      <td>TracIn</td>\n",
       "      <td>no_eos</td>\n",
       "      <td>random</td>\n",
       "      <td>no_accum</td>\n",
       "      <td>0.313894</td>\n",
       "      <td>0.023035</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6125</th>\n",
       "      <td>G.E.1</td>\n",
       "      <td>local</td>\n",
       "      <td>cosine</td>\n",
       "      <td>collapse</td>\n",
       "      <td>mrr</td>\n",
       "      <td>1</td>\n",
       "      <td>TracIn</td>\n",
       "      <td>no_eos</td>\n",
       "      <td>random</td>\n",
       "      <td>no_accum</td>\n",
       "      <td>0.313063</td>\n",
       "      <td>0.021610</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6100</th>\n",
       "      <td>G.E.1</td>\n",
       "      <td>global</td>\n",
       "      <td>cosine</td>\n",
       "      <td>collapse</td>\n",
       "      <td>mrr</td>\n",
       "      <td>1</td>\n",
       "      <td>TracIn</td>\n",
       "      <td>no_eos</td>\n",
       "      <td>random</td>\n",
       "      <td>no_accum</td>\n",
       "      <td>0.297227</td>\n",
       "      <td>0.001729</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4300</th>\n",
       "      <td>G.0,G.E.1</td>\n",
       "      <td>global</td>\n",
       "      <td>cosine</td>\n",
       "      <td>collapse</td>\n",
       "      <td>mrr</td>\n",
       "      <td>1</td>\n",
       "      <td>TracIn</td>\n",
       "      <td>no_eos</td>\n",
       "      <td>random</td>\n",
       "      <td>no_accum</td>\n",
       "      <td>0.297178</td>\n",
       "      <td>0.001694</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6175</th>\n",
       "      <td>G.E.1,G.D.1</td>\n",
       "      <td>global</td>\n",
       "      <td>cosine</td>\n",
       "      <td>collapse</td>\n",
       "      <td>mrr</td>\n",
       "      <td>1</td>\n",
       "      <td>TracIn</td>\n",
       "      <td>no_eos</td>\n",
       "      <td>random</td>\n",
       "      <td>no_accum</td>\n",
       "      <td>0.283794</td>\n",
       "      <td>0.009935</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4375</th>\n",
       "      <td>G.0,G.E.1,G.D.1</td>\n",
       "      <td>global</td>\n",
       "      <td>cosine</td>\n",
       "      <td>collapse</td>\n",
       "      <td>mrr</td>\n",
       "      <td>1</td>\n",
       "      <td>TracIn</td>\n",
       "      <td>no_eos</td>\n",
       "      <td>random</td>\n",
       "      <td>no_accum</td>\n",
       "      <td>0.283710</td>\n",
       "      <td>0.009806</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4925</th>\n",
       "      <td>G.0,G.E.2</td>\n",
       "      <td>local</td>\n",
       "      <td>cosine</td>\n",
       "      <td>collapse</td>\n",
       "      <td>mrr</td>\n",
       "      <td>1</td>\n",
       "      <td>TracIn</td>\n",
       "      <td>no_eos</td>\n",
       "      <td>random</td>\n",
       "      <td>no_accum</td>\n",
       "      <td>0.282829</td>\n",
       "      <td>0.011970</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6725</th>\n",
       "      <td>G.E.2</td>\n",
       "      <td>local</td>\n",
       "      <td>cosine</td>\n",
       "      <td>collapse</td>\n",
       "      <td>mrr</td>\n",
       "      <td>1</td>\n",
       "      <td>TracIn</td>\n",
       "      <td>no_eos</td>\n",
       "      <td>random</td>\n",
       "      <td>no_accum</td>\n",
       "      <td>0.279992</td>\n",
       "      <td>0.012688</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "               layers norm_type normalization      eval metrics  k layer_type  \\\n",
       "                                                                                \n",
       "4225  G.0              global    cosine        collapse  mrr     1  TracIn      \n",
       "4250  G.0              local     cosine        collapse  mrr     1  TracIn      \n",
       "4325  G.0,G.E.1        local     cosine        collapse  mrr     1  TracIn      \n",
       "6125  G.E.1            local     cosine        collapse  mrr     1  TracIn      \n",
       "6100  G.E.1            global    cosine        collapse  mrr     1  TracIn      \n",
       "4300  G.0,G.E.1        global    cosine        collapse  mrr     1  TracIn      \n",
       "6175  G.E.1,G.D.1      global    cosine        collapse  mrr     1  TracIn      \n",
       "4375  G.0,G.E.1,G.D.1  global    cosine        collapse  mrr     1  TracIn      \n",
       "4925  G.0,G.E.2        local     cosine        collapse  mrr     1  TracIn      \n",
       "6725  G.E.2            local     cosine        collapse  mrr     1  TracIn      \n",
       "\n",
       "         eos  subset     accum     score            \n",
       "                                    mean       std  \n",
       "4225  no_eos  random  no_accum  0.319993  0.012945  \n",
       "4250  no_eos  random  no_accum  0.319834  0.013028  \n",
       "4325  no_eos  random  no_accum  0.313894  0.023035  \n",
       "6125  no_eos  random  no_accum  0.313063  0.021610  \n",
       "6100  no_eos  random  no_accum  0.297227  0.001729  \n",
       "4300  no_eos  random  no_accum  0.297178  0.001694  \n",
       "6175  no_eos  random  no_accum  0.283794  0.009935  \n",
       "4375  no_eos  random  no_accum  0.283710  0.009806  \n",
       "4925  no_eos  random  no_accum  0.282829  0.011970  \n",
       "6725  no_eos  random  no_accum  0.279992  0.012688  "
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>metrics</th>\n",
       "      <th>mrr</th>\n",
       "      <th>mrr</th>\n",
       "      <th>mrr_compare_fn_relation</th>\n",
       "      <th>mrr_compare_fn_subject</th>\n",
       "      <th>mrr_compare_fn_object</th>\n",
       "      <th>recall</th>\n",
       "      <th>recall</th>\n",
       "      <th>recall</th>\n",
       "      <th>recall</th>\n",
       "      <th>recall</th>\n",
       "      <th>recall</th>\n",
       "      <th>recall</th>\n",
       "      <th>recall</th>\n",
       "      <th>recall</th>\n",
       "      <th>recall</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>score_text</th>\n",
       "      <th></th>\n",
       "      <td>32.00\\stderr{1.29}</td>\n",
       "      <td>31.64\\stderr{1.39}</td>\n",
       "      <td>64.39\\stderr{1.46}</td>\n",
       "      <td>33.16\\stderr{1.36}</td>\n",
       "      <td>69.10\\stderr{1.82}</td>\n",
       "      <td>8.90\\stderr{2.17}</td>\n",
       "      <td>28.80\\stderr{3.20}</td>\n",
       "      <td>43.56\\stderr{1.98}</td>\n",
       "      <td>16.04\\stderr{3.56}</td>\n",
       "      <td>21.15\\stderr{1.97}</td>\n",
       "      <td>3.41\\stderr{0.93}</td>\n",
       "      <td>10.78\\stderr{1.86}</td>\n",
       "      <td>17.15\\stderr{0.85}</td>\n",
       "      <td>6.05\\stderr{1.49}</td>\n",
       "      <td>7.81\\stderr{1.17}</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>k</th>\n",
       "      <th></th>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>10</td>\n",
       "      <td>25</td>\n",
       "      <td>3</td>\n",
       "      <td>5</td>\n",
       "      <td>1</td>\n",
       "      <td>10</td>\n",
       "      <td>25</td>\n",
       "      <td>3</td>\n",
       "      <td>5</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>eval</th>\n",
       "      <th></th>\n",
       "      <td>collapse</td>\n",
       "      <td>full</td>\n",
       "      <td>full</td>\n",
       "      <td>full</td>\n",
       "      <td>full</td>\n",
       "      <td>collapse</td>\n",
       "      <td>collapse</td>\n",
       "      <td>collapse</td>\n",
       "      <td>collapse</td>\n",
       "      <td>collapse</td>\n",
       "      <td>full</td>\n",
       "      <td>full</td>\n",
       "      <td>full</td>\n",
       "      <td>full</td>\n",
       "      <td>full</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "metrics                     mrr                 mrr mrr_compare_fn_relation  \\\n",
       "score_text   32.00\\stderr{1.29}  31.64\\stderr{1.39}  64.39\\stderr{1.46}       \n",
       "k            1                   1                   1                        \n",
       "eval         collapse            full                full                     \n",
       "\n",
       "metrics     mrr_compare_fn_subject mrr_compare_fn_object             recall  \\\n",
       "score_text   33.16\\stderr{1.36}     69.10\\stderr{1.82}    8.90\\stderr{2.17}   \n",
       "k            1                      1                     1                   \n",
       "eval         full                   full                  collapse            \n",
       "\n",
       "metrics                  recall              recall              recall  \\\n",
       "score_text   28.80\\stderr{3.20}  43.56\\stderr{1.98}  16.04\\stderr{3.56}   \n",
       "k            10                  25                  3                    \n",
       "eval         collapse            collapse            collapse             \n",
       "\n",
       "metrics                  recall             recall              recall  \\\n",
       "score_text   21.15\\stderr{1.97}  3.41\\stderr{0.93}  10.78\\stderr{1.86}   \n",
       "k            5                   1                  10                   \n",
       "eval         collapse            full               full                 \n",
       "\n",
       "metrics                  recall             recall             recall  \n",
       "score_text   17.15\\stderr{0.85}  6.05\\stderr{1.49}  7.81\\stderr{1.17}  \n",
       "k            25                  3                  5                  \n",
       "eval         full                full               full               "
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "largest, dflargest = get_max(scores, 'G', subset='random')\n",
    "display(largest)\n",
    "display(dflargest)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "ename": "IndexError",
     "evalue": "single positional indexer is out-of-bounds",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mIndexError\u001b[0m                                Traceback (most recent call last)",
      "\u001b[0;32m/tmp/ipykernel_2780903/2463032690.py\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0mlargest\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdflargest\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mget_max\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mscores\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m'A'\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mno_prefix\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m'G'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      2\u001b[0m \u001b[0mdisplay\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mlargest\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      3\u001b[0m \u001b[0mdisplay\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdflargest\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/tmp/ipykernel_2780903/3237828414.py\u001b[0m in \u001b[0;36mget_max\u001b[0;34m(df, prefix, no_prefix, subset, metric, eval_type, k)\u001b[0m\n\u001b[1;32m     18\u001b[0m     \u001b[0mnlargest\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mdf2\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mnlargest\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m10\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0;34m'score'\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m'mean'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     19\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 20\u001b[0;31m     \u001b[0mlargest\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mnlargest\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0miloc\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     21\u001b[0m     dflargest = df[(df['layers'] == largest['layers'][0]) &\n\u001b[1;32m     22\u001b[0m                    \u001b[0;34m(\u001b[0m\u001b[0mdf\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m'norm_type'\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m==\u001b[0m \u001b[0mlargest\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m'norm_type'\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m&\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/akyurek/gitother/fewshot_lama/trex/lib/python3.7/site-packages/pandas/core/indexing.py\u001b[0m in \u001b[0;36m__getitem__\u001b[0;34m(self, key)\u001b[0m\n\u001b[1;32m    929\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    930\u001b[0m             \u001b[0mmaybe_callable\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mcom\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mapply_if_callable\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mkey\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mobj\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 931\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_getitem_axis\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmaybe_callable\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0maxis\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0maxis\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    932\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    933\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0m_is_scalar_access\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mkey\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mtuple\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/akyurek/gitother/fewshot_lama/trex/lib/python3.7/site-packages/pandas/core/indexing.py\u001b[0m in \u001b[0;36m_getitem_axis\u001b[0;34m(self, key, axis)\u001b[0m\n\u001b[1;32m   1564\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1565\u001b[0m             \u001b[0;31m# validate the location\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1566\u001b[0;31m             \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_validate_integer\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mkey\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0maxis\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1567\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1568\u001b[0m             \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mobj\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_ixs\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mkey\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0maxis\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0maxis\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/akyurek/gitother/fewshot_lama/trex/lib/python3.7/site-packages/pandas/core/indexing.py\u001b[0m in \u001b[0;36m_validate_integer\u001b[0;34m(self, key, axis)\u001b[0m\n\u001b[1;32m   1498\u001b[0m         \u001b[0mlen_axis\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mlen\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mobj\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_get_axis\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0maxis\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1499\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mkey\u001b[0m \u001b[0;34m>=\u001b[0m \u001b[0mlen_axis\u001b[0m \u001b[0;32mor\u001b[0m \u001b[0mkey\u001b[0m \u001b[0;34m<\u001b[0m \u001b[0;34m-\u001b[0m\u001b[0mlen_axis\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1500\u001b[0;31m             \u001b[0;32mraise\u001b[0m \u001b[0mIndexError\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\"single positional indexer is out-of-bounds\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1501\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1502\u001b[0m     \u001b[0;31m# -------------------------------------------------------------------\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mIndexError\u001b[0m: single positional indexer is out-of-bounds"
     ]
    }
   ],
   "source": [
    "largest, dflargest = get_max(scores, 'A', no_prefix='G')\n",
    "display(largest)\n",
    "display(dflargest)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dfs = {}\n",
    "for eos in (\"no_eos\", ):\n",
    "    for subset in (\"learned\", ):\n",
    "        for accum in (\"accum\", ):\n",
    "            try:\n",
    "                suffix = f\"{eos}+{subset}+{accum}\"\n",
    "                alpha=0.1\n",
    "                df = visualize_one_experiment(\n",
    "                    paths=[\n",
    "                        f\"reranker/sweep_v2/ensemble_bm25/weight_{alpha}/seed_{i}/{subset}/{eos}_{accum}/results_ensemble_arithmetic_reweight_{alpha}_.pickle\"\n",
    "                        for i in range(3)\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\"reranker/sweep_v2/seed_0/{subset}/{eos}_{accum}/results_detailed.pickle\"\n",
    "                )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dfmerged = pd.concat(list(dfs.values()), ignore_index=True)\n",
    "scores = dfmerged.groupby([column for column in dfmerged.columns if column != 'seed' and column != 'score'], axis=0).agg({'score': [np.mean, np.std]}, as_index=False)\n",
    "scores = scores.reset_index()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "largest, dflargest = get_max(scores, 'G')\n",
    "display(largest)\n",
    "display(dflargest)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dfs = {}\n",
    "for eos in (\"no_eos\", ):\n",
    "    for subset in (\"learned\", ):\n",
    "        for accum in (\"accum\", ):\n",
    "            try:\n",
    "                suffix = f\"{eos}+{subset}+{accum}\"\n",
    "                alpha=0.1\n",
    "                df = visualize_one_experiment(\n",
    "                    paths=[\n",
    "                        f\"reranker/sweep_5100/seed_{i}/{subset}/{eos}_{accum}/results_detailed.pickle\"\n",
    "                        for i in range(3)\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\"reranker/sweep_v2/seed_0/{subset}/{eos}_{accum}/results_detailed.pickle\"\n",
    "                )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dfmerged = pd.concat(list(dfs.values()), ignore_index=True)\n",
    "scores = dfmerged.groupby([column for column in dfmerged.columns if column != 'seed' and column != 'score'], axis=0).agg({'score': [np.mean, np.std]}, as_index=False)\n",
    "scores = scores.reset_index()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "largest, dflargest = get_max(scores, 'G')\n",
    "display(largest)\n",
    "display(dflargest)"
   ]
  },
  {
   "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
}