transe.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# TransE\n",
    "\n",
    "Here we will show how to reproduce the TransE model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pydantic import BaseSettings, BaseModel, Field\n",
    "from typing import Optional, Literal, Tuple, Dict, List\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "import torch.nn as nn\n",
    "import torch\n",
    "import numpy as np\n",
    "from abc import ABC, abstractmethod"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Config Class\n",
    "\n",
    "Define out configuration classes that can instantiate configuration items that meet your requirements based on your runtime environment or your own ideas.\n",
    "\n",
    "With the help of configuration classes, you can change the dataset or adjust the hyperparameters of the model without changing the model logic."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DatasetConf(BaseSettings):\n",
    "    \"\"\"\n",
    "    数据集的相关配置信息\n",
    "    \"\"\"\n",
    "    dataset_name: str = Field(title='数据集的名称，方便打印时查看')\n",
    "    base_dir: str = Field(title='数据集的目录')\n",
    "    entity2id_path: str = Field(default='entity2id.txt', title='entity2id 的文件名')\n",
    "    relation2id_path: str = Field(default='relation2id.txt', title='relation2id 的文件名')\n",
    "    train_path: str = Field(default='train.txt', title='training set 的文件')\n",
    "    valid_path: str = Field(default='valid.txt', title='valid set 的文件')\n",
    "    test_path: str = Field(default='test.txt', title='testing set 的目录')\n",
    "\n",
    "\n",
    "class HyperParam(BaseModel):\n",
    "    \"\"\"\n",
    "    超参数\n",
    "    \"\"\"\n",
    "    batch_size: int = 128\n",
    "    valid_batch_size: int = 64\n",
    "    learning_rate: float = 0.001\n",
    "    epoch_size: int = 500\n",
    "    embed_dim: int = 50\n",
    "    norm: int = 1\n",
    "    margin: int = 2.0\n",
    "    valid_freq: int = Field(title='训练过程中，每隔多少次就做一次 valid 来验证是否保存模型')\n",
    "\n",
    "\n",
    "class TrainConf(BaseModel):\n",
    "    \"\"\"\n",
    "    训练的一些配置\n",
    "    \"\"\"\n",
    "    checkpoint_path: str = Field(title='保存模型的路径')\n",
    "    metric_result_path: str = Field(title='运行 test 的 metric 输出位置')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Dataset\n",
    "\n",
    "Defines the classes used to read datasets, including reading entity-to-ID mappings, relationship-to-ID mappings, and triplet collections.\n",
    "\n",
    "+ The `create_mapping` function is used to generate the entity-to-ID mapping dictionary and the relationship-to-ID mapping dictionary.\n",
    "+ `KRLDataset` is a further wrapper around the `Dataset` class in PyTorch, and is similar in usage."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "EntityMapping = Dict[str, int]\n",
    "RelMapping = Dict[str, int]\n",
    "Triple = List[int]\n",
    "\n",
    "def create_mapping(dataset_conf: DatasetConf) -> Tuple[EntityMapping, RelMapping]:\n",
    "    \"\"\"\n",
    "    create mapping of `entity2id` and `relation2id`\n",
    "    \"\"\"\n",
    "    # 读取 entity2id\n",
    "    entity2id = dict()\n",
    "    with open(dataset_conf.base_dir + dataset_conf.entity2id_path) as f:\n",
    "        for line in f:\n",
    "            entity, entity_id = line.split()\n",
    "            entity = entity.strip()\n",
    "            entity_id = int(entity_id.strip())\n",
    "            entity2id[entity] = entity_id\n",
    "    # 读取 relation2id\n",
    "    rel2id = dict()\n",
    "    with open(dataset_conf.base_dir + dataset_conf.relation2id_path) as f:\n",
    "        for line in f:\n",
    "            rel, rel_id = line.split()\n",
    "            rel = rel.strip()\n",
    "            rel_id = int(rel_id.strip())\n",
    "            rel2id[rel] = rel_id\n",
    "    return entity2id, rel2id\n",
    "\n",
    "\n",
    "class KRLDataset(Dataset):\n",
    "    def __init__(self,\n",
    "                 dataset_conf: DatasetConf,\n",
    "                 mode: Literal['train', 'valid', 'test'],\n",
    "                 entity2id: Dict[str, int],\n",
    "                 rel2id: Dict[str, int]) -> None:\n",
    "        super().__init__()\n",
    "        self.conf = dataset_conf\n",
    "        self.mode = mode\n",
    "        self.triples = []\n",
    "        self.entity2id = entity2id\n",
    "        self.rel2id = rel2id\n",
    "        self._read_triples()    # 读取数据集，并获得所有的 triples\n",
    "    \n",
    "    def _split_and_to_id(self, line: str) -> Triple:\n",
    "        \"\"\"将数据集文件中的一行数据进行切分，并将 entity 和 rel 转换成 id\n",
    "\n",
    "        :param line: 数据集的一行数据\n",
    "        :return: [head_id, rel_id, tail_id]\n",
    "        \"\"\"\n",
    "        head, tail, rel = line.split()\n",
    "        head_id = self.entity2id[head.strip()]\n",
    "        rel_id = self.rel2id[rel.strip()]\n",
    "        tail_id = self.entity2id[tail.strip()]\n",
    "        return (head_id, rel_id, tail_id)\n",
    "    \n",
    "    def _read_triples(self):\n",
    "        data_path = {\n",
    "            'train': self.conf.train_path,\n",
    "            'valid': self.conf.valid_path,\n",
    "            'test': self.conf.test_path\n",
    "        }.get(self.mode)\n",
    "        with open(self.conf.base_dir + data_path) as f:\n",
    "            self.triples = [self._split_and_to_id(line) for line in f]\n",
    "    \n",
    "    def __len__(self):\n",
    "        \"\"\"Denotes the total number of samples.\"\"\"\n",
    "        return len(self.triples)\n",
    "    \n",
    "    def __getitem__(self, index) -> Triple:\n",
    "        \"\"\"Returns (head id, relation id, tail id).\"\"\"\n",
    "        triple = self.triples[index]\n",
    "        return triple[0], triple[1], triple[2]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Negative Sampler\n",
    "\n",
    "In order to train the model, we need not only positive samples, but also negative samples. The goal of the negative sampler is to generate negative samples based on the positive samples in the dataset.\n",
    "\n",
    "Since there are multiple negative sampling strategies, we abstract a common abstract class `NegativeSampler`, and all negative samplers that implement different negative sampling strategies should inherit from this abstract base class."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "class NegativeSampler(ABC):\n",
    "    def __init__(self, dataset: KRLDataset, device: torch.device):\n",
    "        self.dataset = dataset\n",
    "        self.device = device\n",
    "    \n",
    "    @abstractmethod\n",
    "    def neg_sample(self, heads, rels, tails):\n",
    "        \"\"\"执行负采样\n",
    "\n",
    "        :param heads: 由 batch_size 个 head idx 组成的 tensor，size: [batch_size]\n",
    "        :param rels: size [batch_size]\n",
    "        :param tails: size [batch_size]\n",
    "        \"\"\"\n",
    "        pass\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The simplest negative sampling strategy is to randomly replace the head entity or tail entity in a triplet to obtain a negative sample."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class RandomNegativeSampler(NegativeSampler):\n",
    "    \"\"\"\n",
    "    随机替换 head 或者 tail 来实现采样\n",
    "    \"\"\"\n",
    "    def __init__(self, dataset: KRLDataset, device: torch.device):\n",
    "        super().__init__(dataset, device)\n",
    "        \n",
    "    def neg_sample(self, heads, rels, tails):\n",
    "        ent_num = len(self.dataset.entity2id)\n",
    "        head_or_tail = torch.randint(high=2, size=heads.size(), device=self.device)\n",
    "        random_entities = torch.randint(high=ent_num, size=heads.size(), device=self.device)\n",
    "        corupted_heads = torch.where(head_or_tail == 1, random_entities, heads)\n",
    "        corupted_tails = torch.where(head_or_tail == 0, random_entities, tails)\n",
    "        return torch.stack([corupted_heads, rels, corupted_tails], dim=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Model\n",
    "\n",
    "Defining the TransE model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "class TransE(nn.Module):\n",
    "    def __init__(\n",
    "        self,\n",
    "        ent_num: int,\n",
    "        rel_num: int,\n",
    "        device: torch.device,\n",
    "        norm: int,\n",
    "        embed_dim: int,\n",
    "        margin: float\n",
    "    ):\n",
    "        super().__init__()\n",
    "        self.ent_num = ent_num\n",
    "        self.rel_num = rel_num\n",
    "        self.device = device\n",
    "        self.norm = norm\n",
    "        self.embed_dim = embed_dim\n",
    "        self.margin = margin\n",
    "\n",
    "        # Initialize ent_embedding\n",
    "        self.ent_embedding = nn.Embedding(self.ent_num, self.embed_dim)\n",
    "        torch.nn.init.xavier_uniform_(self.ent_embedding.weight.data)\n",
    "        #uniform_range = 6 / np.sqrt(self.embed_dim)\n",
    "        #self.ent_embedding.weight.data.uniform_(-uniform_range, uniform_range)\n",
    "        \n",
    "        # Initialize rel_embedding\n",
    "        self.rel_embedding = nn.Embedding(self.rel_num, self.embed_dim)\n",
    "        torch.nn.init.xavier_uniform_(self.rel_embedding.weight.data)\n",
    "        #uniform_range = 6 / np.sqrt(self.embed_dim)\n",
    "        #self.rel_embedding.weight.data.uniform_(-uniform_range, uniform_range)\n",
    "\n",
    "        self.criterion = nn.MarginRankingLoss(margin=self.margin)\n",
    "    \n",
    "    def _distance(self, triples):\n",
    "        \"\"\"Calculate the distance of a batch's triplet\n",
    "\n",
    "        :param triples: triples of a batch，size: [batch, 3]\n",
    "        :return: size: [batch,]\n",
    "        \"\"\"\n",
    "        heads = triples[:, 0]\n",
    "        rels = triples[:, 1]\n",
    "        tails = triples[:, 2]\n",
    "        h_embs = self.ent_embedding(heads)  # h_embs: [batch, embed_dim]\n",
    "        r_embs = self.rel_embedding(rels)\n",
    "        t_embs = self.ent_embedding(tails)\n",
    "        dist = h_embs + r_embs - t_embs  # [batch, embed_dim]\n",
    "        return torch.norm(dist, p=self.norm, dim=1)\n",
    "        \n",
    "    def loss(self, pos_distances, neg_distances):\n",
    "        \"\"\"Calculate the loss of TransE training\n",
    "\n",
    "        :param pos_distances: [batch, ]\n",
    "        :param neg_distances: [batch, ]\n",
    "        :return: loss\n",
    "        \"\"\"\n",
    "        ones = torch.tensor([-1], dtype=torch.long, device=self.device)\n",
    "        return self.criterion(pos_distances, neg_distances, ones)\n",
    "    \n",
    "    def forward(self, pos_triples: torch.Tensor, neg_triples: torch.Tensor):\n",
    "        \"\"\"Return model losses based on the input.\n",
    "\n",
    "        :param pos_triples: triplets of positives in Bx3 shape (B - batch, 3 - head, relation and tail)\n",
    "        :param neg_triples: triplets of negatives in Bx3 shape (B - batch, 3 - head, relation and tail)\n",
    "        :return: tuple of the model loss, positive triplets loss component, negative triples loss component\n",
    "        \"\"\"\n",
    "        assert pos_triples.size()[1] == 3\n",
    "        assert neg_triples.size()[1] == 3\n",
    "        \n",
    "        pos_distances = self._distance(pos_triples)\n",
    "        neg_distances = self._distance(neg_triples)\n",
    "        loss = self.loss(pos_distances, neg_distances)\n",
    "        return loss, pos_distances, neg_distances\n",
    "    \n",
    "    def predict(self, triples: torch.Tensor):\n",
    "        \"\"\"Calculated dissimilarity score for given triplets.\n",
    "\n",
    "        :param triplets: triplets in Bx3 shape (B - batch, 3 - head, relation and tail)\n",
    "        :return: dissimilarity score for given triplets\n",
    "        \"\"\"\n",
    "        return self._distance(triples)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Metric\n",
    "\n",
    "Calculate the metric for measuring the effect of link prediction, i.e. MRR and hits@10."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# metric\n",
    "\n",
    "def cal_hits_at_k(\n",
    "    predictions: torch.Tensor,\n",
    "    ground_truth_idx: torch.Tensor,\n",
    "    device: torch.device,\n",
    "    k: int\n",
    ") -> float:\n",
    "    \"\"\"Calculates number of hits@k.\n",
    "\n",
    "    :param predictions: BxN tensor of prediction values where B is batch size and N number of classes. Predictions\n",
    "    must be sorted in class ids order\n",
    "    :param ground_truth_idx: Bx1 tensor with index of ground truth class\n",
    "    :param k: number of top K results to be considered as hits\n",
    "    :return: Hits@K scoreH\n",
    "    \"\"\"\n",
    "    assert predictions.size()[0] == ground_truth_idx.size()[0]  # has the same batch_size\n",
    "    \n",
    "    zero_tensor = torch.tensor([0], device=device)\n",
    "    one_tensor = torch.tensor([1], device=device)\n",
    "    _, indices = predictions.topk(k, largest=False)  # indices: [batch_size, k]\n",
    "    where_flags = indices == ground_truth_idx  # where_flags: [batch_size, k], type: bool\n",
    "    hits = torch.where(where_flags, one_tensor, zero_tensor).sum().item()\n",
    "    return hits\n",
    "\n",
    "def cal_mrr(predictions: torch.Tensor, ground_truth_idx: torch.Tensor) -> float:\n",
    "    \"\"\"Calculates mean reciprocal rank (MRR) for given predictions and ground truth values.\n",
    "\n",
    "    :param predictions: BxN tensor of prediction values where B is batch size and N number of classes. Predictions\n",
    "    must be sorted in class ids order\n",
    "    :param ground_truth_idx: Bx1 tensor with index of ground truth class\n",
    "    :return: Mean reciprocal rank score\n",
    "    \"\"\"\n",
    "    assert predictions.size(0) == ground_truth_idx.size(0)\n",
    "\n",
    "    indices = predictions.argsort()\n",
    "    return (1.0 / (indices == ground_truth_idx).nonzero()[:, 1].float().add(1.0)).sum().item()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Inference Operation\n",
    "\n",
    "Run the inference process for the model, i.e., iterate through the validation or test set and compute the metric."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def run_testing(\n",
    "    model: TransE,\n",
    "    dataloader: DataLoader,\n",
    "    ent_num: int,\n",
    "    device: torch.device,\n",
    ") -> Tuple[float, float, float, float]:\n",
    "    \"\"\"Run test programs against Trans models\n",
    "\n",
    "    :param model: TransE model\n",
    "    :param ent_num: Number of entities in the dataset\n",
    "    :return: _description_\n",
    "    \"\"\"\n",
    "    hits_at_1 = 0.0\n",
    "    hits_at_3 = 0.0\n",
    "    hits_at_10 = 0.0\n",
    "    mrr = 0.0\n",
    "    examples_count = 0\n",
    "    \n",
    "    # entity_ids = [[0, 1, 2, ..., ent_num]], shape: [1, ent_num]\n",
    "    entitiy_ids = torch.arange(0, ent_num, device=device).unsqueeze(0)\n",
    "    for i, batch in enumerate(dataloader):\n",
    "        # batch: [3, batch_size]\n",
    "        heads, rels, tails = batch[0].to(device), batch[1].to(device), batch[2].to(device)\n",
    "        batch_size = heads.size()[0]\n",
    "        all_entities = entitiy_ids.repeat(batch_size, 1)  # all_entities: [batch_size, ent_num]\n",
    "        # heads: [batch_size,] -> [batch_size, 1] -> [batch_size, ent_num]\n",
    "        heads_expanded = heads.reshape(-1, 1).repeat(1, ent_num)  # _expanded: [batch_size, ent_num]\n",
    "        rels_expanded = rels.reshape(-1, 1).repeat(1, ent_num)\n",
    "        tails_expanded = tails.reshape(-1, 1).repeat(1, ent_num)\n",
    "        # check all possible tails\n",
    "        triplets = torch.stack([heads_expanded, rels_expanded, all_entities], dim=2).reshape(-1, 3)  # triplets: [batch_size * ent_num, 3]\n",
    "        tails_predictions = model.predict(triplets).reshape(batch_size, -1)  # tails_prediction: [batch_size, ent_num]\n",
    "        # check all possible heads\n",
    "        triplets = torch.stack([all_entities, rels_expanded, tails_expanded], dim=2).reshape(-1, 3)\n",
    "        heads_predictions = model.predict(triplets).reshape(batch_size, -1)  # heads_prediction: [batch_size, ent_num]\n",
    "        \n",
    "        # Concept preditions\n",
    "        predictions = torch.cat([tails_predictions, heads_predictions], dim=0)  # predictions: [batch_size * 2, ent_num]\n",
    "        ground_truth_entity_id = torch.cat([tails.reshape(-1, 1), heads.reshape(-1, 1)], dim=0)  # [batch_size * 2, 1]\n",
    "        # calculate metrics\n",
    "        hits_at_1 += cal_hits_at_k(predictions, ground_truth_entity_id, device=device, k=1)\n",
    "        hits_at_3 += cal_hits_at_k(predictions, ground_truth_entity_id, device=device, k=3)\n",
    "        hits_at_10 += cal_hits_at_k(predictions, ground_truth_entity_id, device=device, k=10)\n",
    "        mrr += cal_mrr(predictions, ground_truth_entity_id)\n",
    "        \n",
    "        examples_count += predictions.size()[0]\n",
    "    \n",
    "    hits_at_1_score = hits_at_1 / examples_count * 100\n",
    "    hits_at_3_score = hits_at_3 / examples_count * 100\n",
    "    hits_at_10_score = hits_at_10 / examples_count * 100\n",
    "    mrr_score = mrr / examples_count * 100\n",
    "    \n",
    "    return hits_at_1_score, hits_at_3_score, hits_at_10_score, mrr_score"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Checkpoint\n",
    "\n",
    "During the training process, if the model outperforms the best score on the validation set, the model state at that time should be transformed into a checkpoint and saved to disk.\n",
    "\n",
    "The process of storing and loading checkpoints is simply encapsulated here."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "class CheckpointFormat(BaseModel):\n",
    "    model_state_dict: dict\n",
    "    optim_state_dict: dict\n",
    "    epoch_id: int\n",
    "    best_score: float\n",
    "\n",
    "\n",
    "def save_checkpoint(model: TransE,\n",
    "                    optimzer: torch.optim.Optimizer,\n",
    "                    epoch_id: int,\n",
    "                    best_score: float,\n",
    "                    train_conf: TrainConf):\n",
    "    ckpt = CheckpointFormat(\n",
    "        model_state_dict=model.state_dict(),\n",
    "        optim_state_dict=optimzer.state_dict(),\n",
    "        epoch_id=epoch_id,\n",
    "        best_score=best_score\n",
    "    )\n",
    "    torch.save(ckpt.dict(), train_conf.checkpoint_path)\n",
    "\n",
    "\n",
    "def load_checkpoint(train_conf: TrainConf) -> CheckpointFormat:\n",
    "    ckpt = torch.load(train_conf.checkpoint_path)\n",
    "    return CheckpointFormat.parse_obj(ckpt)\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training Operation\n",
    "\n",
    "The process of training a model using a dataset.\n",
    "\n",
    "In the real library, this part of the functionality is encapsulated in a `Trainer` class."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "def run_training(model: TransE,\n",
    "                 train_conf: TrainConf,\n",
    "                 params: HyperParam,\n",
    "                 device: torch.device,\n",
    "                 dataset_conf: DatasetConf,\n",
    "                 entity2id: Dict[str, int],\n",
    "                 rel2id: Dict[str, int]):\n",
    "    # 准备数据集\n",
    "    train_dataset = KRLDataset(dataset_conf, 'train', entity2id, rel2id)\n",
    "    valid_dataset = KRLDataset(dataset_conf, 'valid', entity2id, rel2id)\n",
    "    # dataset -> dataloader\n",
    "    train_dataloder = DataLoader(train_dataset, params.batch_size)\n",
    "    valid_dataloder = DataLoader(valid_dataset, params.valid_batch_size)\n",
    "    # 负采样器\n",
    "    train_neg_sampler = RandomNegativeSampler(train_dataset, device)\n",
    "    valid_neg_sampler = RandomNegativeSampler(valid_dataset, device)\n",
    "    # 准备训练的工具\n",
    "    optimzer = torch.optim.Adam(model.parameters(), lr=params.learning_rate)\n",
    "    min_valid_loss = 10000.0\n",
    "    best_score = 0.0\n",
    "    # training loop\n",
    "    for epoch_id in range(1, params.epoch_size + 1):\n",
    "        print(\"Starting epoch: \", epoch_id)\n",
    "        loss_sum = 0\n",
    "        model.train()\n",
    "        for i, batch in enumerate(train_dataloder):\n",
    "            # 获取一个 batch 的训练资料\n",
    "            pos_heads, pos_rels, pos_tails = batch[0].to(device), batch[1].to(device), batch[2].to(device)\n",
    "            pos_triples = torch.stack([pos_heads, pos_rels, pos_tails], dim=1)  # pos_triples: [batch_size, 3]\n",
    "            neg_triples = train_neg_sampler.neg_sample(pos_heads, pos_rels, pos_tails)  # neg_triples: [batch_size, 3]\n",
    "            optimzer.zero_grad()\n",
    "            # 计算 loss\n",
    "            loss, pos_dist, neg_dist = model(pos_triples, neg_triples)\n",
    "            loss.backward()\n",
    "            loss_sum += loss.cpu().item()\n",
    "            # update model\n",
    "            optimzer.step()\n",
    "            \n",
    "        if epoch_id % params.valid_freq == 0:\n",
    "            model.eval()\n",
    "            _, _, hits_at_10, _ = run_testing(model, valid_dataloder, len(valid_dataset.entity2id), device)\n",
    "            score = hits_at_10\n",
    "            print('valid hits@10:', score)\n",
    "            if score > best_score:\n",
    "                best_score = score\n",
    "                print('best score of valid: ', best_score)\n",
    "                save_checkpoint(model, optimzer, epoch_id, best_score, train_conf)\n",
    "                "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_device() -> torch.device:\n",
    "    return torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')\n",
    "\n",
    "\n",
    "def main(dataset_conf: DatasetConf,\n",
    "         params: HyperParam,\n",
    "         train_conf: TrainConf,\n",
    "         device: torch.device):\n",
    "    entity2id, rel2id = create_mapping(dataset_conf)\n",
    "    device = get_device()\n",
    "    ent_num = len(entity2id)\n",
    "    rel_num = len(rel2id)\n",
    "    model = TransE(ent_num, rel_num, device,\n",
    "                   norm=params.norm,\n",
    "                   embed_dim=params.embed_dim,\n",
    "                   margin=params.margin)\n",
    "    model = model.to(device)\n",
    "    run_training(model, train_conf, params, device, dataset_conf, entity2id, rel2id)\n",
    "    \n",
    "    # Testing the best checkpoint on test dataset\n",
    "    ckpt = load_checkpoint(train_conf)\n",
    "    model.load_state_dict(ckpt.model_state_dict)\n",
    "    model = model.to(device)\n",
    "    test_dataset = KRLDataset(dataset_conf, 'test', entity2id, rel2id)\n",
    "    test_dataloder = DataLoader(test_dataset, params.valid_batch_size)\n",
    "    hits_at_1, hits_at_3, hits_at_10, mrr = run_testing(model, test_dataloder, ent_num, device)\n",
    "    \n",
    "    # write results\n",
    "    with open(train_conf.metric_result_path, 'w') as f:\n",
    "         f.write(f'dataset: {dataset_conf.dataset_name}\\n')\n",
    "         f.write(f'Hits@1: {hits_at_1}\\n')\n",
    "         f.write(f'Hits@3: {hits_at_3}\\n')\n",
    "         f.write(f'Hits@10: {hits_at_10}\\n')\n",
    "         f.write(f'MRR: {mrr}\\n')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Begin\n",
    "\n",
    "Instantiate the configuration and call the main function.\n",
    "\n",
    "Next, you need the path of the dataset and where to save the checkpoints."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "fb15k_dataset_conf = DatasetConf(\n",
    "    dataset_name='FB15K',\n",
    "    base_dir='/root/yubin/dataset/KRL/master/FB15k/'   # TODO: change it!\n",
    ")\n",
    "\n",
    "fb15k_hyper_params = HyperParam(\n",
    "    valid_freq=5,\n",
    "    batch_size=128,\n",
    "    valid_batch_size=64,\n",
    "    learning_rate=0.001,\n",
    "    epoch_size=500,\n",
    "    embed_dim=50,\n",
    "    norm=1,\n",
    "    margin=2.0\n",
    ")\n",
    "\n",
    "fb15k_train_conf = TrainConf(\n",
    "    checkpoint_path='/root/sharespace/yubin/papers/KRL/scratch/TransX/tmp/transe_fb15k.ckpt',  # TODO: change it!\n",
    "    metric_result_path='/root/sharespace/yubin/papers/KRL/scratch/TransX/tmp/transe_fb15k_metrics.txt'   # TODO: change it!\n",
    ")\n",
    "\n",
    "device = get_device()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "main(fb15k_dataset_conf, fb15k_hyper_params, fb15k_train_conf, device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "ckpt = torch.load(fb15k_train_conf.checkpoint_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'model_state_dict': {'ent_embedding.weight': tensor([[ 0.0126,  1.1495,  0.7716,  ..., -0.8967,  0.3144, -0.2118],\n",
       "          [ 0.6201, -1.0493,  0.7837,  ...,  0.2150,  0.2489,  0.0075],\n",
       "          [-0.2746, -0.2592,  0.1011,  ..., -1.1602,  0.1287,  0.4214],\n",
       "          ...,\n",
       "          [ 0.8894,  0.2856,  0.2504,  ...,  0.8878,  0.9393, -0.2801],\n",
       "          [ 0.8036,  0.7737,  0.1246,  ...,  0.2968, -0.0092, -0.2363],\n",
       "          [-1.0176, -0.0581, -0.5224,  ..., -0.3004, -1.3833, -0.6132]],\n",
       "         device='cuda:0'),\n",
       "  'rel_embedding.weight': tensor([[-0.0909, -0.4988, -0.7200,  ..., -0.0086,  0.0742, -0.1348],\n",
       "          [-0.9219,  0.5674,  0.4704,  ..., -1.6722, -1.1977, -0.0505],\n",
       "          [-0.6558, -0.0297, -0.1181,  ..., -0.2469, -0.2535, -0.6061],\n",
       "          ...,\n",
       "          [ 0.3357,  0.2950, -0.2752,  ...,  0.5626,  0.3279, -0.5322],\n",
       "          [-0.3469, -0.0041, -0.7309,  ...,  0.1166,  0.0848,  0.3135],\n",
       "          [ 0.0369, -0.0189, -0.0132,  ...,  0.1579, -0.0358, -0.0636]],\n",
       "         device='cuda:0')},\n",
       " 'optim_state_dict': {'state': {0: {'step': tensor(1887500.),\n",
       "    'exp_avg': tensor([[-5.6052e-45,  5.6052e-45, -5.6052e-45,  ..., -5.6052e-45,\n",
       "             -5.6052e-45, -5.6052e-45],\n",
       "            [ 5.6052e-45, -5.6052e-45,  5.6052e-45,  ..., -5.6052e-45,\n",
       "             -5.6052e-45, -5.6052e-45],\n",
       "            [ 5.6052e-45,  5.6052e-45,  5.6052e-45,  ..., -5.6052e-45,\n",
       "             -5.6052e-45,  5.6052e-45],\n",
       "            ...,\n",
       "            [ 5.6052e-45,  5.6052e-45,  5.6052e-45,  ..., -5.6052e-45,\n",
       "              5.6052e-45,  5.6052e-45],\n",
       "            [-5.6052e-45,  5.6052e-45, -5.6052e-45,  ...,  5.6052e-45,\n",
       "              5.6052e-45,  5.6052e-45],\n",
       "            [-3.9882e-41,  3.9814e-41,  3.9882e-41,  ..., -3.9882e-41,\n",
       "             -3.9882e-41,  3.9882e-41]], device='cuda:0'),\n",
       "    'exp_avg_sq': tensor([[5.3320e-10, 5.3320e-10, 5.3320e-10,  ..., 5.3320e-10, 5.3320e-10,\n",
       "             5.3320e-10],\n",
       "            [1.1443e-09, 1.1443e-09, 1.1443e-09,  ..., 1.1443e-09, 1.8754e-08,\n",
       "             1.8754e-08],\n",
       "            [3.3255e-08, 3.3255e-08, 2.2694e-08,  ..., 3.3255e-08, 3.3255e-08,\n",
       "             2.2694e-08],\n",
       "            ...,\n",
       "            [3.6491e-13, 3.6491e-13, 3.6491e-13,  ..., 3.6491e-13, 3.6491e-13,\n",
       "             3.6491e-13],\n",
       "            [2.1251e-13, 2.1251e-13, 2.1251e-13,  ..., 2.1251e-13, 2.1251e-13,\n",
       "             2.1251e-13],\n",
       "            [5.5608e-08, 5.5608e-08, 5.5608e-08,  ..., 5.5608e-08, 5.5608e-08,\n",
       "             5.5608e-08]], device='cuda:0')},\n",
       "   1: {'step': tensor(1887500.),\n",
       "    'exp_avg': tensor([[-5.6052e-45, -5.6052e-45,  5.6052e-45,  ..., -5.6052e-45,\n",
       "              5.6052e-45,  5.6052e-45],\n",
       "            [ 5.6052e-45, -5.6052e-45,  5.6052e-45,  ...,  5.6052e-45,\n",
       "              5.6052e-45,  5.6052e-45],\n",
       "            [ 5.6052e-45,  5.6052e-45,  5.6052e-45,  ..., -5.6052e-45,\n",
       "              5.6052e-45, -5.6052e-45],\n",
       "            ...,\n",
       "            [ 5.6052e-45,  5.6052e-45,  5.6052e-45,  ...,  5.6052e-45,\n",
       "              5.6052e-45,  5.6052e-45],\n",
       "            [ 5.6052e-45,  5.6052e-45,  5.6052e-45,  ..., -5.6052e-45,\n",
       "             -5.6052e-45,  5.6052e-45],\n",
       "            [-5.6052e-45,  5.6052e-45,  5.6052e-45,  ..., -5.6052e-45,\n",
       "             -5.6052e-45,  5.6052e-45]], device='cuda:0'),\n",
       "    'exp_avg_sq': tensor([[7.0065e-43, 3.1465e-32, 3.1465e-32,  ..., 4.8395e-37, 7.0065e-43,\n",
       "             7.0065e-43],\n",
       "            [2.6102e-08, 2.8141e-08, 1.2741e-07,  ..., 2.5118e-10, 2.0500e-08,\n",
       "             1.0558e-07],\n",
       "            [2.6078e-24, 7.9196e-29, 7.0065e-43,  ..., 2.6077e-24, 7.9196e-29,\n",
       "             7.0065e-43],\n",
       "            ...,\n",
       "            [6.1684e-23, 1.7540e-34, 7.0065e-43,  ..., 7.0065e-43, 1.7540e-34,\n",
       "             1.7540e-34],\n",
       "            [7.0065e-43, 4.4173e-41, 4.4173e-41,  ..., 7.0065e-43, 7.0065e-43,\n",
       "             7.0065e-43],\n",
       "            [7.0065e-43, 7.0065e-43, 7.0065e-43,  ..., 7.0065e-43, 7.0065e-43,\n",
       "             7.0065e-43]], device='cuda:0')}},\n",
       "  'param_groups': [{'lr': 0.001,\n",
       "    'betas': (0.9, 0.999),\n",
       "    'eps': 1e-08,\n",
       "    'weight_decay': 0,\n",
       "    'amsgrad': False,\n",
       "    'maximize': False,\n",
       "    'foreach': None,\n",
       "    'capturable': False,\n",
       "    'differentiable': False,\n",
       "    'fused': False,\n",
       "    'params': [0, 1]}]},\n",
       " 'epoch_id': 500,\n",
       " 'best_score': 39.291}"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ckpt"
   ]
  }
 ],
 "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.9.15"
  },
  "vscode": {
   "interpreter": {
    "hash": "0418effca45178467ac68c18e34d93809a092be692e0a4443d8690099b71f4bc"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}