files: import itertools
  import os
  import math
  import time
  from collections import defaultdict
  
  from dataclasses import dataclass
  
  import torch
  import torch.utils.data
  
  from kge import Config, Dataset
  from kge.job import Job
  from kge.model import KgeModel
  
  from kge.util import KgeLoss, KgeOptimizer, KgeSampler, KgeLRScheduler
  from typing import Any, Callable, Dict, List, Optional
  import kge.job.util
  
  SLOTS = [0, 1, 2]
  S, P, O = SLOTS
  
  
  class TrainingJob(Job):
      """Abstract base job to train a single model with a fixed set of hyperparameters.
  
      Also used by jobs such as :class:`SearchJob`.
  
      Subclasses for specific training methods need to implement `_prepare` and
      `_process_batch`.
  
      """
  
      def __init__(
          self, config: Config, dataset: Dataset, parent_job: Job = None
      ) -> None:
          from kge.job import EvaluationJob
  
          super().__init__(config, dataset, parent_job)
          self.model: KgeModel = KgeModel.create(config, dataset)
          self.optimizer = KgeOptimizer.create(config, self.model)
          self.kge_lr_scheduler = KgeLRScheduler(config, self.optimizer)
          self.loss = KgeLoss.create(config)
          self.abort_on_nan: bool = config.get("train.abort_on_nan")
          self.batch_size: int = config.get("train.batch_size")
          self.device: str = self.config.get("job.device")
          self.train_split = config.get("train.split")
          valid_conf = config.clone()
          valid_conf.set("job.type", "eval")
          if self.config.get("valid.split") != "":
              valid_conf.set("eval.split", self.config.get("valid.split"))
          valid_conf.set("eval.trace_level", self.config.get("valid.trace_level"))
          self.valid_job = EvaluationJob.create(
              valid_conf, dataset, parent_job=self, model=self.model
          )
          self.config.check("train.trace_level", ["batch", "epoch"])
          self.trace_batch: bool = self.config.get("train.trace_level") == "batch"
          self.epoch: int = 0
          self.valid_trace: List[Dict[str, Any]] = []
          self.is_prepared = False
          self.model.train()
  
          # attributes filled in by implementing classes
          self.loader = None
          self.num_examples = None
          self.type_str: Optional[str] = None
  
          #: Hooks run after training for an epoch.
          #: Signature: job, trace_entry
          self.post_epoch_hooks: List[Callable[[Job, Dict[str, Any]], Any]] = []
  
          #: Hooks run before starting a batch.
          #: Signature: job
          self.pre_batch_hooks: List[Callable[[Job], Any]] = []
  
          #: Hooks run before outputting the trace of a batch. Can modify trace entry.
          #: Signature: job, trace_entry
          self.post_batch_trace_hooks: List[Callable[[Job, Dict[str, Any]], Any]] = []
  
          #: Hooks run before outputting the trace of an epoch. Can modify trace entry.
          #: Signature: job, trace_entry
          self.post_epoch_trace_hooks: List[Callable[[Job, Dict[str, Any]], Any]] = []
  
          #: Hooks run after a validation job.
          #: Signature: job, trace_entry
          self.post_valid_hooks: List[Callable[[Job, Dict[str, Any]], Any]] = []
  
          #: Hooks run after training
          #: Signature: job, trace_entry
          self.post_train_hooks: List[Callable[[Job, Dict[str, Any]], Any]] = []
  
          if self.__class__ == TrainingJob:
              for f in Job.job_created_hooks:
                  f(self)
  
      @staticmethod
      def create(
          config: Config, dataset: Dataset, parent_job: Job = None
      ) -> "TrainingJob":
          """Factory method to create a training job."""
          if config.get("train.type") == "KvsAll":
              return TrainingJobKvsAll(config, dataset, parent_job)
          elif config.get("train.type") == "negative_sampling":
              return TrainingJobNegativeSampling(config, dataset, parent_job)
          elif config.get("train.type") == "1vsAll":
              return TrainingJob1vsAll(config, dataset, parent_job)
          else:
              # perhaps TODO: try class with specified name -> extensibility
              raise ValueError("train.type")
  
      def run(self) -> None:
          """Start/resume the training job and run to completion."""
          self.config.log("Starting training...")
          checkpoint_every = self.config.get("train.checkpoint.every")
          checkpoint_keep = self.config.get("train.checkpoint.keep")
          metric_name = self.config.get("valid.metric")
          patience = self.config.get("valid.early_stopping.patience")
          while True:
              # checking for model improvement according to metric_name
              # and do early stopping and keep the best checkpoint
              if (
                  len(self.valid_trace) > 0
                  and self.valid_trace[-1]["epoch"] == self.epoch
              ):
                  best_index = max(
                      range(len(self.valid_trace)),
                      key=lambda index: self.valid_trace[index][metric_name],
                  )
                  if best_index == len(self.valid_trace) - 1:
                      self.save(self.config.checkpoint_file("best"))
                  if (
                      patience > 0
                      and len(self.valid_trace) > patience
                      and best_index < len(self.valid_trace) - patience
                  ):
                      self.config.log(
                          "Stopping early ({} did not improve over best result ".format(
                              metric_name
                          )
                          + "in the last {} validation runs).".format(patience)
                      )
                      break
                  if self.epoch > self.config.get(
                      "valid.early_stopping.min_threshold.epochs"
                  ) and self.valid_trace[best_index][metric_name] < self.config.get(
                      "valid.early_stopping.min_threshold.metric_value"
                  ):
                      self.config.log(
                          "Stopping early ({} did not achieve min treshold after {} epochs".format(
                              metric_name, self.epoch
                          )
                      )
                      break
  
              # should we stop?
              if self.epoch >= self.config.get("train.max_epochs"):
                  self.config.log("Maximum number
  
  [... the rest is omitted ...]
files: kge/job/train.py
files: import itertools
  import os
  import math
  import time
  from collections import defaultdict
  
  from dataclasses import dataclass
  
  import torch
  import torch.utils.data
  
  from kge import Config, Dataset
  from kge.job import Job
  from kge.model import KgeModel
  
  from kge.util import KgeLoss, KgeOptimizer, KgeSampler, KgeLRScheduler
  from typing import Any, Callable, Dict, List, Optional
  import kge.job.util
  
  SLOTS = [0, 1, 2]
  S, P, O = SLOTS
  
  
  class TrainingJob(Job):
      """Abstract base job to train a single model with a fixed set of hyperparameters.
  
      Also used by jobs such as :class:`SearchJob`.
  
      Subclasses for specific training methods need to implement `_prepare` and
      `_process_batch`.
  
      """
  
      def __init__(
          self, config: Config, dataset: Dataset, parent_job: Job = None
      ) -> None:
          from kge.job import EvaluationJob
  
          super().__init__(config, dataset, parent_job)
          self.model: KgeModel = KgeModel.create(config, dataset)
          self.optimizer = KgeOptimizer.create(config, self.model)
          self.kge_lr_scheduler = KgeLRScheduler(config, self.optimizer)
          self.loss = KgeLoss.create(config)
          self.abort_on_nan: bool = config.get("train.abort_on_nan")
          self.batch_size: int = config.get("train.batch_size")
          self.device: str = self.config.get("job.device")
          self.train_split = config.get("train.split")
          valid_conf = config.clone()
          valid_conf.set("job.type", "eval")
          if self.config.get("valid.split") != "":
              valid_conf.set("eval.split", self.config.get("valid.split"))
          valid_conf.set("eval.trace_level", self.config.get("valid.trace_level"))
          self.valid_job = EvaluationJob.create(
              valid_conf, dataset, parent_job=self, model=self.model
          )
          self.config.check("train.trace_level", ["batch", "epoch"])
          self.trace_batch: bool = self.config.get("train.trace_level") == "batch"
          self.epoch: int = 0
          self.valid_trace: List[Dict[str, Any]] = []
          self.is_prepared = False
          self.model.train()
  
          # attributes filled in by implementing classes
          self.loader = None
          self.num_examples = None
          self.type_str: Optional[str] = None
  
          #: Hooks run after training for an epoch.
          #: Signature: job, trace_entry
          self.post_epoch_hooks: List[Callable[[Job, Dict[str, Any]], Any]] = []
  
          #: Hooks run before starting a batch.
          #: Signature: job
          self.pre_batch_hooks: List[Callable[[Job], Any]] = []
  
          #: Hooks run before outputting the trace of a batch. Can modify trace entry.
          #: Signature: job, trace_entry
          self.post_batch_trace_hooks: List[Callable[[Job, Dict[str, Any]], Any]] = []
  
          #: Hooks run before outputting the trace of an epoch. Can modify trace entry.
          #: Signature: job, trace_entry
          self.post_epoch_trace_hooks: List[Callable[[Job, Dict[str, Any]], Any]] = []
  
          #: Hooks run after a validation job.
          #: Signature: job, trace_entry
          self.post_valid_hooks: List[Callable[[Job, Dict[str, Any]], Any]] = []
  
          #: Hooks run after training
          #: Signature: job, trace_entry
          self.post_train_hooks: List[Callable[[Job, Dict[str, Any]], Any]] = []
  
          if self.__class__ == TrainingJob:
              for f in Job.job_created_hooks:
                  f(self)
  
      @staticmethod
      def create(
          config: Config, dataset: Dataset, parent_job: Job = None
      ) -> "TrainingJob":
          """Factory method to create a training job."""
          if config.get("train.type") == "KvsAll":
              return TrainingJobKvsAll(config, dataset, parent_job)
          elif config.get("train.type") == "negative_sampling":
              return TrainingJobNegativeSampling(config, dataset, parent_job)
          elif config.get("train.type") == "1vsAll":
              return TrainingJob1vsAll(config, dataset, parent_job)
          else:
              # perhaps TODO: try class with specified name -> extensibility
              raise ValueError("train.type")
  
      def run(self) -> None:
          """Start/resume the training job and run to completion."""
          self.config.log("Starting training...")
          checkpoint_every = self.config.get("train.checkpoint.every")
          checkpoint_keep = self.config.get("train.checkpoint.keep")
          metric_name = self.config.get("valid.metric")
          patience = self.config.get("valid.early_stopping.patience")
          while True:
              # checking for model improvement according to metric_name
              # and do early stopping and keep the best checkpoint
              if (
                  len(self.valid_trace) > 0
                  and self.valid_trace[-1]["epoch"] == self.epoch
              ):
                  best_index = max(
                      range(len(self.valid_trace)),
                      key=lambda index: self.valid_trace[index][metric_name],
                  )
                  if best_index == len(self.valid_trace) - 1:
                      self.save(self.config.checkpoint_file("best"))
                  if (
                      patience > 0
                      and len(self.valid_trace) > patience
                      and best_index < len(self.valid_trace) - patience
                  ):
                      self.config.log(
                          "Stopping early ({} did not improve over best result ".format(
                              metric_name
                          )
                          + "in the last {} validation runs).".format(patience)
                      )
                      break
                  if self.epoch > self.config.get(
                      "valid.early_stopping.min_threshold.epochs"
                  ) and self.valid_trace[best_index][metric_name] < self.config.get(
                      "valid.early_stopping.min_threshold.metric_value"
                  ):
                      self.config.log(
                          "Stopping early ({} did not achieve min treshold after {} epochs".format(
                              metric_name, self.epoch
                          )
                      )
                      break
  
              # should we stop?
              if self.epoch >= self.config.get("train.max_epochs"):
                  self.config.log("Maximum number
  
  [... the rest is omitted ...]
files: kge/job/train.py
files: import math
  import torch
  from kge import Config, Dataset
  from kge.model.kge_model import RelationalScorer, KgeModel
  
  
  class ComplExScorer(RelationalScorer):
      r"""Implementation of the ComplEx KGE scorer.
  
      Reference: Théo Trouillon, Johannes Welbl, Sebastian Riedel, Éric Gaussier and
      Guillaume Bouchard: Complex Embeddings for Simple Link Prediction. ICML 2016.
      `<http://proceedings.mlr.press/v48/trouillon16.pdf>`_
  
      """
  
      def __init__(self, config: Config, dataset: Dataset, configuration_key=None):
          super().__init__(config, dataset, configuration_key)
  
      def score_emb(self, s_emb, p_emb, o_emb, combine: str):
          n = p_emb.size(0)
  
          # Here we use a fast implementation of computing the ComplEx scores using
          # Hadamard products, as in Eq. (11) of paper.
          #
          # Split the relation and object embeddings into real part (first half) and
          # imaginary part (second half).
          p_emb_re, p_emb_im = (t.contiguous() for t in p_emb.chunk(2, dim=1))
          o_emb_re, o_emb_im = (t.contiguous() for t in o_emb.chunk(2, dim=1))
  
          # combine them again to create a column block for each required combination
          s_all = torch.cat((s_emb, s_emb), dim=1)  # re, im, re, im
          r_all = torch.cat((p_emb_re, p_emb, -p_emb_im), dim=1)  # re, re, im, -im
          o_all = torch.cat((o_emb, o_emb_im, o_emb_re), dim=1)  # re, im, im, re
  
          if combine == "spo":
              out = (s_all * o_all * r_all).sum(dim=1)
          elif combine == "sp_":
              out = (s_all * r_all).mm(o_all.transpose(0, 1))
          elif combine == "_po":
              out = (r_all * o_all).mm(s_all.transpose(0, 1))
          else:
              return super().score_emb(s_emb, p_emb, o_emb, combine)
  
          return out.view(n, -1)
  
  
  class ComplEx(KgeModel):
      r"""Implementation of the ComplEx KGE model."""
  
      def __init__(self, config: Config, dataset: Dataset, configuration_key=None):
          super().__init__(
              config, dataset, ComplExScorer, configuration_key=configuration_key
          )
  
  
  [... the rest is omitted ...]
files: kge/model/complex.py
files: import math
  import torch
  from kge import Config, Dataset
  from kge.model.kge_model import RelationalScorer, KgeModel
  
  
  class ComplExScorer(RelationalScorer):
      r"""Implementation of the ComplEx KGE scorer.
  
      Reference: Théo Trouillon, Johannes Welbl, Sebastian Riedel, Éric Gaussier and
      Guillaume Bouchard: Complex Embeddings for Simple Link Prediction. ICML 2016.
      `<http://proceedings.mlr.press/v48/trouillon16.pdf>`_
  
      """
  
      def __init__(self, config: Config, dataset: Dataset, configuration_key=None):
          super().__init__(config, dataset, configuration_key)
  
      def score_emb(self, s_emb, p_emb, o_emb, combine: str):
          n = p_emb.size(0)
  
          # Here we use a fast implementation of computing the ComplEx scores using
          # Hadamard products, as in Eq. (11) of paper.
          #
          # Split the relation and object embeddings into real part (first half) and
          # imaginary part (second half).
          p_emb_re, p_emb_im = (t.contiguous() for t in p_emb.chunk(2, dim=1))
          o_emb_re, o_emb_im = (t.contiguous() for t in o_emb.chunk(2, dim=1))
  
          # combine them again to create a column block for each required combination
          s_all = torch.cat((s_emb, s_emb), dim=1)  # re, im, re, im
          r_all = torch.cat((p_emb_re, p_emb, -p_emb_im), dim=1)  # re, re, im, -im
          o_all = torch.cat((o_emb, o_emb_im, o_emb_re), dim=1)  # re, im, im, re
  
          if combine == "spo":
              out = (s_all * o_all * r_all).sum(dim=1)
          elif combine == "sp*":
              out = (s_all * r_all).mm(o_all.transpose(0, 1))
          elif combine == "*po":
              out = (r_all * o_all).mm(s_all.transpose(0, 1))
          else:
              return super().score_emb(s_emb, p_emb, o_emb, combine)
  
          return out.view(n, -1)
  
  
  class ComplEx(KgeModel):
      r"""Implementation of the ComplEx KGE model."""
  
      def __init__(self, config: Config, dataset: Dataset, configuration_key=None):
          super().__init__(
              config, dataset, ComplExScorer, configuration_key=configuration_key
          )
  
  
  [... the rest is omitted ...]
files: kge/model/complex.py
files: import torch
  import math
  
  from kge import Config, Dataset
  from kge.model.kge_model import RelationalScorer, KgeModel
  
  
  class ConvEScorer(RelationalScorer):
      r"""Implementation of the ConvE KGE scorer.
  
      Must be used with ReciprocalRelationsModel."""
  
      def __init__(self, config: Config, dataset: Dataset, configuration_key=None):
          super().__init__(config, dataset, configuration_key)
          # self.configuration_key = configuration_key
          self.emb_dim = self.get_option("entity_embedder.dim") - 1
          aspect_ratio = self.get_option("2D_aspect_ratio")
          self.emb_height = math.sqrt(self.emb_dim / aspect_ratio)
          self.emb_width = self.emb_height * aspect_ratio
  
          # round embedding dimension to match aspect ratio
          rounded_height = math.ceil(self.emb_height)
          if self.get_option("round_dim") and rounded_height != self.emb_height:
              self.emb_height = rounded_height
              self.emb_width = self.emb_height * aspect_ratio
              self.emb_dim = self.emb_height * self.emb_width
              self.set_option("entity_embedder.dim", self.emb_dim + 1)
              self.set_option("relation_embedder.dim", self.emb_dim + 1)
              config.log(
                  "Rounded embedding dimension up to {} to match given aspect ratio.".format(
                      self.emb_dim
                  )
              )
          elif self.emb_dim % self.emb_height or self.emb_dim % self.emb_width:
              raise Exception(
                  "Aspect ratio {} does not produce 2D integers for dimension {}.".format(
                      aspect_ratio, self.emb_dim
                  )
              )
  
          self.filter_size = self.get_option("filter_size")
          self.stride = self.get_option("stride")
          self.padding = self.get_option("padding")
          self.feature_map_dropout = torch.nn.Dropout2d(
              self.get_option("feature_map_dropout")
          )
          self.projection_dropout = torch.nn.Dropout(
              self.get_option("projection_dropout")
          )
          self.convolution = torch.nn.Conv2d(
              in_channels=1,
              out_channels=32,
              kernel_size=(self.filter_size, self.filter_size),
              stride=self.stride,
              padding=self.padding,
              bias=self.get_option("convolution_bias"),
          )
          self.bn1 = torch.nn.BatchNorm2d(32, affine=False)
          self.bn2 = torch.nn.BatchNorm1d(self.emb_dim, affine=False)
          conv_output_height = (
              ((self.emb_height * 2) - self.filter_size + (2 * self.padding))
              / self.stride
          ) + 1
          conv_output_width = (
              (self.emb_width - self.filter_size + (2 * self.padding)) / self.stride
          ) + 1
          self.projection = torch.nn.Linear(
              32 * int(conv_output_height * conv_output_width), int(self.emb_dim)
          )
          self.non_linear = torch.nn.ReLU()
  
      def score_emb(self, s_emb, p_emb, o_emb, combine: str):
          batch_size = p_emb.size(0)
          s_emb_2d = s_emb[:, 1:].view(-1, 1, int(self.emb_height), int(self.emb_width))
          p_emb_2d = p_emb[:, 1:].view(-1, 1, int(self.emb_height), int(self.emb_width))
          stacked_inputs = torch.cat([s_emb_2d, p_emb_2d], 2)
          out = self.convolution(stacked_inputs)
          out = self.bn1(out)
          out = self.non_linear(out)
          out = self.feature_map_dropout(out)
          out = out.view(batch_size, -1)
          out = self.projection(out)
          out = self.projection_dropout(out)
          out = self.bn2(out)
          out = self.non_linear(out)
          if combine == "sp_":
              out = torch.mm(out, o_emb[:, 1:].transpose(1, 0))
          elif combine == "spo":
              out = (out * o_emb[:, 1:]).sum(-1)
          else:
              raise Exception(
                  "Combine {} not supported in ConvE's score function".format(combine)
              )
          out += o_emb[:, 0]
  
          return out.view(batch_size, -1)
  
      # # TODO: move up for all scorers?
      # def get_option(self, name):
      #     return self.config.get_default(self.configuration_key + "." + name)
      #
      # # TODO: move up for all scorers?
      # def set_option(self, name, value):
      #     if self.configuration_key:
      #         self.config.set(self.configuration_key + "." + name, value)
      #     else:
      #         self.config.set(name, value)
  
  
  class ConvE(KgeModel):
      r"""Implementation of the ConvE KGE model."""
  
      def __init__(self, config: Config, dataset: Dataset, configuration_key=None):
          self._init_configuration(config, configuration_key)
          # HACK to add bias terms to embeddings
          self.set_option(
              "entity_embedder.dim", self.get_option("entity_embedder.dim") + 1
          )
          self.set_option(
              "relation_embedder.dim", self.get_option("relation_embedder.dim") + 1
          )
          super().__init__(
              config,
              dataset,
              ConvEScorer(config, dataset, self.configuration_key),
              configuration_key=self.configuration_key,
          )
          # UNDO hack
          self.set_option(
              "entity_embedder.dim", self.get_option("entity_embedder.dim") - 1
          )
          self.set_option(
              "relation_embedder.dim", self.get_option("relation_embedder.dim") - 1
          )
  
  
  [... the rest is omitted ...]
files: kge/model/conve.py
files: import torch
  import math
  
  from kge import Config, Dataset
  from kge.model.kge_model import RelationalScorer, KgeModel
  
  
  class ConvEScorer(RelationalScorer):
      r"""Implementation of the ConvE KGE scorer.
  
      Must be used with ReciprocalRelationsModel."""
  
      def __init__(self, config: Config, dataset: Dataset, configuration_key=None):
          super().__init__(config, dataset, configuration_key)
          # self.configuration_key = configuration_key
          self.emb_dim = self.get_option("entity_embedder.dim") - 1
          aspect_ratio = self.get_option("2D_aspect_ratio")
          self.emb_height = math.sqrt(self.emb_dim / aspect_ratio)
          self.emb_width = self.emb_height * aspect_ratio
  
          # round embedding dimension to match aspect ratio
          rounded_height = math.ceil(self.emb_height)
          if self.get_option("round_dim") and rounded_height != self.emb_height:
              self.emb_height = rounded_height
              self.emb_width = self.emb_height * aspect_ratio
              self.emb_dim = self.emb_height * self.emb_width
              self.set_option("entity_embedder.dim", self.emb_dim + 1)
              self.set_option("relation_embedder.dim", self.emb_dim + 1)
              config.log(
                  "Rounded embedding dimension up to {} to match given aspect ratio.".format(
                      self.emb_dim
                  )
              )
          elif self.emb_dim % self.emb_height or self.emb_dim % self.emb_width:
              raise Exception(
                  "Aspect ratio {} does not produce 2D integers for dimension {}.".format(
                      aspect_ratio, self.emb_dim
                  )
              )
  
          self.filter_size = self.get_option("filter_size")
          self.stride = self.get_option("stride")
          self.padding = self.get_option("padding")
          self.feature_map_dropout = torch.nn.Dropout2d(
              self.get_option("feature_map_dropout")
          )
          self.projection_dropout = torch.nn.Dropout(
              self.get_option("projection_dropout")
          )
          self.convolution = torch.nn.Conv2d(
              in_channels=1,
              out_channels=32,
              kernel_size=(self.filter_size, self.filter_size),
              stride=self.stride,
              padding=self.padding,
              bias=self.get_option("convolution_bias"),
          )
          self.bn1 = torch.nn.BatchNorm2d(32, affine=False)
          self.bn2 = torch.nn.BatchNorm1d(self.emb_dim, affine=False)
          conv_output_height = (
              ((self.emb_height * 2) - self.filter_size + (2 * self.padding))
              / self.stride
          ) + 1
          conv_output_width = (
              (self.emb_width - self.filter_size + (2 * self.padding)) / self.stride
          ) + 1
          self.projection = torch.nn.Linear(
              32 * int(conv_output_height * conv_output_width), int(self.emb_dim)
          )
          self.non_linear = torch.nn.ReLU()
  
      def score_emb(self, s_emb, p_emb, o_emb, combine: str):
          batch_size = p_emb.size(0)
          s_emb_2d = s_emb[:, 1:].view(-1, 1, int(self.emb_height), int(self.emb_width))
          p_emb_2d = p_emb[:, 1:].view(-1, 1, int(self.emb_height), int(self.emb_width))
          stacked_inputs = torch.cat([s_emb_2d, p_emb_2d], 2)
          out = self.convolution(stacked_inputs)
          out = self.bn1(out)
          out = self.non_linear(out)
          out = self.feature_map_dropout(out)
          out = out.view(batch_size, -1)
          out = self.projection(out)
          out = self.projection_dropout(out)
          out = self.bn2(out)
          out = self.non_linear(out)
          if combine == "sp*":
              out = torch.mm(out, o_emb[:, 1:].transpose(1, 0))
          elif combine == "spo":
              out = (out * o_emb[:, 1:]).sum(-1)
          else:
              raise Exception(
                  "Combine {} not supported in ConvE's score function".format(combine)
              )
          out += o_emb[:, 0]
  
          return out.view(batch_size, -1)
  
      # # TODO: move up for all scorers?
      # def get_option(self, name):
      #     return self.config.get_default(self.configuration_key + "." + name)
      #
      # # TODO: move up for all scorers?
      # def set_option(self, name, value):
      #     if self.configuration_key:
      #         self.config.set(self.configuration_key + "." + name, value)
      #     else:
      #         self.config.set(name, value)
  
  
  class ConvE(KgeModel):
      r"""Implementation of the ConvE KGE model."""
  
      def __init__(self, config: Config, dataset: Dataset, configuration_key=None):
          self._init_configuration(config, configuration_key)
          # HACK to add bias terms to embeddings
          self.set_option(
              "entity_embedder.dim", self.get_option("entity_embedder.dim") + 1
          )
          self.set_option(
              "relation_embedder.dim", self.get_option("relation_embedder.dim") + 1
          )
          super().__init__(
              config,
              dataset,
              ConvEScorer(config, dataset, self.configuration_key),
              configuration_key=self.configuration_key,
          )
          # UNDO hack
          self.set_option(
              "entity_embedder.dim", self.get_option("entity_embedder.dim") - 1
          )
          self.set_option(
              "relation_embedder.dim", self.get_option("relation_embedder.dim") - 1
          )
  
  
  [... the rest is omitted ...]
files: kge/model/conve.py
files: import math
  
  import torch
  from kge import Config, Dataset
  from kge.model.kge_model import RelationalScorer, KgeModel
  
  
  class CPScorer(RelationalScorer):
      r"""Implementation of the CP KGE scorer."""
  
      def __init__(self, config: Config, dataset: Dataset, configuration_key=None):
          super().__init__(config, dataset, configuration_key)
  
      def score_emb(self, s_emb, p_emb, o_emb, combine: str):
          n = p_emb.size(0)
  
          # use only first half for subjects and second half for objects
          half_dim = s_emb.shape[1] // 2
          s_emb_h = s_emb[:, :half_dim]
          o_emb_t = o_emb[:, half_dim:]
  
          if combine == "spo":
              out = (s_emb_h * p_emb * o_emb_t).sum(dim=1)
          elif combine == "sp_":
              out = (s_emb_h * p_emb).mm(o_emb_t.transpose(0, 1))
          elif combine == "_po":
              out = (o_emb_t * p_emb).mm(s_emb_h.transpose(0, 1))
          else:
              return super().score_emb(s_emb, p_emb, o_emb, combine)
  
          return out.view(n, -1)
  
  
  class CP(KgeModel):
      r"""Implementation of the CP KGE model."""
  
      def __init__(self, config: Config, dataset: Dataset, configuration_key=None):
          self._init_configuration(config, configuration_key)
          if self.get_option("entity_embedder.dim") % 2 != 0:
              raise ValueError(
                  "CP requires embeddings of even dimensionality"
                  " (got {})".format(self.get_option("entity_embedder.dim"))
              )
          if self.get_option("relation_embedder.dim") < 0:
              self.set_option(
                  "relation_embedder.dim",
                  self.get_option("entity_embedder.dim") // 2,
                  log=True,
              )
          super().__init__(
              config, dataset, CPScorer, configuration_key=self.configuration_key
          )
  
  
  [... the rest is omitted ...]
files: kge/model/cp.py
files: import math
  
  import torch
  from kge import Config, Dataset
  from kge.model.kge_model import RelationalScorer, KgeModel
  
  
  class CPScorer(RelationalScorer):
      r"""Implementation of the CP KGE scorer."""
  
      def __init__(self, config: Config, dataset: Dataset, configuration_key=None):
          super().__init__(config, dataset, configuration_key)
  
      def score_emb(self, s_emb, p_emb, o_emb, combine: str):
          n = p_emb.size(0)
  
          # use only first half for subjects and second half for objects
          half_dim = s_emb.shape[1] // 2
          s_emb_h = s_emb[:, :half_dim]
          o_emb_t = o_emb[:, half_dim:]
  
          if combine == "spo":
              out = (s_emb_h * p_emb * o_emb_t).sum(dim=1)
          elif combine == "sp*":
              out = (s_emb_h * p_emb).mm(o_emb_t.transpose(0, 1))
          elif combine == "*po":
              out = (o_emb_t * p_emb).mm(s_emb_h.transpose(0, 1))
          else:
              return super().score_emb(s_emb, p_emb, o_emb, combine)
  
          return out.view(n, -1)
  
  
  class CP(KgeModel):
      r"""Implementation of the CP KGE model."""
  
      def __init__(self, config: Config, dataset: Dataset, configuration_key=None):
          self._init_configuration(config, configuration_key)
          if self.get_option("entity_embedder.dim") % 2 != 0:
              raise ValueError(
                  "CP requires embeddings of even dimensionality"
                  " (got {})".format(self.get_option("entity_embedder.dim"))
              )
          if self.get_option("relation_embedder.dim") < 0:
              self.set_option(
                  "relation_embedder.dim",
                  self.get_option("entity_embedder.dim") // 2,
                  log=True,
              )
          super().__init__(
              config, dataset, CPScorer, configuration_key=self.configuration_key
          )
  
  
  [... the rest is omitted ...]
files: kge/model/cp.py
files: import math
  
  import torch
  from kge import Config, Dataset
  from kge.model.kge_model import RelationalScorer, KgeModel
  
  
  class DistMultScorer(RelationalScorer):
      r"""Implementation of the DistMult KGE scorer."""
  
      def __init__(self, config: Config, dataset: Dataset, configuration_key=None):
          super().__init__(config, dataset, configuration_key)
  
      def score_emb(self, s_emb, p_emb, o_emb, combine: str):
          n = p_emb.size(0)
  
          if combine == "spo":
              out = (s_emb * p_emb * o_emb).sum(dim=1)
          elif combine == "sp_":
              out = (s_emb * p_emb).mm(o_emb.transpose(0, 1))
          elif combine == "_po":
              out = (o_emb * p_emb).mm(s_emb.transpose(0, 1))
          else:
              return super().score_emb(s_emb, p_emb, o_emb, combine)
  
          return out.view(n, -1)
  
  
  class DistMult(KgeModel):
      r"""Implementation of the DistMult KGE model."""
  
      def __init__(self, config: Config, dataset: Dataset, configuration_key=None):
          super().__init__(
              config, dataset, DistMultScorer, configuration_key=configuration_key
          )
  
  
  [... the rest is omitted ...]
files: kge/model/distmult.py
files: import math
  
  import torch
  from kge import Config, Dataset
  from kge.model.kge_model import RelationalScorer, KgeModel
  
  
  class DistMultScorer(RelationalScorer):
      r"""Implementation of the DistMult KGE scorer."""
  
      def __init__(self, config: Config, dataset: Dataset, configuration_key=None):
          super().__init__(config, dataset, configuration_key)
  
      def score_emb(self, s_emb, p_emb, o_emb, combine: str):
          n = p_emb.size(0)
  
          if combine == "spo":
              out = (s_emb * p_emb * o_emb).sum(dim=1)
          elif combine == "sp*":
              out = (s_emb * p_emb).mm(o_emb.transpose(0, 1))
          elif combine == "*po":
              out = (o_emb * p_emb).mm(s_emb.transpose(0, 1))
          else:
              return super().score_emb(s_emb, p_emb, o_emb, combine)
  
          return out.view(n, -1)
  
  
  class DistMult(KgeModel):
      r"""Implementation of the DistMult KGE model."""
  
      def __init__(self, config: Config, dataset: Dataset, configuration_key=None):
          super().__init__(
              config, dataset, DistMultScorer, configuration_key=configuration_key
          )
  
  
  [... the rest is omitted ...]
files: kge/model/distmult.py
files: import importlib
  import tempfile
  
  from torch import Tensor
  import torch.nn
  import os
  
  import kge
  from kge import Config, Configurable, Dataset
  from kge.misc import filename_in_module
  from typing import Any, Dict, List, Optional, Union
  
  from typing import TYPE_CHECKING
  
  if TYPE_CHECKING:
      from kge.job import Job
  
  
  SLOTS = [0, 1, 2]
  S, P, O = SLOTS
  
  
  class KgeBase(torch.nn.Module, Configurable):
      r"""Base class for all KGE models, scorers, and embedders."""
  
      def __init__(self, config: Config, dataset: Dataset, configuration_key=None):
          Configurable.__init__(self, config, configuration_key)
          torch.nn.Module.__init__(self)
          self.dataset = dataset
          self.meta: Dict[str, Any] = dict()  #: meta-data stored with this module
  
      def initialize(self, what: Tensor, initialize: str, initialize_args):
          try:
              getattr(torch.nn.init, initialize)(what, **initialize_args)
          except:
              raise ValueError(
                  "invalid initialization options: {} with args {}".format(
                      initialize, initialize_args
                  )
              )
  
      def prepare_job(self, job: "Job", **kwargs):
          r"""Prepares the given job to work with this model.
  
          If this model does not support the specified job type, this function may raise
          an error.
  
          This function commonly registers hooks specific to this model. For a list of
          available hooks during training or evaluation, see :class:`TrainingJob` or
          :class:`EvaluationJob`:, respectively.
  
          """
  
      def penalty(self, **kwargs) -> List[Tensor]:
          r"""Returns additional penalty terms that are added to the loss during training.
  
          This method is called once per batch during training. The arguments being passed
          depend on the trainer being used.
  
          Returns a (possibly empty) list of penalty terms.
  
          """
  
          return []
  
      def save(self):
          "Returns data structure to save state"
          return (self.state_dict(), self.meta)
  
      def load(self, savepoint):
          "Loads state from a saved data structure"
          self.load_state_dict(savepoint[0])
          self.meta = savepoint[1]
  
  
  class RelationalScorer(KgeBase):
      r"""Base class for all relational scorers.
  
      Relational scorers take as input the embeddings of (subject, predicate,
      object)-triple and produce a score.
  
      Implementations of this class should either implement
      :func:`~RelationalScorer.score_emb_spo` (the quick way, but potentially inefficient)
      or :func:`~RelationalScorer.score_emb` (the hard way, potentially more efficient).
  
      """
  
      def __init__(self, config: Config, dataset: Dataset, configuration_key: str):
          super().__init__(config, dataset, configuration_key)
  
      def score_emb_spo(self, s_emb: Tensor, p_emb: Tensor, o_emb: Tensor) -> Tensor:
          r"""Scores a set of triples specified by their embeddings.
  
          `s_emb`, `p_emb`, and `o_emb` are tensors of size :math:`n\times d_e`,
          :math:`n\times d_r`, and :math:`n\times d_e`, where :math:`d_e` and
          :math:`d_r` are the sizes of the entity and relation embeddings, respectively.
  
          The embeddings are combined row-wise. The output is a :math`n\times 1` tensor,
          in which the :math:`i`-th entry holds the score of the embedding triple
          :math:`(s_i, p_i, o_i)`.
  
          """
          return self.score_emb(s_emb, p_emb, o_emb, "spo")
  
      def score_emb(
          self, s_emb: Tensor, p_emb: Tensor, o_emb: Tensor, combine: str
      ) -> Tensor:
          r"""Scores a set of triples specified by their embeddings.
  
          `s_emb`, `p_emb`, and `o_emb` are tensors of size :math:`n_s\times d_e`,
          :math:`n_p\times d_r`, and :math:`n_o\times d_e`, where :math:`d_e` and
          :math:`d_r` are the sizes of the entity and relation embeddings, respectively.
  
          The provided embeddings are combined based on the value of `combine`. Common
          values are :code:`"spo"`, :code:`"sp_"`, and :code:`"_po"`. Not all models may
          support all combinations.
  
          When `combine` is :code:`"spo"`, then embeddings are combined row-wise. In this
          case, it is required that :math:`n_s=n_p=n_o=n`. The output is identical to
          :func:`~RelationalScorer.score_emb_spo`, i.e., a :math`n\times 1` tensor, in
          which the :math:`i`-th entry holds the score of the embedding triple
          :math:`(s_i, p_i, o_i)`.
  
          When `combine` is :code:`"sp_"`, the subjects and predicates are taken row-wise
          and subsequently combined with all objects. In this case, it is required that
          :math:`n_s=n_p=n`. The output is a :math`n\times n_o` tensor, in which the
          :math:`(i,j)`-th entry holds the score of the embedding triple :math:`(s_i, p_i,
          o_j)`.
  
          When `combine` is :code:`"_po"`, predicates and objects are taken row-wise and
          subsequently combined with all subjects. In this case, it is required that
          :math:`n_p=n_o=n`. The output is a :math`n\times n_s` tensor, in which the
          :math:`(i,j)`-th entry holds the score of the embedding triple :math:`(s_j, p_i,
          o_i)`.
  
          """
          n = p_emb.size(0)
  
          if combine == "spo":
              assert s_emb.size(0) == n and o_emb.size(0) == n
              out = self.score_emb_spo(s_emb, p_emb, o_emb)
          elif combine == "sp_":
              assert s_emb.size(0) == n
              n_o = o_emb.size(0)
              s_embs = s_emb.repeat_interleave(n_o, 0)
              p_embs = p_emb
  
  [... the rest is omitted ...]
files: kge/model/kge_model.py
files: import importlib
  import tempfile
  
  from torch import Tensor
  import torch.nn
  import os
  
  import kge
  from kge import Config, Configurable, Dataset
  from kge.misc import filename_in_module
  from typing import Any, Dict, List, Optional, Union
  
  from typing import TYPE_CHECKING
  
  if TYPE_CHECKING:
      from kge.job import Job
  
  
  SLOTS = [0, 1, 2]
  S, P, O = SLOTS
  
  
  class KgeBase(torch.nn.Module, Configurable):
      r"""Base class for all KGE models, scorers, and embedders."""
  
      def __init__(self, config: Config, dataset: Dataset, configuration_key=None):
          Configurable.__init__(self, config, configuration_key)
          torch.nn.Module.__init__(self)
          self.dataset = dataset
          self.meta: Dict[str, Any] = dict()  #: meta-data stored with this module
  
      def initialize(self, what: Tensor, initialize: str, initialize_args):
          try:
              getattr(torch.nn.init, initialize)(what, **initialize_args)
          except:
              raise ValueError(
                  "invalid initialization options: {} with args {}".format(
                      initialize, initialize_args
                  )
              )
  
      def prepare_job(self, job: "Job", **kwargs):
          r"""Prepares the given job to work with this model.
  
          If this model does not support the specified job type, this function may raise
          an error.
  
          This function commonly registers hooks specific to this model. For a list of
          available hooks during training or evaluation, see :class:`TrainingJob` or
          :class:`EvaluationJob`:, respectively.
  
          """
  
      def penalty(self, **kwargs) -> List[Tensor]:
          r"""Returns additional penalty terms that are added to the loss during training.
  
          This method is called once per batch during training. The arguments being passed
          depend on the trainer being used.
  
          Returns a (possibly empty) list of penalty terms.
  
          """
  
          return []
  
      def save(self):
          "Returns data structure to save state"
          return (self.state_dict(), self.meta)
  
      def load(self, savepoint):
          "Loads state from a saved data structure"
          self.load_state_dict(savepoint[0])
          self.meta = savepoint[1]
  
  
  class RelationalScorer(KgeBase):
      r"""Base class for all relational scorers.
  
      Relational scorers take as input the embeddings of (subject, predicate,
      object)-triple and produce a score.
  
      Implementations of this class should either implement
      :func:`~RelationalScorer.score_emb_spo` (the quick way, but potentially inefficient)
      or :func:`~RelationalScorer.score_emb` (the hard way, potentially more efficient).
  
      """
  
      def __init__(self, config: Config, dataset: Dataset, configuration_key: str):
          super().__init__(config, dataset, configuration_key)
  
      def score_emb_spo(self, s_emb: Tensor, p_emb: Tensor, o_emb: Tensor) -> Tensor:
          r"""Scores a set of triples specified by their embeddings.
  
          `s_emb`, `p_emb`, and `o_emb` are tensors of size :math:`n\times d_e`,
          :math:`n\times d_r`, and :math:`n\times d_e`, where :math:`d_e` and
          :math:`d_r` are the sizes of the entity and relation embeddings, respectively.
  
          The embeddings are combined row-wise. The output is a :math`n\times 1` tensor,
          in which the :math:`i`-th entry holds the score of the embedding triple
          :math:`(s_i, p_i, o_i)`.
  
          """
          return self.score_emb(s_emb, p_emb, o_emb, "spo")
  
      def score_emb(
          self, s_emb: Tensor, p_emb: Tensor, o_emb: Tensor, combine: str
      ) -> Tensor:
          r"""Scores a set of triples specified by their embeddings.
  
          `s_emb`, `p_emb`, and `o_emb` are tensors of size :math:`n_s\times d_e`,
          :math:`n_p\times d_r`, and :math:`n_o\times d_e`, where :math:`d_e` and
          :math:`d_r` are the sizes of the entity and relation embeddings, respectively.
  
          The provided embeddings are combined based on the value of `combine`. Common
          values are :code:`"spo"`, :code:`"sp*"`, and :code:`"*po"`. Not all models may
          support all combinations.
  
          When `combine` is :code:`"spo"`, then embeddings are combined row-wise. In this
          case, it is required that :math:`n_s=n_p=n_o=n`. The output is identical to
          :func:`~RelationalScorer.score_emb_spo`, i.e., a :math`n\times 1` tensor, in
          which the :math:`i`-th entry holds the score of the embedding triple
          :math:`(s_i, p_i, o_i)`.
  
          When `combine` is :code:`"sp*"`, the subjects and predicates are taken row-wise
          and subsequently combined with all objects. In this case, it is required that
          :math:`n_s=n_p=n`. The output is a :math`n\times n_o` tensor, in which the
          :math:`(i,j)`-th entry holds the score of the embedding triple :math:`(s_i, p_i,
          o_j)`.
  
          When `combine` is :code:`"*po"`, predicates and objects are taken row-wise and
          subsequently combined with all subjects. In this case, it is required that
          :math:`n_p=n_o=n`. The output is a :math`n\times n_s` tensor, in which the
          :math:`(i,j)`-th entry holds the score of the embedding triple :math:`(s_j, p_i,
          o_i)`.
  
          """
          n = p_emb.size(0)
  
          if combine == "spo":
              assert s_emb.size(0) == n and o_emb.size(0) == n
              out = self.score_emb_spo(s_emb, p_emb, o_emb)
          elif combine == "sp*":
              assert s_emb.size(0) == n
              n_o = o_emb.size(0)
              s_embs = s_emb.repeat_interleave(n_o, 0)
              p_embs
  
  [... the rest is omitted ...]
files: kge/model/kge_model.py
files: import torch
  from torch import Tensor
  from kge import Config, Dataset
  from kge.model.kge_model import KgeModel
  
  
  class ReciprocalRelationsModel(KgeModel):
      """Modifies a base model to use different relation embeddings for predicting subject and object.
  
      This implements the reciprocal relations training procedure of [TODO cite ConvE].
      Note that this model cannot be used to score a single triple, but only to rank sp_
      or _po questions.
  
      """
  
      def __init__(self, config: Config, dataset: Dataset, configuration_key=None):
          self._init_configuration(config, configuration_key)
  
          # Initialize base model
          # Using a dataset with twice the number of relations to initialize base model
          alt_dataset = dataset.shallow_copy()
          alt_dataset._num_relations = dataset.num_relations() * 2
          base_model = KgeModel.create(
              config, alt_dataset, self.configuration_key + ".base_model"
          )
  
          # Initialize this model
          super().__init__(
              config, dataset, base_model.get_scorer(), initialize_embedders=False
          )
          self._base_model = base_model
          # TODO change entity_embedder assignment to sub and obj embedders when support
          # for that is added
          self._entity_embedder = self._base_model.get_s_embedder()
          self._relation_embedder = self._base_model.get_p_embedder()
  
      def prepare_job(self, job, **kwargs):
          self._base_model.prepare_job(job, **kwargs)
  
      def penalty(self, **kwargs):
          return super().penalty(**kwargs) + self._base_model.penalty(**kwargs)
  
      def score_spo(self, s: Tensor, p: Tensor, o: Tensor, direction=None) -> Tensor:
          if direction == "o":
              return super().score_spo(s, p, o, "o")
          elif direction == "s":
              return super().score_spo(o, p + self.dataset.num_relations(), s, "o")
          else:
              raise Exception(
                  "The reciprocal relations model cannot compute "
                  "undirected spo scores."
              )
  
      def score_po(self, p, o, s=None):
          if s is None:
              s = self.get_s_embedder().embed_all()
          else:
              s = self.get_s_embedder().embed(s)
          p = self.get_p_embedder().embed(p + self.dataset.num_relations())
          o = self.get_o_embedder().embed(o)
          return self._scorer.score_emb(o, p, s, combine="sp_")
  
      def score_so(self, s, o, p=None):
          raise Exception("The reciprocal relations model cannot score relations.")
  
      def score_sp_po(
          self,
          s: torch.Tensor,
          p: torch.Tensor,
          o: torch.Tensor,
          entity_subset: torch.Tensor = None,
      ) -> torch.Tensor:
          s = self.get_s_embedder().embed(s)
          p_inv = self.get_p_embedder().embed(p + self.dataset.num_relations())
          p = self.get_p_embedder().embed(p)
          o = self.get_o_embedder().embed(o)
          if self.get_s_embedder() is self.get_o_embedder():
              if entity_subset is not None:
                  all_entities = self.get_s_embedder().embed(entity_subset)
              else:
                  all_entities = self.get_s_embedder().embed_all()
              sp_scores = self._scorer.score_emb(s, p, all_entities, combine="sp_")
              po_scores = self._scorer.score_emb(o, p_inv, all_entities, combine="sp_")
          else:
              if entity_subset is not None:
                  all_objects = self.get_o_embedder().embed(entity_subset)
                  all_subjects = self.get_s_embedder().embed(entity_subset)
              else:
                  all_objects = self.get_o_embedder().embed_all()
                  all_subjects = self.get_s_embedder().embed_all()
              sp_scores = self._scorer.score_emb(s, p, all_objects, combine="sp_")
              po_scores = self._scorer.score_emb(o, p_inv, all_subjects, combine="sp_")
          return torch.cat((sp_scores, po_scores), dim=1)
  
  
  [... the rest is omitted ...]
files: kge/model/reciprocal_relations_model.py
files: import torch
  from torch import Tensor
  from kge import Config, Dataset
  from kge.model.kge_model import KgeModel
  
  
  class ReciprocalRelationsModel(KgeModel):
      """Modifies a base model to use different relation embeddings for predicting subject and object.
  
      This implements the reciprocal relations training procedure of [TODO cite ConvE].
      Note that this model cannot be used to score a single triple, but only to rank sp*
      or *po questions.
  
      """
  
      def __init__(self, config: Config, dataset: Dataset, configuration_key=None):
          self._init_configuration(config, configuration_key)
  
          # Initialize base model
          # Using a dataset with twice the number of relations to initialize base model
          alt_dataset = dataset.shallow_copy()
          alt_dataset._num_relations = dataset.num_relations() * 2
          base_model = KgeModel.create(
              config, alt_dataset, self.configuration_key + ".base_model"
          )
  
          # Initialize this model
          super().__init__(
              config, dataset, base_model.get_scorer(), initialize_embedders=False
          )
          self._base_model = base_model
          # TODO change entity_embedder assignment to sub and obj embedders when support
          # for that is added
          self._entity_embedder = self._base_model.get_s_embedder()
          self._relation_embedder = self._base_model.get_p_embedder()
  
      def prepare_job(self, job, **kwargs):
          self._base_model.prepare_job(job, **kwargs)
  
      def penalty(self, **kwargs):
          return super().penalty(**kwargs) + self._base_model.penalty(**kwargs)
  
      def score_spo(self, s: Tensor, p: Tensor, o: Tensor, direction=None) -> Tensor:
          if direction == "o":
              return super().score_spo(s, p, o, "o")
          elif direction == "s":
              return super().score_spo(o, p + self.dataset.num_relations(), s, "o")
          else:
              raise Exception(
                  "The reciprocal relations model cannot compute "
                  "undirected spo scores."
              )
  
      def score_po(self, p, o, s=None):
          if s is None:
              s = self.get_s_embedder().embed_all()
          else:
              s = self.get_s_embedder().embed(s)
          p = self.get_p_embedder().embed(p + self.dataset.num_relations())
          o = self.get_o_embedder().embed(o)
          return self._scorer.score_emb(o, p, s, combine="sp*")
  
      def score_so(self, s, o, p=None):
          raise Exception("The reciprocal relations model cannot score relations.")
  
      def score_sp_po(
          self,
          s: torch.Tensor,
          p: torch.Tensor,
          o: torch.Tensor,
          entity_subset: torch.Tensor = None,
      ) -> torch.Tensor:
          s = self.get_s_embedder().embed(s)
          p_inv = self.get_p_embedder().embed(p + self.dataset.num_relations())
          p = self.get_p_embedder().embed(p)
          o = self.get_o_embedder().embed(o)
          if self.get_s_embedder() is self.get_o_embedder():
              if entity_subset is not None:
                  all_entities = self.get_s_embedder().embed(entity_subset)
              else:
                  all_entities = self.get_s_embedder().embed_all()
              sp_scores = self._scorer.score_emb(s, p, all_entities, combine="sp*")
              po_scores = self._scorer.score_emb(o, p_inv, all_entities, combine="sp*")
          else:
              if entity_subset is not None:
                  all_objects = self.get_o_embedder().embed(entity_subset)
                  all_subjects = self.get_s_embedder().embed(entity_subset)
              else:
                  all_objects = self.get_o_embedder().embed_all()
                  all_subjects = self.get_s_embedder().embed_all()
              sp_scores = self._scorer.score_emb(s, p, all_objects, combine="sp*")
              po_scores = self._scorer.score_emb(o, p_inv, all_subjects, combine="sp*")
          return torch.cat((sp_scores, po_scores), dim=1)
  
  
  [... the rest is omitted ...]
files: kge/model/reciprocal_relations_model.py
files: import math
  
  import torch
  
  from kge import Config, Dataset
  from kge.model.kge_model import KgeEmbedder, KgeModel, RelationalScorer
  
  
  class RescalScorer(RelationalScorer):
      r"""Implementation of the RESCAL KGE scorer."""
  
      def __init__(self, config: Config, dataset: Dataset, configuration_key=None):
          super().__init__(config, dataset, configuration_key)
  
      def score_emb(
          self,
          s_emb: torch.Tensor,
          p_emb: torch.Tensor,
          o_emb: torch.Tensor,
          combine: str,
      ):
          batch_size = p_emb.size(0)
          entity_size = s_emb.size(-1)
  
          # reshape relation embeddings to obtain mixing matrices for RESCAL
          p_mixmat = p_emb.view(-1, entity_size, entity_size)
  
          if combine == "spo":
              out = (
                  s_emb.unsqueeze(1)  # [batch x 1 x entity_size]
                  .bmm(p_mixmat)  # apply mixing matrices
                  .view(batch_size, entity_size)  # drop dim 1
                  * o_emb  # apply object embeddings
              ).sum(
                  dim=-1
              )  # and sum to obtain predictions
          elif combine == "sp_":
              out = (
                  s_emb.unsqueeze(1)
                  .bmm(p_mixmat)
                  .view(batch_size, entity_size)
                  .mm(o_emb.transpose(0, 1))
              )
          elif combine == "_po":
              out = (
                  p_mixmat.bmm(o_emb.unsqueeze(2))
                  .view(batch_size, entity_size)
                  .mm(s_emb.transpose(0, 1))
              )
          else:
              super().score_emb(s_emb, p_emb, o_emb, combine)
  
          return out.view(batch_size, -1)
  
  
  class Rescal(KgeModel):
      r"""Implementation of the RÃ‰SCAL KGE model."""
  
      def __init__(self, config: Config, dataset: Dataset, configuration_key=None):
          self._init_configuration(config, configuration_key)
          rescal_set_relation_embedder_dim(
              config, dataset, self.configuration_key + ".relation_embedder"
          )
          super().__init__(
              config, dataset, RescalScorer, configuration_key=self.configuration_key
          )
  
  
  def rescal_set_relation_embedder_dim(config, dataset, rel_emb_conf_key):
      """Set the relation embedder dimensionality for RESCAL in the config.
  
      If <0, set it to the square of the size of the entity embedder. Else leave
      unchanged.
  
      """
      dim = config.get_default(rel_emb_conf_key + ".dim")
      if dim < 0:  # autodetect relation embedding dimensionality
          ent_emb_conf_key = rel_emb_conf_key.replace(
              "relation_embedder", "entity_embedder"
          )
          if ent_emb_conf_key == rel_emb_conf_key:
              raise ValueError(
                  "Cannot determine relation embedding size; please set manually."
              )
          dim = config.get_default(ent_emb_conf_key + ".dim") ** 2
          config.set(rel_emb_conf_key + ".dim", dim, log=True)
  
  
  [... the rest is omitted ...]
files: kge/model/rescal.py
files: import math
  
  import torch
  
  from kge import Config, Dataset
  from kge.model.kge_model import KgeEmbedder, KgeModel, RelationalScorer
  
  
  class RescalScorer(RelationalScorer):
      r"""Implementation of the RESCAL KGE scorer."""
  
      def __init__(self, config: Config, dataset: Dataset, configuration_key=None):
          super().__init__(config, dataset, configuration_key)
  
      def score_emb(
          self,
          s_emb: torch.Tensor,
          p_emb: torch.Tensor,
          o_emb: torch.Tensor,
          combine: str,
      ):
          batch_size = p_emb.size(0)
          entity_size = s_emb.size(-1)
  
          # reshape relation embeddings to obtain mixing matrices for RESCAL
          p_mixmat = p_emb.view(-1, entity_size, entity_size)
  
          if combine == "spo":
              out = (
                  s_emb.unsqueeze(1)  # [batch x 1 x entity_size]
                  .bmm(p_mixmat)  # apply mixing matrices
                  .view(batch_size, entity_size)  # drop dim 1
                  * o_emb  # apply object embeddings
              ).sum(
                  dim=-1
              )  # and sum to obtain predictions
          elif combine == "sp*":
              out = (
                  s_emb.unsqueeze(1)
                  .bmm(p_mixmat)
                  .view(batch_size, entity_size)
                  .mm(o_emb.transpose(0, 1))
              )
          elif combine == "*po":
              out = (
                  p_mixmat.bmm(o_emb.unsqueeze(2))
                  .view(batch_size, entity_size)
                  .mm(s_emb.transpose(0, 1))
              )
          else:
              super().score_emb(s_emb, p_emb, o_emb, combine)
  
          return out.view(batch_size, -1)
  
  
  class Rescal(KgeModel):
      r"""Implementation of the RÃ‰SCAL KGE model."""
  
      def __init__(self, config: Config, dataset: Dataset, configuration_key=None):
          self._init_configuration(config, configuration_key)
          rescal_set_relation_embedder_dim(
              config, dataset, self.configuration_key + ".relation_embedder"
          )
          super().__init__(
              config, dataset, RescalScorer, configuration_key=self.configuration_key
          )
  
  
  def rescal_set_relation_embedder_dim(config, dataset, rel_emb_conf_key):
      """Set the relation embedder dimensionality for RESCAL in the config.
  
      If <0, set it to the square of the size of the entity embedder. Else leave
      unchanged.
  
      """
      dim = config.get_default(rel_emb_conf_key + ".dim")
      if dim < 0:  # autodetect relation embedding dimensionality
          ent_emb_conf_key = rel_emb_conf_key.replace(
              "relation_embedder", "entity_embedder"
          )
          if ent_emb_conf_key == rel_emb_conf_key:
              raise ValueError(
                  "Cannot determine relation embedding size; please set manually."
              )
          dim = config.get_default(ent_emb_conf_key + ".dim") ** 2
          config.set(rel_emb_conf_key + ".dim", dim, log=True)
  
  
  [... the rest is omitted ...]
files: kge/model/rescal.py
files: import torch
  from kge import Config, Dataset
  from kge.model.kge_model import RelationalScorer, KgeModel
  from torch.nn import functional as F
  
  
  # TODO sp_ and _po scoring with RotatE leads to *large* intermediate results. It's
  # unclear whether this can be fixed. Expect out-of-memory errors when using RotatE with
  # 1vsAll or KvsAll training. To do validation/evaluation, you may want to set
  # eval.chunk_size.
  class RotatEScorer(RelationalScorer):
      r"""Implementation of the RotatE KGE scorer."""
  
      def __init__(self, config: Config, dataset: Dataset, configuration_key=None):
          super().__init__(config, dataset, configuration_key)
          self._norm = self.get_option("l_norm")
  
      def score_emb(self, s_emb, p_emb, o_emb, combine: str):
          n = p_emb.size(0)
  
          # determine real and imaginary part
          s_emb_re, s_emb_im = torch.chunk(s_emb, 2, dim=1)
          o_emb_re, o_emb_im = torch.chunk(o_emb, 2, dim=1)
  
          # convert from radians to points on complex unix ball
          p_emb_re, p_emb_im = torch.cos(p_emb), torch.sin(p_emb)
  
          if combine == "spo":
              # compute the difference vector (s*p-t)
              sp_emb_re, sp_emb_im = hadamard_complex(
                  s_emb_re, s_emb_im, p_emb_re, p_emb_im
              )
              diff_re, diff_im = diff_complex(sp_emb_re, sp_emb_im, o_emb_re, o_emb_im)
  
              # compute the absolute values for each (complex) element of the difference
              # vector
              diff_abs = norm_complex(diff_re, diff_im)
  
              # now take the norm of the absolute values of the difference vector
              out = torch.norm(diff_abs, dim=1, p=self._norm)
          elif combine == "sp_":
              # as above, but pair each sp-pair with each object
              sp_emb_re, sp_emb_im = hadamard_complex(
                  s_emb_re, s_emb_im, p_emb_re, p_emb_im
              )  # sp x dim
              diff_re, diff_im = pairwise_diff_complex(
                  sp_emb_re, sp_emb_im, o_emb_re, o_emb_im
              )  # sp x o x dim
              diff_abs = norm_complex(diff_re, diff_im)  # sp x o x dim
              out = torch.norm(diff_abs, dim=2, p=self._norm)
          elif combine == "_po":
              # as above, but pair each subject with each po-pair
              sp_emb_re, sp_emb_im = pairwise_hadamard_complex(
                  s_emb_re, s_emb_im, p_emb_re, p_emb_im
              )  # s x p x dim
              diff_re, diff_im = diff_complex(
                  sp_emb_re, sp_emb_im, o_emb_re, o_emb_im
              )  # s x po x dim
              diff_abs = norm_complex(diff_re, diff_im)  # s x po x dim
              out = torch.norm(diff_abs, dim=2, p=self._norm).t()
          else:
              return super().score_emb(s_emb, p_emb, o_emb, combine)
          return out.view(n, -1)
  
  
  class RotatE(KgeModel):
      r"""Implementation of the RotatE KGE model."""
  
      def __init__(self, config: Config, dataset: Dataset, configuration_key=None):
          self._init_configuration(config, configuration_key)
          if self.get_option("entity_embedder.dim") % 2 != 0:
              raise ValueError(
                  "RotatE requires embeddings of even dimensionality"
                  " (got {})".format(self.get_option("entity_embedder.dim"))
              )
          if self.get_option("relation_embedder.dim") < 0:
              self.set_option(
                  "relation_embedder.dim",
                  self.get_option("entity_embedder.dim") // 2,
                  log=True,
              )
          super().__init__(
              config, dataset, RotatEScorer, configuration_key=self.configuration_key
          )
  
  
  def pairwise_sum(X, Y):
      """Compute pairwise sum of rows of X and Y.
  
      Returns tensor of shape len(X) x len(Y) x dim."""
      return X.unsqueeze(1) + Y
  
  
  def pairwise_diff(X, Y):
      """Compute pairwise difference of rows of X and Y.
  
      Returns tensor of shape len(X) x len(Y) x dim."""
      return X.unsqueeze(1) - Y
  
  
  def pairwise_hadamard(X, Y):
      """Compute pairwise Hadamard product of rows of X and Y.
  
      Returns tensor of shape len(X) x len(Y) x dim."""
      return X.unsqueeze(1) * Y
  
  
  def hadamard_complex(x_re, x_im, y_re, y_im):
      "Hadamard product for complex vectors"
      result_re = x_re * y_re - x_im * y_im
      result_im = x_re * y_im + x_im * y_re
      return result_re, result_im
  
  
  def pairwise_hadamard_complex(x_re, x_im, y_re, y_im):
      "Pairwise Hadamard product for complex vectors"
      result_re = pairwise_hadamard(x_re, y_re) - pairwise_hadamard(x_im, y_im)
      result_im = pairwise_hadamard(x_re, y_im) + pairwise_hadamard(x_im, y_re)
      return result_re, result_im
  
  
  def diff_complex(x_re, x_im, y_re, y_im):
      "Difference of complex vectors"
      return x_re - y_re, x_im - y_im
  
  
  def pairwise_diff_complex(x_re, x_im, y_re, y_im):
      "Pairwise difference of complex vectors"
      return pairwise_diff(x_re, y_re), pairwise_diff(x_im, y_im)
  
  
  def norm_complex(x_re, x_im):
      "Compute magnitude of given complex numbers"
      x_re
  
  [... the rest is omitted ...]
files: kge/model/rotate.py
files: import torch
  from kge import Config, Dataset
  from kge.model.kge_model import RelationalScorer, KgeModel
  from torch.nn import functional as F
  
  
  # TODO sp* and *po scoring with RotatE leads to *large* intermediate results. It's
  # unclear whether this can be fixed. Expect out-of-memory errors when using RotatE with
  # 1vsAll or KvsAll training. To do validation/evaluation, you may want to set
  # eval.chunk_size.
  class RotatEScorer(RelationalScorer):
      r"""Implementation of the RotatE KGE scorer."""
  
      def __init__(self, config: Config, dataset: Dataset, configuration_key=None):
          super().__init__(config, dataset, configuration_key)
          self._norm = self.get_option("l_norm")
  
      def score_emb(self, s_emb, p_emb, o_emb, combine: str):
          n = p_emb.size(0)
  
          # determine real and imaginary part
          s_emb_re, s_emb_im = torch.chunk(s_emb, 2, dim=1)
          o_emb_re, o_emb_im = torch.chunk(o_emb, 2, dim=1)
  
          # convert from radians to points on complex unix ball
          p_emb_re, p_emb_im = torch.cos(p_emb), torch.sin(p_emb)
  
          if combine == "spo":
              # compute the difference vector (s*p-t)
              sp_emb_re, sp_emb_im = hadamard_complex(
                  s_emb_re, s_emb_im, p_emb_re, p_emb_im
              )
              diff_re, diff_im = diff_complex(sp_emb_re, sp_emb_im, o_emb_re, o_emb_im)
  
              # compute the absolute values for each (complex) element of the difference
              # vector
              diff_abs = norm_complex(diff_re, diff_im)
  
              # now take the norm of the absolute values of the difference vector
              out = torch.norm(diff_abs, dim=1, p=self._norm)
          elif combine == "sp*":
              # as above, but pair each sp-pair with each object
              sp_emb_re, sp_emb_im = hadamard_complex(
                  s_emb_re, s_emb_im, p_emb_re, p_emb_im
              )  # sp x dim
              diff_re, diff_im = pairwise_diff_complex(
                  sp_emb_re, sp_emb_im, o_emb_re, o_emb_im
              )  # sp x o x dim
              diff_abs = norm_complex(diff_re, diff_im)  # sp x o x dim
              out = torch.norm(diff_abs, dim=2, p=self._norm)
          elif combine == "*po":
              # as above, but pair each subject with each po-pair
              sp_emb_re, sp_emb_im = pairwise_hadamard_complex(
                  s_emb_re, s_emb_im, p_emb_re, p_emb_im
              )  # s x p x dim
              diff_re, diff_im = diff_complex(
                  sp_emb_re, sp_emb_im, o_emb_re, o_emb_im
              )  # s x po x dim
              diff_abs = norm_complex(diff_re, diff_im)  # s x po x dim
              out = torch.norm(diff_abs, dim=2, p=self._norm).t()
          else:
              return super().score_emb(s_emb, p_emb, o_emb, combine)
          return out.view(n, -1)
  
  
  class RotatE(KgeModel):
      r"""Implementation of the RotatE KGE model."""
  
      def __init__(self, config: Config, dataset: Dataset, configuration_key=None):
          self._init_configuration(config, configuration_key)
          if self.get_option("entity_embedder.dim") % 2 != 0:
              raise ValueError(
                  "RotatE requires embeddings of even dimensionality"
                  " (got {})".format(self.get_option("entity_embedder.dim"))
              )
          if self.get_option("relation_embedder.dim") < 0:
              self.set_option(
                  "relation_embedder.dim",
                  self.get_option("entity_embedder.dim") // 2,
                  log=True,
              )
          super().__init__(
              config, dataset, RotatEScorer, configuration_key=self.configuration_key
          )
  
  
  def pairwise_sum(X, Y):
      """Compute pairwise sum of rows of X and Y.
  
      Returns tensor of shape len(X) x len(Y) x dim."""
      return X.unsqueeze(1) + Y
  
  
  def pairwise_diff(X, Y):
      """Compute pairwise difference of rows of X and Y.
  
      Returns tensor of shape len(X) x len(Y) x dim."""
      return X.unsqueeze(1) - Y
  
  
  def pairwise_hadamard(X, Y):
      """Compute pairwise Hadamard product of rows of X and Y.
  
      Returns tensor of shape len(X) x len(Y) x dim."""
      return X.unsqueeze(1) * Y
  
  
  def hadamard_complex(x_re, x_im, y_re, y_im):
      "Hadamard product for complex vectors"
      result_re = x_re * y_re - x_im * y_im
      result_im = x_re * y_im + x_im * y_re
      return result_re, result_im
  
  
  def pairwise_hadamard_complex(x_re, x_im, y_re, y_im):
      "Pairwise Hadamard product for complex vectors"
      result_re = pairwise_hadamard(x_re, y_re) - pairwise_hadamard(x_im, y_im)
      result_im = pairwise_hadamard(x_re, y_im) + pairwise_hadamard(x_im, y_re)
      return result_re, result_im
  
  
  def diff_complex(x_re, x_im, y_re, y_im):
      "Difference of complex vectors"
      return x_re - y_re, x_im - y_im
  
  
  def pairwise_diff_complex(x_re, x_im, y_re, y_im):
      "Pairwise difference of complex vectors"
      return pairwise_diff(x_re, y_re), pairwise_diff(x_im, y_im)
  
  
  def norm_complex(x_re, x_im):
      "Compute magnitude of given complex numbers"
      x_re
  
  [... the rest is omitted ...]
files: kge/model/rotate.py
files: import math
  
  import torch
  from kge import Config, Dataset
  from kge.model.kge_model import RelationalScorer, KgeModel
  
  
  class SimplEScorer(RelationalScorer):
      r"""Implementation of the SimplE KGE scorer."""
  
      def __init__(self, config: Config, dataset: Dataset, configuration_key=None):
          super().__init__(config, dataset, configuration_key)
  
      def score_emb(self, s_emb, p_emb, o_emb, combine: str):
          n = p_emb.size(0)
  
          # split left/right
          s_emb_h, s_emb_t = torch.chunk(s_emb, 2, dim=1)
          p_emb_forward, p_emb_backward = torch.chunk(p_emb, 2, dim=1)
          o_emb_h, o_emb_t = torch.chunk(o_emb, 2, dim=1)
  
          if combine == "spo":
              out1 = (s_emb_h * p_emb_forward * o_emb_t).sum(dim=1)
              out2 = (s_emb_t * p_emb_backward * o_emb_h).sum(dim=1)
          elif combine == "sp_":
              out1 = (s_emb_h * p_emb_forward).mm(o_emb_t.transpose(0, 1))
              out2 = (s_emb_t * p_emb_backward).mm(o_emb_h.transpose(0, 1))
          elif combine == "_po":
              out1 = (o_emb_t * p_emb_forward).mm(s_emb_h.transpose(0, 1))
              out2 = (o_emb_h * p_emb_backward).mm(s_emb_t.transpose(0, 1))
          else:
              return super().score_emb(s_emb, p_emb, o_emb, combine)
  
          return (out1 + out2).view(n, -1) / 2.0
  
  
  class SimplE(KgeModel):
      r"""Implementation of the SimplE KGE model."""
  
      def __init__(self, config: Config, dataset: Dataset, configuration_key=None):
          self._init_configuration(config, configuration_key)
          if self.get_option("entity_embedder.dim") % 2 != 0:
              raise ValueError(
                  "SimplE requires embeddings of even dimensionality"
                  " (got {})".format(self.get_option("entity_embedder.dim"))
              )
          super().__init__(
              config, dataset, SimplEScorer, configuration_key=self.configuration_key
          )
  
  
  [... the rest is omitted ...]
files: kge/model/simple.py
files: import math
  
  import torch
  from kge import Config, Dataset
  from kge.model.kge_model import RelationalScorer, KgeModel
  
  
  class SimplEScorer(RelationalScorer):
      r"""Implementation of the SimplE KGE scorer."""
  
      def __init__(self, config: Config, dataset: Dataset, configuration_key=None):
          super().__init__(config, dataset, configuration_key)
  
      def score_emb(self, s_emb, p_emb, o_emb, combine: str):
          n = p_emb.size(0)
  
          # split left/right
          s_emb_h, s_emb_t = torch.chunk(s_emb, 2, dim=1)
          p_emb_forward, p_emb_backward = torch.chunk(p_emb, 2, dim=1)
          o_emb_h, o_emb_t = torch.chunk(o_emb, 2, dim=1)
  
          if combine == "spo":
              out1 = (s_emb_h * p_emb_forward * o_emb_t).sum(dim=1)
              out2 = (s_emb_t * p_emb_backward * o_emb_h).sum(dim=1)
          elif combine == "sp*":
              out1 = (s_emb_h * p_emb_forward).mm(o_emb_t.transpose(0, 1))
              out2 = (s_emb_t * p_emb_backward).mm(o_emb_h.transpose(0, 1))
          elif combine == "*po":
              out1 = (o_emb_t * p_emb_forward).mm(s_emb_h.transpose(0, 1))
              out2 = (o_emb_h * p_emb_backward).mm(s_emb_t.transpose(0, 1))
          else:
              return super().score_emb(s_emb, p_emb, o_emb, combine)
  
          return (out1 + out2).view(n, -1) / 2.0
  
  
  class SimplE(KgeModel):
      r"""Implementation of the SimplE KGE model."""
  
      def __init__(self, config: Config, dataset: Dataset, configuration_key=None):
          self._init_configuration(config, configuration_key)
          if self.get_option("entity_embedder.dim") % 2 != 0:
              raise ValueError(
                  "SimplE requires embeddings of even dimensionality"
                  " (got {})".format(self.get_option("entity_embedder.dim"))
              )
          super().__init__(
              config, dataset, SimplEScorer, configuration_key=self.configuration_key
          )
  
  
  [... the rest is omitted ...]
files: kge/model/simple.py
files: import torch
  from kge import Config, Dataset
  from kge.model.kge_model import RelationalScorer, KgeModel
  from torch.nn import functional as F
  
  
  class TransEScorer(RelationalScorer):
      r"""Implementation of the TransE KGE scorer."""
  
      def __init__(self, config: Config, dataset: Dataset, configuration_key=None):
          super().__init__(config, dataset, configuration_key)
          self._norm = self.get_option("l_norm")
  
      def score_emb(self, s_emb, p_emb, o_emb, combine: str):
          n = p_emb.size(0)
          if combine == "spo":
              out = -F.pairwise_distance(s_emb + p_emb, o_emb, p=self._norm)
          elif combine == "sp_":
              out = -torch.cdist(s_emb + p_emb, o_emb, p=self._norm)
          elif combine == "_po":
              out = -torch.cdist(o_emb - p_emb, s_emb, p=self._norm)
          else:
              super().score_emb(s_emb, p_emb, o_emb, combine)
          return out.view(n, -1)
  
  
  class TransE(KgeModel):
      r"""Implementation of the TransE KGE model."""
  
      def __init__(self, config: Config, dataset: Dataset, configuration_key=None):
          super().__init__(
              config, dataset, TransEScorer, configuration_key=configuration_key
          )
  
  
  [... the rest is omitted ...]
files: kge/model/transe.py
files: import torch
  from kge import Config, Dataset
  from kge.model.kge_model import RelationalScorer, KgeModel
  from torch.nn import functional as F
  
  
  class TransEScorer(RelationalScorer):
      r"""Implementation of the TransE KGE scorer."""
  
      def __init__(self, config: Config, dataset: Dataset, configuration_key=None):
          super().__init__(config, dataset, configuration_key)
          self._norm = self.get_option("l_norm")
  
      def score_emb(self, s_emb, p_emb, o_emb, combine: str):
          n = p_emb.size(0)
          if combine == "spo":
              out = -F.pairwise_distance(s_emb + p_emb, o_emb, p=self._norm)
          elif combine == "sp*":
              out = -torch.cdist(s_emb + p_emb, o_emb, p=self._norm)
          elif combine == "*po":
              out = -torch.cdist(o_emb - p_emb, s_emb, p=self._norm)
          else:
              super().score_emb(s_emb, p_emb, o_emb, combine)
          return out.view(n, -1)
  
  
  class TransE(KgeModel):
      r"""Implementation of the TransE KGE model."""
  
      def __init__(self, config: Config, dataset: Dataset, configuration_key=None):
          super().__init__(
              config, dataset, TransEScorer, configuration_key=configuration_key
          )
  
  
  [... the rest is omitted ...]
files: kge/model/transe.py