main.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import torch as th\n",
    "import numpy as np\n",
    "from tqdm import tqdm\n",
    "from ray import tune\n",
    "from ray.tune.search.hyperopt import HyperOptSearch\n",
    "from ray.air.config import RunConfig\n",
    "from hyperopt import hp\n",
    "\n",
    "import setproctitle\n",
    "setproctitle.setproctitle(\"RelEns@yueling\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_name = 'wn18rr' # wn18rr, fb15k_237, nell\n",
    "model_names = ['TransE', 'RotatE', 'ComplEx', 'ConvE', 'compgcn', 'house']\n",
    "\n",
    "rank_dir = \"ranks\"\n",
    "checkpoint_path = f\"weights/{dataset_name}_rel_weights.npy\"\n",
    "\n",
    "max_concurrent_trials = 8\n",
    "n_samples = 300\n",
    "n_initial_points = 64"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loading dataset: wn18rr\n",
      "(6268, 3)\n"
     ]
    }
   ],
   "source": [
    "# (N, 3) Columns: head, realtion, tail\n",
    "print(f\"Loading dataset: {dataset_name}\")\n",
    "triplets = {\n",
    "    'test': np.load(f\"dataset/{dataset_name}/test_triples.npy\"),\n",
    "    'valid': np.load(f\"dataset/{dataset_name}/valid_triples.npy\")\n",
    "}\n",
    "print(triplets['test'].shape)\n",
    "\n",
    "pos_idx = {\n",
    "    'test': triplets['test'][:, 2],\n",
    "    'valid': triplets['valid'][:, 2]\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(6268, 40943)\n",
      "6\n"
     ]
    }
   ],
   "source": [
    "all_ranks = {\n",
    "    'test': [np.load(f\"{rank_dir}/{dataset_name}/{m}_test_ranks.npy\") for m in model_names],\n",
    "    'valid': [np.load(f\"{rank_dir}/{dataset_name}/{m}_valid_ranks.npy\") for m in model_names]\n",
    "}\n",
    "print(all_ranks['test'][0].shape)\n",
    "\n",
    "n_model = len(all_ranks['test'])\n",
    "print(n_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def eval_model(pos_idx, ranks):\n",
    "    ''' \n",
    "        pos_idx: np.ndarray \n",
    "        ranks: np.ndarray\n",
    "    '''\n",
    "    pos_idx = th.from_numpy(pos_idx)\n",
    "    ranks = th.from_numpy(ranks)\n",
    "    argsort = th.argsort(ranks , dim=1, descending=False)\n",
    "    pos_ranking = th.nonzero(argsort == pos_idx.unsqueeze(1))[:, 1].cpu().numpy() + 1\n",
    "\n",
    "    # calculate metrics\n",
    "    all_ranking        = np.array(pos_ranking)\n",
    "    metrics            = {}\n",
    "    metrics['mrr']     = round(np.mean(1/all_ranking), 4)\n",
    "    metrics['mr']      = round(np.mean(all_ranking), 4)\n",
    "    metrics['hits@1']  = round(np.mean(all_ranking<=1), 4)\n",
    "    metrics['hits@3']  = round(np.mean(all_ranking<=3), 4)\n",
    "    metrics['hits@10'] = round(np.mean(all_ranking<=10), 4)\n",
    "\n",
    "    return metrics\n",
    "\n",
    "def objective(config, data):\n",
    "    sub_pos_idx, sub_ranks = data\n",
    "\n",
    "    n_model = len(sub_ranks)\n",
    "\n",
    "    weights = [config[f\"w_{i}\"] for i in range(n_model)]\n",
    "\n",
    "    ranks_avg = np.average(sub_ranks, weights=weights, axis=0)\n",
    "\n",
    "    mrr = eval_model(sub_pos_idx, ranks_avg)\n",
    "\n",
    "    return mrr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# get rel indexes\n",
    "relations = {\n",
    "    'test': triplets['test'][:, 1],\n",
    "    'valid': triplets['valid'][:, 1]\n",
    "}\n",
    "\n",
    "num_relation = relations['valid'].max()+1\n",
    "\n",
    "rel_indexes = {\n",
    "    'valid': {},\n",
    "    'test': {}\n",
    "} # relation_id -> np array\n",
    "\n",
    "# relation_id [0, 534]\n",
    "for relation_id in range(num_relation):\n",
    "    rel_indexes['test'][relation_id] = np.where(relations['test'] == relation_id)[0]\n",
    "    rel_indexes['valid'][relation_id] = np.where(relations['valid'] == relation_id)[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Load existing models\n"
     ]
    }
   ],
   "source": [
    "default_config = {f'w_{i}': 1/n_model for i in range(n_model)}\n",
    "\n",
    "if checkpoint_path and os.path.exists(checkpoint_path):\n",
    "    print(\"Load existing models\")\n",
    "    rel_weights = np.load(checkpoint_path)\n",
    "else:\n",
    "    print(\"Searching for ensemble weights\")\n",
    "    rel_weights = np.zeros((num_relation, n_model))\n",
    "\n",
    "    search_space = {f\"w_{i}\": hp.uniform(f\"w_{i}\", 0, 1) for i in range(n_model)}\n",
    "    hyperopt_search = HyperOptSearch(search_space, metric=\"mrr\", mode=\"max\", n_initial_points=n_initial_points)\n",
    "\n",
    "    for rel_id in range(num_relation):\n",
    "        if len(rel_indexes['valid'][rel_id]) == 0:\n",
    "            # default weights\n",
    "            rel_weights[rel_id] = np.fromiter(default_config.values(), dtype=np.float32)\n",
    "            continue\n",
    "\n",
    "        subranks = [model_rank[rel_indexes['valid'][rel_id]] for model_rank in all_ranks['valid']]\n",
    "        sub_pos_idx = pos_idx['valid'][rel_indexes['valid'][rel_id]]\n",
    "\n",
    "        tuner = tune.Tuner(tune.with_parameters(objective, data=(sub_pos_idx, subranks)), param_space=search_space,\n",
    "                tune_config=tune.TuneConfig(num_samples=n_samples, search_alg=hyperopt_search, max_concurrent_trials=max_concurrent_trials))\n",
    "        results = tuner.fit()\n",
    "\n",
    "        best_weights = np.fromiter(results.get_best_result(metric=\"mrr\", mode=\"max\").config.values(), dtype=float)\n",
    "        rel_weights[rel_id] = best_weights\n",
    "\n",
    "    # np.save(f\"{dataset_name}_rel_weights.npy\", rel_weights)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Calculate relation performance"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Evaluating\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 22/22 [00:30<00:00,  1.41s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "rel_ensemble\n",
      "mrr: 0.5201\n",
      "hits@1: 0.4770\n",
      "hits@3: 0.5375\n",
      "hits@10: 0.6039\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "print(\"Evaluating\")\n",
    "rel_res = {\n",
    "    'rel_ensemble': [{'mrr': 0, 'mr': 0, 'hits@1': 0, 'hits@3': 0, 'hits@10': 0} for _ in range(num_relation)]\n",
    "}\n",
    "\n",
    "mode = 'test'\n",
    "for rel_id in tqdm(range(num_relation)):\n",
    "    if len(rel_indexes[mode][rel_id]) == 0:\n",
    "        continue\n",
    "\n",
    "    sub_ranks = [model_rank[rel_indexes[mode][rel_id]] for model_rank in all_ranks[mode]]\n",
    "    sub_pos_idx = pos_idx[mode][rel_indexes[mode][rel_id]]\n",
    "\n",
    "    config = {f\"w_{i}\": rel_weights[rel_id][i] for i in range(n_model)}\n",
    "    metrics = objective(config, (sub_pos_idx, sub_ranks))\n",
    "    rel_res['rel_ensemble'][rel_id] = metrics\n",
    "\n",
    "final_res = {\n",
    "    'rel_ensemble': {'mrr': 0, 'hits@1': 0, 'hits@3': 0, 'hits@10': 0}\n",
    "}\n",
    "\n",
    "for rel_id in range(num_relation):\n",
    "    for metric in ['mrr', 'hits@1', 'hits@3', 'hits@10']:\n",
    "        final_res['rel_ensemble'][metric] += rel_res['rel_ensemble'][rel_id][metric] * len(rel_indexes[mode][rel_id])\n",
    "\n",
    "print(\"rel_ensemble\")\n",
    "for k, v in final_res['rel_ensemble'].items():\n",
    "    v_mean = v / len(relations[mode])\n",
    "    print(f\"{k}: {v_mean:.4f}\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.8.16 ('th1110')",
   "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.8.16"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "c5f194073aaff85b697eed6685e89bdb344177f148991fb6fdbaca70a347ec53"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}