graph_data/prepare_dataset.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prepare training, validation and test datasets"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We need to prepare the training, validation and external test datasets. We will use the training dataset to train the model and the test dataset to evaluate the model for all KGE models."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### [Require to Modify According to Your Situation] Prepare all relation files"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Merging the following files:\n",
      "/Users/jy006/Documents/Code/BioMedGPS/biomedgps-data/graph_data/formatted_relations/drkg/unformatted_drkg.tsv\n",
      "/Users/jy006/Documents/Code/BioMedGPS/biomedgps-data/graph_data/formatted_relations/drkg/formatted_drkg.tsv\n",
      "/Users/jy006/Documents/Code/BioMedGPS/biomedgps-data/graph_data/formatted_relations/hsdn/formatted_hsdn.tsv\n",
      "/Users/jy006/Documents/Code/BioMedGPS/biomedgps-data/graph_data/custom_relations/formatted_malacards_mecfs.tsv\n",
      "/Users/jy006/Documents/Code/BioMedGPS/biomedgps-data/graph_data/custom_relations/formatted_custom_all_v20240119.tsv\n",
      "Number of entities: 815043\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "\n",
    "root_dir = os.path.dirname(os.getcwd())\n",
    "\n",
    "dataset_name = \"biomedgps-full-v20240127\"\n",
    "outputdir = os.path.join(root_dir, \"datasets\", dataset_name)\n",
    "graph_data_dir = os.path.join(root_dir, \"graph_data\")\n",
    "\n",
    "formatted_ctd = os.path.join(\n",
    "    graph_data_dir, \"formatted_relations\", \"drkg\", \"unformatted_drkg.tsv\"\n",
    ")\n",
    "formatted_drkg = os.path.join(\n",
    "    graph_data_dir, \"formatted_relations\", \"drkg\", \"formatted_drkg.tsv\"\n",
    ")\n",
    "formatted_hsdn = os.path.join(\n",
    "    graph_data_dir, \"formatted_relations\", \"hsdn\", \"formatted_hsdn.tsv\"\n",
    ")\n",
    "formatted_malacards_mecfs = os.path.join(\n",
    "    graph_data_dir, \"custom_relations\", \"formatted_malacards_mecfs.tsv\"\n",
    ")\n",
    "formatted_custom = os.path.join(\n",
    "    graph_data_dir, \"custom_relations\", \"formatted_custom_all_v20240119.tsv\"\n",
    ")\n",
    "\n",
    "files = [\n",
    "    formatted_ctd,\n",
    "    formatted_drkg,\n",
    "    formatted_hsdn,\n",
    "    formatted_malacards_mecfs,\n",
    "    formatted_custom,\n",
    "]\n",
    "\n",
    "print(\"Merging the following files:\")\n",
    "print(\"\\n\".join(files))\n",
    "\n",
    "entity_file = os.path.join(graph_data_dir, \"entities.tsv\")\n",
    "print(\"Number of entities: {}\".format(len(open(entity_file).readlines())))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Dependencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Adding /Users/jy006/Documents/Code/BioMedGPS/biomedgps-data/lib to sys.path\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import sys\n",
    "\n",
    "lib_dir = os.path.join(os.path.dirname(os.getcwd()), \"lib\")\n",
    "\n",
    "print(\"Adding {} to sys.path\".format(lib_dir))\n",
    "sys.path.append(lib_dir)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Merge all relation files into one file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Running: python3 /Users/jy006/Documents/Code/BioMedGPS/biomedgps-data/lib/data.py merge-files --input /Users/jy006/Documents/Code/BioMedGPS/biomedgps-data/graph_data/formatted_relations/drkg/unformatted_drkg.tsv --input /Users/jy006/Documents/Code/BioMedGPS/biomedgps-data/graph_data/formatted_relations/drkg/formatted_drkg.tsv --input /Users/jy006/Documents/Code/BioMedGPS/biomedgps-data/graph_data/formatted_relations/hsdn/formatted_hsdn.tsv --input /Users/jy006/Documents/Code/BioMedGPS/biomedgps-data/graph_data/custom_relations/formatted_malacards_mecfs.tsv --input /Users/jy006/Documents/Code/BioMedGPS/biomedgps-data/graph_data/custom_relations/formatted_custom_all_v20240119.tsv --output /var/folders/4s/d4nr1sg91ps1k3qz00h28w_r0000gp/T/tmpnluxqzag/knowledge_graph.tsv\n",
      "Merging the following files:  ('/Users/jy006/Documents/Code/BioMedGPS/biomedgps-data/graph_data/formatted_relations/drkg/unformatted_drkg.tsv', '/Users/jy006/Documents/Code/BioMedGPS/biomedgps-data/graph_data/formatted_relations/drkg/formatted_drkg.tsv', '/Users/jy006/Documents/Code/BioMedGPS/biomedgps-data/graph_data/formatted_relations/hsdn/formatted_hsdn.tsv', '/Users/jy006/Documents/Code/BioMedGPS/biomedgps-data/graph_data/custom_relations/formatted_malacards_mecfs.tsv', '/Users/jy006/Documents/Code/BioMedGPS/biomedgps-data/graph_data/custom_relations/formatted_custom_all_v20240119.tsv')\n",
      "The following columns are in the /Users/jy006/Documents/Code/BioMedGPS/biomedgps-data/graph_data/formatted_relations/drkg/unformatted_drkg.tsv: ['raw_source_id', 'raw_target_id', 'raw_source_type', 'raw_target_type', 'relation_type', 'resource', 'pmids', 'key_sentence', 'source_id', 'source_type', 'target_id', 'target_type']\n",
      "The following columns are in the /Users/jy006/Documents/Code/BioMedGPS/biomedgps-data/graph_data/formatted_relations/drkg/formatted_drkg.tsv: ['raw_source_id', 'raw_target_id', 'raw_source_type', 'raw_target_type', 'relation_type', 'resource', 'pmids', 'key_sentence', 'source_id', 'source_type', 'target_id', 'target_type']\n",
      "The following columns are in the /Users/jy006/Documents/Code/BioMedGPS/biomedgps-data/graph_data/formatted_relations/hsdn/formatted_hsdn.tsv: ['raw_source_id', 'raw_target_id', 'raw_source_type', 'raw_target_type', 'relation_type', 'resource', 'pmids', 'key_sentence', 'source_id', 'source_type', 'target_id', 'target_type']\n",
      "The following columns are in the /Users/jy006/Documents/Code/BioMedGPS/biomedgps-data/graph_data/custom_relations/formatted_malacards_mecfs.tsv: ['source_id', 'source_name', 'source_type', 'target_id', 'target_name', 'target_type', 'relation_type', 'resource']\n",
      "The following columns are in the /Users/jy006/Documents/Code/BioMedGPS/biomedgps-data/graph_data/custom_relations/formatted_custom_all_v20240119.tsv: ['relation_type', 'source_name', 'source_type', 'source_id', 'target_name', 'target_type', 'target_id', 'key_sentence', 'pmids', 'resource']\n",
      "The following columns will be used:  ['target_type', 'source_id', 'relation_type', 'source_type', 'resource', 'target_id']\n",
      "Number of unique entity ids: 138000\n",
      "Number of relations: 6005665\n",
      "Running: python3 /Users/jy006/Documents/Code/BioMedGPS/biomedgps-data/graph_data/scripts/annotate_relations.py --entity-file /Users/jy006/Documents/Code/BioMedGPS/biomedgps-data/graph_data/entities.tsv --relation-file /var/folders/4s/d4nr1sg91ps1k3qz00h28w_r0000gp/T/tmpnluxqzag/knowledge_graph.tsv --output-dir /var/folders/4s/d4nr1sg91ps1k3qz00h28w_r0000gp/T/tmpnluxqzag\n",
      "File written to: /var/folders/4s/d4nr1sg91ps1k3qz00h28w_r0000gp/T/tmpnluxqzag/knowledge_graph_annotated.tsv\n"
     ]
    }
   ],
   "source": [
    "import subprocess\n",
    "import pandas as pd\n",
    "import tempfile\n",
    "\n",
    "temp_dir = tempfile.mkdtemp()\n",
    "\n",
    "args = [\"python3\", os.path.join(lib_dir, \"data.py\"), \"merge-files\"]\n",
    "\n",
    "for f in files:\n",
    "    args.extend([\"--input\", f])\n",
    "\n",
    "kg_file = os.path.join(temp_dir, \"knowledge_graph.tsv\")\n",
    "annotated_kg_file = os.path.join(temp_dir, \"knowledge_graph_annotated.tsv\")\n",
    "args.extend([\"--output\", kg_file])\n",
    "\n",
    "print(\"Running: {}\".format(\" \".join(args)))\n",
    "args_str = \" \".join(args)\n",
    "!{args_str}\n",
    "\n",
    "if os.path.exists(kg_file):\n",
    "    df = pd.read_csv(kg_file, sep=\"\\t\")\n",
    "    source_ids = df[[\"source_id\", \"source_type\"]].drop_duplicates()\n",
    "    target_ids = df[[\"target_id\", \"target_type\"]].drop_duplicates()\n",
    "    ids = pd.concat([source_ids, target_ids]).drop_duplicates()\n",
    "    print(\"Number of unique entity ids: {}\".format(len(ids)))\n",
    "    print(\"Number of relations: {}\".format(len(df.drop_duplicates())))\n",
    "\n",
    "# Annotate the knowledge graph with the entities\n",
    "args = [\n",
    "    \"python3\",\n",
    "    os.path.join(os.path.dirname(lib_dir), \"graph_data\", \"scripts\", \"annotate_relations.py\"),\n",
    "    \"--entity-file\",\n",
    "    entity_file,\n",
    "    \"--relation-file\",\n",
    "    kg_file,\n",
    "    \"--output-dir\",\n",
    "    os.path.dirname(kg_file),\n",
    "]\n",
    "\n",
    "print(\"Running: {}\".format(\" \".join(args)))\n",
    "args_str = \" \".join(args)\n",
    "!{args_str}\n",
    "print(\"File written to: {}\".format(annotated_kg_file))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Split the merged relation file into training, validation and test files"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Running: python3 /Users/jy006/Documents/Code/BioMedGPS/biomedgps-data/lib/data.py split --input /var/folders/4s/d4nr1sg91ps1k3qz00h28w_r0000gp/T/tmpnluxqzag/knowledge_graph.tsv --output-1 /var/folders/4s/d4nr1sg91ps1k3qz00h28w_r0000gp/T/tmpnluxqzag/train_validation.tsv --output-2 /var/folders/4s/d4nr1sg91ps1k3qz00h28w_r0000gp/T/tmpnluxqzag/test.tsv --ratio 0.95\n",
      "Split files created: /var/folders/4s/d4nr1sg91ps1k3qz00h28w_r0000gp/T/tmpnluxqzag/train_validation.tsv and /var/folders/4s/d4nr1sg91ps1k3qz00h28w_r0000gp/T/tmpnluxqzag/test.tsv.\n",
      "Running: python3 /Users/jy006/Documents/Code/BioMedGPS/biomedgps-data/lib/data.py split --input /var/folders/4s/d4nr1sg91ps1k3qz00h28w_r0000gp/T/tmpnluxqzag/train_validation.tsv --output-1 /var/folders/4s/d4nr1sg91ps1k3qz00h28w_r0000gp/T/tmpnluxqzag/train.tsv --output-2 /var/folders/4s/d4nr1sg91ps1k3qz00h28w_r0000gp/T/tmpnluxqzag/valid.tsv --ratio 0.95\n",
      "Split files created: /var/folders/4s/d4nr1sg91ps1k3qz00h28w_r0000gp/T/tmpnluxqzag/train.tsv and /var/folders/4s/d4nr1sg91ps1k3qz00h28w_r0000gp/T/tmpnluxqzag/valid.tsv.\n"
     ]
    }
   ],
   "source": [
    "train_validation_file = os.path.join(temp_dir, \"train_validation.tsv\")\n",
    "train_file = os.path.join(temp_dir, \"train.tsv\")\n",
    "test_file = os.path.join(temp_dir, \"test.tsv\")\n",
    "valid_file = os.path.join(temp_dir, \"valid.tsv\")\n",
    "\n",
    "split_cmd = [\n",
    "    \"python3\",\n",
    "    os.path.join(lib_dir, \"data.py\"),\n",
    "    \"split\",\n",
    "    \"--input\",\n",
    "    kg_file,\n",
    "    \"--output-1\",\n",
    "    train_validation_file,\n",
    "    \"--output-2\",\n",
    "    test_file,\n",
    "    \"--ratio\",\n",
    "    \"0.95\",\n",
    "]\n",
    "\n",
    "print(\"Running: {}\".format(\" \".join(split_cmd)))\n",
    "split_cmd_str = \" \".join(split_cmd)\n",
    "!{split_cmd_str}\n",
    "print(f\"Split files created: {train_validation_file} and {test_file}.\")\n",
    "\n",
    "split_cmd = [\n",
    "    \"python3\",\n",
    "    os.path.join(lib_dir, \"data.py\"),\n",
    "    \"split\",\n",
    "    \"--input\",\n",
    "    train_validation_file,\n",
    "    \"--output-1\",\n",
    "    train_file,\n",
    "    \"--output-2\",\n",
    "    valid_file,\n",
    "    \"--ratio\",\n",
    "    \"0.95\",\n",
    "]\n",
    "\n",
    "print(\"Running: {}\".format(\" \".join(split_cmd)))\n",
    "split_cmd_str = \" \".join(split_cmd)\n",
    "!{split_cmd_str}\n",
    "print(\"Split files created: {} and {}.\".format(train_file, valid_file))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Check whether number of ids in train, validation, and test are the same.\n",
    "If you see the following message, you need to run the section \"Keep the same number of ids in train, validation, and test\" in the notebook.\n",
    "\n",
    "```\n",
    "ValueError: You need to keep the entity ids and relation types in the test and validation files the same as the ones in the train file.\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Running: python3 /Users/jy006/Documents/Code/BioMedGPS/biomedgps-data/lib/data.py check-ids --input /var/folders/4s/d4nr1sg91ps1k3qz00h28w_r0000gp/T/tmpnluxqzag/train.tsv --input /var/folders/4s/d4nr1sg91ps1k3qz00h28w_r0000gp/T/tmpnluxqzag/valid.tsv --input /var/folders/4s/d4nr1sg91ps1k3qz00h28w_r0000gp/T/tmpnluxqzag/test.tsv --output /var/folders/4s/d4nr1sg91ps1k3qz00h28w_r0000gp/T/tmpnluxqzag/id_checked.tsv\n",
      "Number of ids in train, validation, and test: \n",
      "       file1      file2  num_intersection  num_file1  num_file2        category\n",
      "0  train.tsv  valid.tsv             37719      69247      38234             ids\n",
      "1  train.tsv  valid.tsv               114        123        114  relation_types\n",
      "2  train.tsv   test.tsv             38080      69247      38587             ids\n",
      "3  train.tsv   test.tsv               110        123        110  relation_types\n",
      "4  valid.tsv   test.tsv             31890      38234      38587             ids\n",
      "5  valid.tsv   test.tsv               110        114        110  relation_types\n",
      "Traceback (most recent call last):\n",
      "  File \"/Users/jy006/Documents/Code/BioMedGPS/biomedgps-data/lib/data.py\", line 452, in <module>\n",
      "    cli()\n",
      "  File \"/Users/jy006/miniconda3/envs/biomedgps-data/lib/python3.10/site-packages/click/core.py\", line 1157, in __call__\n",
      "    return self.main(*args, **kwargs)\n",
      "  File \"/Users/jy006/miniconda3/envs/biomedgps-data/lib/python3.10/site-packages/click/core.py\", line 1078, in main\n",
      "    rv = self.invoke(ctx)\n",
      "  File \"/Users/jy006/miniconda3/envs/biomedgps-data/lib/python3.10/site-packages/click/core.py\", line 1688, in invoke\n",
      "    return _process_result(sub_ctx.command.invoke(sub_ctx))\n",
      "  File \"/Users/jy006/miniconda3/envs/biomedgps-data/lib/python3.10/site-packages/click/core.py\", line 1434, in invoke\n",
      "    return ctx.invoke(self.callback, **ctx.params)\n",
      "  File \"/Users/jy006/miniconda3/envs/biomedgps-data/lib/python3.10/site-packages/click/core.py\", line 783, in invoke\n",
      "    return __callback(*args, **kwargs)\n",
      "  File \"/Users/jy006/Documents/Code/BioMedGPS/biomedgps-data/lib/data.py\", line 286, in check_ids\n",
      "    raise ValueError(\n",
      "ValueError: You need to keep the entity ids and relation types in the test and validation files the same as the ones in the train file.\n",
      "Checked files created: /var/folders/4s/d4nr1sg91ps1k3qz00h28w_r0000gp/T/tmpnluxqzag/id_checked.tsv.\n"
     ]
    }
   ],
   "source": [
    "id_checked_file = os.path.join(temp_dir, \"id_checked.tsv\")\n",
    "check_ids_cmd = [\n",
    "    \"python3\",\n",
    "    os.path.join(lib_dir, \"data.py\"),\n",
    "    \"check-ids\",\n",
    "    \"--input\",\n",
    "    train_file,\n",
    "    \"--input\",\n",
    "    valid_file,\n",
    "    \"--input\",\n",
    "    test_file,\n",
    "    \"--output\",\n",
    "    id_checked_file,\n",
    "]\n",
    "\n",
    "print(\"Running: {}\".format(\" \".join(check_ids_cmd)))\n",
    "check_ids_cmd_str = \" \".join(check_ids_cmd)\n",
    "# Catch the error\n",
    "check_ids_cmd_str += \" || true\"\n",
    "!{check_ids_cmd_str}\n",
    "print(\"Checked files created: {}.\".format(id_checked_file))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Keep the entity id and relation type same among validation, test and training datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Running: python3 /Users/jy006/Documents/Code/BioMedGPS/biomedgps-data/lib/data.py keep-valid --train-file /var/folders/4s/d4nr1sg91ps1k3qz00h28w_r0000gp/T/tmpnluxqzag/train.tsv --valid-file /var/folders/4s/d4nr1sg91ps1k3qz00h28w_r0000gp/T/tmpnluxqzag/valid.tsv --test-file /var/folders/4s/d4nr1sg91ps1k3qz00h28w_r0000gp/T/tmpnluxqzag/test.tsv --output-dir /var/folders/4s/d4nr1sg91ps1k3qz00h28w_r0000gp/T/tmpnluxqzag/keep_valid\n",
      "Reading the train, test and validation dataframes...\n",
      "Found 5244530 rows in train\n",
      "Found 290571 rows in test\n",
      "Found 276043 rows in validation\n",
      "Checking the format of the dataframes...\n",
      "Found 69247 entities in train\n",
      "Found 38587 entities in test\n",
      "Found 38234 entities in validation\n",
      "Found 31842 entities in interaction\n",
      "Found 110 rows in relation intersection\n",
      "Found 281325 rows in test intersection\n",
      "Found 31712 entities in test intersection\n",
      "Found 267395 rows in validation intersection\n",
      "Found 31709 entities in validation intersection\n",
      "Found 9246 rows in test not in intersection\n",
      "Found 8648 rows in validation not in intersection\n",
      "Found 5262424 rows in extended train\n",
      "Files created: ['test_valid.tsv', 'train_valid.tsv', 'valid_valid.tsv'].\n"
     ]
    }
   ],
   "source": [
    "keep_valid_dir = os.path.join(temp_dir, \"keep_valid\")\n",
    "os.makedirs(keep_valid_dir, exist_ok=True)\n",
    "\n",
    "keep_valid_cmd = [\n",
    "    \"python3\",\n",
    "    os.path.join(lib_dir, \"data.py\"),\n",
    "    \"keep-valid\",\n",
    "    \"--train-file\",\n",
    "    train_file,\n",
    "    \"--valid-file\",\n",
    "    valid_file,\n",
    "    \"--test-file\",\n",
    "    test_file,\n",
    "    \"--output-dir\",\n",
    "    keep_valid_dir,\n",
    "]\n",
    "\n",
    "print(\"Running: {}\".format(\" \".join(keep_valid_cmd)))\n",
    "keep_valid_cmd_str = \" \".join(keep_valid_cmd)\n",
    "!{keep_valid_cmd_str}\n",
    "print(\"Files created: {}.\".format(os.listdir(keep_valid_dir)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### [Again] Check whether number of ids in train, validation, and test are the same."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Running: python3 /Users/jy006/Documents/Code/BioMedGPS/biomedgps-data/lib/data.py check-ids --input /var/folders/4s/d4nr1sg91ps1k3qz00h28w_r0000gp/T/tmpnluxqzag/keep_valid/train_valid.tsv --input /var/folders/4s/d4nr1sg91ps1k3qz00h28w_r0000gp/T/tmpnluxqzag/keep_valid/valid_valid.tsv --input /var/folders/4s/d4nr1sg91ps1k3qz00h28w_r0000gp/T/tmpnluxqzag/keep_valid/test_valid.tsv --output /var/folders/4s/d4nr1sg91ps1k3qz00h28w_r0000gp/T/tmpnluxqzag/keep_valid/id_checked.tsv\n",
      "Number of ids in train, validation, and test: \n",
      "             file1            file2  ...  num_file2        category\n",
      "0  train_valid.tsv  valid_valid.tsv  ...      31709             ids\n",
      "1  train_valid.tsv  valid_valid.tsv  ...        109  relation_types\n",
      "2  train_valid.tsv   test_valid.tsv  ...      31712             ids\n",
      "3  train_valid.tsv   test_valid.tsv  ...        110  relation_types\n",
      "4  valid_valid.tsv   test_valid.tsv  ...      31712             ids\n",
      "5  valid_valid.tsv   test_valid.tsv  ...        110  relation_types\n",
      "\n",
      "[6 rows x 6 columns]\n",
      "The files are valid\n",
      "Checked files created: /var/folders/4s/d4nr1sg91ps1k3qz00h28w_r0000gp/T/tmpnluxqzag/keep_valid/id_checked.tsv.\n"
     ]
    }
   ],
   "source": [
    "id_checked_file = os.path.join(keep_valid_dir, \"id_checked.tsv\")\n",
    "train_valid_file = os.path.join(keep_valid_dir, \"train_valid.tsv\")\n",
    "valid_valid_file = os.path.join(keep_valid_dir, \"valid_valid.tsv\")\n",
    "test_valid_file = os.path.join(keep_valid_dir, \"test_valid.tsv\")\n",
    "\n",
    "check_ids_cmd = [\n",
    "    \"python3\",\n",
    "    os.path.join(lib_dir, \"data.py\"),\n",
    "    \"check-ids\",\n",
    "    \"--input\",\n",
    "    train_valid_file,\n",
    "    \"--input\",\n",
    "    valid_valid_file,\n",
    "    \"--input\",\n",
    "    test_valid_file,\n",
    "    \"--output\",\n",
    "    id_checked_file,\n",
    "]\n",
    "\n",
    "print(\"Running: {}\".format(\" \".join(check_ids_cmd)))\n",
    "check_ids_cmd_str = \" \".join(check_ids_cmd)\n",
    "!{check_ids_cmd_str}\n",
    "print(\"Checked files created: {}.\".format(id_checked_file))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Running: python3 /Users/jy006/Documents/Code/BioMedGPS/biomedgps-data/lib/data.py hrt --input /var/folders/4s/d4nr1sg91ps1k3qz00h28w_r0000gp/T/tmpnluxqzag/keep_valid/train_valid.tsv --output /var/folders/4s/d4nr1sg91ps1k3qz00h28w_r0000gp/T/tmpnluxqzag/hrt/train.tsv\n",
      "HRT files created: /var/folders/4s/d4nr1sg91ps1k3qz00h28w_r0000gp/T/tmpnluxqzag/hrt/train.tsv.\n",
      "Running: python3 /Users/jy006/Documents/Code/BioMedGPS/biomedgps-data/lib/data.py hrt --input /var/folders/4s/d4nr1sg91ps1k3qz00h28w_r0000gp/T/tmpnluxqzag/keep_valid/valid_valid.tsv --output /var/folders/4s/d4nr1sg91ps1k3qz00h28w_r0000gp/T/tmpnluxqzag/hrt/valid.tsv\n",
      "HRT files created: /var/folders/4s/d4nr1sg91ps1k3qz00h28w_r0000gp/T/tmpnluxqzag/hrt/valid.tsv.\n",
      "Running: python3 /Users/jy006/Documents/Code/BioMedGPS/biomedgps-data/lib/data.py hrt --input /var/folders/4s/d4nr1sg91ps1k3qz00h28w_r0000gp/T/tmpnluxqzag/keep_valid/test_valid.tsv --output /var/folders/4s/d4nr1sg91ps1k3qz00h28w_r0000gp/T/tmpnluxqzag/hrt/test.tsv\n",
      "HRT files created: /var/folders/4s/d4nr1sg91ps1k3qz00h28w_r0000gp/T/tmpnluxqzag/hrt/test.tsv.\n"
     ]
    }
   ],
   "source": [
    "hrt_dir = os.path.join(temp_dir, \"hrt\")\n",
    "os.makedirs(hrt_dir, exist_ok=True)\n",
    "\n",
    "files = [\n",
    "    (os.path.join(keep_valid_dir, \"train_valid.tsv\"), os.path.join(hrt_dir, \"train.tsv\")),\n",
    "    (os.path.join(keep_valid_dir, \"valid_valid.tsv\"), os.path.join(hrt_dir, \"valid.tsv\")),\n",
    "    (os.path.join(keep_valid_dir, \"test_valid.tsv\"), os.path.join(hrt_dir, \"test.tsv\")),\n",
    "]\n",
    "\n",
    "for input_file, output_file in files:\n",
    "    hrt_cmd = [\n",
    "        \"python3\",\n",
    "        os.path.join(lib_dir, \"data.py\"),\n",
    "        \"hrt\",\n",
    "        \"--input\",\n",
    "        input_file,\n",
    "        \"--output\",\n",
    "        output_file,\n",
    "    ]\n",
    "\n",
    "    print(\"Running: {}\".format(\" \".join(hrt_cmd)))\n",
    "    hrt_cmd_str = \" \".join(hrt_cmd)\n",
    "    !{hrt_cmd_str}\n",
    "    print(\"HRT files created: {}.\".format(output_file))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Copy all files to the dataset folder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Please found all files in /Users/jy006/Documents/Code/BioMedGPS/biomedgps-data/datasets/biomedgps-full-v20240127\n"
     ]
    }
   ],
   "source": [
    "os.makedirs(outputdir, exist_ok=True)\n",
    "\n",
    "files = [\n",
    "    (os.path.join(hrt_dir, \"train.tsv\"), os.path.join(outputdir, \"train.tsv\")),\n",
    "    (os.path.join(hrt_dir, \"valid.tsv\"), os.path.join(outputdir, \"valid.tsv\")),\n",
    "    (os.path.join(hrt_dir, \"test.tsv\"), os.path.join(outputdir, \"test.tsv\")),\n",
    "    (\n",
    "        os.path.join(keep_valid_dir, \"id_checked.tsv\"),\n",
    "        os.path.join(outputdir, \"id_checked.tsv\"),\n",
    "    ),\n",
    "    (entity_file, os.path.join(outputdir, \"annotated_entities.tsv\")),\n",
    "    (kg_file, os.path.join(outputdir, \"knowledge_graph.tsv\")),\n",
    "    (annotated_kg_file, os.path.join(outputdir, \"knowledge_graph_annotated.tsv\")),\n",
    "]\n",
    "\n",
    "for f, output_file in files:\n",
    "    subprocess.check_output([\"cp\", f, output_file])\n",
    "\n",
    "print(\"Please found all files in {}\".format(outputdir))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "biomedgps",
   "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.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}