relation-eval.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_macro_PRF(predicted_idx, gold_idx, i=-1, empty_label=None):\n",
    "    '''\n",
    "    This evaluation function follows work from Sorokin and Gurevych(https://www.aclweb.org/anthology/D17-1188.pdf)\n",
    "    code borrowed from the following link:\n",
    "    https://github.com/UKPLab/emnlp2017-relation-extraction/blob/master/relation_extraction/evaluation/metrics.py\n",
    "    '''\n",
    "    if i == -1:\n",
    "        i = len(predicted_idx)\n",
    "\n",
    "    complete_rel_set = set(gold_idx) - {empty_label}\n",
    "    avg_prec = 0.0\n",
    "    avg_rec = 0.0\n",
    "\n",
    "    for r in complete_rel_set:\n",
    "        r_indices = (predicted_idx[:i] == r)\n",
    "        tp = len((predicted_idx[:i][r_indices] == gold_idx[:i][r_indices]).nonzero()[0])\n",
    "        tp_fp = len(r_indices.nonzero()[0])\n",
    "        tp_fn = len((gold_idx == r).nonzero()[0])\n",
    "        prec = (tp / tp_fp) if tp_fp > 0 else 0\n",
    "        rec = tp / tp_fn\n",
    "        #print(id_to_labels[r], prec, rec, 2.0 * prec * rec / (prec + rec))\n",
    "        avg_prec += prec\n",
    "        avg_rec += rec\n",
    "    f1 = 0\n",
    "    avg_prec = avg_prec / len(set(predicted_idx[:i]))\n",
    "    avg_rec = avg_rec / len(complete_rel_set)\n",
    "    if (avg_rec+avg_prec) > 0:\n",
    "        f1 = 2.0 * avg_prec * avg_rec / (avg_prec + avg_rec)\n",
    "\n",
    "    return avg_prec, avg_rec, f1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "import numpy as np\n",
    "from sklearn.metrics import classification_report\n",
    "import collections\n",
    "from sklearn.metrics import normalized_mutual_info_score, adjusted_rand_score, homogeneity_score, completeness_score, v_measure_score\n",
    "\n",
    "\n",
    "class ClusterEvaluation:\n",
    "    '''\n",
    "    groundtruthlabels and predicted_clusters should be two list, for example:\n",
    "    groundtruthlabels = [0, 0, 1, 1], that means the 0th and 1th data is in cluster 0,\n",
    "    and the 2th and 3th data is in cluster 1\n",
    "    '''\n",
    "    def __init__(self, groundtruthlabels, predicted_clusters):\n",
    "        self.relations = {}\n",
    "        self.groundtruthsets, self.assessableElemSet = self.createGroundTruthSets(groundtruthlabels)\n",
    "        self.predictedsets = self.createPredictedSets(predicted_clusters)\n",
    "\n",
    "    def createGroundTruthSets(self, labels):\n",
    "\n",
    "        groundtruthsets= {}\n",
    "        assessableElems = set()\n",
    "\n",
    "        for i, c in enumerate(labels):\n",
    "            assessableElems.add(i)\n",
    "            groundtruthsets.setdefault(c, set()).add(i)\n",
    "\n",
    "        return groundtruthsets, assessableElems\n",
    "\n",
    "    def createPredictedSets(self, cs):\n",
    "\n",
    "        predictedsets = {}\n",
    "        for i, c in enumerate(cs):\n",
    "            predictedsets.setdefault(c, set()).add(i)\n",
    "\n",
    "        return predictedsets\n",
    "\n",
    "    def b3precision(self, response_a, reference_a):\n",
    "        return len(response_a.intersection(reference_a)) / float(len(response_a.intersection(self.assessableElemSet)))\n",
    "\n",
    "    def b3recall(self, response_a, reference_a):\n",
    "        return len(response_a.intersection(reference_a)) / float(len(reference_a))\n",
    "\n",
    "    def b3TotalElementPrecision(self):\n",
    "        totalPrecision = 0.0\n",
    "        for c in self.predictedsets:\n",
    "            for r in self.predictedsets[c]:\n",
    "                totalPrecision += self.b3precision(self.predictedsets[c],\n",
    "                                                   self.findCluster(r, self.groundtruthsets))\n",
    "\n",
    "        return totalPrecision / float(len(self.assessableElemSet))\n",
    "\n",
    "    def b3TotalElementRecall(self):\n",
    "        totalRecall = 0.0\n",
    "        for c in self.predictedsets:\n",
    "            for r in self.predictedsets[c]:\n",
    "                totalRecall += self.b3recall(self.predictedsets[c], self.findCluster(r, self.groundtruthsets))\n",
    "\n",
    "        return totalRecall / float(len(self.assessableElemSet))\n",
    "\n",
    "    def findCluster(self, a, setsDictionary):\n",
    "        for c in setsDictionary:\n",
    "            if a in setsDictionary[c]:\n",
    "                return setsDictionary[c]\n",
    "\n",
    "    def printEvaluation(self):\n",
    "\n",
    "        recB3 = self.b3TotalElementRecall()\n",
    "        precB3 = self.b3TotalElementPrecision()\n",
    "        betasquare = math.pow(0.5, 2)\n",
    "        if recB3 == 0.0 and precB3 == 0.0:\n",
    "            F1B3 = 0.0\n",
    "            F05B3 = 0.0\n",
    "        else:\n",
    "            betasquare = math.pow(0.5, 2)\n",
    "            F1B3 = (2 * recB3 * precB3) / (recB3 + precB3)\n",
    "            F05B3 = ((1+betasquare) * recB3 * precB3)/((betasquare*precB3)+recB3)\n",
    "\n",
    "        m = {'F1': F1B3, 'F0.5': F05B3, 'precision': precB3, 'recall': recB3}\n",
    "        return m\n",
    "\n",
    "    def getF05(self):\n",
    "        recB3 = self.b3TotalElementRecall()\n",
    "        precB3 = self.b3TotalElementPrecision()\n",
    "        betasquare = math.pow(0.5, 2)\n",
    "        if recB3 == 0.0 and precB3 == 0.0:\n",
    "            F05B3 = 0.0\n",
    "        else:\n",
    "            F05B3 = ((1+betasquare) * recB3 * precB3)/((betasquare*precB3)+recB3)\n",
    "        return F05B3\n",
    "\n",
    "    def getF1(self):\n",
    "        recB3 = self.b3TotalElementRecall()\n",
    "        precB3 = self.b3TotalElementPrecision()\n",
    "\n",
    "        if recB3 == 0.0 and precB3 == 0.0:\n",
    "            F1B3 = 0.0\n",
    "        else:\n",
    "            F1B3 = (2 * recB3 * precB3) / (recB3 + precB3)\n",
    "        return F1B3\n",
    "\n",
    "def compute_cluster_metrics(gold_label_ids, predict_labels):\n",
    "    unseen_label_ids = gold_label_ids\n",
    "    cluster_eval = ClusterEvaluation(unseen_label_ids, predict_labels).printEvaluation()\n",
    "    print('B3', cluster_eval)\n",
    "    # NMI, ARI, V_measure\n",
    "    nmi = normalized_mutual_info_score\n",
    "    print('NMI', nmi(unseen_label_ids, predict_labels))\n",
    "    print('ARI', adjusted_rand_score(unseen_label_ids, predict_labels))\n",
    "    print('Homogeneity', homogeneity_score(unseen_label_ids, predict_labels))\n",
    "    print('Completeness', completeness_score(unseen_label_ids, predict_labels))\n",
    "    print('V_measure', v_measure_score(unseen_label_ids, predict_labels))\n",
    "\n",
    "    B3_F1 = cluster_eval['F1']\n",
    "    B3_precision = cluster_eval['precision']\n",
    "    B3_recall = cluster_eval['recall']\n",
    "    NMI = normalized_mutual_info_score(unseen_label_ids, predict_labels)\n",
    "    ARI = adjusted_rand_score(unseen_label_ids, predict_labels)\n",
    "    Homogeneity = homogeneity_score(unseen_label_ids, predict_labels)\n",
    "    Completeness = completeness_score(unseen_label_ids, predict_labels)\n",
    "    V_measure = v_measure_score(unseen_label_ids, predict_labels)\n",
    "    return B3_F1, NMI, ARI"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "RelationPrompt p: 0.4776472720786848, r: 0.3602666666666667, f1: 0.4071264051681733\n",
      "RelationPrompt with added data p: 0.5564498521345045, r: 0.5548761904761905, f1: 0.5552867829940056\n"
     ]
    }
   ],
   "source": [
    "#RelationPrompt Results.\n",
    "\n",
    "results_no_added_data = {\n",
    "    \"12321\": {\n",
    "        \"precision\": 0.563572016705731,\n",
    "        \"recall\": 0.3225714285714286,\n",
    "        \"score\": 0.41029978047129695,\n",
    "\n",
    "    },\n",
    "    \"111\": {\n",
    "        \"precision\": 0.3782919141331947,\n",
    "        \"recall\": 0.2439047619047619,\n",
    "        \"score\": 0.2965853171530759,\n",
    "    },\n",
    "    \"943\": {\n",
    "        \"precision\": 0.5670379080898744,\n",
    "        \"recall\": 0.424,\n",
    "        \"score\": 0.48519652188380946,\n",
    "    },\n",
    "    \"300\": {\n",
    "        \"precision\": 0.4160565573101449,\n",
    "        \"recall\": 0.3947619047619048,\n",
    "        \"score\": 0.4051295986347155,\n",
    "    },\n",
    "    \"1300\": {\n",
    "        \"precision\": 0.46327796415447897,\n",
    "        \"recall\": 0.4160952380952381,\n",
    "        \"score\": 0.43842080769796876,\n",
    "    }\n",
    "}\n",
    "\n",
    "results_with_added_data = {\n",
    "    \"12321\": {\n",
    "        \"precision\": 0.6335737530573299,\n",
    "        \"recall\": 0.6580952380952382,\n",
    "        \"score\": 0.6456017334551125,\n",
    "    },\n",
    "    \"943\": {\n",
    "        \"precision\": 0.6568560306481417,\n",
    "        \"recall\": 0.6818095238095239,\n",
    "        \"score\": 0.6691002035217998,\n",
    "    },\n",
    "    \"111\": {\n",
    "        \"precision\": 0.5521802938349817,\n",
    "        \"recall\": 0.534952380952381,\n",
    "        \"score\": 0.5434298310641532,\n",
    "    },\n",
    "    \"300\": {\n",
    "        \"precision\": 0.42557058135412607,\n",
    "        \"recall\": 0.3803809523809524,\n",
    "        \"score\": 0.4017088776805001,\n",
    "    },\n",
    "    \"1300\": {\n",
    "        \"precision\": 0.5140686017779429,\n",
    "        \"recall\": 0.5191428571428571,\n",
    "        \"score\": 0.5165932692484626,\n",
    "    }\n",
    "}\n",
    "\n",
    "no_added_avg_f1 = 0.0\n",
    "no_added_avg_p = 0.0\n",
    "no_added_avg_r = 0.0\n",
    "for seed, scores in results_no_added_data.items():\n",
    "    no_added_avg_r += scores[\"recall\"]\n",
    "    no_added_avg_p += scores[\"precision\"]\n",
    "    no_added_avg_f1 += scores[\"score\"]\n",
    "\n",
    "no_added_avg_f1 = no_added_avg_f1 / len(results_no_added_data.keys())\n",
    "no_added_avg_p = no_added_avg_p / len(results_no_added_data.keys())\n",
    "no_added_avg_r = no_added_avg_r / len(results_no_added_data.keys())\n",
    "print(\"RelationPrompt p: {}, r: {}, f1: {}\".format(no_added_avg_p, no_added_avg_r, no_added_avg_f1))\n",
    "\n",
    "with_added_avg_f1 = 0.0\n",
    "with_added_avg_p = 0.0\n",
    "with_added_avg_r = 0.0\n",
    "for seed, scores in results_with_added_data.items():\n",
    "    with_added_avg_r += scores[\"recall\"]\n",
    "    with_added_avg_p += scores[\"precision\"]\n",
    "    with_added_avg_f1 += scores[\"score\"]\n",
    "\n",
    "with_added_avg_f1 = with_added_avg_f1 / len(results_with_added_data.keys())\n",
    "with_added_avg_p = with_added_avg_p / len(results_with_added_data.keys())\n",
    "with_added_avg_r = with_added_avg_r / len(results_with_added_data.keys())\n",
    "print(\"RelationPrompt with added data p: {}, r: {}, f1: {}\".format(with_added_avg_p, with_added_avg_r, with_added_avg_f1))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "RelationPrompt p: 0.5317085909855976, r: 0.451352380952381, f1: 0.48336916871825614\n",
      "RelationPrompt with added data p: 0.5965900376944065, r: 0.6104571428571429, f1: 0.602583640622186\n"
     ]
    }
   ],
   "source": [
    "# Fixed relation classification approach using RelationPrompt on the fewrel data.\n",
    "\n",
    "results_no_added_data = {\n",
    "    \"12321\": {\n",
    "        \"precision\": 0.575983824836923,\n",
    "        \"recall\": 0.48914285714285716,\n",
    "        \"score\": 0.5290232204588634,\n",
    "\n",
    "    },\n",
    "    \"111\": {\n",
    "       \"precision\": 0.5198497623076235,\n",
    "       \"recall\": 0.3634285714285714,\n",
    "       \"score\": 0.42778872583411154,\n",
    "    },\n",
    "    \"943\": {\n",
    "        \"precision\": 0.6549421430758446,                                                                                         \n",
    "        \"recall\": 0.4341904761904762,                                                                                            \n",
    "        \"score\": 0.5221947005331152,\n",
    "       \n",
    "    },\n",
    "    \"300\": {\n",
    "        \"precision\": 0.41382154126798826,\n",
    "        \"recall\": 0.43857142857142856,\n",
    "        \"score\": 0.425837168886372,\n",
    "        \n",
    "    },\n",
    "    \"1300\": {\n",
    "        \"precision\": 0.49394568343960854,\n",
    "        \"recall\": 0.5314285714285715,\n",
    "        \"score\": 0.5120020278788189,\n",
    "    }\n",
    "}\n",
    "\n",
    "results_with_added_data = {\n",
    "    \"12321\": {\n",
    "        \"precision\": 0.5846474695919295,\n",
    "        \"recall\": 0.6363809523809524,\n",
    "        \"score\": 0.6094182687490122,\n",
    "    },\n",
    "    \"943\": {\n",
    "       \"precision\": 0.7871256865879325,\n",
    "       \"recall\": 0.8087619047619047,\n",
    "       \"score\": 0.7977971293497309,\n",
    "    },\n",
    "    \"111\": {\n",
    "        \"precision\": 0.6450093534005406,\n",
    "        \"recall\": 0.5740952380952381,\n",
    "        \"score\": 0.6074897935702198,\n",
    "    },\n",
    "    \"300\": {\n",
    "        \"precision\": 0.4314928287593796,\n",
    "        \"recall\": 0.4784761904761904,\n",
    "        \"score\": 0.4537715912482854,\n",
    "    },\n",
    "    \"1300\": {\n",
    "        \"precision\": 0.5346748501322504,\n",
    "        \"recall\": 0.5545714285714285,\n",
    "        \"score\": 0.544441420193681,\n",
    "    }\n",
    "}\n",
    "\n",
    "no_added_avg_f1 = 0.0\n",
    "no_added_avg_p = 0.0\n",
    "no_added_avg_r = 0.0\n",
    "for seed, scores in results_no_added_data.items():\n",
    "    no_added_avg_r += scores[\"recall\"]\n",
    "    no_added_avg_p += scores[\"precision\"]\n",
    "    no_added_avg_f1 += scores[\"score\"]\n",
    "\n",
    "no_added_avg_f1 = no_added_avg_f1 / len(results_no_added_data.keys())\n",
    "no_added_avg_p = no_added_avg_p / len(results_no_added_data.keys())\n",
    "no_added_avg_r = no_added_avg_r / len(results_no_added_data.keys())\n",
    "print(\"RelationPrompt p: {}, r: {}, f1: {}\".format(no_added_avg_p, no_added_avg_r, no_added_avg_f1))\n",
    "\n",
    "with_added_avg_f1 = 0.0\n",
    "with_added_avg_p = 0.0\n",
    "with_added_avg_r = 0.0\n",
    "for seed, scores in results_with_added_data.items():\n",
    "    with_added_avg_r += scores[\"recall\"]\n",
    "    with_added_avg_p += scores[\"precision\"]\n",
    "    with_added_avg_f1 += scores[\"score\"]\n",
    "\n",
    "with_added_avg_f1 = with_added_avg_f1 / len(results_with_added_data.keys())\n",
    "with_added_avg_p = with_added_avg_p / len(results_with_added_data.keys())\n",
    "with_added_avg_r = with_added_avg_r / len(results_with_added_data.keys())\n",
    "print(\"RelationPrompt with added data p: {}, r: {}, f1: {}\".format(with_added_avg_p, with_added_avg_r, with_added_avg_f1))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "RelationPrompt p: 0.5235653148614057, r: 0.3876612070058003, f1: 0.4451737262895918\n",
      "RelationPrompt with added data p: 0.5970622565406768, r: 0.6046564388097418, f1: 0.5999932087544672\n"
     ]
    }
   ],
   "source": [
    "# RelationPrompt Results on the wikizsl dataset.\n",
    "\n",
    "results_no_added_data = {\n",
    "    \"12321\": {\n",
    "        \"precision\": 0.5165223978964989,                                                                                         \n",
    "        \"recall\": 0.411289554607959,                                                                                             \n",
    "        \"score\": 0.45793819836539623, \n",
    "    },\n",
    "    \"111\": {\n",
    "       \"precision\": 0.5523308999611009,                                                                                         \n",
    "        \"recall\": 0.4217834340695948,                                                                                            \n",
    "        \"score\": 0.47830940494301755,\n",
    "    },\n",
    "    \"943\": {\n",
    "        \"precision\": 0.4520375943184734,                                                                                         \n",
    "        \"recall\": 0.3400661179949058,                                                                                            \n",
    "        \"score\": 0.3881377337285404\n",
    "    },\n",
    "    \"300\": {\n",
    "        \"precision\": 0.5620754472446287,                                                                                         \n",
    "        \"recall\": 0.3949786312161496,                                                                                            \n",
    "        \"score\": 0.46393990849491334,\n",
    "    },\n",
    "    \"1300\": {\n",
    "        \"precision\": 0.5348602348863266,\n",
    "        \"recall\": 0.3701882971403922,\n",
    "        \"score\": 0.43754338591609165,\n",
    "    }\n",
    "}\n",
    "\n",
    "results_with_added_data = {\n",
    "    \"12321\": {\n",
    "        \"precision\": 0.4418063455810748,\n",
    "        \"recall\": 0.4008679015415297,\n",
    "        \"score\": 0.42034269647034644,\n",
    "    },\n",
    "    \"943\": {\n",
    "       \"precision\": 0.6822823985760695,                                                                                         \n",
    "       \"recall\": 0.6623873885906933,                                                                                            \n",
    "       \"score\": 0.6721877156568392,\n",
    "    },\n",
    "    \"111\": {\n",
    "        \"precision\": 0.5678342725708723,                                                                                         \n",
    "        \"recall\": 0.5645445692592689,                                                                                            \n",
    "        \"score\": 0.5661846424136192,\n",
    "    },\n",
    "    \"300\": {\n",
    "        \"precision\": 0.6104572470893014,                                                                                         \n",
    "        \"recall\": 0.616261447562777,                                                                                             \n",
    "        \"score\": 0.6133456160837898,\n",
    "    },\n",
    "    \"1300\": {\n",
    "        \"precision\": 0.682931018886066,\n",
    "        \"recall\": 0.7792208870944407,\n",
    "        \"score\": 0.7279053731477408,\n",
    "    }\n",
    "}\n",
    "\n",
    "no_added_avg_f1 = 0.0\n",
    "no_added_avg_p = 0.0\n",
    "no_added_avg_r = 0.0\n",
    "for seed, scores in results_no_added_data.items():\n",
    "    no_added_avg_r += scores[\"recall\"]\n",
    "    no_added_avg_p += scores[\"precision\"]\n",
    "    no_added_avg_f1 += scores[\"score\"]\n",
    "\n",
    "no_added_avg_f1 = no_added_avg_f1 / len(results_no_added_data.keys())\n",
    "no_added_avg_p = no_added_avg_p / len(results_no_added_data.keys())\n",
    "no_added_avg_r = no_added_avg_r / len(results_no_added_data.keys())\n",
    "print(\"RelationPrompt p: {}, r: {}, f1: {}\".format(no_added_avg_p, no_added_avg_r, no_added_avg_f1))\n",
    "\n",
    "with_added_avg_f1 = 0.0\n",
    "with_added_avg_p = 0.0\n",
    "with_added_avg_r = 0.0\n",
    "for seed, scores in results_with_added_data.items():\n",
    "    with_added_avg_r += scores[\"recall\"]\n",
    "    with_added_avg_p += scores[\"precision\"]\n",
    "    with_added_avg_f1 += scores[\"score\"]\n",
    "\n",
    "with_added_avg_f1 = with_added_avg_f1 / len(results_with_added_data.keys())\n",
    "with_added_avg_p = with_added_avg_p / len(results_with_added_data.keys())\n",
    "with_added_avg_r = with_added_avg_r / len(results_with_added_data.keys())\n",
    "print(\"RelationPrompt with added data p: {}, r: {}, f1: {}\".format(with_added_avg_p, with_added_avg_r, with_added_avg_f1))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "RelationPrompt p: 0.4809365190519602, r: 0.3596027819736719, f1: 0.40810433973533466\n",
      "RelationPrompt with added data p: 0.5553768576109043, r: 0.5543951363546602, f1: 0.5545221177068979\n"
     ]
    }
   ],
   "source": [
    "# RelationPrompt Results on the fewrel test data without sentences having multiple triples.\n",
    "\"\"\" num examples \"\"\"\n",
    "results_no_added_data = {\n",
    "    \"12321\": {\n",
    "        \"precision\": 0.5648831705030012,\n",
    "        \"recall\": 0.32112829305955065,\n",
    "        \"score\": 0.40947544311065043,\n",
    "    },\n",
    "    \"111\": {\n",
    "        \"precision\": 0.37927777403601654,                                                                                        \n",
    "        \"recall\": 0.2433876050126119,                                                                                            \n",
    "        \"score\": 0.2965043895589107,\n",
    "    },\n",
    "    \"943\": {\n",
    "      \"precision\": 0.567304509224951,                                                                                        \n",
    "      \"recall\": 0.4237896547829373,                                                                                          \n",
    "      \"score\": 0.48515628656114684,  \n",
    "    },\n",
    "    \"300\": {\n",
    "       \"precision\": 0.43034837269703086,\n",
    "        \"recall\": 0.3948975359486661,\n",
    "        \"score\": 0.4118615074541075,\n",
    "    },\n",
    "    \"1300\": {\n",
    "       \"precision\": 0.46286876879880123,\n",
    "        \"recall\": 0.41481082106459366,\n",
    "        \"score\": 0.4375240719918582,\n",
    "    }\n",
    "}\n",
    "\n",
    "results_with_added_data = {\n",
    "    \"12321\": {\n",
    "        \"precision\": 0.6317241484841962,\n",
    "        \"recall\": 0.6561757053352106,\n",
    "        \"score\": 0.6437178131196967,\n",
    "    },\n",
    "    \"943\": {\n",
    "        \"precision\": 0.6540826555485862,\n",
    "        \"recall\": 0.6823299628017044,\n",
    "        \"score\": 0.6679077822246751,\n",
    "    },\n",
    "    \"111\": {\n",
    "        \"precision\": 0.5528180890880586,\n",
    "        \"recall\": 0.5362529911578472,\n",
    "        \"score\": 0.5444095600678331\n",
    "    },\n",
    "    \"300\": {\n",
    "        \"precision\": 0.42326730897659876,\n",
    "        \"recall\": 0.3800854262665035,\n",
    "        \"score\": 0.40051580955497273,\n",
    "    },\n",
    "    \"1300\": {\n",
    "        \"precision\": 0.5149920859570816,\n",
    "        \"recall\": 0.517131596212035,\n",
    "        \"score\": 0.5160596235673119,\n",
    "    }\n",
    "}\n",
    "\n",
    "no_added_avg_f1 = 0.0\n",
    "no_added_avg_p = 0.0\n",
    "no_added_avg_r = 0.0\n",
    "for seed, scores in results_no_added_data.items():\n",
    "    no_added_avg_r += scores[\"recall\"]\n",
    "    no_added_avg_p += scores[\"precision\"]\n",
    "    no_added_avg_f1 += scores[\"score\"]\n",
    "\n",
    "no_added_avg_f1 = no_added_avg_f1 / len(results_no_added_data.keys())\n",
    "no_added_avg_p = no_added_avg_p / len(results_no_added_data.keys())\n",
    "no_added_avg_r = no_added_avg_r / len(results_no_added_data.keys())\n",
    "print(\"RelationPrompt p: {}, r: {}, f1: {}\".format(no_added_avg_p, no_added_avg_r, no_added_avg_f1))\n",
    "\n",
    "with_added_avg_f1 = 0.0\n",
    "with_added_avg_p = 0.0\n",
    "with_added_avg_r = 0.0\n",
    "for seed, scores in results_with_added_data.items():\n",
    "    with_added_avg_r += scores[\"recall\"]\n",
    "    with_added_avg_p += scores[\"precision\"]\n",
    "    with_added_avg_f1 += scores[\"score\"]\n",
    "\n",
    "with_added_avg_f1 = with_added_avg_f1 / len(results_with_added_data.keys())\n",
    "with_added_avg_p = with_added_avg_p / len(results_with_added_data.keys())\n",
    "with_added_avg_r = with_added_avg_r / len(results_with_added_data.keys())\n",
    "print(\"RelationPrompt with added data p: {}, r: {}, f1: {}\".format(with_added_avg_p, with_added_avg_r, with_added_avg_f1))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "RelationPrompt p: 0.46809274436460635, r: 0.20921904761904758, f1: 0.2850566304723004\n",
      "RelationPrompt with added data p: 0.49878761344938083, r: 0.48481904761904754, f1: 0.4904729285339803\n"
     ]
    }
   ],
   "source": [
    "# RelationPrompt Results on the test data with negs data.\n",
    "\n",
    "results_no_added_data = {\n",
    "    \"12321\": {\n",
    "       \"precision\": 0.459251382294596,\n",
    "        \"recall\": 0.11276190476190477,\n",
    "        \"score\": 0.18106593606788426,\n",
    "    },\n",
    "    \"111\": {\n",
    "       \"precision\": 0.3305721717039897,\n",
    "        \"recall\": 0.17371428571428568,\n",
    "        \"score\": 0.22774797078061565,\n",
    "    },\n",
    "    \"943\": {\n",
    "       \"precision\": 0.48672764408023866,\n",
    "        \"recall\": 0.25685714285714284,\n",
    "        \"score\": 0.3362615110052447,\n",
    "    },\n",
    "    \"300\": {\n",
    "       \"precision\": 0.5159171136238571,                                                                                               \n",
    "        \"recall\": 0.27723809523809523,                                                                                                 \n",
    "        \"score\": 0.3606655451132733,         \n",
    "    },\n",
    "    \"1300\": {\n",
    "        \"precision\": 0.5479954101203502,\n",
    "        \"recall\": 0.22552380952380954,\n",
    "        \"score\": 0.3195421893944839,\n",
    "    }\n",
    "}\n",
    "\n",
    "results_with_added_data = {\n",
    "    \"12321\": {\n",
    "        \"precision\": 0.5501224678769095,\n",
    "        \"recall\": 0.5651428571428572,\n",
    "        \"score\": 0.5575315152363876, \n",
    "    },\n",
    "    \"943\": {\n",
    "       \"precision\": 0.543395939080254,\n",
    "        \"recall\": 0.4959047619047619,\n",
    "        \"score\": 0.518565288244702,\n",
    "    },\n",
    "    \"111\": {\n",
    "       \"precision\": 0.4962629947119975,\n",
    "        \"recall\": 0.523142857142857,\n",
    "        \"score\": 0.5093485395939629,\n",
    "    },\n",
    "    \"300\": {\n",
    "        \"precision\": 0.40499570471570406,\n",
    "         \"recall\": 0.3205714285714285,\n",
    "        \"score\": 0.35787192023932907,\n",
    "    },\n",
    "    \"1300\": {\n",
    "        \"precision\": 0.49916096086203887,\n",
    "        \"recall\": 0.5193333333333333,\n",
    "        \"score\": 0.5090473793555202,\n",
    "    }\n",
    "}\n",
    "\n",
    "no_added_avg_f1 = 0.0\n",
    "no_added_avg_p = 0.0\n",
    "no_added_avg_r = 0.0\n",
    "for seed, scores in results_no_added_data.items():\n",
    "    no_added_avg_r += scores[\"recall\"]\n",
    "    no_added_avg_p += scores[\"precision\"]\n",
    "    no_added_avg_f1 += scores[\"score\"]\n",
    "\n",
    "no_added_avg_f1 = no_added_avg_f1 / len(results_no_added_data.keys())\n",
    "no_added_avg_p = no_added_avg_p / len(results_no_added_data.keys())\n",
    "no_added_avg_r = no_added_avg_r / len(results_no_added_data.keys())\n",
    "print(\"RelationPrompt p: {}, r: {}, f1: {}\".format(no_added_avg_p, no_added_avg_r, no_added_avg_f1))\n",
    "\n",
    "with_added_avg_f1 = 0.0\n",
    "with_added_avg_p = 0.0\n",
    "with_added_avg_r = 0.0\n",
    "for seed, scores in results_with_added_data.items():\n",
    "    with_added_avg_r += scores[\"recall\"]\n",
    "    with_added_avg_p += scores[\"precision\"]\n",
    "    with_added_avg_f1 += scores[\"score\"]\n",
    "\n",
    "with_added_avg_f1 = with_added_avg_f1 / len(results_with_added_data.keys())\n",
    "with_added_avg_p = with_added_avg_p / len(results_with_added_data.keys())\n",
    "with_added_avg_r = with_added_avg_r / len(results_with_added_data.keys())\n",
    "print(\"RelationPrompt with added data p: {}, r: {}, f1: {}\".format(with_added_avg_p, with_added_avg_r, with_added_avg_f1))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "12321 ~/sep-1/fewrel/concat_run_12321/relation.concat.run.12321.epoch.0.dev.predictions.step.9900.csv 0.6225757946913466\n",
      "943 ~/sep-1/fewrel/concat_run_943/relation.concat.run.943.epoch.0.dev.predictions.step.1300.csv 0.5099570516870796\n",
      "111 ~/sep-1/fewrel/concat_run_111/relation.concat.run.111.epoch.0.dev.predictions.step.5000.csv 0.5992184520328142\n",
      "300 ~/sep-1/fewrel/concat_run_300/relation.concat.run.300.epoch.0.dev.predictions.step.5500.csv 0.7290633687525613\n",
      "1300 ~/sep-1/fewrel/concat_run_1300/relation.concat.run.1300.epoch.0.dev.predictions.step.1300.csv 0.6430064368535299\n"
     ]
    }
   ],
   "source": [
    "# Dev prediction for the model on the fewrel dataset using the concat model without the negative examples.\n",
    "from re import I\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "\n",
    "gold_files = {\n",
    "    12321: \"~/codes/QA-ZRE/fewrl_data/val_data_12321.csv\",\n",
    "    943: \"~/codes/QA-ZRE/fewrl_data/val_data_943.csv\",\n",
    "    111: \"~/codes/QA-ZRE/fewrl_data/val_data_111.csv\",\n",
    "    300: \"~/codes/QA-ZRE/fewrl_data/val_data_300.csv\",\n",
    "    1300: \"~/codes/QA-ZRE/fewrl_data/val_data_1300.csv\",\n",
    "}\n",
    "id_files = {\n",
    "    12321: \"~/codes/QA-ZRE/fewrl_data/val_ids_12321.csv\",\n",
    "    943: \"~/codes/QA-ZRE/fewrl_data/val_ids_943.csv\",\n",
    "    111: \"~/codes/QA-ZRE/fewrl_data/val_ids_111.csv\",\n",
    "    300: \"~/codes/QA-ZRE/fewrl_data/val_ids_300.csv\",\n",
    "    1300: \"~/codes/QA-ZRE/fewrl_data/val_ids_1300.csv\",\n",
    "}\n",
    "\n",
    "seeds = [12321, 943, 111, 300, 1300]\n",
    "\n",
    "for seed in seeds:\n",
    "    predictions = [\"~/sep-1/fewrel/concat_run_{}/relation.concat.run.{}.epoch.0.dev.predictions.step.{}.csv\".format(seed, seed, step * 100) for step in range(1, 106, 1)]\n",
    "    max_f1 = 0.0\n",
    "    max_file = None\n",
    "    prediction_files = predictions\n",
    "    df = pd.read_csv(gold_files[seed], sep=',')\n",
    "    ids = {val:i for i, val in enumerate(pd.read_csv(id_files[seed], sep=',')[\"relation_ids\"].tolist())}\n",
    "    actual_ids = df[\"actual_ids\"].tolist()\n",
    "    num_examples = len(actual_ids) // 5\n",
    "\n",
    "    gold_indices = []\n",
    "    for each_relation_id in actual_ids:\n",
    "        gold_indices.append(ids[each_relation_id])\n",
    "\n",
    "    gold_indices = np.max(np.reshape(np.array(gold_indices), (num_examples, 5)), axis=1)\n",
    "\n",
    "    max_f1 = 0.0\n",
    "    max_file = \"None\"\n",
    "    for prediction_file in prediction_files:\n",
    "        pred_log_ps = pd.read_csv(prediction_file, sep=',')[\"relation_log_p\"].tolist()\n",
    "        pred_ids = np.argmax(np.reshape(np.array(pred_log_ps), (num_examples, 5)), axis=1)\n",
    "        avg_prec, avg_rec, f1 = compute_macro_PRF(pred_ids, gold_indices)\n",
    "        if max_f1 <= f1:\n",
    "            max_f1 = f1\n",
    "            max_file = prediction_file\n",
    "\n",
    "    print(seed, max_file, max_f1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "12321 ~/sep-28/wikizsl/concat_run_12321/relation.concat.run.epoch.0.dev.predictions.step.1700.csv 0.4742412814446165\n",
      "943 ~/sep-28/wikizsl/concat_run_943/relation.concat.run.epoch.0.dev.predictions.step.4000.csv 0.6433634224082675\n",
      "111 ~/sep-28/wikizsl/concat_run_111/relation.concat.run.epoch.0.dev.predictions.step.2000.csv 0.462240973169985\n",
      "300 ~/sep-28/wikizsl/concat_run_300/relation.concat.run.epoch.0.dev.predictions.step.4700.csv 0.6777896183831733\n",
      "1300 ~/sep-28/wikizsl/concat_run_1300/relation.concat.run.epoch.0.dev.predictions.step.4400.csv 0.48692508873389795\n"
     ]
    }
   ],
   "source": [
    "# Dev prediction for the model on the wikizsl dataset using the concat model without the negative examples.\n",
    "from re import I\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import json\n",
    "\n",
    "gold_files = {\n",
    "    12321: \"~/codes/QA-ZRE/wikizsl_data/val_data_12321.csv.sampled.csv\",\n",
    "    943: \"~/codes/QA-ZRE/wikizsl_data/val_data_943.csv.sampled.csv\",\n",
    "    111: \"~/codes/QA-ZRE/wikizsl_data/val_data_111.csv.sampled.csv\",\n",
    "    300: \"~/codes/QA-ZRE/wikizsl_data/val_data_300.csv.sampled.csv\",\n",
    "    1300: \"~/codes/QA-ZRE/wikizsl_data/val_data_1300.csv.sampled.csv\",\n",
    "}\n",
    "id_files = {\n",
    "    12321: \"~/codes/QA-ZRE/wikizsl_data/val_ids_12321.csv\",\n",
    "    943: \"~/codes/QA-ZRE/wikizsl_data/val_ids_943.csv\",\n",
    "    111: \"~/codes/QA-ZRE/wikizsl_data/val_ids_111.csv\",\n",
    "    300: \"~/codes/QA-ZRE/wikizsl_data/val_ids_300.csv\",\n",
    "    1300: \"~/codes/QA-ZRE/wikizsl_data/val_ids_1300.csv\",\n",
    "}\n",
    "\n",
    "seeds = [12321, 943, 111, 300, 1300]\n",
    "\n",
    "for seed in seeds:\n",
    "    predictions = [\"~/sep-28/wikizsl/concat_run_{}/relation.concat.run.epoch.0.dev.predictions.step.{}.csv\".format(seed, step * 100) for step in range(1, 48, 1)]\n",
    "    max_f1 = 0.0\n",
    "    max_file = None\n",
    "    prediction_files = predictions\n",
    "    df = pd.read_csv(gold_files[seed], sep=',')\n",
    "    gold_entity_relations = df[\"entity_relations\"].tolist()\n",
    "    gold_relations = [row.split(\"<SEP>\")[1].strip() for row in gold_entity_relations][:5]\n",
    "\n",
    "    label_to_id = {}\n",
    "    with open(\"./relation_descriptions.json\", \"r\") as fd:\n",
    "        re_desc_data = json.load(fd)\n",
    "        for row in re_desc_data:\n",
    "            re_label = row[\"relation_label\"]\n",
    "            re_id = row[\"relation_id\"]\n",
    "            label_to_id[re_label] = re_id\n",
    "\n",
    "    ids = {label_to_id[rel_label]: i for i, rel_label in enumerate(gold_relations)}\n",
    "\n",
    "    # ids = {val:i for i, val in enumerate(pd.read_csv(id_files[seed], sep=',')[\"relation_ids\"].tolist())}\n",
    "\n",
    "    actual_ids = df[\"actual_ids\"].tolist()\n",
    "    num_examples = len(actual_ids) // 5\n",
    "\n",
    "    gold_indices = []\n",
    "    for each_relation_id in actual_ids:\n",
    "        gold_indices.append(ids[each_relation_id])\n",
    "\n",
    "    gold_indices = np.max(np.reshape(np.array(gold_indices), (num_examples, 5)), axis=1)\n",
    "\n",
    "    max_f1 = 0.0\n",
    "    max_file = \"None\"\n",
    "    for prediction_file in prediction_files:\n",
    "        pred_log_ps = pd.read_csv(prediction_file, sep=',')[\"relation_log_p\"].tolist()\n",
    "        pred_ids = np.argmax(np.reshape(np.array(pred_log_ps), (num_examples, 5)), axis=1)\n",
    "        avg_prec, avg_rec, f1 = compute_macro_PRF(pred_ids, gold_indices)\n",
    "        if max_f1 <= f1:\n",
    "            max_f1 = f1\n",
    "            max_file = prediction_file\n",
    "\n",
    "    print(seed, max_file, max_f1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "12321 ~/sep-28/wikizsl/concat_run_12321_with_unks/relation.concat.run.epoch.0.dev.predictions.step.8100.csv 0.7789526506917566\n",
      "943 ~/sep-28/wikizsl/concat_run_943_with_unks/relation.concat.run.epoch.0.dev.predictions.step.3900.csv 0.7846069028486728\n",
      "111 ~/sep-28/wikizsl/concat_run_111_with_unks/relation.concat.run.epoch.0.dev.predictions.step.7700.csv 0.5954701697186493\n",
      "300 ~/sep-28/wikizsl/concat_run_300_with_unks/relation.concat.run.epoch.0.dev.predictions.step.8200.csv 0.8075851218740074\n",
      "1300 ~/sep-28/wikizsl/concat_run_1300_with_unks/relation.concat.run.epoch.0.dev.predictions.step.4200.csv 0.710696564194337\n"
     ]
    }
   ],
   "source": [
    "# Dev prediction for the model on the wikizsl dataset using the concat model with the negative examples.\n",
    "from re import I\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "\n",
    "gold_files = {\n",
    "    12321: \"~/codes/QA-ZRE/wikizsl_data_unks/val_data_12321.csv.sampled.csv\",\n",
    "    943: \"~/codes/QA-ZRE/wikizsl_data_unks/val_data_943.csv.sampled.csv\",\n",
    "    111: \"~/codes/QA-ZRE/wikizsl_data_unks/val_data_111.csv.sampled.csv\",\n",
    "    300: \"~/codes/QA-ZRE/wikizsl_data_unks/val_data_300.csv.sampled.csv\",\n",
    "    1300: \"~/codes/QA-ZRE/wikizsl_data_unks/val_data_1300.csv.sampled.csv\",\n",
    "}\n",
    "id_files = {\n",
    "    12321: \"~/codes/QA-ZRE/wikizsl_data_unks/val_ids_12321.csv\",\n",
    "    943: \"~/codes/QA-ZRE/wikizsl_data_unks/val_ids_943.csv\",\n",
    "    111: \"~/codes/QA-ZRE/wikizsl_data_unks/val_ids_111.csv\",\n",
    "    300: \"~/codes/QA-ZRE/wikizsl_data_unks/val_ids_300.csv\",\n",
    "    1300: \"~/codes/QA-ZRE/wikizsl_data_unks/val_ids_1300.csv\",\n",
    "}\n",
    "\n",
    "seeds = [12321, 943, 111, 300, 1300]\n",
    "\n",
    "for seed in seeds:\n",
    "    predictions = [\"~/sep-28/wikizsl/concat_run_{}_with_unks/relation.concat.run.epoch.0.dev.predictions.step.{}.csv\".format(seed, step * 100) for step in range(1, 94, 1)]\n",
    "    max_f1 = 0.0\n",
    "    max_file = None\n",
    "    prediction_files = predictions\n",
    "    df = pd.read_csv(gold_files[seed], sep=',')\n",
    "\n",
    "    gold_entity_relations = df[\"entity_relations\"].tolist()\n",
    "    gold_relations = [row.split(\"<SEP>\")[1].strip() for row in gold_entity_relations][:5]\n",
    "\n",
    "    label_to_id = {}\n",
    "    with open(\"./relation_descriptions.json\", \"r\") as fd:\n",
    "        re_desc_data = json.load(fd)\n",
    "        for row in re_desc_data:\n",
    "            re_label = row[\"relation_label\"]\n",
    "            re_id = row[\"relation_id\"]\n",
    "            label_to_id[re_label] = re_id\n",
    "\n",
    "    ids = {label_to_id[rel_label]: i for i, rel_label in enumerate(gold_relations)}\n",
    "\n",
    "    #ids = {val:i for i, val in enumerate(pd.read_csv(id_files[seed], sep=',')[\"relation_ids\"].tolist())}\n",
    "\n",
    "    actual_ids = df[\"actual_ids\"].tolist()\n",
    "    num_examples = len(actual_ids) // 5\n",
    "\n",
    "    gold_indices = []\n",
    "    for each_relation_id in actual_ids:\n",
    "        gold_indices.append(ids[each_relation_id])\n",
    "\n",
    "    gold_indices = np.max(np.reshape(np.array(gold_indices), (num_examples, 5)), axis=1)\n",
    "\n",
    "    max_f1 = 0.0\n",
    "    max_file = \"None\"\n",
    "    for prediction_file in prediction_files:\n",
    "        pred_log_ps = pd.read_csv(prediction_file, sep=',')[\"relation_log_p\"].tolist()\n",
    "        pred_ids = np.argmax(np.reshape(np.array(pred_log_ps), (num_examples, 5)), axis=1)\n",
    "        avg_prec, avg_rec, f1 = compute_macro_PRF(pred_ids, gold_indices)\n",
    "        if max_f1 <= f1:\n",
    "            max_f1 = f1\n",
    "            max_file = prediction_file\n",
    "\n",
    "    print(seed, max_file, max_f1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "12321 ~/sep-1/fewrel/concat_run_12321_with_unks/relation.concat.run.epoch.0.dev.predictions.step.8600.csv 0.8560795626010362\n",
      "943 ~/sep-1/fewrel/concat_run_943_with_unks/relation.concat.run.epoch.0.dev.predictions.step.3400.csv 0.7936263560920236\n",
      "111 ~/sep-1/fewrel/concat_run_111_with_unks/relation.concat.run.epoch.0.dev.predictions.step.4600.csv 0.8865212316090652\n",
      "300 ~/sep-1/fewrel/concat_run_300_with_unks/relation.concat.run.epoch.0.dev.predictions.step.6100.csv 0.8899964022223363\n",
      "1300 ~/sep-1/fewrel/concat_run_1300_with_unks/relation.concat.run.epoch.0.dev.predictions.step.13700.csv 0.8507511193598049\n",
      "0.8553949343768531\n"
     ]
    }
   ],
   "source": [
    "# Dev prediction for the model on the fewrel dataset using the concat model with the negative examples.\n",
    "from re import I\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "\n",
    "gold_files = {\n",
    "    12321: \"~/codes/QA-ZRE/fewrl_data_unks/val_data_12321.csv\",\n",
    "    943: \"~/codes/QA-ZRE/fewrl_data_unks/val_data_943.csv\",\n",
    "    111: \"~/codes/QA-ZRE/fewrl_data_unks/val_data_111.csv\",\n",
    "    300: \"~/codes/QA-ZRE/fewrl_data_unks/val_data_300.csv\",\n",
    "    1300: \"~/codes/QA-ZRE/fewrl_data_unks/val_data_1300.csv\",\n",
    "}\n",
    "id_files = {\n",
    "    12321: \"~/codes/QA-ZRE/fewrl_data_unks/val_ids_12321.csv\",\n",
    "    943: \"~/codes/QA-ZRE/fewrl_data_unks/val_ids_943.csv\",\n",
    "    111: \"~/codes/QA-ZRE/fewrl_data_unks/val_ids_111.csv\",\n",
    "    300: \"~/codes/QA-ZRE/fewrl_data_unks/val_ids_300.csv\",\n",
    "    1300: \"~/codes/QA-ZRE/fewrl_data_unks/val_ids_1300.csv\",\n",
    "}\n",
    "\n",
    "seeds = [12321, 943, 111, 300, 1300]\n",
    "\n",
    "avg_f1_dev = 0.0\n",
    "for seed in seeds:\n",
    "    predictions = [\"~/sep-1/fewrel/concat_run_{}_with_unks/relation.concat.run.epoch.0.dev.predictions.step.{}.csv\".format(seed, step * 100) for step in range(1, 210, 1)]\n",
    "    max_f1 = 0.0\n",
    "    max_file = None\n",
    "    prediction_files = predictions\n",
    "    df = pd.read_csv(gold_files[seed], sep=',')\n",
    "    ids = {val:i for i, val in enumerate(pd.read_csv(id_files[seed], sep=',')[\"relation_ids\"].tolist())}\n",
    "    actual_ids = df[\"actual_ids\"].tolist()\n",
    "    num_examples = len(actual_ids) // 5\n",
    "\n",
    "    gold_indices = []\n",
    "    for each_relation_id in actual_ids:\n",
    "        gold_indices.append(ids[each_relation_id])\n",
    "\n",
    "    gold_indices = np.max(np.reshape(np.array(gold_indices), (num_examples, 5)), axis=1)\n",
    "\n",
    "    max_f1 = 0.0\n",
    "    max_file = \"None\"\n",
    "    for prediction_file in prediction_files:\n",
    "        pred_log_ps = pd.read_csv(prediction_file, sep=',')[\"relation_log_p\"].tolist()\n",
    "        pred_ids = np.argmax(np.reshape(np.array(pred_log_ps), (num_examples, 5)), axis=1)\n",
    "        avg_prec, avg_rec, f1 = compute_macro_PRF(pred_ids, gold_indices)\n",
    "        if max_f1 <= f1:\n",
    "            max_f1 = f1\n",
    "            max_file = prediction_file\n",
    "\n",
    "    avg_f1_dev += max_f1\n",
    "    print(seed, max_file, max_f1)\n",
    "\n",
    "print(avg_f1_dev/5.0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "12321 ~/sep-1/fewrel/concat_run_12321/relation.concat.run.12321.epoch.0.test.predictions.step.9900.csv 0.34701733359825704\n",
      "943 ~/sep-1/fewrel/concat_run_943/relation.concat.run.943.epoch.0.test.predictions.step.1300.csv 0.3671554049998069\n",
      "111 ~/sep-1/fewrel/concat_run_111/relation.concat.run.111.epoch.0.test.predictions.step.5000.csv 0.256033773197929\n",
      "300 ~/sep-1/fewrel/concat_run_300/relation.concat.run.300.epoch.0.test.predictions.step.5500.csv 0.3014579520878535\n",
      "1300 ~/sep-1/fewrel/concat_run_1300/relation.concat.run.1300.epoch.0.test.predictions.step.1300.csv 0.3067879888834918\n",
      "0.31569049055346765\n",
      "0.3082683125936674\n",
      "0.32371428571428573\n"
     ]
    }
   ],
   "source": [
    "# Test prediction for the model on the fewrel dataset using the concat model without the negative examples.\n",
    "from re import I\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "\n",
    "gold_files = {\n",
    "    12321: \"~/codes/QA-ZRE/fewrl_data/test_data_12321.csv\",\n",
    "    943: \"~/codes/QA-ZRE/fewrl_data/test_data_943.csv\",\n",
    "    111: \"~/codes/QA-ZRE/fewrl_data/test_data_111.csv\",\n",
    "    300: \"~/codes/QA-ZRE/fewrl_data/test_data_300.csv\",\n",
    "    1300: \"~/codes/QA-ZRE/fewrl_data/test_data_1300.csv\",\n",
    "}\n",
    "id_files = {\n",
    "    12321: \"~/codes/QA-ZRE/fewrl_data/test_ids_12321.csv\",\n",
    "    943: \"~/codes/QA-ZRE/fewrl_data/test_ids_943.csv\",\n",
    "    111: \"~/codes/QA-ZRE/fewrl_data/test_ids_111.csv\",\n",
    "    300: \"~/codes/QA-ZRE/fewrl_data/test_ids_300.csv\",\n",
    "    1300: \"~/codes/QA-ZRE/fewrl_data/test_ids_1300.csv\",\n",
    "}\n",
    "\n",
    "prediction_files = {\n",
    "    12321: \"~/sep-1/fewrel/concat_run_12321/relation.concat.run.12321.epoch.0.test.predictions.step.9900.csv\",\n",
    "    943: \"~/sep-1/fewrel/concat_run_943/relation.concat.run.943.epoch.0.test.predictions.step.1300.csv\",\n",
    "    111: \"~/sep-1/fewrel/concat_run_111/relation.concat.run.111.epoch.0.test.predictions.step.5000.csv\",\n",
    "    300: \"~/sep-1/fewrel/concat_run_300/relation.concat.run.300.epoch.0.test.predictions.step.5500.csv\",\n",
    "    1300: \"~/sep-1/fewrel/concat_run_1300/relation.concat.run.1300.epoch.0.test.predictions.step.1300.csv\",\n",
    "}\n",
    "seeds = [12321, 943, 111, 300, 1300]\n",
    "\n",
    "avg_f1 = 0.0\n",
    "avg_p = 0.0\n",
    "avg_r = 0.0\n",
    "for seed in seeds:\n",
    "    predictions = [prediction_files[seed]]\n",
    "    max_f1 = 0.0\n",
    "    max_file = None\n",
    "    df = pd.read_csv(gold_files[seed], sep=',')\n",
    "    ids = {val:i for i, val in enumerate(pd.read_csv(id_files[seed], sep=',')[\"relation_ids\"].tolist())}\n",
    "    actual_ids = df[\"actual_ids\"].tolist()\n",
    "    num_examples = len(actual_ids) // 15\n",
    "\n",
    "    gold_indices = []\n",
    "    for each_relation_id in actual_ids:\n",
    "        gold_indices.append(ids[each_relation_id])\n",
    "\n",
    "    gold_indices = np.max(np.reshape(np.array(gold_indices), (num_examples, 15)), axis=1)\n",
    "\n",
    "    max_f1 = 0.0\n",
    "    max_file = \"None\"\n",
    "    for prediction_file in predictions:\n",
    "        pred_log_ps = pd.read_csv(prediction_file, sep=',')[\"relation_log_p\"].tolist()\n",
    "        pred_ids = np.argmax(np.reshape(np.array(pred_log_ps), (num_examples, 15)), axis=1)\n",
    "        prec, rec, f1 = compute_macro_PRF(pred_ids, gold_indices)\n",
    "        avg_f1 += f1\n",
    "        avg_p += prec\n",
    "        avg_r += rec\n",
    "        if max_f1 <= f1:\n",
    "            max_f1 = f1\n",
    "            max_file = prediction_file\n",
    "\n",
    "    print(seed, max_file, max_f1)\n",
    "\n",
    "print(avg_f1/5.0)\n",
    "print(avg_p/5.0)\n",
    "print(avg_r/5.0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "12321 ~/sep-1/fewrel/concat_run_12321_with_unks/relation.concat.run.12321.epoch.0.test.predictions.step.8600.csv 0.6744073305090348\n",
      "943 ~/sep-1/fewrel/concat_run_943_with_unks/relation.concat.run.943.epoch.0.test.predictions.step.3400.csv 0.649950240610028\n",
      "111 ~/sep-1/fewrel/concat_run_111_with_unks/relation.concat.run.111.epoch.0.test.predictions.step.4600.csv 0.5539922978879085\n",
      "300 ~/sep-1/fewrel/concat_run_300_with_unks/relation.concat.run.300.epoch.0.test.predictions.step.6100.csv 0.4930803340711556\n",
      "1300 ~/sep-1/fewrel/concat_run_1300_with_unks/relation.concat.run.1300.epoch.0.test.predictions.step.13700.csv 0.6530094576640766\n",
      "f1:  0.6048879321484406\n",
      "p:  0.6280573304427337\n",
      "r:  0.5836190476190477\n"
     ]
    }
   ],
   "source": [
    "# Test prediction for the model on the fewrel dataset using the concat model with the negative examples.\n",
    "from re import I\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "\n",
    "gold_files = {\n",
    "    12321: \"~/codes/QA-ZRE/fewrl_data_unks/test_data_12321.csv\",\n",
    "    943: \"~/codes/QA-ZRE/fewrl_data_unks/test_data_943.csv\",\n",
    "    111: \"~/codes/QA-ZRE/fewrl_data_unks/test_data_111.csv\",\n",
    "    300: \"~/codes/QA-ZRE/fewrl_data_unks/test_data_300.csv\",\n",
    "    1300: \"~/codes/QA-ZRE/fewrl_data_unks/test_data_1300.csv\",\n",
    "}\n",
    "id_files = {\n",
    "    12321: \"~/codes/QA-ZRE/fewrl_data_unks/test_ids_12321.csv\",\n",
    "    943: \"~/codes/QA-ZRE/fewrl_data_unks/test_ids_943.csv\",\n",
    "    111: \"~/codes/QA-ZRE/fewrl_data_unks/test_ids_111.csv\",\n",
    "    300: \"~/codes/QA-ZRE/fewrl_data_unks/test_ids_300.csv\",\n",
    "    1300: \"~/codes/QA-ZRE/fewrl_data_unks/test_ids_1300.csv\",\n",
    "}\n",
    "\n",
    "prediction_files = {\n",
    "    12321: \"~/sep-1/fewrel/concat_run_12321_with_unks/relation.concat.run.12321.epoch.0.test.predictions.step.8600.csv\",\n",
    "    943: \"~/sep-1/fewrel/concat_run_943_with_unks/relation.concat.run.943.epoch.0.test.predictions.step.3400.csv\",\n",
    "    111: \"~/sep-1/fewrel/concat_run_111_with_unks/relation.concat.run.111.epoch.0.test.predictions.step.4600.csv\",\n",
    "    300: \"~/sep-1/fewrel/concat_run_300_with_unks/relation.concat.run.300.epoch.0.test.predictions.step.6100.csv\",\n",
    "    1300: \"~/sep-1/fewrel/concat_run_1300_with_unks/relation.concat.run.1300.epoch.0.test.predictions.step.13700.csv\",\n",
    "}\n",
    "seeds = [12321, 943, 111, 300, 1300]\n",
    "\n",
    "avg_f1 = 0.0\n",
    "avg_p = 0.0\n",
    "avg_r = 0.0\n",
    "for seed in seeds:\n",
    "    predictions = [prediction_files[seed]]\n",
    "    max_f1 = 0.0\n",
    "    max_file = None\n",
    "    df = pd.read_csv(gold_files[seed], sep=',')\n",
    "    ids = {val:i for i, val in enumerate(pd.read_csv(id_files[seed], sep=',')[\"relation_ids\"].tolist())}\n",
    "    actual_ids = df[\"actual_ids\"].tolist()\n",
    "    num_examples = len(actual_ids) // 15\n",
    "\n",
    "    gold_indices = []\n",
    "    for each_relation_id in actual_ids:\n",
    "        gold_indices.append(ids[each_relation_id])\n",
    "\n",
    "    gold_indices = np.max(np.reshape(np.array(gold_indices), (num_examples, 15)), axis=1)\n",
    "\n",
    "    max_f1 = 0.0\n",
    "    max_file = \"None\"\n",
    "    for prediction_file in predictions:\n",
    "        pred_log_ps = pd.read_csv(prediction_file, sep=',')[\"relation_log_p\"].tolist()\n",
    "        pred_ids = np.argmax(np.reshape(np.array(pred_log_ps), (num_examples, 15)), axis=1)\n",
    "        prec, rec, f1 = compute_macro_PRF(pred_ids, gold_indices)\n",
    "        avg_f1 += f1\n",
    "        avg_p += prec\n",
    "        avg_r += rec\n",
    "        if max_f1 <= f1:\n",
    "            max_f1 = f1\n",
    "            max_file = prediction_file\n",
    "\n",
    "    print(seed, max_file, max_f1)\n",
    "\n",
    "print(\"f1: \", avg_f1/5.0)\n",
    "print(\"p: \", avg_p/5.0)\n",
    "print(\"r: \", avg_r/5.0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "12321 ~/sep-28/wikizsl/concat_run_12321/relation.concat.run.12321.epoch.0.test.predictions.step.1700.csv 0.287491358493333\n",
      "943 ~/sep-28/wikizsl/concat_run_943/relation.concat.run.943.epoch.0.test.predictions.step.4000.csv 0.33523016814229356\n",
      "111 ~/sep-28/wikizsl/concat_run_111/relation.concat.run.111.epoch.0.test.predictions.step.2000.csv 0.25736755828754043\n",
      "300 ~/sep-28/wikizsl/concat_run_300/relation.concat.run.300.epoch.0.test.predictions.step.4700.csv 0.46728152259303596\n",
      "1300 ~/sep-28/wikizsl/concat_run_1300/relation.concat.run.1300.epoch.0.test.predictions.step.4400.csv 0.34425972150396533\n",
      "f1:  0.3383260658040336\n",
      "p:  0.3310861160390882\n",
      "r:  0.3464293234111697\n"
     ]
    }
   ],
   "source": [
    "# Test prediction for the model on the wikizsl dataset using the concat model without the negative examples.\n",
    "from re import I\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import json\n",
    "\n",
    "gold_files = {\n",
    "    12321: \"~/codes/QA-ZRE/wikizsl_data/test_data_12321.csv\",\n",
    "    943: \"~/codes/QA-ZRE/wikizsl_data/test_data_943.csv\",\n",
    "    111: \"~/codes/QA-ZRE/wikizsl_data/test_data_111.csv\",\n",
    "    300: \"~/codes/QA-ZRE/wikizsl_data/test_data_300.csv\",\n",
    "    1300: \"~/codes/QA-ZRE/wikizsl_data/test_data_1300.csv\",\n",
    "}\n",
    "id_files = {\n",
    "    12321: \"~/codes/QA-ZRE/wikizsl_data/test_ids_12321.csv\",\n",
    "    943: \"~/codes/QA-ZRE/wikizsl_data/test_ids_943.csv\",\n",
    "    111: \"~/codes/QA-ZRE/wikizsl_data/test_ids_111.csv\",\n",
    "    300: \"~/codes/QA-ZRE/wikizsl_data/test_ids_300.csv\",\n",
    "    1300: \"~/codes/QA-ZRE/wikizsl_data/test_ids_1300.csv\",\n",
    "}\n",
    "\n",
    "prediction_files = {\n",
    "    12321: \"~/sep-28/wikizsl/concat_run_12321/relation.concat.run.12321.epoch.0.test.predictions.step.1700.csv\",\n",
    "    943: \"~/sep-28/wikizsl/concat_run_943/relation.concat.run.943.epoch.0.test.predictions.step.4000.csv\",\n",
    "    111: \"~/sep-28/wikizsl/concat_run_111/relation.concat.run.111.epoch.0.test.predictions.step.2000.csv\",\n",
    "    300: \"~/sep-28/wikizsl/concat_run_300/relation.concat.run.300.epoch.0.test.predictions.step.4700.csv\",\n",
    "    1300: \"~/sep-28/wikizsl/concat_run_1300/relation.concat.run.1300.epoch.0.test.predictions.step.4400.csv\",\n",
    "}\n",
    "seeds = [12321, 943, 111, 300, 1300]\n",
    "\n",
    "avg_f1 = 0.0\n",
    "avg_p = 0.0\n",
    "avg_r = 0.0\n",
    "for seed in seeds:\n",
    "    predictions = [prediction_files[seed]]\n",
    "    max_f1 = 0.0\n",
    "    max_file = None\n",
    "    df = pd.read_csv(gold_files[seed], sep=',')\n",
    "\n",
    "    gold_entity_relations = df[\"entity_relations\"].tolist()\n",
    "    gold_relations = [row.split(\"<SEP>\")[1].strip() for row in gold_entity_relations][:15]\n",
    "\n",
    "    label_to_id = {}\n",
    "    with open(\"./relation_descriptions.json\", \"r\") as fd:\n",
    "        re_desc_data = json.load(fd)\n",
    "        for row in re_desc_data:\n",
    "            re_label = row[\"relation_label\"]\n",
    "            re_id = row[\"relation_id\"]\n",
    "            label_to_id[re_label] = re_id\n",
    "\n",
    "    ids = {label_to_id[rel_label]: i for i, rel_label in enumerate(gold_relations)}\n",
    "\n",
    "\n",
    "    #ids = {val:i for i, val in enumerate(pd.read_csv(id_files[seed], sep=',')[\"relation_ids\"].tolist())}\n",
    "    actual_ids = df[\"actual_ids\"].tolist()\n",
    "    num_examples = len(actual_ids) // 15\n",
    "\n",
    "    gold_indices = []\n",
    "    for each_relation_id in actual_ids:\n",
    "        gold_indices.append(ids[each_relation_id])\n",
    "\n",
    "    gold_indices = np.max(np.reshape(np.array(gold_indices), (num_examples, 15)), axis=1)\n",
    "\n",
    "    max_f1 = 0.0\n",
    "    max_file = \"None\"\n",
    "    for prediction_file in predictions:\n",
    "        pred_log_ps = pd.read_csv(prediction_file, sep=',')[\"relation_log_p\"].tolist()\n",
    "        pred_ids = np.argmax(np.reshape(np.array(pred_log_ps), (num_examples, 15)), axis=1)\n",
    "        prec, rec, f1 = compute_macro_PRF(pred_ids, gold_indices)\n",
    "        avg_f1 += f1\n",
    "        avg_p += prec\n",
    "        avg_r += rec\n",
    "        if max_f1 <= f1:\n",
    "            max_f1 = f1\n",
    "            max_file = prediction_file\n",
    "\n",
    "    print(seed, max_file, max_f1)\n",
    "\n",
    "print(\"f1: \", avg_f1/5.0)\n",
    "print(\"p: \", avg_p/5.0)\n",
    "print(\"r: \", avg_r/5.0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "12321 ~/sep-28/wikizsl/concat_run_12321_with_unks/relation.concat.run.12321.epoch.0.test.predictions.step.8100.csv 0.6049281662236237\n",
      "943 ~/sep-28/wikizsl/concat_run_943_with_unks/relation.concat.run.943.epoch.0.test.predictions.step.3900.csv 0.5855951376072791\n",
      "111 ~/sep-28/wikizsl/concat_run_111_with_unks/relation.concat.run.111.epoch.0.test.predictions.step.7700.csv 0.5588637367774965\n",
      "300 ~/sep-28/wikizsl/concat_run_300_with_unks/relation.concat.run.300.epoch.0.test.predictions.step.8200.csv 0.5861322571377339\n",
      "1300 ~/sep-28/wikizsl/concat_run_1300_with_unks/relation.concat.run.1300.epoch.0.test.predictions.step.4200.csv 0.7149613469058409\n",
      "f1:  0.6100961289303948\n",
      "p:  0.625870654679628\n",
      "r:  0.5954022669100214\n"
     ]
    }
   ],
   "source": [
    "# Test prediction for the model on the wikizsl dataset using the concat model with the negative examples.\n",
    "from re import I\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import json\n",
    "\n",
    "gold_files = {\n",
    "    12321: \"~/codes/QA-ZRE/wikizsl_data_unks/test_data_12321.csv\",\n",
    "    943: \"~/codes/QA-ZRE/wikizsl_data_unks/test_data_943.csv\",\n",
    "    111: \"~/codes/QA-ZRE/wikizsl_data_unks/test_data_111.csv\",\n",
    "    300: \"~/codes/QA-ZRE/wikizsl_data_unks/test_data_300.csv\",\n",
    "    1300: \"~/codes/QA-ZRE/wikizsl_data_unks/test_data_1300.csv\",\n",
    "}\n",
    "\n",
    "prediction_files = {\n",
    "    12321: \"~/sep-28/wikizsl/concat_run_12321_with_unks/relation.concat.run.12321.epoch.0.test.predictions.step.8100.csv\",\n",
    "    943: \"~/sep-28/wikizsl/concat_run_943_with_unks/relation.concat.run.943.epoch.0.test.predictions.step.3900.csv\",\n",
    "    111: \"~/sep-28/wikizsl/concat_run_111_with_unks/relation.concat.run.111.epoch.0.test.predictions.step.7700.csv\",\n",
    "    300: \"~/sep-28/wikizsl/concat_run_300_with_unks/relation.concat.run.300.epoch.0.test.predictions.step.8200.csv\",\n",
    "    1300: \"~/sep-28/wikizsl/concat_run_1300_with_unks/relation.concat.run.1300.epoch.0.test.predictions.step.4200.csv\",\n",
    "}\n",
    "seeds = [12321, 943, 111, 300, 1300]\n",
    "\n",
    "avg_f1 = 0.0\n",
    "avg_p = 0.0\n",
    "avg_r = 0.0\n",
    "for seed in seeds:\n",
    "    predictions = [prediction_files[seed]]\n",
    "    max_f1 = 0.0\n",
    "    max_file = None\n",
    "    df = pd.read_csv(gold_files[seed], sep=',')\n",
    "\n",
    "    gold_entity_relations = df[\"entity_relations\"].tolist()\n",
    "    gold_relations = [row.split(\"<SEP>\")[1].strip() for row in gold_entity_relations][:15]\n",
    "\n",
    "    label_to_id = {}\n",
    "    with open(\"./relation_descriptions.json\", \"r\") as fd:\n",
    "        re_desc_data = json.load(fd)\n",
    "        for row in re_desc_data:\n",
    "            re_label = row[\"relation_label\"]\n",
    "            re_id = row[\"relation_id\"]\n",
    "            label_to_id[re_label] = re_id\n",
    "\n",
    "    ids = {label_to_id[rel_label]: i for i, rel_label in enumerate(gold_relations)}\n",
    "\n",
    "\n",
    "    actual_ids = df[\"actual_ids\"].tolist()\n",
    "    num_examples = len(actual_ids) // 15\n",
    "\n",
    "    gold_indices = []\n",
    "    for each_relation_id in actual_ids:\n",
    "        gold_indices.append(ids[each_relation_id])\n",
    "\n",
    "    gold_indices = np.max(np.reshape(np.array(gold_indices), (num_examples, 15)), axis=1)\n",
    "\n",
    "    max_f1 = 0.0\n",
    "    max_file = \"None\"\n",
    "    for prediction_file in predictions:\n",
    "        pred_log_ps = pd.read_csv(prediction_file, sep=',')[\"relation_log_p\"].tolist()\n",
    "        pred_ids = np.argmax(np.reshape(np.array(pred_log_ps), (num_examples, 15)), axis=1)\n",
    "        prec, rec, f1 = compute_macro_PRF(pred_ids, gold_indices)\n",
    "        avg_f1 += f1\n",
    "        avg_p += prec\n",
    "        avg_r += rec\n",
    "        if max_f1 <= f1:\n",
    "            max_f1 = f1\n",
    "            max_file = prediction_file\n",
    "\n",
    "    print(seed, max_file, max_f1)\n",
    "\n",
    "print(\"f1: \", avg_f1/5.0)\n",
    "print(\"p: \", avg_p/5.0)\n",
    "print(\"r: \", avg_r/5.0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "12321 ~/sep-1/fewrel/run_12321/relation.offmml-pgg.run.epoch.0.dev.predictions.step.1800.csv 0.6498331456500441\n",
      "943 ~/sep-1/fewrel/run_943/relation.offmml-pgg.run.epoch.0.dev.predictions.step.1900.csv 0.5659540894916784\n",
      "111 ~/sep-1/fewrel/run_111/relation.offmml-pgg.run.epoch.0.dev.predictions.step.8800.csv 0.6217527243186586\n",
      "300 ~/sep-1/fewrel/run_300/relation.offmml-pgg.run.epoch.0.dev.predictions.step.1900.csv 0.7017740927559872\n",
      "1300 ~/sep-1/fewrel/run_1300/relation.offmml-pgg.run.epoch.0.dev.predictions.step.7900.csv 0.6402843822685743\n"
     ]
    }
   ],
   "source": [
    "# Dev prediction for the model on the fewrel dataset using the offmml-g model without the negative examples.\n",
    "from re import I\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "\n",
    "gold_files = {\n",
    "    12321: \"~/codes/QA-ZRE/fewrl_data/val_data_12321.csv\",\n",
    "    943: \"~/codes/QA-ZRE/fewrl_data/val_data_943.csv\",\n",
    "    111: \"~/codes/QA-ZRE/fewrl_data/val_data_111.csv\",\n",
    "    300: \"~/codes/QA-ZRE/fewrl_data/val_data_300.csv\",\n",
    "    1300: \"~/codes/QA-ZRE/fewrl_data/val_data_1300.csv\",\n",
    "}\n",
    "id_files = {\n",
    "    12321: \"~/codes/QA-ZRE/fewrl_data/val_ids_12321.csv\",\n",
    "    943: \"~/codes/QA-ZRE/fewrl_data/val_ids_943.csv\",\n",
    "    111: \"~/codes/QA-ZRE/fewrl_data/val_ids_111.csv\",\n",
    "    300: \"~/codes/QA-ZRE/fewrl_data/val_ids_300.csv\",\n",
    "    1300: \"~/codes/QA-ZRE/fewrl_data/val_ids_1300.csv\",\n",
    "}\n",
    "\n",
    "seeds = [12321, 943, 111, 300, 1300]\n",
    "\n",
    "for seed in seeds:\n",
    "    predictions = [\"~/sep-1/fewrel/run_{}/relation.offmml-pgg.run.epoch.0.dev.predictions.step.{}.csv\".format(seed, step * 100) for step in range(1, 106, 1)]\n",
    "    max_f1 = 0.0\n",
    "    max_file = None\n",
    "    prediction_files = predictions\n",
    "    df = pd.read_csv(gold_files[seed], sep=',')\n",
    "    ids = {val:i for i, val in enumerate(pd.read_csv(id_files[seed], sep=',')[\"relation_ids\"].tolist())}\n",
    "    actual_ids = df[\"actual_ids\"].tolist()\n",
    "    num_examples = len(actual_ids) // 5\n",
    "\n",
    "    gold_indices = []\n",
    "    for each_relation_id in actual_ids:\n",
    "        gold_indices.append(ids[each_relation_id])\n",
    "\n",
    "    gold_indices = np.max(np.reshape(np.array(gold_indices), (num_examples, 5)), axis=1)\n",
    "\n",
    "    max_f1 = 0.0\n",
    "    max_file = \"None\"\n",
    "    for prediction_file in prediction_files:\n",
    "        pred_log_ps = pd.read_csv(prediction_file, sep=',')[\"answer_log_p\"].tolist()\n",
    "        pred_log_ps = np.log(np.mean(np.reshape(np.exp(np.array(pred_log_ps)), (num_examples, 5, 8)), axis=2))\n",
    "        pred_ids = np.argmax(pred_log_ps, axis=1)\n",
    "        avg_prec, avg_rec, f1 = compute_macro_PRF(pred_ids, gold_indices)\n",
    "        if max_f1 <= f1:\n",
    "            max_f1 = f1\n",
    "            max_file = prediction_file\n",
    "\n",
    "    print(seed, max_file, max_f1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "12321 ~/sep-28/wikizsl/run_12321/relation.offmml-pgg.run.epoch.0.dev.predictions.step.800.csv 0.4203692048623931\n",
      "943 ~/sep-28/wikizsl/run_943/relation.offmml-pgg.run.epoch.0.dev.predictions.step.3000.csv 0.6375243450837228\n",
      "111 ~/sep-28/wikizsl/run_111/relation.offmml-pgg.run.epoch.0.dev.predictions.step.500.csv 0.5169302461492905\n",
      "300 ~/sep-28/wikizsl/run_300/relation.offmml-pgg.run.epoch.0.dev.predictions.step.1500.csv 0.6790189264600169\n",
      "1300 ~/sep-28/wikizsl/run_1300/relation.offmml-pgg.run.epoch.0.dev.predictions.step.2400.csv 0.4979349584311499\n"
     ]
    }
   ],
   "source": [
    "# Dev prediction for the model on the wikizsl dataset using the offmml-g model without the negative examples.\n",
    "from re import I\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import json\n",
    "\n",
    "gold_files = {\n",
    "    12321: \"~/codes/QA-ZRE/wikizsl_data/val_data_12321.csv.sampled.csv\",\n",
    "    943: \"~/codes/QA-ZRE/wikizsl_data/val_data_943.csv.sampled.csv\",\n",
    "    111: \"~/codes/QA-ZRE/wikizsl_data/val_data_111.csv.sampled.csv\",\n",
    "    300: \"~/codes/QA-ZRE/wikizsl_data/val_data_300.csv.sampled.csv\",\n",
    "    1300: \"~/codes/QA-ZRE/wikizsl_data/val_data_1300.csv.sampled.csv\",\n",
    "}\n",
    "\n",
    "seeds = [12321, 943, 111, 300, 1300]\n",
    "\n",
    "for seed in seeds:\n",
    "    predictions = [\"~/sep-28/wikizsl/run_{}/relation.offmml-pgg.run.epoch.0.dev.predictions.step.{}.csv\".format(seed, step * 100) for step in range(1, 46, 1)]\n",
    "    max_f1 = 0.0\n",
    "    max_file = None\n",
    "    prediction_files = predictions\n",
    "    df = pd.read_csv(gold_files[seed], sep=',')\n",
    "    gold_entity_relations = df[\"entity_relations\"].tolist()\n",
    "    gold_relations = [row.split(\"<SEP>\")[1].strip() for row in gold_entity_relations][:5]\n",
    "\n",
    "    label_to_id = {}\n",
    "    with open(\"./relation_descriptions.json\", \"r\") as fd:\n",
    "        re_desc_data = json.load(fd)\n",
    "        for row in re_desc_data:\n",
    "            re_label = row[\"relation_label\"]\n",
    "            re_id = row[\"relation_id\"]\n",
    "            label_to_id[re_label] = re_id\n",
    "\n",
    "    ids = {label_to_id[rel_label]: i for i, rel_label in enumerate(gold_relations)}\n",
    "\n",
    "    actual_ids = df[\"actual_ids\"].tolist()\n",
    "    num_examples = len(actual_ids) // 5\n",
    "\n",
    "    gold_indices = []\n",
    "    for each_relation_id in actual_ids:\n",
    "        gold_indices.append(ids[each_relation_id])\n",
    "\n",
    "    gold_indices = np.max(np.reshape(np.array(gold_indices), (num_examples, 5)), axis=1)\n",
    "\n",
    "    max_f1 = 0.0\n",
    "    max_file = \"None\"\n",
    "    for prediction_file in prediction_files:\n",
    "        pred_log_ps = pd.read_csv(prediction_file, sep=',')[\"answer_log_p\"].tolist()\n",
    "        pred_log_ps = np.log(np.mean(np.reshape(np.exp(np.array(pred_log_ps)), (num_examples, 5, 8)), axis=2))\n",
    "        pred_ids = np.argmax(pred_log_ps, axis=1)\n",
    "        avg_prec, avg_rec, f1 = compute_macro_PRF(pred_ids, gold_indices)\n",
    "        if max_f1 <= f1:\n",
    "            max_f1 = f1\n",
    "            max_file = prediction_file\n",
    "\n",
    "    print(seed, max_file, max_f1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "12321 ~/sep-28/wikizsl/run_12321_with_unks/relation.offmml-pgg.run.epoch.0.dev.predictions.step.8900.csv 0.8487020087551942\n",
      "943 ~/sep-28/wikizsl/run_943_with_unks/relation.offmml-pgg.run.epoch.0.dev.predictions.step.3700.csv 0.7728171239157858\n",
      "111 ~/sep-28/wikizsl/run_111_with_unks/relation.offmml-pgg.run.epoch.0.dev.predictions.step.6900.csv 0.7598997604137664\n",
      "300 ~/sep-28/wikizsl/run_300_with_unks/relation.offmml-pgg.run.epoch.0.dev.predictions.step.600.csv 0.7626723708239196\n",
      "1300 ~/sep-28/wikizsl/run_1300_with_unks/relation.offmml-pgg.run.epoch.0.dev.predictions.step.4800.csv 0.7242507938718629\n",
      "f1 63.21830747186563\n",
      "p 64.67936565168954\n",
      "r 61.92106181417089\n"
     ]
    }
   ],
   "source": [
    "# Dev prediction for the model on the wikizsl dataset using the offmml-g model with the negative examples.\n",
    "from re import I\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import json\n",
    "\n",
    "gold_files = {\n",
    "    12321: \"~/codes/QA-ZRE/wikizsl_data_unks/val_data_12321.csv.sampled.csv\",\n",
    "    943: \"~/codes/QA-ZRE/wikizsl_data_unks/val_data_943.csv.sampled.csv\",\n",
    "    111: \"~/codes/QA-ZRE/wikizsl_data_unks/val_data_111.csv.sampled.csv\",\n",
    "    300: \"~/codes/QA-ZRE/wikizsl_data_unks/val_data_300.csv.sampled.csv\",\n",
    "    1300: \"~/codes/QA-ZRE/wikizsl_data_unks/val_data_1300.csv.sampled.csv\",\n",
    "}\n",
    "\n",
    "\n",
    "seeds = [12321, 943, 111, 300, 1300]\n",
    "\n",
    "avg_f1 = 0.0\n",
    "avg_p = 0.0\n",
    "avg_r = 0.0\n",
    "for seed in seeds:\n",
    "    predictions = [\"~/sep-28/wikizsl/run_{}_with_unks/relation.offmml-pgg.run.epoch.0.dev.predictions.step.{}.csv\".format(seed, step * 100) for step in range(1, 93, 1)]\n",
    "\n",
    "    max_f1 = 0.0\n",
    "    max_file = None\n",
    "    df = pd.read_csv(gold_files[seed], sep=',')\n",
    "    gold_entity_relations = df[\"entity_relations\"].tolist()\n",
    "    gold_relations = [row.split(\"<SEP>\")[1].strip() for row in gold_entity_relations][:5]\n",
    "\n",
    "    label_to_id = {}\n",
    "    with open(\"./relation_descriptions.json\", \"r\") as fd:\n",
    "        re_desc_data = json.load(fd)\n",
    "        for row in re_desc_data:\n",
    "            re_label = row[\"relation_label\"]\n",
    "            re_id = row[\"relation_id\"]\n",
    "            label_to_id[re_label] = re_id\n",
    "\n",
    "    ids = {label_to_id[rel_label]: i for i, rel_label in enumerate(gold_relations)}\n",
    "\n",
    "    actual_ids = df[\"actual_ids\"].tolist()\n",
    "    num_examples = len(actual_ids) // 5\n",
    "\n",
    "    gold_indices = []\n",
    "    for each_relation_id in actual_ids:\n",
    "        gold_indices.append(ids[each_relation_id])\n",
    "\n",
    "    gold_indices = np.max(np.reshape(np.array(gold_indices), (num_examples, 5)), axis=1)\n",
    "\n",
    "    max_f1 = 0.0\n",
    "    max_file = \"None\"\n",
    "    for prediction_file in predictions:\n",
    "        pred_log_ps = pd.read_csv(prediction_file, sep=',')[\"answer_log_p\"].tolist()\n",
    "        pred_log_ps = np.log(np.mean(np.reshape(np.exp(np.array(pred_log_ps)), (num_examples, 5, 8)), axis=2))\n",
    "        pred_ids = np.argmax(pred_log_ps, axis=1)\n",
    "        prec, rec, f1 = compute_macro_PRF(pred_ids, gold_indices)\n",
    "        avg_f1 += f1\n",
    "        avg_p += prec\n",
    "        avg_r += rec\n",
    "        if max_f1 <= f1:\n",
    "            max_f1 = f1\n",
    "            max_file = prediction_file\n",
    "\n",
    "    print(seed, max_file, max_f1)\n",
    "\n",
    "print(\"f1\", avg_f1/5.0)\n",
    "print(\"p\", avg_p/5.0)\n",
    "print(\"r\", avg_r/5.0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "11956 11955\n",
      "12321 ~/sep-28/wikizsl/run_12321_with_unks/relation.offmml-pgg.run.epoch.0.test.predictions.step.8900.csv 0.668822746571222\n",
      "12393 12391\n",
      "943 ~/sep-28/wikizsl/run_943_with_unks/relation.offmml-pgg.run.epoch.0.test.predictions.step.3700.csv 0.6003950701125171\n",
      "12503 12501\n",
      "111 ~/sep-28/wikizsl/run_111_with_unks/relation.offmml-pgg.run.epoch.0.test.predictions.step.6900.csv 0.5870916969967358\n",
      "14677 14676\n",
      "300 ~/sep-28/wikizsl/run_300_with_unks/relation.offmml-pgg.run.epoch.0.test.predictions.step.600.csv 0.527820373518701\n",
      "11418 11417\n",
      "1300 ~/sep-28/wikizsl/run_1300_with_unks/relation.offmml-pgg.run.epoch.0.test.predictions.step.4800.csv 0.7116625838241473\n",
      "f1 0.6191584942046646\n",
      "p 0.63602764660785\n",
      "r 0.6035599258415971\n"
     ]
    }
   ],
   "source": [
    "# Test prediction for the model on the wikizsl dataset using the offmml-g model with the negative examples.\n",
    "from re import I\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import json\n",
    "\n",
    "def hash_tokens(tokens):\n",
    "    return \"\".join(\"\".join(tokens).split()).replace('\"', '').replace(\"'\", \"\").strip().lower()\n",
    "\n",
    "gold_files = {\n",
    "    12321: \"~/codes/QA-ZRE/wikizsl_data_unks/test_data_12321.csv\",\n",
    "    943: \"~/codes/QA-ZRE/wikizsl_data_unks/test_data_943.csv\",\n",
    "    111: \"~/codes/QA-ZRE/wikizsl_data_unks/test_data_111.csv\",\n",
    "    300: \"~/codes/QA-ZRE/wikizsl_data_unks/test_data_300.csv\",\n",
    "    1300: \"~/codes/QA-ZRE/wikizsl_data_unks/test_data_1300.csv\",\n",
    "}\n",
    "\n",
    "prediction_files = {\n",
    "    12321: \"~/sep-28/wikizsl/run_12321_with_unks/relation.offmml-pgg.run.epoch.0.test.predictions.step.8900.csv\",\n",
    "    943: \"~/sep-28/wikizsl/run_943_with_unks/relation.offmml-pgg.run.epoch.0.test.predictions.step.3700.csv\",\n",
    "    111: \"~/sep-28/wikizsl/run_111_with_unks/relation.offmml-pgg.run.epoch.0.test.predictions.step.6900.csv\",\n",
    "    300: \"~/sep-28/wikizsl/run_300_with_unks/relation.offmml-pgg.run.epoch.0.test.predictions.step.600.csv\",\n",
    "    1300: \"~/sep-28/wikizsl/run_1300_with_unks/relation.offmml-pgg.run.epoch.0.test.predictions.step.4800.csv\",\n",
    "}\n",
    "\n",
    "seeds = [12321, 943, 111, 300, 1300]\n",
    "\n",
    "avg_f1 = 0.0\n",
    "avg_p = 0.0\n",
    "avg_r = 0.0\n",
    "for seed in seeds:\n",
    "    predictions = [prediction_files[seed]]\n",
    "    max_f1 = 0.0\n",
    "    max_file = None\n",
    "    df = pd.read_csv(gold_files[seed], sep=',')\n",
    "    gold_entity_relations = df[\"entity_relations\"].tolist()\n",
    "    gold_relations = [row.split(\"<SEP>\")[1].strip() for row in gold_entity_relations][:15]\n",
    "\n",
    "    label_to_id = {}\n",
    "    with open(\"./relation_descriptions.json\", \"r\") as fd:\n",
    "        re_desc_data = json.load(fd)\n",
    "        for row in re_desc_data:\n",
    "            re_label = row[\"relation_label\"]\n",
    "            re_id = row[\"relation_id\"]\n",
    "            label_to_id[re_label] = re_id\n",
    "\n",
    "    ids = {label_to_id[rel_label]: i for i, rel_label in enumerate(gold_relations)}\n",
    "\n",
    "    sentences = []\n",
    "    sent_hash = {}\n",
    "    for sent in df[\"passages\"].tolist():\n",
    "        hash = hash_tokens(sent.split())\n",
    "        sent_hash[hash] = sent_hash.get(hash, 0) + 1\n",
    "        sentences.append(hash)\n",
    "    \n",
    "    print(len(sentences)//15, len(sent_hash))\n",
    "    actual_ids = df[\"actual_ids\"].tolist()\n",
    "    num_examples = len(actual_ids) // 15\n",
    "\n",
    "    gold_indices = []\n",
    "    valid_indices = set()\n",
    "    for index, each_relation_id in enumerate(actual_ids):\n",
    "        if sent_hash[sentences[index]] == 15:\n",
    "            valid_indices.add(index//15)\n",
    "        gold_indices.append(ids[each_relation_id])\n",
    "\n",
    "    gold_indices = np.max(np.reshape(np.array(gold_indices), (num_examples, 15)), axis=1)\n",
    "\n",
    "    max_f1 = 0.0\n",
    "    max_file = \"None\"\n",
    "    for prediction_file in predictions:\n",
    "        pred_log_ps = pd.read_csv(prediction_file, sep=',')[\"answer_log_p\"].tolist()\n",
    "        pred_log_ps = np.log(np.mean(np.reshape(np.exp(np.array(pred_log_ps)), (num_examples, 15, 8)), axis=2))\n",
    "        pred_ids = np.argmax(pred_log_ps, axis=1)\n",
    "        #valid_pred_ids = np.take(pred_ids, list(valid_indices))\n",
    "        #valid_gold_ids = np.take(gold_indices, list(valid_indices))\n",
    "        valid_pred_ids = pred_ids\n",
    "        valid_gold_ids = gold_indices\n",
    "        prec, rec, f1 = compute_macro_PRF(valid_pred_ids, valid_gold_ids)\n",
    "        avg_f1 += f1\n",
    "        avg_p += prec\n",
    "        avg_r += rec\n",
    "        if max_f1 <= f1:\n",
    "            max_f1 = f1\n",
    "            max_file = prediction_file\n",
    "\n",
    "    print(seed, max_file, max_f1)\n",
    "\n",
    "print(\"f1\", avg_f1/5.0)\n",
    "print(\"p\", avg_p/5.0)\n",
    "print(\"r\", avg_r/5.0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "12321 ~/sep-1/fewrel/run_12321/relation.offmml-pgg.run.epoch.0.test.predictions.step.1800.csv 0.32358858387038925\n",
      "943 ~/sep-1/fewrel/run_943/relation.offmml-pgg.run.epoch.0.test.predictions.step.1900.csv 0.3172183349253462\n",
      "111 ~/sep-1/fewrel/run_111/relation.offmml-pgg.run.epoch.0.test.predictions.step.8800.csv 0.284561011851722\n",
      "300 ~/sep-1/fewrel/run_300/relation.offmml-pgg.run.epoch.0.test.predictions.step.1900.csv 0.2525921583604101\n",
      "1300 ~/sep-1/fewrel/run_1300/relation.offmml-pgg.run.epoch.0.test.predictions.step.7900.csv 0.3175353309188262\n",
      "f: 0.2990990839853388\n",
      "p: 0.2939819684227859\n",
      "r: 0.3044952380952381\n"
     ]
    }
   ],
   "source": [
    "# Test prediction for the model on the fewrel dataset using the offmml-g model without the negative examples.\n",
    "from re import I\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "\n",
    "gold_files = {\n",
    "    12321: \"~/codes/QA-ZRE/fewrl_data/test_data_12321.csv\",\n",
    "    943: \"~/codes/QA-ZRE/fewrl_data/test_data_943.csv\",\n",
    "    111: \"~/codes/QA-ZRE/fewrl_data/test_data_111.csv\",\n",
    "    300: \"~/codes/QA-ZRE/fewrl_data/test_data_300.csv\",\n",
    "    1300: \"~/codes/QA-ZRE/fewrl_data/test_data_1300.csv\",\n",
    "}\n",
    "id_files = {\n",
    "    12321: \"~/codes/QA-ZRE/fewrl_data/test_ids_12321.csv\",\n",
    "    943: \"~/codes/QA-ZRE/fewrl_data/test_ids_943.csv\",\n",
    "    111: \"~/codes/QA-ZRE/fewrl_data/test_ids_111.csv\",\n",
    "    300: \"~/codes/QA-ZRE/fewrl_data/test_ids_300.csv\",\n",
    "    1300: \"~/codes/QA-ZRE/fewrl_data/test_ids_1300.csv\",\n",
    "}\n",
    "\n",
    "prediction_files = {\n",
    "    12321: \"~/sep-1/fewrel/run_12321/relation.offmml-pgg.run.epoch.0.test.predictions.step.1800.csv\",\n",
    "    943: \"~/sep-1/fewrel/run_943/relation.offmml-pgg.run.epoch.0.test.predictions.step.1900.csv\",\n",
    "    111: \"~/sep-1/fewrel/run_111/relation.offmml-pgg.run.epoch.0.test.predictions.step.8800.csv\",\n",
    "    300: \"~/sep-1/fewrel/run_300/relation.offmml-pgg.run.epoch.0.test.predictions.step.1900.csv\",\n",
    "    1300: \"~/sep-1/fewrel/run_1300/relation.offmml-pgg.run.epoch.0.test.predictions.step.7900.csv\",\n",
    "}\n",
    "seeds = [12321, 943, 111, 300, 1300]\n",
    "\n",
    "avg_f1 = 0.0\n",
    "avg_p = 0.0\n",
    "avg_r = 0.0\n",
    "for seed in seeds:\n",
    "    predictions = [prediction_files[seed]]\n",
    "    max_f1 = 0.0\n",
    "    max_file = None\n",
    "    df = pd.read_csv(gold_files[seed], sep=',')\n",
    "    ids = {val:i for i, val in enumerate(pd.read_csv(id_files[seed], sep=',')[\"relation_ids\"].tolist())}\n",
    "    actual_ids = df[\"actual_ids\"].tolist()\n",
    "    num_examples = len(actual_ids) // 15\n",
    "\n",
    "    gold_indices = []\n",
    "    for each_relation_id in actual_ids:\n",
    "        gold_indices.append(ids[each_relation_id])\n",
    "\n",
    "    gold_indices = np.max(np.reshape(np.array(gold_indices), (num_examples, 15)), axis=1)\n",
    "\n",
    "    max_f1 = 0.0\n",
    "    max_file = \"None\"\n",
    "    for prediction_file in predictions:\n",
    "        pred_log_ps = pd.read_csv(prediction_file, sep=',')[\"answer_log_p\"].tolist()\n",
    "        pred_log_ps = np.log(np.mean(np.reshape(np.exp(np.array(pred_log_ps)), (num_examples, 15, 8)), axis=2))\n",
    "        pred_ids = np.argmax(pred_log_ps, axis=1)\n",
    "        prec, rec, f1 = compute_macro_PRF(pred_ids, gold_indices)\n",
    "        avg_f1 += f1\n",
    "        avg_p += prec\n",
    "        avg_r += rec\n",
    "        if max_f1 <= f1:\n",
    "            max_f1 = f1\n",
    "            max_file = prediction_file\n",
    "\n",
    "    print(seed, max_file, max_f1)\n",
    "\n",
    "print(\"f:\", avg_f1/5.0)\n",
    "print(\"p:\", avg_p/5.0)\n",
    "print(\"r:\", avg_r/5.0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "12321 ~/sep-1/fewrel/run_12321_with_unks/relation.supp_data.offmml-pgg.run.epoch.0.dev.predictions.step.1500.csv 0.849902924915118\n",
      "943 ~/sep-1/fewrel/run_943_with_unks/relation.supp_data.offmml-pgg.run.epoch.0.dev.predictions.step.2200.csv 0.7716361725216287\n",
      "111 ~/sep-1/fewrel/run_111_with_unks/relation.supp_data.offmml-pgg.run.epoch.0.dev.predictions.step.1900.csv 0.9133416796371262\n",
      "300 ~/sep-1/fewrel/run_300_with_unks/relation.supp_data.offmml-pgg.run.epoch.0.dev.predictions.step.100.csv 0.8243735676508117\n",
      "1300 ~/sep-1/fewrel/run_1300_with_unks/relation.supp_data.offmml-pgg.run.epoch.0.dev.predictions.step.1300.csv 0.9242521139813203\n",
      "0.8567012917412009\n"
     ]
    }
   ],
   "source": [
    "# Dev prediction for the model on the fewrel dataset using the offmml-g model with the negative examples.\n",
    "from re import I\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "\n",
    "gold_files = {\n",
    "    12321: \"~/codes/QA-ZRE/fewrl_data_unks/val_data_12321.csv\",\n",
    "    943: \"~/codes/QA-ZRE/fewrl_data_unks/val_data_943.csv\",\n",
    "    111: \"~/codes/QA-ZRE/fewrl_data_unks/val_data_111.csv\",\n",
    "    300: \"~/codes/QA-ZRE/fewrl_data_unks/val_data_300.csv\",\n",
    "    1300: \"~/codes/QA-ZRE/fewrl_data_unks/val_data_1300.csv\",\n",
    "}\n",
    "id_files = {\n",
    "    12321: \"~/codes/QA-ZRE/fewrl_data_unks/val_ids_12321.csv\",\n",
    "    943: \"~/codes/QA-ZRE/fewrl_data_unks/val_ids_943.csv\",\n",
    "    111: \"~/codes/QA-ZRE/fewrl_data_unks/val_ids_111.csv\",\n",
    "    300: \"~/codes/QA-ZRE/fewrl_data_unks/val_ids_300.csv\",\n",
    "    1300: \"~/codes/QA-ZRE/fewrl_data_unks/val_ids_1300.csv\",\n",
    "}\n",
    "\n",
    "seeds = [12321, 943, 111, 300, 1300]\n",
    "\n",
    "avg_f1_dev = 0.0\n",
    "for seed in seeds:\n",
    "    predictions = [\"~/sep-1/fewrel/run_{}_with_unks/relation.supp_data.offmml-pgg.run.epoch.0.dev.predictions.step.{}.csv\".format(seed, step * 100) for step in range(1, 25, 1)]\n",
    "    max_f1 = 0.0\n",
    "    max_file = None\n",
    "    prediction_files = predictions\n",
    "    df = pd.read_csv(gold_files[seed], sep=',')\n",
    "    ids = {val:i for i, val in enumerate(pd.read_csv(id_files[seed], sep=',')[\"relation_ids\"].tolist())}\n",
    "    actual_ids = df[\"actual_ids\"].tolist()\n",
    "    num_examples = len(actual_ids) // 5\n",
    "\n",
    "    gold_indices = []\n",
    "    for each_relation_id in actual_ids:\n",
    "        gold_indices.append(ids[each_relation_id])\n",
    "\n",
    "    gold_indices = np.max(np.reshape(np.array(gold_indices), (num_examples, 5)), axis=1)\n",
    "\n",
    "    max_f1 = 0.0\n",
    "    max_file = \"None\"\n",
    "    for prediction_file in prediction_files:\n",
    "        pred_log_ps = pd.read_csv(prediction_file, sep=',')[\"answer_log_p\"].tolist()\n",
    "        pred_log_ps = np.log(np.mean(np.reshape(np.exp(np.array(pred_log_ps)), (num_examples, 5, 8)), axis=2))\n",
    "        pred_ids = np.argmax(pred_log_ps, axis=1)\n",
    "        avg_prec, avg_rec, f1 = compute_macro_PRF(pred_ids, gold_indices)\n",
    "        if max_f1 <= f1:\n",
    "            max_f1 = f1\n",
    "            max_file = prediction_file\n",
    "\n",
    "    avg_f1_dev += max_f1\n",
    "    print(seed, max_file, max_f1)\n",
    "\n",
    "print(avg_f1_dev / 5.0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "12321 ~/sep-1/fewrel/run_12321_with_unks/relation.offmml-pgg.run.epoch.0.dev.predictions.step.7700.csv 0.8984147957805454\n",
      "943 ~/sep-1/fewrel/run_943_with_unks/relation.offmml-pgg.run.epoch.0.dev.predictions.step.5900.csv 0.8066627814635562\n",
      "111 ~/sep-1/fewrel/run_111_with_unks/relation.offmml-pgg.run.epoch.0.dev.predictions.step.11000.csv 0.9040280878278296\n",
      "300 ~/sep-1/fewrel/run_300_with_unks/relation.offmml-pgg.run.epoch.0.dev.predictions.step.10700.csv 0.912398403495852\n",
      "1300 ~/sep-1/fewrel/run_1300_with_unks/relation.offmml-pgg.run.epoch.0.dev.predictions.step.15200.csv 0.8778187449916541\n",
      "0.8798645627118875\n"
     ]
    }
   ],
   "source": [
    "# Dev prediction for the model on the fewrel dataset using the offmml-g model with the negative examples.\n",
    "from re import I\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "\n",
    "gold_files = {\n",
    "    12321: \"~/codes/QA-ZRE/fewrl_data_unks/val_data_12321.csv\",\n",
    "    943: \"~/codes/QA-ZRE/fewrl_data_unks/val_data_943.csv\",\n",
    "    111: \"~/codes/QA-ZRE/fewrl_data_unks/val_data_111.csv\",\n",
    "    300: \"~/codes/QA-ZRE/fewrl_data_unks/val_data_300.csv\",\n",
    "    1300: \"~/codes/QA-ZRE/fewrl_data_unks/val_data_1300.csv\",\n",
    "}\n",
    "id_files = {\n",
    "    12321: \"~/codes/QA-ZRE/fewrl_data_unks/val_ids_12321.csv\",\n",
    "    943: \"~/codes/QA-ZRE/fewrl_data_unks/val_ids_943.csv\",\n",
    "    111: \"~/codes/QA-ZRE/fewrl_data_unks/val_ids_111.csv\",\n",
    "    300: \"~/codes/QA-ZRE/fewrl_data_unks/val_ids_300.csv\",\n",
    "    1300: \"~/codes/QA-ZRE/fewrl_data_unks/val_ids_1300.csv\",\n",
    "}\n",
    "\n",
    "seeds = [12321, 943, 111, 300, 1300]\n",
    "\n",
    "avg_f1_dev = 0.0\n",
    "for seed in seeds:\n",
    "    predictions = [\"~/sep-1/fewrel/run_{}_with_unks/relation.offmml-pgg.run.epoch.0.dev.predictions.step.{}.csv\".format(seed, step * 100) for step in range(1, 202, 1)]\n",
    "    max_f1 = 0.0\n",
    "    max_file = None\n",
    "    prediction_files = predictions\n",
    "    df = pd.read_csv(gold_files[seed], sep=',')\n",
    "    ids = {val:i for i, val in enumerate(pd.read_csv(id_files[seed], sep=',')[\"relation_ids\"].tolist())}\n",
    "    actual_ids = df[\"actual_ids\"].tolist()\n",
    "    num_examples = len(actual_ids) // 5\n",
    "\n",
    "    gold_indices = []\n",
    "    for each_relation_id in actual_ids:\n",
    "        gold_indices.append(ids[each_relation_id])\n",
    "\n",
    "    gold_indices = np.max(np.reshape(np.array(gold_indices), (num_examples, 5)), axis=1)\n",
    "\n",
    "    max_f1 = 0.0\n",
    "    max_file = \"None\"\n",
    "    for prediction_file in prediction_files:\n",
    "        pred_log_ps = pd.read_csv(prediction_file, sep=',')[\"answer_log_p\"].tolist()\n",
    "        pred_log_ps = np.log(np.mean(np.reshape(np.exp(np.array(pred_log_ps)), (num_examples, 5, 8)), axis=2))\n",
    "        pred_ids = np.argmax(pred_log_ps, axis=1)\n",
    "        avg_prec, avg_rec, f1 = compute_macro_PRF(pred_ids, gold_indices)\n",
    "        if max_f1 <= f1:\n",
    "            max_f1 = f1\n",
    "            max_file = prediction_file\n",
    "\n",
    "    avg_f1_dev += max_f1\n",
    "    print(seed, max_file, max_f1)\n",
    "\n",
    "print(avg_f1_dev / 5.0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "10082\n",
      "B3 {'F1': 0.5505287079830417, 'F0.5': 0.5243468625428549, 'precision': 0.5082332793930865, 'recall': 0.6005028571428087}\n",
      "NMI 0.6335432549568847\n",
      "ARI 0.4743914269765192\n",
      "Homogeneity 0.6130224342776839\n",
      "Completeness 0.6554855201408198\n",
      "V_measure 0.6335432549568847\n",
      "12321 ~/sep-1/fewrel/run_12321_with_unks/relation.offmml-pgg.run.epoch.0.test.predictions.step.7700.csv 0.6539281780314811\n",
      "10109\n",
      "B3 {'F1': 0.5322902649849284, 'F0.5': 0.5197443785830872, 'precision': 0.5117039292394131, 'recall': 0.5546024489795853}\n",
      "NMI 0.6036032631756123\n",
      "ARI 0.4560230689508492\n",
      "Homogeneity 0.5930270687859815\n",
      "Completeness 0.6145635445077966\n",
      "V_measure 0.6036032631756123\n",
      "943 ~/sep-1/fewrel/run_943_with_unks/relation.offmml-pgg.run.epoch.0.test.predictions.step.5900.csv 0.6459036663934058\n",
      "10243\n",
      "B3 {'F1': 0.477282043621157, 'F0.5': 0.47680422230773106, 'precision': 0.47648620592337787, 'recall': 0.47808054421767665}\n",
      "NMI 0.5772413120731947\n",
      "ARI 0.4030604655629921\n",
      "Homogeneity 0.5695324125087463\n",
      "Completeness 0.5851617625290726\n",
      "V_measure 0.5772413120731946\n",
      "111 ~/sep-1/fewrel/run_111_with_unks/relation.offmml-pgg.run.epoch.0.test.predictions.step.11000.csv 0.6397429356638897\n",
      "9681\n",
      "B3 {'F1': 0.3714598996493775, 'F0.5': 0.35647308858768817, 'precision': 0.3471361211863297, 'recall': 0.39944925170064816}\n",
      "NMI 0.43846589519754875\n",
      "ARI 0.2685729998792066\n",
      "Homogeneity 0.4196070968363362\n",
      "Completeness 0.45909964445428425\n",
      "V_measure 0.4384658951975487\n",
      "300 ~/sep-1/fewrel/run_300_with_unks/relation.offmml-pgg.run.epoch.0.test.predictions.step.10700.csv 0.5094849804267173\n",
      "10005\n",
      "B3 {'F1': 0.5125657613828108, 'F0.5': 0.4884719949081218, 'precision': 0.47362964312083455, 'recall': 0.5584770068027213}\n",
      "NMI 0.5931828183052745\n",
      "ARI 0.45604755402296077\n",
      "Homogeneity 0.5785942549329224\n",
      "Completeness 0.6085260768321393\n",
      "V_measure 0.5931828183052745\n",
      "1300 ~/sep-1/fewrel/run_1300_with_unks/relation.offmml-pgg.run.epoch.0.test.predictions.step.15200.csv 0.6165207370401674\n",
      "f: 0.6131160995111322\n",
      "p: 0.637325315840368\n",
      "r: 0.5915238095238096\n",
      "b f1: 0.48882533552426305\n",
      "nmi : 0.5692073087417031\n",
      "arr: 0.4116191030785056\n"
     ]
    }
   ],
   "source": [
    "# Test prediction for the model on the fewrel dataset using the offmml-g model with the negative examples.\n",
    "from re import I\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import json\n",
    "from pathlib import Path\n",
    "\n",
    "def hash_tokens(tokens):\n",
    "    return \"\".join(\"\".join(tokens).split()).replace('\"', '').replace(\"'\", \"\").strip().lower()\n",
    "\n",
    "def read_jsonl(file_path):\n",
    "    \"\"\"Read a .jsonl file and yield its contents line by line.\n",
    "    file_path (unicode / Path): The file path.\n",
    "    YIELDS: The loaded JSON contents of each line.\n",
    "    \"\"\"\n",
    "    with Path(file_path).open('r', encoding='utf8') as f:\n",
    "        for line in f:\n",
    "            try:  # hack to handle broken jsonl\n",
    "                yield json.loads(line.strip())\n",
    "            except ValueError:\n",
    "                continue\n",
    "\n",
    "gold_files = {\n",
    "    12321: \"~/codes/QA-ZRE/fewrl_data_unks/test_data_12321.csv\",\n",
    "    943: \"~/codes/QA-ZRE/fewrl_data_unks/test_data_943.csv\",\n",
    "    111: \"~/codes/QA-ZRE/fewrl_data_unks/test_data_111.csv\",\n",
    "    300: \"~/codes/QA-ZRE/fewrl_data_unks/test_data_300.csv\",\n",
    "    1300: \"~/codes/QA-ZRE/fewrl_data_unks/test_data_1300.csv\",\n",
    "}\n",
    "id_files = {\n",
    "    12321: \"~/codes/QA-ZRE/fewrl_data_unks/test_ids_12321.csv\",\n",
    "    943: \"~/codes/QA-ZRE/fewrl_data_unks/test_ids_943.csv\",\n",
    "    111: \"~/codes/QA-ZRE/fewrl_data_unks/test_ids_111.csv\",\n",
    "    300: \"~/codes/QA-ZRE/fewrl_data_unks/test_ids_300.csv\",\n",
    "    1300: \"~/codes/QA-ZRE/fewrl_data_unks/test_ids_1300.csv\",\n",
    "}\n",
    "\n",
    "prediction_files = {\n",
    "    #12321: \"~/sep-1/fewrel/run_12321_with_unks/relation.supp_data.offmml-pgg.run.epoch.0.test.predictions.step.full.csv\",\n",
    "    12321: \"~/sep-1/fewrel/run_12321_with_unks/relation.offmml-pgg.run.epoch.0.test.predictions.step.7700.csv\",\n",
    "    943: \"~/sep-1/fewrel/run_943_with_unks/relation.offmml-pgg.run.epoch.0.test.predictions.step.5900.csv\",\n",
    "    111: \"~/sep-1/fewrel/run_111_with_unks/relation.offmml-pgg.run.epoch.0.test.predictions.step.11000.csv\",\n",
    "    300: \"~/sep-1/fewrel/run_300_with_unks/relation.offmml-pgg.run.epoch.0.test.predictions.step.10700.csv\",\n",
    "    1300: \"~/sep-1/fewrel/run_1300_with_unks/relation.offmml-pgg.run.epoch.0.test.predictions.step.15200.csv\",\n",
    "    #943: \"~/sep-1/fewrel/run_943_with_unks/relation.supp_data.offmml-pgg.run.epoch.0.test.predictions.step..csv\",\n",
    "    #111: \"~/sep-1/fewrel/run_111_with_unks/relation.supp_data.offmml-pgg.run.epoch.0.test.predictions.step..csv\",\n",
    "    #300: \"~/sep-1/fewrel/run_300_with_unks/relation.supp_data.offmml-pgg.run.epoch.0.test.predictions.step..csv\",\n",
    "    #1300: \"~/sep-1/fewrel/run_1300_with_unks/relation.supp_data.offmml-pgg.run.epoch.0.test.predictions.step..csv\",\n",
    "}\n",
    "seeds = [12321, 943, 111, 300, 1300]\n",
    "\n",
    "avg_f1 = 0.0\n",
    "avg_p = 0.0\n",
    "avg_r = 0.0\n",
    "avg_nmi = 0.0\n",
    "avg_b_f1 = 0.0\n",
    "avg_arr = 0.0\n",
    "for seed in seeds:\n",
    "    sent_hash = {}\n",
    "    predictions = [prediction_files[seed]]\n",
    "    max_f1 = 0.0\n",
    "    max_file = None\n",
    "    df = pd.read_csv(gold_files[seed], sep=',')\n",
    "    ids = {val:i for i, val in enumerate(pd.read_csv(id_files[seed], sep=',')[\"relation_ids\"].tolist())}\n",
    "    actual_ids = df[\"actual_ids\"].tolist()\n",
    "    sentences = []\n",
    "    for sent in df[\"passages\"].tolist():\n",
    "        hash = hash_tokens(sent.split())\n",
    "        sent_hash[hash] = sent_hash.get(hash, 0) + 1\n",
    "        sentences.append(hash)\n",
    "    num_examples = len(actual_ids) // 15\n",
    "\n",
    "    gold_indices = []\n",
    "    valid_indices = set()\n",
    "    for index, each_relation_id in enumerate(actual_ids):\n",
    "        if sent_hash[sentences[index]] == 15:\n",
    "            valid_indices.add(index//15)\n",
    "        gold_indices.append(ids[each_relation_id])\n",
    "\n",
    "    print(len(valid_indices))\n",
    "    gold_indices = np.max(np.reshape(np.array(gold_indices), (num_examples, 15)), axis=1)\n",
    "\n",
    "    max_f1 = 0.0\n",
    "    max_file = \"None\"\n",
    "    for prediction_file in predictions:\n",
    "        pred_log_ps = pd.read_csv(prediction_file, sep=',')[\"answer_log_p\"].tolist()\n",
    "        pred_log_ps = np.log(np.mean(np.reshape(np.exp(np.array(pred_log_ps)), (num_examples, 15, 8)), axis=2))\n",
    "        pred_ids = np.argmax(pred_log_ps, axis=1)\n",
    "        #valid_pred_ids = np.take(pred_ids, list(valid_indices))\n",
    "        #valid_gold_ids = np.take(gold_indices, list(valid_indices))\n",
    "        valid_pred_ids = pred_ids\n",
    "        valid_gold_ids = gold_indices\n",
    "        b_f1, nmi, arr = compute_cluster_metrics(valid_gold_ids, valid_pred_ids)\n",
    "        avg_b_f1 += b_f1\n",
    "        avg_nmi += nmi\n",
    "        avg_arr += arr\n",
    "        prec, rec, f1 = compute_macro_PRF(valid_pred_ids, valid_gold_ids)\n",
    "        avg_f1 += f1\n",
    "        avg_p += prec\n",
    "        avg_r += rec\n",
    "        if max_f1 <= f1:\n",
    "            max_f1 = f1\n",
    "            max_file = prediction_file\n",
    "\n",
    "    print(seed, max_file, max_f1)\n",
    "\n",
    "print(\"f:\", avg_f1/5.0)\n",
    "print(\"p:\", avg_p/5.0)\n",
    "print(\"r:\", avg_r/5.0)\n",
    "print(\"b f1:\", avg_b_f1/5.0)\n",
    "print(\"nmi :\", avg_nmi/5.0)\n",
    "print(\"arr:\", avg_arr/5.0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Eval of the RE-QA for relation extraction using the concat model.\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "\n",
    "mean_f1 = 0.0\n",
    "for fold_i in range(1, 11, 1):\n",
    "    gold_file = \"./zero-shot-extraction/relation_splits/dev.{}.concat.relation_data.csv\".format(str(fold_i-1))\n",
    "    gold_indices = []\n",
    "    df = pd.read_csv(gold_file, sep=',')\n",
    "    correct_indices = df[\"correct_indices\"].tolist()\n",
    "    for i, index in enumerate(correct_indices):\n",
    "        if index:\n",
    "            gold_indices.append(int(i % 12))\n",
    "\n",
    "    num_examples = len(correct_indices) // 12\n",
    "    gold_indices = np.array(gold_indices)\n",
    "\n",
    "    max_file = None\n",
    "    max_f1 = 0.0\n",
    "    for checkpoint_i in range(1, 200, 1):\n",
    "        prediction_file = \"~/reqa-predictions/fold_{}/concat/relation.concat.dev.predictions.fold.{}.step.{}.csv\".format(str(fold_i), str(fold_i), str(100 * checkpoint_i))\n",
    "        pred_log_ps = pd.read_csv(prediction_file, sep=',')[\"relation_log_p\"].tolist()\n",
    "        pred_ids = np.argmax(np.reshape(np.array(pred_log_ps), (num_examples, 12)), axis=1)\n",
    "        avg_prec, avg_rec, f1 = compute_macro_PRF(pred_ids, gold_indices)\n",
    "        if f1 >= max_f1:\n",
    "            max_f1 = f1\n",
    "            max_file = prediction_file\n",
    "\n",
    "    print(fold_i)\n",
    "    print(max_f1)\n",
    "    print(max_file)\n",
    "    print(\"\\r\\n\")\n",
    "    mean_f1 += max_f1\n",
    "\n",
    "print(mean_f1/10.0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Test prediction for the RelationPrompt on the RE-QA dataset.\n",
    "from re import I\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "\n",
    "gold_files = {\n",
    "    1: \"~/codes/RelationPrompt/train_reqa_models/fold_1/extractor/pred_in_single.jsonl\",\n",
    "}\n",
    "\n",
    "prediction_arrs = {\n",
    "    1: [\"~/codes/RelationPrompt/train_reqa_models/fold_1/extractor/pred_out_single.jsonl\"]\n",
    "}\n",
    "\n",
    "for fold_id in range(1, 2, 1):\n",
    "    prediction_files = prediction_arrs[fold_id]\n",
    "    df = pd.read_csv(gold_files[fold_id], sep=',')\n",
    "    answers = [ans.replace(\"</s>\", \"\").strip() for ans in df[\"answers\"].tolist()]\n",
    "    all_classes = set(answers)\n",
    "    ids = {val:i for i, val in enumerate(list(all_classes))}\n",
    "    actual_ids = [ids[ans] for ans in answers]\n",
    "    gold_indices = np.array(actual_ids)\n",
    "    for prediction_file in prediction_files:\n",
    "        prediction_ids = []\n",
    "        for pred_class in pd.read_csv(prediction_file, sep=',')[\"predictions_str\"].tolist():\n",
    "            if pred_class.strip() in ids:\n",
    "                prediction_ids.append(ids[pred_class.strip()])\n",
    "            else:\n",
    "                prediction_ids.append(-1)\n",
    "        pred_ids = np.array(prediction_ids)\n",
    "        avg_prec, avg_rec, f1 = compute_macro_PRF(pred_ids, gold_indices)\n",
    "        if f1 > max_f1:\n",
    "            max_f1 = f1\n",
    "            max_file = prediction_file\n",
    "        print(prediction_file, avg_prec, avg_rec, f1)\n",
    "\n",
    "print(max_f1, max_file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1\n",
      "0.6645079529607614\n",
      "~/reqa-predictions/fold_1/concat/relation.concat.dev.predictions.fold.1.step.3600.csv\n",
      "\n",
      "\n",
      "2\n",
      "0.7355692801363015\n",
      "~/reqa-predictions/fold_2/concat/relation.concat.dev.predictions.fold.2.step.4300.csv\n",
      "\n",
      "\n",
      "3\n",
      "0.816466070295921\n",
      "~/reqa-predictions/fold_3/concat/relation.concat.dev.predictions.fold.3.step.5200.csv\n",
      "\n",
      "\n",
      "4\n",
      "0.820538067780762\n",
      "~/reqa-predictions/fold_4/concat/relation.concat.dev.predictions.fold.4.step.1600.csv\n",
      "\n",
      "\n",
      "5\n",
      "0.7970665456384882\n",
      "~/reqa-predictions/fold_5/concat/relation.concat.dev.predictions.fold.5.step.2900.csv\n",
      "\n",
      "\n",
      "6\n",
      "0.9100498471715361\n",
      "~/reqa-predictions/fold_6/concat/relation.concat.dev.predictions.fold.6.step.1400.csv\n",
      "\n",
      "\n",
      "7\n",
      "0.7789862082105365\n",
      "~/reqa-predictions/fold_7/concat/relation.concat.dev.predictions.fold.7.step.2500.csv\n",
      "\n",
      "\n",
      "8\n",
      "0.7585498509710629\n",
      "~/reqa-predictions/fold_8/concat/relation.concat.dev.predictions.fold.8.step.400.csv\n",
      "\n",
      "\n",
      "9\n",
      "0.7690369496615103\n",
      "~/reqa-predictions/fold_9/concat/relation.concat.dev.predictions.fold.9.step.2600.csv\n",
      "\n",
      "\n",
      "10\n",
      "0.7394406760740089\n",
      "~/reqa-predictions/fold_10/concat/relation.concat.dev.predictions.fold.10.step.800.csv\n",
      "\n",
      "\n",
      "0.7790211448900889\n"
     ]
    }
   ],
   "source": [
    "# Eval of the RE-QA for relation extraction using the concat model.\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "\n",
    "mean_f1 = 0.0\n",
    "for fold_i in range(1, 11, 1):\n",
    "    gold_file = \"./zero-shot-extraction/relation_splits/dev.{}.concat.relation_data.csv\".format(str(fold_i-1))\n",
    "    gold_indices = []\n",
    "    df = pd.read_csv(gold_file, sep=',')\n",
    "    correct_indices = df[\"correct_indices\"].tolist()\n",
    "    for i, index in enumerate(correct_indices):\n",
    "        if index:\n",
    "            gold_indices.append(int(i % 12))\n",
    "\n",
    "    num_examples = len(correct_indices) // 12\n",
    "    gold_indices = np.array(gold_indices)\n",
    "\n",
    "    max_file = None\n",
    "    max_f1 = 0.0\n",
    "    for checkpoint_i in range(1, 200, 1):\n",
    "        prediction_file = \"~/reqa-predictions/fold_{}/concat/relation.concat.dev.predictions.fold.{}.step.{}.csv\".format(str(fold_i), str(fold_i), str(100 * checkpoint_i))\n",
    "        pred_log_ps = pd.read_csv(prediction_file, sep=',')[\"relation_log_p\"].tolist()\n",
    "        pred_ids = np.argmax(np.reshape(np.array(pred_log_ps), (num_examples, 12)), axis=1)\n",
    "        avg_prec, avg_rec, f1 = compute_macro_PRF(pred_ids, gold_indices)\n",
    "        if f1 >= max_f1:\n",
    "            max_f1 = f1\n",
    "            max_file = prediction_file\n",
    "\n",
    "    print(fold_i)\n",
    "    print(max_f1)\n",
    "    print(max_file)\n",
    "    print(\"\\r\\n\")\n",
    "    mean_f1 += max_f1\n",
    "\n",
    "print(mean_f1/10.0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "157 1\n",
      "159 1\n",
      "162 1\n",
      "163 1\n",
      "164 1\n",
      "165 1\n",
      "167 1\n",
      "168 1\n",
      "169 1\n",
      "171 1\n",
      "1\n",
      "0.7958029874558785\n",
      "~/reqa-predictions/fold_1/gold/relation.gold.dev.predictions.fold.1.step.600.csv\n",
      "\n",
      "\n",
      "2\n",
      "0.8055480874453981\n",
      "~/reqa-predictions/fold_2/gold/relation.gold.dev.predictions.fold.2.step.1900.csv\n",
      "\n",
      "\n",
      "3\n",
      "0.8088784230714176\n",
      "~/reqa-predictions/fold_3/gold/relation.gold.dev.predictions.fold.3.step.200.csv\n",
      "\n",
      "\n",
      "4\n",
      "0.7950547270773886\n",
      "~/reqa-predictions/fold_4/gold/relation.gold.dev.predictions.fold.4.step.9500.csv\n",
      "\n",
      "\n",
      "5\n",
      "0.8218222460881007\n",
      "~/reqa-predictions/fold_5/gold/relation.gold.dev.predictions.fold.5.step.15300.csv\n",
      "\n",
      "\n",
      "6\n",
      "0.9343882793208368\n",
      "~/reqa-predictions/fold_6/gold/relation.gold.dev.predictions.fold.6.step.1100.csv\n",
      "\n",
      "\n",
      "7\n",
      "0.7977715930389874\n",
      "~/reqa-predictions/fold_7/gold/relation.gold.dev.predictions.fold.7.step.2600.csv\n",
      "\n",
      "\n",
      "8\n",
      "0.8869445616734918\n",
      "~/reqa-predictions/fold_8/gold/relation.gold.dev.predictions.fold.8.step.1000.csv\n",
      "\n",
      "\n",
      "9\n",
      "0.8353253359450881\n",
      "~/reqa-predictions/fold_9/gold/relation.gold.dev.predictions.fold.9.step.1900.csv\n",
      "\n",
      "\n",
      "10\n",
      "0.8742971188094225\n",
      "~/reqa-predictions/fold_10/gold/relation.gold.dev.predictions.fold.10.step.4000.csv\n",
      "\n",
      "\n",
      "0.8355833359926009\n"
     ]
    }
   ],
   "source": [
    "# Eval of the RE-QA using the Gold Templates on the dev data over all the folds.\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "\n",
    "mean_f1 = 0.0\n",
    "for fold_i in range(1, 11, 1):\n",
    "    gold_file = \"./zero-shot-extraction/relation_splits/dev.{}.relation_data.csv\".format(str(fold_i-1))\n",
    "    gold_indices = []\n",
    "    df = pd.read_csv(gold_file, sep=',')\n",
    "    correct_indices = df[\"correct_indices\"].tolist()\n",
    "    for i, index in enumerate(correct_indices):\n",
    "        if index:\n",
    "            gold_indices.append(int(i % 12))\n",
    "\n",
    "    num_examples = len(correct_indices) // 12\n",
    "    gold_indices = np.array(gold_indices)\n",
    "\n",
    "    max_file = None\n",
    "    max_f1 = 0.0\n",
    "    for checkpoint_i in range(1, 200, 1):\n",
    "        try:\n",
    "            prediction_file = \"~/reqa-predictions/fold_{}/gold/relation.gold.dev.predictions.fold.{}.step.{}.csv\".format(str(fold_i), str(fold_i), str(100 * checkpoint_i))\n",
    "            pred_log_ps = pd.read_csv(prediction_file, sep=',')[\"relation_log_p\"].tolist()\n",
    "            pred_ids = np.argmax(np.reshape(np.array(pred_log_ps), (num_examples, 12)), axis=1)\n",
    "            avg_prec, avg_rec, f1 = compute_macro_PRF(pred_ids, gold_indices)\n",
    "            if f1 >= max_f1:\n",
    "                max_f1 = f1\n",
    "                max_file = prediction_file\n",
    "        except:\n",
    "            print(checkpoint_i, fold_i)\n",
    "\n",
    "    print(fold_i)\n",
    "    print(max_f1)\n",
    "    print(max_file)\n",
    "    print(\"\\r\\n\")\n",
    "    mean_f1 += max_f1\n",
    "\n",
    "print(mean_f1/10.0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1\n",
      "0.7021862131500355\n",
      "~/reqa-predictions/fold_1/mml-pgg-off-sim/relation.mml-pgg-off-sim.run.fold_1.dev.predictions.step.4700.csv\n",
      "\n",
      "\n",
      "2\n",
      "0.7318462713898469\n",
      "~/reqa-predictions/fold_2/mml-pgg-off-sim/relation.mml-pgg-off-sim.run.fold_2.dev.predictions.step.400.csv\n",
      "\n",
      "\n",
      "3\n",
      "0.7766533200558716\n",
      "~/reqa-predictions/fold_3/mml-pgg-off-sim/relation.mml-pgg-off-sim.run.fold_3.dev.predictions.step.3600.csv\n",
      "\n",
      "\n",
      "4\n",
      "0.8437707696480834\n",
      "~/reqa-predictions/fold_4/mml-pgg-off-sim/relation.mml-pgg-off-sim.run.fold_4.dev.predictions.step.800.csv\n",
      "\n",
      "\n",
      "5\n",
      "0.8300206299665337\n",
      "~/reqa-predictions/fold_5/mml-pgg-off-sim/relation.mml-pgg-off-sim.run.fold_5.dev.predictions.step.7900.csv\n",
      "\n",
      "\n",
      "6\n",
      "0.8906375171815566\n",
      "~/reqa-predictions/fold_6/mml-pgg-off-sim/relation.mml-pgg-off-sim.run.fold_6.dev.predictions.step.700.csv\n",
      "\n",
      "\n",
      "197\n",
      "198\n",
      "199\n",
      "7\n",
      "0.7827607798234402\n",
      "~/reqa-predictions/fold_7/mml-pgg-off-sim/relation.mml-pgg-off-sim.run.fold_7.dev.predictions.step.2100.csv\n",
      "\n",
      "\n",
      "198\n",
      "199\n",
      "8\n",
      "0.795102231532206\n",
      "~/reqa-predictions/fold_8/mml-pgg-off-sim/relation.mml-pgg-off-sim.run.fold_8.dev.predictions.step.6800.csv\n",
      "\n",
      "\n",
      "185\n",
      "186\n",
      "187\n",
      "188\n",
      "189\n",
      "190\n",
      "191\n",
      "192\n",
      "193\n",
      "194\n",
      "195\n",
      "196\n",
      "197\n",
      "198\n",
      "199\n",
      "9\n",
      "0.7819016572177454\n",
      "~/reqa-predictions/fold_9/mml-pgg-off-sim/relation.mml-pgg-off-sim.run.fold_9.dev.predictions.step.4300.csv\n",
      "\n",
      "\n",
      "10\n",
      "0.7755543602531488\n",
      "~/reqa-predictions/fold_10/mml-pgg-off-sim/relation.mml-pgg-off-sim.run.fold_10.dev.predictions.step.1600.csv\n",
      "\n",
      "\n",
      "0.7910433750218469\n"
     ]
    }
   ],
   "source": [
    "# MML-OFF-PGG performance for Relation Extraction on all the dev folds.\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "\n",
    "mean_f1 = 0.0\n",
    "for fold_i in range(1, 11, 1):\n",
    "    gold_file = \"./zero-shot-extraction/relation_splits/dev.{}.qq.relation_data.csv\".format(str(fold_i-1))\n",
    "    gold_indices = []\n",
    "    df = pd.read_csv(gold_file, sep=',')\n",
    "    correct_indices = df[\"correct_indices\"].tolist()\n",
    "    for i, index in enumerate(correct_indices):\n",
    "        if index:\n",
    "            gold_indices.append(int(i % 12))\n",
    "\n",
    "    num_examples = len(correct_indices) // 12\n",
    "    gold_indices = np.array(gold_indices)\n",
    "\n",
    "    max_file = None\n",
    "    max_f1 = 0.0\n",
    "    for checkpoint_i in range(1, 200, 1):\n",
    "        try:\n",
    "            prediction_file = \"~/reqa-predictions/fold_{}/mml-pgg-off-sim/relation.mml-pgg-off-sim.run.fold_{}.dev.predictions.step.{}.csv\".format(str(fold_i), str(fold_i), str(100 * checkpoint_i))\n",
    "            pred_log_ps = pd.read_csv(prediction_file, sep=',')[\"answer_log_p\"].tolist()\n",
    "            pred_log_ps = np.log(np.mean(np.reshape(np.exp(np.array(pred_log_ps)), (num_examples, 12, 8)), axis=2))\n",
    "            pred_ids = np.argmax(pred_log_ps, axis=1)\n",
    "            avg_prec, avg_rec, f1 = compute_macro_PRF(pred_ids, gold_indices)\n",
    "            if f1 >= max_f1:\n",
    "                max_f1 = f1\n",
    "                max_file = prediction_file\n",
    "        except:\n",
    "            print(checkpoint_i)\n",
    "\n",
    "    print(fold_i)\n",
    "    print(max_f1)\n",
    "    print(max_file)\n",
    "    print(\"\\r\\n\")\n",
    "    mean_f1 += max_f1\n",
    "\n",
    "print(mean_f1/10.0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1 concat alone 0.6868200485458184 0.6926748847477643 0.6810633587961235\n",
      "1 gold alone 0.7404285767936459 0.7531358825049224 0.728142964768295\n",
      "\n",
      "\n",
      "2 concat alone 0.5913023110477719 0.5975156340730365 0.585216877806697\n",
      "2 gold alone 0.6586987557834644 0.6757540315314027 0.6424831986356715\n",
      "\n",
      "\n",
      "3 concat alone 0.6324106523285894 0.6578301488771372 0.60888256050775\n",
      "3 gold alone 0.676931710052006 0.6806536959107133 0.6732502083282599\n",
      "\n",
      "\n",
      "4 concat alone 0.5794697498339418 0.6063706954690022 0.5548542716952589\n",
      "4 gold alone 0.6728400404740734 0.6895175863911674 0.656950210374016\n",
      "\n",
      "\n",
      "5 concat alone 0.6191377271586156 0.6268067958412001 0.6116540543838874\n",
      "5 gold alone 0.49349253129168735 0.4986497121305824 0.48844093262952004\n",
      "\n",
      "\n",
      "6 concat alone 0.5782347137534999 0.6001076131310151 0.5579002025330327\n",
      "6 gold alone 0.6607776835391415 0.6691998506324075 0.6525648747509227\n",
      "\n",
      "\n",
      "7 concat alone 0.6735910798575685 0.6934384410876097 0.6548482345074617\n",
      "7 gold alone 0.6474800154241827 0.6618318212248525 0.633737435642919\n",
      "\n",
      "\n",
      "8 concat alone 0.6046606629785962 0.6215315803731405 0.58868143336638\n",
      "8 gold alone 0.7714364531887481 0.780457536718944 0.7626215309046493\n",
      "\n",
      "\n",
      "9 concat alone 0.6017628265979826 0.6153436850776661 0.5887684922695174\n",
      "9 gold alone 0.6243678988994422 0.6462041634257861 0.6039591600991271\n",
      "\n",
      "\n",
      "10 concat alone 0.6167540970127698 0.6150992932792283 0.6184178286132765\n",
      "10 gold alone 0.6636083208337981 0.6717276147811493 0.6556829615307239\n",
      "\n",
      "\n",
      "gold alone p: 0.6727131895251928\n",
      "gold alone r: 0.6497833477664104\n",
      "gold alone f: 0.661006198628019\n",
      "concat alone p: 0.6326718771956801\n",
      "concat alone r: 0.6050287314479386\n",
      "concat alone f: 0.6184143869115154\n"
     ]
    }
   ],
   "source": [
    "# Test set performance over the 10 folds of the RE-QA dataset for the concat and gold models.\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import json\n",
    "\n",
    "gold_files = {\n",
    "    1: \"relation.gold.test.predictions.fold.1.step.600.csv\",\n",
    "    2: \"relation.gold.test.predictions.fold.2.step.1900.csv\",\n",
    "    3: \"relation.gold.test.predictions.fold.3.step.200.csv\",\n",
    "    4: \"relation.gold.test.predictions.fold.4.step.9500.csv\",\n",
    "    5: \"relation.gold.test.predictions.fold.5.step.15300.csv\",\n",
    "    6: \"relation.gold.test.predictions.fold.6.step.1100.csv\",\n",
    "    7: \"relation.gold.test.predictions.fold.7.step.2600.csv\",\n",
    "    8: \"relation.gold.test.predictions.fold.8.step.1000.csv\",\n",
    "    9: \"relation.gold.test.predictions.fold.9.step.1900.csv\",\n",
    "    10: \"relation.gold.test.predictions.fold.10.step.4000.csv\"\n",
    "}\n",
    "\n",
    "concat_files = {\n",
    "    1: \"relation.concat.test.predictions.fold.1.step.3600.csv\",\n",
    "    2: \"relation.concat.test.predictions.fold.2.step.4300.csv\",\n",
    "    3: \"relation.concat.test.predictions.fold.3.step.5200.csv\",\n",
    "    4: \"relation.concat.test.predictions.fold.4.step.1600.csv\",\n",
    "    5: \"relation.concat.test.predictions.fold.5.step.2900.csv\",\n",
    "    6: \"relation.concat.test.predictions.fold.6.step.1400.csv\",\n",
    "    7: \"relation.concat.test.predictions.fold.7.step.2500.csv\",\n",
    "    8: \"relation.concat.test.predictions.fold.8.step.400.csv\",\n",
    "    9: \"relation.concat.test.predictions.fold.9.step.2600.csv\",\n",
    "    10: \"relation.concat.test.predictions.fold.10.step.800.csv\"\n",
    "}\n",
    "\n",
    "gold_alone_p_r_f = {'f': [], 'r': [], 'p': []}\n",
    "concat_alone_p_r_f = {'f': [], 'r': [], 'p': []}\n",
    "\n",
    "for fold_i in range(1, 11, 1):\n",
    "    gold_file = \"./zero-shot-extraction/relation_splits/test.{}.relation_data.csv\".format(str(fold_i-1))\n",
    "    fewrel_file = \"./zero-shot-extraction/relation_splits/test.{}.fewrel_format.json\".format(str(fold_i-1))\n",
    "\n",
    "    example_indices_to_consider = set()\n",
    "    with open(fewrel_file, 'r') as fin:\n",
    "        short_examples = json.load(fin)\n",
    "        for key, val in short_examples.items():\n",
    "            for row in val:\n",
    "                example_indices_to_consider.add(row[\"example_index\"])\n",
    "\n",
    "\n",
    "    gold_indices = []\n",
    "    df = pd.read_csv(gold_file, sep=',')\n",
    "    correct_indices = df[\"correct_indices\"].tolist()\n",
    "    for i, index in enumerate(correct_indices):\n",
    "        if index:\n",
    "            gold_indices.append(int(i % 24))\n",
    "\n",
    "    gold_indices_to_consider = []\n",
    "    for i, index in enumerate(gold_indices):\n",
    "        if i in example_indices_to_consider:\n",
    "            gold_indices_to_consider.append(index)\n",
    "\n",
    "    num_examples = len(correct_indices) // 24\n",
    "    gold_indices_to_consider = np.array(gold_indices_to_consider)\n",
    "    \n",
    "    concat_prediction_file = \"~/may-20/fold_{}/concat/{}\".format(fold_i, concat_files[fold_i])\n",
    "    concat_pred_log_ps = pd.read_csv(concat_prediction_file, sep=',')[\"relation_log_p\"].tolist()\n",
    "    concat_pred_log_ps = np.reshape(np.array(concat_pred_log_ps), (num_examples, 24))\n",
    "    concat_pred_ids = np.argmax(concat_pred_log_ps, axis=1)\n",
    "    concat_pred_ids_to_consider = []\n",
    "    for i, index in enumerate(concat_pred_ids):\n",
    "        if i in example_indices_to_consider:\n",
    "            concat_pred_ids_to_consider.append(index)\n",
    "    \n",
    "    concat_pred_ids_to_consider = np.array(concat_pred_ids_to_consider)\n",
    "    avg_prec, avg_rec, f1 = compute_macro_PRF(concat_pred_ids_to_consider, gold_indices_to_consider)\n",
    "    concat_alone_p_r_f[\"f\"].append(f1)\n",
    "    concat_alone_p_r_f[\"r\"].append(avg_rec)\n",
    "    concat_alone_p_r_f[\"p\"].append(avg_prec)\n",
    "    print(fold_i, \"concat alone\", f1, avg_prec, avg_rec)\n",
    "    \n",
    "    gold_prediction_file = \"~/may-20/fold_{}/gold/{}\".format(fold_i, gold_files[fold_i])\n",
    "    gold_pred_log_ps = pd.read_csv(gold_prediction_file, sep=',')[\"relation_log_p\"].tolist()\n",
    "    gold_pred_log_ps = np.reshape(np.array(gold_pred_log_ps), (num_examples, 24))\n",
    "    gold_pred_ids = np.argmax(gold_pred_log_ps, axis=1)\n",
    "\n",
    "    gold_pred_ids_to_consider = []\n",
    "    for i, index in enumerate(gold_pred_ids):\n",
    "        if i in example_indices_to_consider:\n",
    "            gold_pred_ids_to_consider.append(index)\n",
    "    \n",
    "    gold_pred_ids_to_consider = np.array(gold_pred_ids_to_consider)\n",
    "    avg_prec, avg_rec, f1 = compute_macro_PRF(gold_pred_ids_to_consider, gold_indices_to_consider)\n",
    "    print(fold_i, \"gold alone\", f1, avg_prec, avg_rec)\n",
    "    gold_alone_p_r_f[\"f\"].append(f1)\n",
    "    gold_alone_p_r_f[\"r\"].append(avg_rec)\n",
    "    gold_alone_p_r_f[\"p\"].append(avg_prec)\n",
    "    print(\"\\n\")\n",
    "\n",
    "print(\"gold alone p:\", np.mean(np.array(gold_alone_p_r_f[\"p\"])))\n",
    "print(\"gold alone r:\", np.mean(np.array(gold_alone_p_r_f[\"r\"])))\n",
    "print(\"gold alone f:\", np.mean(np.array(gold_alone_p_r_f[\"f\"])))\n",
    "\n",
    "print(\"concat alone p:\", np.mean(np.array(concat_alone_p_r_f[\"p\"])))\n",
    "print(\"concat alone r:\", np.mean(np.array(concat_alone_p_r_f[\"r\"])))\n",
    "print(\"concat alone f:\", np.mean(np.array(concat_alone_p_r_f[\"f\"])))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1 mml 0.7324313594407584 0.7347101443139104 0.7301666666666667\n",
      "\n",
      "\n"
     ]
    },
    {
     "ename": "ValueError",
     "evalue": "cannot reshape array of size 692139 into shape (6000,24,8)",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mValueError\u001b[0m                                Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-4-6217db65b1f8>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m     32\u001b[0m     \u001b[0mmml_pred_log_ps\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mpd\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mread_csv\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmml_prediction_file\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0msep\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\"answer_log_p\"\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtolist\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     33\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 34\u001b[0;31m     \u001b[0mpred_log_ps\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mnp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mlog\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mnp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmean\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mnp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mreshape\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mnp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mexp\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mnp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0marray\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmml_pred_log_ps\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0mnum_examples\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m24\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m8\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0maxis\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m2\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[0m\u001b[1;32m     35\u001b[0m     \u001b[0mpred_ids\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mnp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0margmax\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mpred_log_ps\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0maxis\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     36\u001b[0m     \u001b[0mavg_prec\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mavg_rec\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mf1\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mcompute_macro_PRF\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mpred_ids\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mgold_indices\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m<__array_function__ internals>\u001b[0m in \u001b[0;36mreshape\u001b[0;34m(*args, **kwargs)\u001b[0m\n",
      "\u001b[0;32m~/.local/lib/python3.6/site-packages/numpy/core/fromnumeric.py\u001b[0m in \u001b[0;36mreshape\u001b[0;34m(a, newshape, order)\u001b[0m\n\u001b[1;32m    297\u001b[0m            [5, 6]])\n\u001b[1;32m    298\u001b[0m     \"\"\"\n\u001b[0;32m--> 299\u001b[0;31m     \u001b[0;32mreturn\u001b[0m \u001b[0m_wrapfunc\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0ma\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m'reshape'\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnewshape\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0morder\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0morder\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    300\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    301\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.local/lib/python3.6/site-packages/numpy/core/fromnumeric.py\u001b[0m in \u001b[0;36m_wrapfunc\u001b[0;34m(obj, method, *args, **kwds)\u001b[0m\n\u001b[1;32m     56\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     57\u001b[0m     \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 58\u001b[0;31m         \u001b[0;32mreturn\u001b[0m \u001b[0mbound\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwds\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     59\u001b[0m     \u001b[0;32mexcept\u001b[0m \u001b[0mTypeError\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     60\u001b[0m         \u001b[0;31m# A TypeError occurs if the object does have such a method in its\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mValueError\u001b[0m: cannot reshape array of size 692139 into shape (6000,24,8)"
     ]
    }
   ],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "\n",
    "mml_files = {\n",
    "    1: \"relation.mml-pgg-off-sim.run.fold_1.test.predictions.step.4700.csv\",\n",
    "    2: \"relation.mml-pgg-off-sim.run.fold_2.test.predictions.step.400.csv\",\n",
    "    3: \"relation.mml-pgg-off-sim.run.fold_3.test.predictions.step.3600.csv\",\n",
    "    4: \"relation.mml-pgg-off-sim.run.fold_4.test.predictions.step.800.csv\",\n",
    "    5: \"relation.mml-pgg-off-sim.run.fold_5.test.predictions.step.7900.csv\",\n",
    "    6: \"relation.mml-pgg-off-sim.run.fold_6.test.predictions.step.700.csv\",\n",
    "    7: \"relation.mml-pgg-off-sim.run.fold_7.test.predictions.step.2100.csv\",\n",
    "    8: \"relation.mml-pgg-off-sim.run.fold_8.test.predictions.step.6800.csv\",\n",
    "    9: \"relation.mml-pgg-off-sim.run.fold_9.test.predictions.step.4300.csv\",\n",
    "    10: \"relation.mml-pgg-off-sim.run.fold_10.test.predictions.step.1600.csv\"\n",
    "}\n",
    "\n",
    "mml_pgg_p_r_f = {'f': [], 'r': [], 'p': []}\n",
    "\n",
    "for fold_i in range(1, 11, 1):\n",
    "    gold_file = \"./zero-shot-extraction/relation_splits/test.{}.qq.relation_data.csv\".format(str(fold_i-1))\n",
    "    gold_indices = []\n",
    "    df = pd.read_csv(gold_file, sep=',')\n",
    "    correct_indices = df[\"correct_indices\"].tolist()\n",
    "    for i, index in enumerate(correct_indices):\n",
    "        if index:\n",
    "            gold_indices.append(int(i % 24))\n",
    "\n",
    "    num_examples = len(correct_indices) // 24\n",
    "    gold_indices = np.array(gold_indices)\n",
    "    \n",
    "    mml_prediction_file = \"~/may-20/fold_{}/{}\".format(fold_i, mml_files[fold_i])\n",
    "    mml_pred_log_ps = pd.read_csv(mml_prediction_file, sep=',')[\"answer_log_p\"].tolist()\n",
    "\n",
    "    pred_log_ps = np.log(np.mean(np.reshape(np.exp(np.array(mml_pred_log_ps)), (num_examples, 24, 8)), axis=2))\n",
    "    pred_ids = np.argmax(pred_log_ps, axis=1)\n",
    "    avg_prec, avg_rec, f1 = compute_macro_PRF(pred_ids, gold_indices)\n",
    "    mml_pgg_p_r_f[\"f\"].append(f1)\n",
    "    mml_pgg_p_r_f[\"r\"].append(avg_rec)\n",
    "    mml_pgg_p_r_f[\"p\"].append(avg_prec)\n",
    "    print(fold_i, \"mml pgg off\", f1, avg_prec, avg_rec)\n",
    "    print(\"\\n\")\n",
    "\n",
    "print(\"mml pgg off p:\", np.mean(np.array(mml_pgg_p_r_f[\"p\"])))\n",
    "print(\"mml pgg off r:\", np.mean(np.array(mml_pgg_p_r_f[\"r\"])))\n",
    "print(\"mml pgg off f:\", np.mean(np.array(mml_pgg_p_r_f[\"f\"])))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.6.9 64-bit",
   "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.6.9"
  },
  "vscode": {
   "interpreter": {
    "hash": "31f2aee4e71d21fbe5cf8b01ff0e069b9275f58929596ceb00d14d90e3e16cd6"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}