mrr-eval-from-scores.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "f37bd916",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pickle\n",
    "import os\n",
    "from tqdm import tqdm \n",
    "from typing import Dict\n",
    "from collections import defaultdict\n",
    "import numpy as np\n",
    "from pymongo import MongoClient"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "de252fac",
   "metadata": {},
   "outputs": [],
   "source": [
    "def softmax(x):\n",
    "    \"\"\"Compute softmax values for each sets of scores in x.\"\"\"\n",
    "    return np.exp(x) / np.sum(np.exp(x), axis=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "ce535663",
   "metadata": {},
   "outputs": [],
   "source": [
    "entities = []\n",
    "id2ent = {}\n",
    "with open(\"wd5m_aliases_entities_v3.txt\", 'r') as f:\n",
    "    for line in f:\n",
    "        entities.append(line.split('\\t')[1][:-1])\n",
    "        id_, ent = line.split('\\t')\n",
    "        id2ent[id_] = ent[:-1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "af06e662",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "('Wikimedia list article', 'genus')"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "id2ent['Q13406463'], id2ent['Q34740']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "65da8184",
   "metadata": {},
   "outputs": [],
   "source": [
    "fname = 'scores/baseline_model.pickle'\n",
    "# fname = 'scores_500_base_trie.pickle'\n",
    "scores_data = pickle.load(open(fname, 'rb'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "4936372f",
   "metadata": {},
   "outputs": [],
   "source": [
    "entity_strings_set = set(entities)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "95842d8a",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Exception ignored in: <function tqdm.__del__ at 0x7f570032a790>\n",
      "Traceback (most recent call last):\n",
      "  File \"/data/home/admin/t5-experiments-from-yura/venv/lib/python3.8/site-packages/tqdm/std.py\", line 1162, in __del__\n",
      "    self.close()\n",
      "  File \"/data/home/admin/t5-experiments-from-yura/venv/lib/python3.8/site-packages/tqdm/notebook.py\", line 288, in close\n",
      "    self.disp(bar_style='danger', check_delay=False)\n",
      "AttributeError: 'tqdm_notebook' object has no attribute 'disp'\n"
     ]
    }
   ],
   "source": [
    "def create_filter_dict(collection_name) -> Dict[str, int]:\n",
    "    \n",
    "    client = MongoClient('localhost', 27017)\n",
    "    collection = client['KGLM'][collection_name]\n",
    "    \n",
    "    filter_dict = defaultdict(list)\n",
    "#     tot = client['KGLM'].command(\"collstats\", collection_name)['count']\n",
    "#     for doc in tqdm(collection.find({}), total=tot):\n",
    "    for doc in collection.find({}):\n",
    "        inp = id2ent[doc['head']]\n",
    "        filter_dict[inp].append(doc['tail'])\n",
    "    return filter_dict\n",
    "\n",
    "def getAllFilteringEntities(inp, filter_dicts):\n",
    "    entities = []\n",
    "    splits = ['train', 'test', 'valid']\n",
    "    for s in splits:\n",
    "        entities.extend(filter_dicts[s][inp])\n",
    "    return list(set(entities))\n",
    "\n",
    "filter_dicts = {}\n",
    "splits = ['train', 'valid', 'test']\n",
    "for split in splits:\n",
    "    filter_dicts[split] = create_filter_dict(split)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "546f8cc6",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "5133it [00:00, 6603.81it/s]\n"
     ]
    }
   ],
   "source": [
    "predictions_scores_dicts = []\n",
    "for string_arr, score_arr in tqdm(zip(scores_data['prediction_strings'], scores_data['scores'])):\n",
    "    ps_pairs = [(p,s) for p,s in zip(string_arr, score_arr)]\n",
    "    ps_pairs = list(set(ps_pairs)) # while sampling, duplicates are created\n",
    "    # remove predictions that are not entities\n",
    "    ps_dict_only_entities = defaultdict(list)\n",
    "    for ps in ps_pairs:\n",
    "        if ps[0] in entity_strings_set:\n",
    "            ps_dict_only_entities[ps[0]] = ps[1]\n",
    "    predictions_scores_dicts.append(ps_dict_only_entities)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "58f99693",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 5133/5133 [00:00<00:00, 10286.67it/s]\n"
     ]
    }
   ],
   "source": [
    "predictions_filtered = []\n",
    "head_num_filter = 0\n",
    "tail_num_filter = 0\n",
    "\n",
    "for i in tqdm(range(len(predictions_scores_dicts))):\n",
    "    ps_dict = predictions_scores_dicts[i].copy()\n",
    "    target = scores_data['target_strings'][i]\n",
    "    inputs = scores_data['input_strings'][i]\n",
    "    prediction_strings = ps_dict.keys()\n",
    "    \n",
    "    if target in prediction_strings:\n",
    "        original_score = ps_dict[target]\n",
    "    # get filtering entities\n",
    "    filtering_entities = getAllFilteringEntities(inputs, filter_dicts)\n",
    "\n",
    "    for ent in filtering_entities:\n",
    "        if ent in ps_dict:\n",
    "            ps_dict[ent] = -float(\"inf\")\n",
    "    if target in prediction_strings:\n",
    "        ps_dict[target] = original_score\n",
    "    # softmax for scores\n",
    "    names_arr = []\n",
    "    scores_arr = []\n",
    "    for k, v in ps_dict.items():\n",
    "        names_arr.append(k)\n",
    "        scores_arr.append(v)\n",
    "    scores_arr = np.array(scores_arr)\n",
    "#     scores_arr = softmax(scores_arr)\n",
    "    for name, score in zip(names_arr, scores_arr):\n",
    "        ps_dict[name] = score\n",
    "    predictions_filtered.append(ps_dict)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "9d04c870",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 5133/5133 [00:00<00:00, 62631.05it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "hits@1 0.11805961426066627\n",
      "hits@3 0.19286966686148452\n",
      "hits@10 0.2577440093512566\n",
      "mrr 0.1655579625790256\n",
      "0.13208649912331968 were <10 length preds array without answer\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "count = {}\n",
    "reciprocal_ranks = 0.0\n",
    "k_list = [1, 3, 10]\n",
    "\n",
    "for k in k_list:\n",
    "    count[k] = 0\n",
    "    \n",
    "num_small_arrs = 0\n",
    "total_count = 0\n",
    "\n",
    "for i in tqdm(range(len(predictions_filtered))):\n",
    "    target = scores_data['target_strings'][i]\n",
    "    ps_dict = predictions_filtered[i]\n",
    "    ps_sorted = sorted(ps_dict.items(), key=lambda item: -item[1])\n",
    "    inputs = scores_data['input_strings'][i]\n",
    "\n",
    "    if len(ps_dict) == 0:\n",
    "        preds = []\n",
    "    else:\n",
    "        preds = [x[0] for x in ps_sorted]\n",
    "    if len(filtering_entities) > 1 and target in preds[:1]:\n",
    "        print(i)\n",
    "    if target in preds:\n",
    "        rank = preds.index(target) + 1\n",
    "        reciprocal_ranks += 1./rank\n",
    "    for k in k_list:\n",
    "        if target in preds[:k]:\n",
    "            count[k] += 1\n",
    "    if len(preds) < 10 and target not in preds:\n",
    "        num_small_arrs += 1\n",
    "        \n",
    "total_count = len(predictions_filtered)\n",
    "for k in k_list:\n",
    "    hits_at_k = count[k]/total_count\n",
    "    print('hits@{}'.format(k), hits_at_k)\n",
    "print('mrr', reciprocal_ranks/total_count)\n",
    "print(num_small_arrs/total_count, 'were <10 length preds array without answer')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "0adddd91",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 5133/5133 [00:00<00:00, 398308.34it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "hits@1 0.0\n",
      "hits@3 0.0\n",
      "hits@10 0.0\n",
      "mrr 0.0\n",
      "1.0 were <10 length preds array without answer\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "count = {}\n",
    "reciprocal_ranks = 0.0\n",
    "k_list = [1, 3, 10]\n",
    "\n",
    "for k in k_list:\n",
    "    count[k] = 0\n",
    "    \n",
    "num_small_arrs = 0\n",
    "total_count = 0\n",
    "\n",
    "for i in tqdm(range(len(predictions_filtered))):\n",
    "    target = scores_data['target_strings'][i]\n",
    "    ps_dict = predictions_filtered[i]\n",
    "    ps_sorted = sorted(ps_dict.items(), key=lambda item: -item[1])\n",
    "    inputs = scores_data['input_strings'][i]\n",
    "\n",
    "    if len(ps_dict) == 0:\n",
    "        preds = []\n",
    "    else:\n",
    "        preds = [x[0] for x in ps_sorted]\n",
    "    if len(filtering_entities) > 1 and target in preds[:1]:\n",
    "        print(i)\n",
    "    if target in preds:\n",
    "        rank = preds.index(target) + 1\n",
    "        reciprocal_ranks += 1./rank\n",
    "    for k in k_list:\n",
    "        if target in preds[:k]:\n",
    "            count[k] += 1\n",
    "    if len(preds) < 10 and target not in preds:\n",
    "        num_small_arrs += 1\n",
    "        \n",
    "total_count = len(predictions_filtered)\n",
    "for k in k_list:\n",
    "    hits_at_k = count[k]/total_count\n",
    "    print('hits@{}'.format(k), hits_at_k)\n",
    "print('mrr', reciprocal_ranks/total_count)\n",
    "print(num_small_arrs/total_count, 'were <10 length preds array without answer')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d7e3606a",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "venv",
   "language": "python",
   "name": "venv"
  },
  "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.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}