kubernetes_EARec.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "kubernetes_EARec.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/qwiksilva/cs224w-github-rec/blob/master/kubernetes_EARec.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "NvL9QnTMV2yp",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/gdrive')\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "STwRc7AmuyMj",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import pandas as pd\n",
        "import numpy as np\n",
        "import networkx as nx\n",
        "from collections import Counter\n",
        "from datetime import datetime\n",
        "import time\n",
        "from tqdm.auto import tqdm\n",
        "from collections import defaultdict\n",
        "from scipy import sparse\n",
        "import time\n",
        "import functools\n",
        "import numpy as np\n",
        "from scipy.optimize import fmin_bfgs, fmin_l_bfgs_b\n",
        "import os\n",
        "import array\n",
        "from datetime import datetime as dt\n",
        "from pprint import pprint\n",
        "import pytz\n",
        "import random\n",
        "import joblib\n",
        "import itertools\n",
        "import matplotlib.pyplot as plt\n",
        "from matplotlib.ticker import FormatStrFormatter\n",
        "\n",
        "home = \"/gdrive/My Drive/Colab Notebooks/cs224w-data/final-data-12062019\"\n",
        "os.listdir(home)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "AKH49QaWWVNL",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def load_data(suffix=\"v4_2018_only\", path=\"/gdrive/My Drive/Colab Notebooks/cs224w-data/final-data-12062019\"):\n",
        "  \"Load the required files\"\n",
        "\n",
        "  # train_srw_output = joblib.load(os.path.join(path, f\"train_srw_output_{suffix}.joblib\"))\n",
        "  # train_node_mapping = joblib.load(os.path.join(path, f\"train_node_mapping_{suffix}.joblib\"))\n",
        "  # test_srw_output= joblib.load(os.path.join(path, f\"test_srw_output_{suffix}.joblib\"))\n",
        "  # test_node_mapping = joblib.load(os.path.join(path, f\"test_node_mapping_{suffix}.joblib\"))\n",
        "  test_cosine_sim_matrix = joblib.load(os.path.join(path, f\"test_cosine_sim_matrix_{suffix}.joblib\"))\n",
        "  test_pr_cosine_index_map = joblib.load(os.path.join(path, f\"test_pr_cosine_index_map_{suffix}.joblib\"))\n",
        "  train_pr_cosine_index_map = joblib.load(os.path.join(path, f\"train_pr_cosine_index_map_{suffix}.joblib\"))\n",
        "\n",
        "  dtypes = {\n",
        "    'comment': str,\n",
        "    'commenter_id': int,\n",
        "    'commenter_login': str, \n",
        "    'pr_id': int,\n",
        "    'user_id': int,\n",
        "    'username':str,\n",
        "    'title': str,\n",
        "    'description': str,\n",
        "    'pr_created_at': str,\n",
        "    'comment_created_at': str,\n",
        "  }\n",
        "\n",
        "  train_comments = joblib.load(os.path.join(path, 'train_comments_df_v4_2018_only.joblib'))\n",
        "  test_comments = joblib.load(os.path.join(path, 'train_comments_df_v4_2018_only.joblib'))\n",
        "\n",
        "  comments = pd.concat([train_comments,test_comments])\n",
        "  comments = comments.reset_index()\n",
        "\n",
        "  data = {\n",
        "      'test_cosine_sim_matrix':test_cosine_sim_matrix,\n",
        "      'test_pr_cosine_index_map':test_pr_cosine_index_map,\n",
        "      'train_pr_cosine_index_map':train_pr_cosine_index_map,\n",
        "      'train_comments':train_comments,\n",
        "      'comments':comments\n",
        "  }\n",
        "  return data\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "NVMFMYY7u48W",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def convert_to_tic(s):\n",
        "    #Jan, 1, 2014\n",
        "    start_time = time.mktime(datetime(2018, 1, 1).timetuple())\n",
        "    #Jan, 1, 2019\n",
        "    end_time = time.mktime(datetime(2019, 1, 1).timetuple())\n",
        "    time_delta = end_time-start_time\n",
        "    return float(time.mktime(pd.to_datetime(s).timetuple()) - start_time) / time_delta\n",
        "\n",
        "def path2List(fileString):\n",
        "    return fileString.split(\"/\")\n",
        "\n",
        "def LCP(f1,f2):\n",
        "    f1 = path2List(f1)\n",
        "    f2 = path2List(f2)\n",
        "    common_path = 0\n",
        "    min_length = min(len(f1),len(f2))\n",
        "    for i in range(min_length):\n",
        "        if f1[i] == f2[i]:\n",
        "            common_path += 1\n",
        "        else:\n",
        "            break\n",
        "    common_path = common_path / max(len(f1), len(f2))\n",
        "    return common_path\n",
        "\n",
        "def get_pr_similarities(pr_obj1, pr_obj2, file_similarities):\n",
        "    similarity = 0.0\n",
        "    for filename1 in pr_obj1:\n",
        "        for filename2 in pr_obj2: \n",
        "            if (filename1, filename2) in file_similarities:\n",
        "                similarity += file_similarities[(filename1, filename2)]\n",
        "            else:\n",
        "                file_similarity = LCP(filename1, filename2)\n",
        "                file_similarities[(filename1, filename2)] = file_similarity\n",
        "                similarity += file_similarity\n",
        "\n",
        "    similarity = similarity / (len(pr_obj1)*len(pr_obj2))\n",
        "    return similarity\n",
        "\n",
        "def create_user_to_user_graph(graph, data, pr_id='pr_id', commenter_id='commenter_id', owner_id='user_id'):\n",
        "    print(\"converting times to tics\")\n",
        "    comment_times = data['comment_created_at'].map(convert_to_tic)\n",
        "    data['weight'] = comment_times\n",
        "    data = data.sort_values(by='weight')\n",
        "    decay = 0.8\n",
        "\n",
        "    print(f\"finished times to tics, starting graph creation on {len(data)} rows\")\n",
        "    counts = Counter()\n",
        "    for index, row in data.iterrows():\n",
        "        if row[commenter_id] == row[owner_id]:\n",
        "            continue\n",
        "            \n",
        "        row['weight'] *= decay**counts[(row[pr_id], row[commenter_id])]\n",
        "        counts[(row[pr_id], row[commenter_id])] += 1\n",
        "\n",
        "    print(\"finished counting, starting weight calculations\")\n",
        "    edge_weights = Counter()\n",
        "    for index, row in data.iterrows():\n",
        "        if row[commenter_id] == row[owner_id]:\n",
        "              continue\n",
        "              \n",
        "        edge_weights[(row[commenter_id], row[owner_id])] += row['weight']\n",
        "\n",
        "    print(\"finished weights, creating graph\")\n",
        "    for edge in edge_weights:\n",
        "        graph.add_edge(edge[1], edge[0], weight=edge_weights[edge])\n",
        "\n",
        "    return graph\n",
        "\n",
        "def add_pr(graph, new_pr_id, pr_to_files, pr_to_users, comment_counts, file_similarities={}):\n",
        "    if int(new_pr_id) in graph.nodes:\n",
        "        print(f\"Warning: found node already in graph with id {new_pr_id}\")\n",
        "        return graph\n",
        "    \n",
        "    edge_weights = Counter()\n",
        "    new_pr_files = pr_to_files[new_pr_id]\n",
        "    for pr_id in pr_to_users: \n",
        "        if pr_id == new_pr_id:\n",
        "            continue\n",
        "            \n",
        "        pr_files = pr_to_files[pr_id]\n",
        "        similarity = get_pr_similarities(new_pr_files, pr_files, file_similarities)\n",
        "        for user_id in pr_to_users[pr_id]:\n",
        "            edge_weights[user_id] += comment_counts[user_id][pr_id] * similarity\n",
        "            \n",
        "    for user_id in edge_weights:\n",
        "        weight = edge_weights[user_id] / sum(comment_counts[user_id].values())\n",
        "        graph.add_edge(new_pr_id, user_id, weight=weight)\n",
        "\n",
        "    return graph\n",
        "\n",
        "def add_pr_precomputed(graph, new_pr_id, new_pr_idx, pr_to_users, comment_counts, pr_similarities, train_pr_cosine_index_map):\n",
        "    # graph.add_node(new_pr_id)\n",
        "    \n",
        "    edge_weights = Counter()\n",
        "    for pr_id in pr_to_users:\n",
        "            \n",
        "        similarity = pr_similarities[new_pr_idx][train_pr_cosine_index_map[pr_id]]\n",
        "        for user_id in pr_to_users[pr_id]:\n",
        "            edge_weights[user_id] += comment_counts[user_id][pr_id] * similarity\n",
        "            \n",
        "    for user_id in edge_weights:\n",
        "        weight = edge_weights[user_id] / sum(comment_counts[user_id].values())\n",
        "        edge_weights[user_id] = weight\n",
        "\n",
        "    # edges = sorted([edge for edge in graph.out_edges(new_pr_id, data='weight')], key=lambda x:x[2], reverse=True)\n",
        "    # print(\"max edges\", edges[:5])\n",
        "\n",
        "    return edge_weights\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "GJyyAf9EveEL",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def setup_add_pr(train_comments, file_paths, commenter_id='commenter_id'):\n",
        "    t = time.time()\n",
        "    grouped_comments = train_comments.groupby([commenter_id, 'pr_id']).size()\n",
        "    grouped_files = file_paths.groupby(['pr_number'])\n",
        "\n",
        "    comment_counts = {}\n",
        "    for commenter_group in grouped_comments.index:\n",
        "        if commenter_group[0] not in comment_counts:\n",
        "            comment_counts[commenter_group[0]] = Counter()\n",
        "        comment_counts[commenter_group[0]][commenter_group[1]] = grouped_comments[commenter_group]\n",
        "\n",
        "    pr_to_users = {}\n",
        "    grouped_prs = train_comments.groupby(['pr_id'])\n",
        "    for pr_id, group in grouped_prs:\n",
        "        pr_to_users[pr_id] = group[commenter_id].unique()\n",
        "\n",
        "    pr_to_files = {}\n",
        "    for name, group in grouped_files:\n",
        "        pr_to_files[name] = group.filename.values\n",
        "\n",
        "    elapsed = time.time() - t\n",
        "    print(f\"setup took {elapsed} sec\")\n",
        "\n",
        "    return comment_counts, pr_to_users, pr_to_files\n",
        "\n",
        "def setup_add_pr_precomputed(train_comments, commenter_id='commenter_id'):\n",
        "    t = time.time()\n",
        "    grouped_comments = train_comments.groupby([commenter_id, 'pr_id']).size()\n",
        "    # grouped_files = file_paths.groupby(['pr_number'])\n",
        "\n",
        "    comment_counts = {}\n",
        "    for commenter_group in grouped_comments.index:\n",
        "        if commenter_group[0] not in comment_counts:\n",
        "            comment_counts[commenter_group[0]] = Counter()\n",
        "        comment_counts[commenter_group[0]][commenter_group[1]] = grouped_comments[commenter_group]\n",
        "\n",
        "    pr_to_users = {}\n",
        "    grouped_prs = train_comments.groupby(['pr_id'])\n",
        "    for pr_id, group in grouped_prs:\n",
        "        pr_to_users[pr_id] = group[commenter_id].unique()\n",
        "\n",
        "    # pr_to_files = {}\n",
        "    # for name, group in grouped_files:\n",
        "    #     pr_to_files[name] = group.filename.values\n",
        "\n",
        "    elapsed = time.time() - t\n",
        "    print(f\"setup took {elapsed} sec\")\n",
        "\n",
        "    return comment_counts, pr_to_users"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "vlEChEgZfZak",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "data = load_data()\n",
        "test_cosine_sim_matrix = data['test_cosine_sim_matrix']\n",
        "test_pr_cosine_index_map = data['test_pr_cosine_index_map']\n",
        "train_pr_cosine_index_map = data['train_pr_cosine_index_map']\n",
        "train_comments = data['train_comments']\n",
        "\n",
        "graph = nx.DiGraph()\n",
        "graph = create_user_to_user_graph(graph, train_comments)\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "VuYc2kEcdKcV",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "commenters = set(train_comments['user_id'].unique())\n",
        "authors = set(train_comments['commenter_id'].unique())\n",
        "both = commenters.union(authors)\n",
        "print(len(commenters - authors)/len(both))\n",
        "print(len(authors - commenters)/len(both))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "AxPFl-8iIXmd",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "earec = graph\n",
        "[len(c) for c in sorted(nx.weakly_connected_components(earec), key=len, reverse=True)]"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ux0sNRB0IU93",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "print('Nodes', len(earec.nodes()))\n",
        "print('Edges', len(earec.edges()))\n",
        "print('average_shortest_path_length', nx.average_shortest_path_length(earec))\n",
        "print('average_clustering', nx.algorithms.cluster.average_clustering(earec))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "cmKdmgz4IFMR",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "\n",
        "\n",
        "degree_sequence = sorted([d for n, d in earec.degree()], reverse=True)\n",
        "degreeCount = Counter(degree_sequence)\n",
        "deg, cnt = zip(*degreeCount.items())\n",
        "\n",
        "fig, (ax1, ax2) = plt.subplots(1, 2)\n",
        "ax1.plot(deg, cnt, '+', color='b')\n",
        "ax1.set_yscale('log')\n",
        "ax1.set_xscale('log')\n",
        "\n",
        "ax1.set_title(\"Degree Histogram\")\n",
        "ax1.set_ylabel(\"Count\")\n",
        "ax1.set_label(\"Degree\")\n",
        "\n",
        "degree_sequence = sorted([d for n, d in earec.in_degree()], reverse=True)\n",
        "degreeCount = Counter(degree_sequence)\n",
        "deg_in, cnt_in = zip(*degreeCount.items())\n",
        "\n",
        "degree_sequence = sorted([d for n, d in earec.out_degree()], reverse=True)\n",
        "degreeCount = Counter(degree_sequence)\n",
        "deg_out, cnt_out = zip(*degreeCount.items())\n",
        "\n",
        "ax2.plot(deg_in, cnt_in, '+', color='b', label='in degree')\n",
        "ax2.plot(deg_out, cnt_out, '+', color='r', label='out degree')\n",
        "ax2.set_yscale('log')\n",
        "ax2.set_xscale('log')\n",
        "\n",
        "ax2.set_title(\"Degree Histogram\")\n",
        "ax2.set_ylabel(\"Count\")\n",
        "ax2.set_xlabel(\"Degree\")\n",
        "ax2.legend()\n",
        "\n",
        "fig.savefig(\"degree_hist.png\", format=\"PNG\")"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "XyB3JESiIOOh",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "degree_sequence = sorted([d for n, d in earec.in_degree()], reverse=True)\n",
        "degreeCount = Counter(degree_sequence)\n",
        "deg_in, cnt_in = zip(*degreeCount.items())\n",
        "\n",
        "degree_sequence = sorted([d for n, d in earec.out_degree()], reverse=True)\n",
        "degreeCount = Counter(degree_sequence)\n",
        "deg_out, cnt_out = zip(*degreeCount.items())\n",
        "\n",
        "fig, ax = plt.subplots()\n",
        "plt.plot(deg_in, cnt_in, '+', color='b', label='in degree')\n",
        "plt.plot(deg_out, cnt_out, '+', color='r', label='out degree')\n",
        "ax.set_yscale('log')\n",
        "ax.set_xscale('log')\n",
        "\n",
        "plt.title(\"Degree Histogram\")\n",
        "plt.ylabel(\"Count\")\n",
        "plt.xlabel(\"Degree\")\n",
        "plt.legend()\n",
        "plt.savefig(\"dir_degree_hist.png\", format=\"PNG\")"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "PB_Ozd7oIQKV",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "fig = plt.figure(figsize=(12,12))\n",
        "ax = plt.subplot(111)\n",
        "ax.set_title('EARec', fontsize=10)\n",
        "pos = nx.spring_layout(earec)\n",
        "nx.draw(earec, pos, node_size=1, width=0.5)\n",
        "plt.savefig(\"earec_vis.png\", format=\"PNG\")"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "jT1fN242PRD5",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "comment_counts, pr_to_users = setup_add_pr_precomputed(train_comments)\n",
        "weights = {}\n",
        "for new_pr_id in tqdm(test_pr_cosine_index_map):\n",
        "  weights[new_pr_id] = add_pr_precomputed(graph, new_pr_id, test_pr_cosine_index_map[new_pr_id], pr_to_users, comment_counts, test_cosine_sim_matrix, train_pr_cosine_index_map)\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "y1Ou20k1IX5r",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "test_prs = test_pr_cosine_index_map\n",
        "pageranks = {}\n",
        "pr_scores = {}\n",
        "sim_scores = {}\n",
        "alpha = 0.05\n",
        "\n",
        "i = 0\n",
        "for new_pr_id in tqdm(test_pr_cosine_index_map):\n",
        "    new_pr_node = \"p\"+str(new_pr_id)\n",
        "    pagerank_dict = nx.pagerank(graph, personalization=weights[new_pr_id], alpha=alpha, weight='weight')\n",
        "    scores = sorted([(k, v) for k, v in pagerank_dict.items()], key= lambda x: x[1], reverse=True)\n",
        "    pageranks[new_pr_id] = scores\n",
        "\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "P1hNXaauUvw5",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "_ = generate_results(pageranks, test_dict, top_k=10, verbose=True)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "cAGpOxPZrZwk",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "joblib.dump(pageranks, os.path.join(home, 'EARec_best_result_12102019.joblib'))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "5KObT6I5Xs0P",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "sim_scores = {}\n",
        "for k, v in weights.items():\n",
        "  sim_scores[k] = sorted([(k_, v_) for k_, v_ in v.items()], key= lambda x: x[1], reverse=True)\n",
        "# sim_scores = sorted([(k, v) for k, v in weights.items()], key= lambda x: x[1], reverse=True)\n",
        "_ = generate_results(sim_scores, test_dict, top_k=10, verbose=True)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "cqpRw8xY0Kyn",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "alphas = [0, .025, .05, .1, .25, .5, .75, .9, 1.0]\n",
        "aps = [0.003, 0.056, 0.062, 0.060, 0.060, 0.056, 0.056, .056, 0.055]\n",
        "ars = [0.010, 0.224, 0.249, 0.239, 0.238, 0.219, 0.216, .221, 0.214]\n",
        "f1s = [0.004, 0.086, 0.094, 0.091, 0.091, 0.085, 0.084, 0.085, 0.083]\n",
        "mrrs = [0.011, 0.119, 0.127, 0.126, 0.124, 0.124, 0.122, 0.119, 0.122]"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "kmDetCSd2jt6",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "fig, (ax1, ax2) = plt.subplots(2)\n",
        "fig.suptitle('Alpha hyperparameter search on EARec network ')\n",
        "ax1.plot(alphas[1:], f1s[1:])\n",
        "ax1.set(ylabel='F-1 score')\n",
        "ax1.yaxis.set_major_formatter(FormatStrFormatter('%.3f'))\n",
        "\n",
        "ax2.plot(alphas[1:], mrrs[1:])\n",
        "ax2.set(xlabel='alpha', ylabel='MRR score')\n",
        "ax2.yaxis.set_major_formatter(FormatStrFormatter('%.3f'))\n",
        "\n",
        "fig.savefig('EARec_param_graph.PNG')"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "5LRF0zKGYpfB",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "pagerank_dict = nx.pagerank(graph, alpha=.75, weight='weight')\n",
        "pagerank_items = sorted(pagerank_dict.items(), key=lambda x:x[1], reverse=True)\n",
        "baseline = {pr_id:pagerank_items for pr_id in test_pr_cosine_index_map}\n",
        "\n",
        "_ = generate_results(baseline, test_dict, top_k=10, verbose=True)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ur53FHbqUZ0d",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "train_comments_df = joblib.load(os.path.join(home, 'train_comments_df_v4_2018_only.joblib'))\n",
        "test_comments_df = joblib.load(os.path.join(home, 'test_comments_df_v4_2018_only.joblib'))\n",
        "filepath_summary_df = joblib.load(os.path.join(home, 'filepath_feat_df_v4_2018_only.joblib'))\n",
        "train_dict = joblib.load(os.path.join(home, 'train_dict_v4_2018_only.joblib'))\n",
        "test_dict = joblib.load(os.path.join(home, 'test_dict_v4_2018_only.joblib'))\n",
        "\n",
        "def generate_results(pred_dict, test_dict, top_k=10, verbose=False):\n",
        "    ap = 0.0\n",
        "    ar = 0.0\n",
        "    amrr = 0.0\n",
        "    af1 = 0.0\n",
        "    for test_pr_id, gt_test_reviewers in test_dict.items():\n",
        "        # if test_pr_id not in pred_dict:\n",
        "        #   continue\n",
        "          \n",
        "        pred = pred_dict[test_pr_id]\n",
        "        ranked_candidates = np.array([cand for cand, score in sorted(pred, key=lambda x: x[1], reverse=True)])\n",
        "        pred_set = set(ranked_candidates[:top_k])\n",
        "        actual_set = set(gt_test_reviewers)\n",
        "        print(pred_set, actual_set)\n",
        "        precision = len(pred_set & actual_set) / len(pred_set)\n",
        "        recall = len(pred_set & actual_set) / len(actual_set)\n",
        "        f1 = (2 * precision * recall) / (precision + recall + 10**(-10))\n",
        "        gt_rank = []\n",
        "        mrr = 0.0\n",
        "        for gt_reviewer in gt_test_reviewers:\n",
        "            lookup = np.argwhere(ranked_candidates == gt_reviewer).tolist()\n",
        "            if len(lookup) == 0:\n",
        "                rank = 0\n",
        "            else:\n",
        "                rank = lookup[0][0] + 1\n",
        "                mrr += 1/rank\n",
        "            gt_rank.append(rank)\n",
        "        amrr += mrr/len(gt_rank)\n",
        "        ap += precision\n",
        "        ar += recall\n",
        "        af1 += f1\n",
        "    if verbose:\n",
        "        print(\"Top-%d-Average-Precision = %0.3f\\nTop-%d-Average-Recall = %0.3f\\nTop-%d-Average-F1 = %0.3f\\nAverage-MRR = %0.3f\" \n",
        "            %(top_k, ap/len(test_dict), top_k, ar/len(test_dict), top_k, af1/len(test_dict), amrr/len(test_dict)))\n",
        "    return ap/len(test_dict), ar/len(test_dict), af1/len(test_dict), amrr/len(test_dict)"
      ],
      "execution_count": 0,
      "outputs": []
    }
  ]
}