{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3d02d392-75f2-402e-be93-5032fd879b3c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Copyright (c) 2022 Graphcore Ltd. All rights reserved."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1ebc7833-d4f4-44c8-8b64-007e1fa99573",
   "metadata": {},
   "source": [
    "# Build entity mapping database\n",
    "\n",
    "_Note: this notebook is included to document the process and so will not run out of the box._\n",
    "\n",
    "In order to provide an interface for querying the model, we build a full text searchable mapping from string to entity and relation ID. To do this, we:\n",
    " - Download a full list of entity and relation labels from wikidata.\n",
    " - Use OGB's `data/ogbl_wikikg2/mapping/` metadata to filter entities of interest, and map them to contiguous OGB dataset IDs.\n",
    " - Build a SQLite database with FTS3 indicies for efficient local retrieval.\n",
    " \n",
    "Also contains the command to build a faster-loading `.npz` file containing ogbl-wikikg2."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f3e2dcd4-da30-4703-9315-73f5f62b4413",
   "metadata": {},
   "outputs": [],
   "source": [
    "import bz2\n",
    "import json\n",
    "import tqdm\n",
    "import itertools as it\n",
    "from pathlib import Path\n",
    "import csv\n",
    "import gzip\n",
    "import sys\n",
    "\n",
    "import kge_mapping\n",
    "import kge_training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f63157dc",
   "metadata": {},
   "outputs": [],
   "source": [
    "from examples_utils import notebook_logging\n",
    "%load_ext gc_logger"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "33addfcb-8a72-4150-8c8b-37c269c00fd2",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "# Source:\n",
    "# https://dumps.wikimedia.org/wikidatawiki/entities/latest-all.json.bz2\n",
    "\n",
    "labels = {}\n",
    "path = Path(\"/localdata/research/scratch/douglaso/latest-all.json.bz2\")\n",
    "\n",
    "f0 = open(path, \"rb\")\n",
    "f = bz2.BZ2File(f0)\n",
    "f.readline()  # opening \"[\"\n",
    "tq = tqdm.tqdm(it.count())\n",
    "for n in tq:\n",
    "    line = f.readline().decode().rstrip(\"\\n ,\")\n",
    "    if line == \"]\":\n",
    "        break\n",
    "    e = json.loads()\n",
    "    labels[e[\"id\"]] = e[\"labels\"].get(\"en\", dict(value=\"\"))[\"value\"]\n",
    "    if n % int(1e3) == 0:\n",
    "        tq.set_description(f\"{f0.tell() / path.stat().st_size:.0%}, {f0.tell() / 2**30:.1f} GiB\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e2d35fe3-fe10-4232-867f-93c135797a2b",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "# Source:\n",
    "# https://hay.toolforge.org/propbrowse/props.json\n",
    "\n",
    "with Path(\"/localdata/research/scratch/douglaso/props.json\").open() as f:\n",
    "    props = {item[\"id\"]: item[\"label\"] for item in json.load(f)}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "edab16f6-72ef-4f4b-9b2c-17688e1fb190",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "records = []\n",
    "\n",
    "with gzip.open(\"data/ogbl_wikikg2/mapping/nodeidx2entityid.csv.gz\", \"rt\") as f:\n",
    "    for item in csv.DictReader(f):\n",
    "        records.append(dict(\n",
    "            type=\"entity\",\n",
    "            idx=int(item[\"node idx\"]),\n",
    "            wikidata_id=item[\"entity id\"],\n",
    "            wikidata_label=labels.get(item[\"entity id\"], \"\"),\n",
    "        ))\n",
    "\n",
    "with gzip.open(\"data/ogbl_wikikg2/mapping/reltype2relid.csv.gz\", \"rt\") as f:\n",
    "    for item in csv.DictReader(f):\n",
    "        records.append(dict(\n",
    "            type=\"relation\",\n",
    "            idx=int(item[\"reltype\"]),\n",
    "            wikidata_id=item[\"rel id\"],\n",
    "            wikidata_label=props.get(item[\"rel id\"], \"\"),\n",
    "        ))\n",
    "\n",
    "print(f\"Missing entity labels: {sum(not r['wikidata_label'] for r in records if r['type'] == 'entity') / (sum(1 for r in records if r['type'] == 'entity')):.1%}\")\n",
    "print(f\"Missing relation labels: {sum(not r['wikidata_label'] for r in records if r['type'] == 'relation') / (sum(1 for r in records if r['type'] == 'relation')):.1%}\")\n",
    "\n",
    "with gzip.open(\"data/ogbl_wikikg2_mapping.jsonl.gz\", \"wt\") as f:\n",
    "    for record in records:\n",
    "        print(json.dumps(record), file=f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "392c9edd-d861-4efb-8fc9-f84a27763c31",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "kge_mapping.Database.build(\n",
    "    Path(\"data/ogbl_wikikg2_mapping.sqlite\"),\n",
    "    kge_mapping.RawData.load(Path(\"data/ogbl_wikikg2_mapping.jsonl.gz\"), Path(\"data\")),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e1b3d8ea-fb6f-40ea-9694-9f44f9b13483",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "kge_training.Dataset.build_wikikg2(\n",
    "    Path(\"data\"),\n",
    "    Path(\"data/ogbl_wikikg2.npz\"),\n",
    "    seed=1000,\n",
    ")"
   ]
  }
 ],
 "metadata": {
  "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.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}