KRED-Data_Preparation.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "source": [
    "## Data creation Notebook for KRED: Knowledge-Aware Document Representation for News Recommendations\n",
    "\n",
    "This notebook is useful to create a pickle file that is useful to train and test the architecture of KRED.\n",
    "\n",
    "> **Note:** The file generated name is \"data_dict_{timestamp}.pkl\"\n"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import time\n",
    "import torch\n",
    "import os\n",
    "from utils import *\n",
    "from train_test import *\n",
    "import argparse\n",
    "from parse_config import ConfigParser"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "WFqGbb7AN84n",
    "outputId": "4c5214cb-559b-48a0-c7ec-b8f07921a243",
    "tags": []
   },
   "outputs": [],
   "source": [
    "# The download part has been removed since it is deprecated\n",
    "\n",
    "MIND_type = 'demo'\n",
    "data_path = \"/datasets/\"\n",
    "\n",
    "train_news_file = os.path.join(data_path, 'mind_train', r'news.tsv')\n",
    "train_behaviors_file = os.path.join(data_path, 'mind_train', r'behaviors.tsv')\n",
    "valid_news_file = os.path.join(data_path, 'mind_val', r'news.tsv')\n",
    "valid_behaviors_file = os.path.join(data_path, 'mind_val', r'behaviors.tsv')\n",
    "knowledge_graph_file = os.path.join(data_path, 'mind_kg', r'wikidata-graph.tsv')\n",
    "entity_embedding_file = os.path.join(data_path, 'mind_kg', r'entity2vecd100.vec')\n",
    "relation_embedding_file = os.path.join(data_path, 'mind_kg', r'relation2vecd100.vec')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "id": "g_wFBRK2Ye7S",
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Config Generation\n",
    "\n",
    "parser = argparse.ArgumentParser(description='KRED')\n",
    "parser.add_argument('-f')\n",
    "parser.add_argument('-c', '--config', default=\"./config.json\", type=str,\n",
    "                    help='config file path (default: None)')\n",
    "parser.add_argument('-r', '--resume', default=None, type=str,\n",
    "                    help='path to latest checkpoint (default: None)')\n",
    "parser.add_argument('-d', '--device', default=None, type=str,\n",
    "                    help='indices of GPUs to enable (default: all)')\n",
    "\n",
    "config = ConfigParser.from_args(parser)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "outputs": [],
   "source": [
    "from utils.cleaner import *\n",
    "config = cleaner(config)"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "entity2id_dict = entity_to_id(config, entities_news(config))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "constructing embedding ...\n"
     ]
    }
   ],
   "source": [
    "entity_embedding = []\n",
    "entity_embedding.append(np.zeros(config['model']['entity_embedding_dim']))\n",
    "entity2embedding_dict = {}\n",
    "entity2embedding_dict, entity_embedding, relation_embedding = construct_embedding_mind(config, entity2id_dict, entity_embedding, entity2embedding_dict)"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "constructing adjacency matrix ...\n",
      "construct_adj_mind finish\n"
     ]
    }
   ],
   "source": [
    "entity_adj, relation_adj = construct_adj_mind(config, entity2id_dict, entity2embedding_dict)"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "constructing embedding ...\n"
     ]
    }
   ],
   "source": [
    "entities_not_embedded = set([item for items in entity_adj for item in items]).difference(set(entity2id_dict.values()))\n",
    "entity2id_dict_not_embedded = id_to_entity(config, entities_not_embedded)\n",
    "entity2embedding_dict, entity_embedding, relation_embedding = construct_embedding_mind(config, entity2id_dict_not_embedded, entity_embedding, entity2embedding_dict)\n",
    "\n",
    "# Add the new entities to the dictionary\n",
    "entity2id_dict.update(entity2id_dict_not_embedded)\n",
    "# Invert the dictionary\n",
    "id2entity_dict = {v: k for k, v in entity2id_dict.items()}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "for i in range(1, len(entity_adj)):\n",
    "    for j in range(0, len(entity_adj[i])):\n",
    "        entity_adj[i][j] = entity2embedding_dict[id2entity_dict[entity_adj[i][j]]]\n",
    "entity_embedding = torch.FloatTensor(np.array(entity_embedding))\n",
    "relation_embedding = torch.FloatTensor(np.array(relation_embedding))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "embedding_folder=None\n",
    "news_feature, max_entity_freq, max_entity_pos, max_entity_type = build_news_features_mind(config, entity2embedding_dict, embedding_folder)\n",
    "\n",
    "# Load the user history\n",
    "user_history = build_user_history(config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'lifestyle': 0, 'news': 1, 'health': 2, 'weather': 3, 'entertainment': 4, 'autos': 5, 'travel': 6, 'sports': 7, 'foodanddrink': 8, 'tv': 9, 'movies': 10, 'finance': 11, 'video': 12, 'music': 13, 'middleeast': 14, 'kids': 15, 'northamerica': 16}\n"
     ]
    }
   ],
   "source": [
    "train_data, dev_data = get_user2item_data(config)\n",
    "vert_train, vert_test = build_vert_data(config)\n",
    "pop_train, pop_test = build_pop_data(config)\n",
    "item2item_train, item2item_test = build_item2item_data(config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "data_dict = {\n",
    "    'user_history' : user_history,\n",
    "    'entity_embedding' : entity_embedding,\n",
    "    'relation_embedding' : relation_embedding,\n",
    "    'entity_adj' : entity_adj,\n",
    "    'relation_adj' : relation_adj,\n",
    "    'news_feature': news_feature,\n",
    "    'max_entity_freq':max_entity_freq,\n",
    "    'max_entity_pos': max_entity_pos,\n",
    "    'max_entity_type':max_entity_type,\n",
    "    'train_data': train_data,\n",
    "    'dev_data':dev_data,\n",
    "    'vert_train': vert_train,\n",
    "    'vert_test' : vert_test,\n",
    "    'pop_train': pop_train,\n",
    "    'pop_test':pop_test,\n",
    "    'item2item_train': item2item_train,\n",
    "    'item2item_test': item2item_test\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "import gzip\n",
    "import pickle\n",
    "timestamp = int(time.time()*1000)\n",
    "def save_compressed_pickle(filename, obj):\n",
    "    with gzip.open(filename, 'wb') as f:\n",
    "        pickle.dump(obj, f)\n",
    "save_compressed_pickle(f'data_dict_{timestamp}.pkl', data_dict)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [],
   "metadata": {
    "collapsed": false
   }
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "provenance": []
  },
  "gpuClass": "standard",
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}