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 of epochs reached.")
                  break
  
              # start a new epoch
              self.epoch += 1
              self.config.log("Starting epoch {}...".format(self.epoch))
              trace_entry = self.run_epoch()
              for f in self.post_epoch_hooks:
                  f(self, trace_entry)
              self.config.log("Finished epoch {}.".format(self.epoch))
  
              # update model metadata
              self.model.meta["train_job_trace_entry"] = self.trace_entry
              self.model.meta["train_epoch"] = self.epoch
              self.model.meta["train_config"] = self.config
              self.model.meta["train_trace_entry"] = trace_entry
  
              # validate and update learning rate
              if (
                  self.config.get("valid.every") > 0
                  and self.epoch % self.config.get("valid.every") == 0
              ):
                  self.valid_job.epoch = self.epoch
                  trace_entry = self.valid_job.run()
                  self.valid_trace.append(trace_entry)
                  for f in self.post_valid_hooks:
                      f(self, trace_entry)
                  self.model.meta["valid_trace_entry"] = trace_entry
  
                  # metric-based scheduler step
                  self.kge_lr_scheduler.step(trace_entry[metric_name])
              else:
                  self.kge_lr_scheduler.step()
  
              # create checkpoint and delete old one, if necessary
              self.save(self.config.checkpoint_file(self.epoch))
              if self.epoch > 1:
                  delete_checkpoint_epoch = -1
                  if checkpoint_every == 0:
                      # do not keep any old checkpoints
                      delete_checkpoint_epoch = self.epoch - 1
                  elif (self.epoch - 1) % checkpoint_every != 0:
                      # delete checkpoints that are not in the checkpoint.every schedule
                      delete_checkpoint_epoch = self.epoch - 1
                  elif checkpoint_keep > 0:
                      # keep a maximum number of checkpoint_keep checkpoints
                      delete_checkpoint_epoch = (
                          self.epoch - 1 - checkpoint_every * checkpoint_keep
                      )
                  if delete_checkpoint_epoch > 0:
                      if os.path.exists(
                          self.config.checkpoint_file(delete_checkpoint_epoch)
                      ):
                          self.config.log(
                              "Removing old checkpoint {}...".format(
                                  self.config.checkpoint_file(delete_checkpoint_epoch)
                              )
                          )
                          os.remove(self.config.checkpoint_file(delete_checkpoint_epoch))
                      else:
                          self.config.log(
                              "Could not delete old checkpoint {}, does not exits.".format(
                                  self.config.checkpoint_file(delete_checkpoint_epoch)
                              )
                          )
  
          for f in self.post_train_hooks:
              f(self, trace_entry)
          self.trace(event="train_completed")
  
      def save(self, filename) -> None:
          """Save current state to specified file"""
          self.config.log("Saving checkpoint to {}...".format(filename))
          torch.save(
              {
                  "type": "train",
                  "config": self.config,
                  "epoch": self.epoch,
                  "valid_trace": self.valid_trace,
                  "model": self.model.save(),
                  "optimizer_state_dict": self.optimizer.state_dict(),
                  "lr_scheduler_state_dict": self.kge_lr_scheduler.state_dict(),
                  "job_id": self.job_id,
              },
              filename,
          )
  
      def load(self, filename: str) -> str:
          """Load job state from specified file.
  
          Returns job id of the job that created the checkpoint."""
          self.config.log("Loading checkpoint from {}...".format(filename))
          checkpoint = torch.load(filename, map_location="cpu")
          if "model" in checkpoint:
              # new format
              self.model.load(checkpoint["model"])
          else:
              # old format (deprecated, will eventually be removed)
              self.model.load_state_dict(checkpoint["model_state_dict"])
          self.optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
          if "lr_scheduler_state_dict" in checkpoint:
              # new format
              self.kge_lr_scheduler.load_state_dict(checkpoint["lr_scheduler_state_dict"])
          self.epoch = checkpoint["epoch"]
          self.valid_trace = checkpoint["valid_trace"]
          self.model.train()
          return checkpoint.get("job_id")
  
      def resume(self, checkpoint_file: str = None) -> None:
          if checkpoint_file is None:
              last_checkpoint = self.config.last_checkpoint()
              if last_checkpoint is not None:
                  checkpoint_file = self.config.checkpoint_file(last_checkpoint)
  
          if checkpoint_file is not None:
              self.resumed_from_job_id = self.load(checkpoint_file)
              self.trace(
                  event="job_resumed", epoch=self.epoch, checkpoint_file=checkpoint_file
              )
              self.config.log(
                  "Resumed from {} of job {}".format(
                      checkpoint_file, self.resumed_from_job_id
                  )
              )
          else:
              self.config.log("No checkpoint found, starting from scratch...")
  
      def run_epoch(self) -> Dict[str, Any]:
          "Runs an epoch and returns a trace entry."
  
          # prepare the job is not done already
          if not self.is_prepared:
              self._prepare()
              self.model.prepare_job(self)  # let the model add some hooks
              self.is_prepared = True
  
          # variables that record various statitics
          sum_loss = 0.0
          sum_penalty = 0.0
          sum_penalties = defaultdict(lambda: 0.0)
          epoch_time = -time.time()
          prepare_time = 0.0
          forward_time = 0.0
          backward_time = 0.0
          optimizer_time = 0.0
  
          # process each batch
          for batch_index, batch in enumerate(self.loader):
              for f in self.pre_batch_hooks:
                  f(self)
  
              # process batch (preprocessing + forward pass + backward pass on loss)
              self.optimizer.zero_grad()
              batch_result: TrainingJob._ProcessBatchResult = self._process_batch(
                  batch_index, batch
              )
              sum_loss += batch_result.avg_loss * batch_result.size
  
              # determine penalty terms (forward pass)
              batch_forward_time = batch_result.forward_time - time.time()
              penalties_torch = self.model.penalty(
                  epoch=self.epoch,
                  batch_index=batch_index,
                  num_batches=len(self.loader),
                  batch=batch,
              )
              batch_forward_time += time.time()
  
              # backward pass on penalties
              batch_backward_time = batch_result.backward_time - time.time()
              penalty = 0.0
              for index, (penalty_key, penalty_value_torch) in enumerate(penalties_torch):
                  penalty_value_torch.backward()
                  penalty += penalty_value_torch.item()
                  sum_penalties[penalty_key] += penalty_value_torch.item()
              sum_penalty += penalty
              batch_backward_time += time.time()
  
              # determine full cost
              cost_value = batch_result.avg_loss + penalty
  
              # abort on nan
              if self.abort_on_nan and math.isnan(cost_value):
                  raise FloatingPointError("Cost became nan, aborting training job")
  
              # TODO # visualize graph
              # if (
              #     self.epoch == 1
              #     and batch_index == 0
              #     and self.config.get("train.visualize_graph")
              # ):
              #     from torchviz import make_dot
  
              #     f = os.path.join(self.config.folder, "cost_value")
              #     graph = make_dot(cost_value, params=dict(self.model.named_parameters()))
              #     graph.save(f"{f}.gv")
              #     graph.render(f)  # needs graphviz installed
              #     self.config.log("Exported compute graph to " + f + ".{gv,pdf}")
  
              # print memory stats
              if self.epoch == 1 and batch_index == 0:
                  if self.device.startswith("cuda"):
                      self.config.log(
                          "CUDA memory after first batch: allocated={:14,} "
                          "cached={:14,} max_allocated={:14,}".format(
                              torch.cuda.memory_allocated(self.device),
                              torch.cuda.memory_cached(self.device),
                              torch.cuda.max_memory_allocated(self.device),
                          )
                      )
  
              # update parameters
              batch_optimizer_time = -time.time()
              self.optimizer.step()
              batch_optimizer_time += time.time()
  
              # tracing/logging
              if self.trace_batch:
                  batch_trace = {
                      "type": self.type_str,
                      "scope": "batch",
                      "epoch": self.epoch,
                      "split": self.train_split,
                      "batch": batch_index,
                      "size": batch_result.size,
                      "batches": len(self.loader),
                      "lr": [group["lr"] for group in self.optimizer.param_groups],
                      "avg_loss": batch_result.avg_loss,
                      "penalties": [p.item() for k, p in penalties_torch],
                      "penalty": penalty,
                      "cost": cost_value,
                      "prepare_time": batch_result.prepare_time,
                      "forward_time": batch_forward_time,
                      "backward_time": batch_backward_time,
                      "optimizer_time": batch_optimizer_time,
                  }
                  for f in self.post_batch_trace_hooks:
                      f(self, batch_trace)
                  self.trace(**batch_trace, event="batch_completed")
              print(
                  (
                      "\r"  # go back
                      + "{}  batch{: "
                      + str(1 + int(math.ceil(math.log10(len(self.loader)))))
                      + "d}/{}"
                      + ", avg_loss {:.4E}, penalty {:.4E}, cost {:.4E}, time {:6.2f}s"
                      + "\033[K"  # clear to right
                  ).format(
                      self.config.log_prefix,
                      batch_index,
                      len(self.loader) - 1,
                      batch_result.avg_loss,
                      penalty,
                      cost_value,
                      batch_result.prepare_time
                      + batch_forward_time
                      + batch_backward_time
                      + batch_optimizer_time,
                  ),
                  end="",
                  flush=True,
              )
  
              # update times
              prepare_time += batch_result.prepare_time
              forward_time += batch_forward_time
              backward_time += batch_backward_time
              optimizer_time += batch_optimizer_time
  
          # all done; now trace and log
          epoch_time += time.time()
          print("\033[2K\r", end="", flush=True)  # clear line and go back
  
          other_time = (
              epoch_time - prepare_time - forward_time - backward_time - optimizer_time
          )
          trace_entry = dict(
              type=self.type_str,
              scope="epoch",
              epoch=self.epoch,
              split=self.train_split,
              batches=len(self.loader),
              size=self.num_examples,
              lr = [group["lr"] for group in self.optimizer.param_groups],
              avg_loss=sum_loss / self.num_examples,
              avg_penalty=sum_penalty / len(self.loader),
              avg_penalties={k: p / len(self.loader) for k, p in sum_penalties.items()},
              avg_cost=sum_loss / self.num_examples + sum_penalty / len(self.loader),
              epoch_time=epoch_time,
              prepare_time=prepare_time,
              forward_time=forward_time,
              backward_time=backward_time,
              optimizer_time=optimizer_time,
              other_time=other_time,
              event="epoch_completed",
          )
          for f in self.post_epoch_trace_hooks:
              f(self, trace_entry)
          trace_entry = self.trace(**trace_entry, echo=True, echo_prefix="  ", log=True)
          return trace_entry
  
      def _prepare(self):
          """Prepare this job for running.
  
          Sets (at least) the `loader`, `num_examples`, and `type_str` attributes of this
          job to a data loader, number of examples per epoch, and a name for the trainer,
          repectively.
  
          Guaranteed to be called exactly once before running the first epoch.
  
          """
          raise NotImplementedError
  
      @dataclass
      class _ProcessBatchResult:
          """Result of running forward+backward pass on a batch."""
  
          avg_loss: float
          size: int
          prepare_time: float
          forward_time: float
          backward_time: float
  
      def _process_batch(
          self, batch_index: int, batch
      ) -> "TrainingJob._ProcessBatchResult":
          "Run forward and backward pass on batch and return results."
          raise NotImplementedError
  
  
  class TrainingJobKvsAll(TrainingJob):
      def __init__(self, config, dataset, parent_job=None):
          super().__init__(config, dataset, parent_job)
          self.label_smoothing = config.check_range(
              "KvsAll.label_smoothing", float("-inf"), 1.0, max_inclusive=False
          )
          if self.label_smoothing < 0:
              if config.get("train.auto_correct"):
                  config.log(
                      "Setting label_smoothing to 0, "
                      "was set to {}.".format(self.label_smoothing)
                  )
                  self.label_smoothing = 0
              else:
                  raise Exception(
                      "Label_smoothing was set to {}, "
                      "should be at least 0.".format(self.label_smoothing)
                  )
          elif self.label_smoothing > 0 and self.label_smoothing <= (
              1.0 / dataset.num_entities()
          ):
              if config.get("train.auto_correct"):
                  # just to be sure it's used correctly
                  config.log(
                      "Setting label_smoothing to 1/num_entities = {}, "
                      "was set to {}.".format(
                          1.0 / dataset.num_entities(), self.label_smoothing
                      )
                  )
                  self.label_smoothing = 1.0 / dataset.num_entities()
              else:
                  raise Exception(
                      "Label_smoothing was set to {}, "
                      "should be at least {}.".format(
                          self.label_smoothing, 1.0 / dataset.num_entities()
                      )
                  )
  
          config.log("Initializing 1-to-N training job...")
          self.type_str = "KvsAll"
  
          if self.__class__ == TrainingJobKvsAll:
              for f in Job.job_created_hooks:
                  f(self)
  
      def _prepare(self):
          # create sp and po label_coords (if not done before)
          train_sp = self.dataset.index(f"{self.train_split}_sp_to_o")
          train_po = self.dataset.index(f"{self.train_split}_po_to_s")
  
          # convert indexes to pytoch tensors: a nx2 keys tensor (rows = keys),
          # an offset vector (row = starting offset in values for corresponding
          # key), a values vector (entries correspond to values of original
          # index)
          #
          # Afterwards, it holds:
          # index[keys[i]] = values[offsets[i]:offsets[i+1]]
  
          (
              self.train_sp_keys,
              self.train_sp_values,
              self.train_sp_offsets,
          ) = kge.indexing.prepare_index(train_sp)
          (
              self.train_po_keys,
              self.train_po_values,
              self.train_po_offsets,
          ) = kge.indexing.prepare_index(train_po)
  
          # create dataloader
          self.loader = torch.utils.data.DataLoader(
              range(len(train_sp) + len(train_po)),
              collate_fn=self._get_collate_fun(),
              shuffle=True,
              batch_size=self.batch_size,
              num_workers=self.config.get("train.num_workers"),
              pin_memory=self.config.get("train.pin_memory"),
          )
          self.num_examples = len(train_sp) + len(train_po)
  
      def _get_collate_fun(self):
          num_sp = len(self.train_sp_keys)
  
          # create the collate function
          def collate(batch):
              """For a batch of size n, returns a triple of:
  
              - pairs (nx2 tensor, row = sp or po indexes),
              - label coordinates (position of ones in a batch_size x num_entities tensor)
              - is_sp (vector of size n, 1 if corresponding example_index is sp, 0 if po)
              - triples (all true triples in the batch: needed for weighted penalties only)
  
              """
              # count how many labels we have
              num_ones = 0
              for example_index in batch:
                  if example_index < num_sp:
                      num_ones += self.train_sp_offsets[example_index + 1]
                      num_ones -= self.train_sp_offsets[example_index]
                  else:
                      example_index -= num_sp
                      num_ones += self.train_po_offsets[example_index + 1]
                      num_ones -= self.train_po_offsets[example_index]
  
              # now create the results
              sp_po_batch = torch.zeros([len(batch), 2], dtype=torch.long)
              is_sp = torch.zeros([len(batch)], dtype=torch.long)
              label_coords = torch.zeros([num_ones, 2], dtype=torch.int)
              current_index = 0
              triples = torch.zeros([num_ones, 3], dtype=torch.long)
              for batch_index, example_index in enumerate(batch):
                  is_sp[batch_index] = 1 if example_index < num_sp else 0
                  if is_sp[batch_index]:
                      keys = self.train_sp_keys
                      offsets = self.train_sp_offsets
                      values = self.train_sp_values
                      sp_po_col_1, sp_po_col_2, o_s_col = S, P, O
                  else:
                      example_index -= num_sp
                      keys = self.train_po_keys
                      offsets = self.train_po_offsets
                      values = self.train_po_values
                      o_s_col, sp_po_col_1, sp_po_col_2 = S, P, O
  
                  sp_po_batch[batch_index,] = keys[example_index]
                  start = offsets[example_index]
                  end = offsets[example_index + 1]
                  size = end - start
                  label_coords[current_index : (current_index + size), 0] = batch_index
                  label_coords[current_index : (current_index + size), 1] = values[
                      start:end
                  ]
                  triples[current_index : (current_index + size), sp_po_col_1] = keys[
                      example_index
                  ][0]
                  triples[current_index : (current_index + size), sp_po_col_2] = keys[
                      example_index
                  ][1]
                  triples[current_index : (current_index + size), o_s_col] = values[
                      start:end
                  ]
                  current_index += size
  
              # all done
              return {
                  "sp_po_batch": sp_po_batch,
                  "label_coords": label_coords,
                  "is_sp": is_sp,
                  "triples": triples,
              }
  
          return collate
  
      def _process_batch(self, batch_index, batch) -> TrainingJob._ProcessBatchResult:
          # prepare
          prepare_time = -time.time()
          sp_po_batch = batch["sp_po_batch"].to(self.device)
          batch_size = len(sp_po_batch)
          label_coords = batch["label_coords"].to(self.device)
          is_sp = batch["is_sp"]
          sp_indexes = is_sp.nonzero().to(self.device).view(-1)
          po_indexes = (is_sp == 0).nonzero().to(self.device).view(-1)
          labels = kge.job.util.coord_to_sparse_tensor(
              batch_size, self.dataset.num_entities(), label_coords, self.device
          ).to_dense()
          if self.label_smoothing > 0.0:
              # as in ConvE: https://github.com/TimDettmers/ConvE
              labels = (1.0 - self.label_smoothing) * labels + 1.0 / labels.size(1)
          prepare_time += time.time()
  
          # forward/backward pass (sp)
          loss_value = 0.0
          if len(sp_indexes) > 0:
              forward_time = -time.time()
              scores_sp = self.model.score_sp(
                  sp_po_batch[sp_indexes, 0], sp_po_batch[sp_indexes, 1]
              )
              loss_value_sp = self.loss(scores_sp, labels[sp_indexes,]) / batch_size
              loss_value = loss_value_sp.item()
              forward_time += time.time()
              backward_time = -time.time()
              loss_value_sp.backward()
              backward_time += time.time()
  
          # forward/backward pass (po)
          if len(po_indexes) > 0:
              forward_time = -time.time()
              scores_po = self.model.score_po(
                  sp_po_batch[po_indexes, 0], sp_po_batch[po_indexes, 1]
              )
              loss_value_po = self.loss(scores_po, labels[po_indexes,]) / batch_size
              loss_value += loss_value_po.item()
              forward_time += time.time()
              backward_time = -time.time()
              loss_value_po.backward()
              backward_time += time.time()
  
          # all done
          return TrainingJob._ProcessBatchResult(
              loss_value, batch_size, prepare_time, forward_time, backward_time
          )
  
  
  class TrainingJobNegativeSampling(TrainingJob):
      def __init__(self, config, dataset, parent_job=None):
          super().__init__(config, dataset, parent_job)
          self._sampler = KgeSampler.create(config, "negative_sampling", dataset)
          self.is_prepared = False
          self._implementation = self.config.check(
              "negative_sampling.implementation", ["triple", "all", "batch", "auto"],
          )
          if self._implementation == "auto":
              max_nr_of_negs = max(self._sampler.num_samples)
              if self._sampler.shared:
                  self._implementation = "batch"
              elif max_nr_of_negs <= 30:
                  self._implementation = "triple"
              elif max_nr_of_negs > 30:
                  self._implementation = "batch"
          self._max_chunk_size = self.config.get("negative_sampling.chunk_size")
  
          config.log(
              "Initializing negative sampling training job with "
              "'{}' scoring function ...".format(self._implementation)
          )
          self.type_str = "negative_sampling"
  
          if self.__class__ == TrainingJobNegativeSampling:
              for f in Job.job_created_hooks:
                  f(self)
  
      def _prepare(self):
          """Construct dataloader"""
  
          if self.is_prepared:
              return
  
          self.num_examples = self.dataset.split(self.train_split).size(0)
          self.loader = torch.utils.data.DataLoader(
              range(self.num_examples),
              collate_fn=self._get_collate_fun(),
              shuffle=True,
              batch_size=self.batch_size,
              num_workers=self.config.get("train.num_workers"),
              pin_memory=self.config.get("train.pin_memory"),
          )
  
          self.is_prepared = True
  
      def _get_collate_fun(self):
          # create the collate function
          def collate(batch):
              """For a batch of size n, returns a tuple of:
  
              - triples (tensor of shape [n,3], ),
              - negative_samples (list of tensors of shape [n,num_samples]; 3 elements
                in order S,P,O)
              """
  
              triples = self.dataset.split(self.train_split)[batch, :].long()
              # labels = torch.zeros((len(batch), self._sampler.num_negatives_total + 1))
              # labels[:, 0] = 1
              # labels = labels.view(-1)
  
              negative_samples = list()
              for slot in [S, P, O]:
                  negative_samples.append(self._sampler.sample(triples, slot))
              return {"triples": triples, "negative_samples": negative_samples}
  
          return collate
  
      def _process_batch(self, batch_index, batch) -> TrainingJob._ProcessBatchResult:
          # prepare
          prepare_time = -time.time()
          batch_triples = batch["triples"].to(self.device)
          batch_negative_samples = [
              ns.to(self.device) for ns in batch["negative_samples"]
          ]
          batch_size = len(batch_triples)
          prepare_time += time.time()
  
          loss_value = 0.0
          forward_time = 0.0
          backward_time = 0.0
          labels = None
  
          # perform processing of batch in smaller chunks to save memory
          max_chunk_size = (
              self._max_chunk_size if self._max_chunk_size > 0 else batch_size
          )
          for chunk_number in range(math.ceil(batch_size / max_chunk_size)):
              # determine data used for this chunk
              chunk_start = max_chunk_size * chunk_number
              chunk_end = min(max_chunk_size * (chunk_number + 1), batch_size)
              negative_samples = [
                  ns[chunk_start:chunk_end, :] for ns in batch_negative_samples
              ]
              triples = batch_triples[chunk_start:chunk_end, :]
              chunk_size = chunk_end - chunk_start
  
              # process the chunk
              for slot in [S, P, O]:
                  num_samples = self._sampler.num_samples[slot]
                  if num_samples <= 0:
                      continue
  
                  # construct gold labels: first column corresponds to positives,
                  # remaining columns to negatives
                  if labels is None or labels.shape != torch.Size(
                      [chunk_size, 1 + num_samples]
                  ):
                      prepare_time -= time.time()
                      labels = torch.zeros(
                          (chunk_size, 1 + num_samples), device=self.device
                      )
                      labels[:, 0] = 1
                      prepare_time += time.time()
  
                  # compute corresponding scores
                  scores = None
                  if self._implementation == "triple":
                      # construct triples
                      prepare_time -= time.time()
                      triples_to_score = triples.repeat(1, 1 + num_samples).view(-1, 3)
                      triples_to_score[:, slot] = torch.cat(
                          (
                              triples[:, [slot]],  # positives
                              negative_samples[slot],  # negatives
                          ),
                          1,
                      ).view(-1)
                      prepare_time += time.time()
  
                      # and score them
                      forward_time -= time.time()
                      scores = self.model.score_spo(
                          triples_to_score[:, 0],
                          triples_to_score[:, 1],
                          triples_to_score[:, 2],
                          direction="s" if slot == S else ("o" if slot == O else "p"),
                      ).view(chunk_size, -1)
                      forward_time += time.time()
                  elif self._implementation == "all":
                      # Score against all possible targets. Creates a score matrix of size
                      # [chunk_size, num_entities] or [chunk_size, num_relations]. All
                      # scores relevant for positive and negative triples are contained in
                      # this score matrix.
  
                      # compute all scores for slot
                      forward_time -= time.time()
                      if slot == S:
                          all_scores = self.model.score_po(triples[:, P], triples[:, O])
                      elif slot == P:
                          all_scores = self.model.score_so(triples[:, S], triples[:, O])
                      elif slot == O:
                          all_scores = self.model.score_sp(triples[:, S], triples[:, P])
                      else:
                          raise NotImplementedError
                      forward_time += time.time()
  
                      # determine indexes of relevant scores in scoring matrix
                      prepare_time -= time.time()
                      row_indexes = (
                          torch.arange(chunk_size, device=self.device)
                          .unsqueeze(1)
                          .repeat(1, 1 + num_samples)
                          .view(-1)
                      )  # 000 111 222; each 1+num_negative times (here: 3)
                      column_indexes = torch.cat(
                          (
                              triples[:, [slot]],  # positives
                              negative_samples[slot],  # negatives
                          ),
                          1,
                      ).view(-1)
                      prepare_time += time.time()
  
                      # now pick the scores we need
                      forward_time -= time.time()
                      scores = all_scores[row_indexes, column_indexes].view(
                          chunk_size, -1
                      )
                      forward_time += time.time()
                  elif self._implementation == "batch":
                      # Score against all targets contained in the chunk. Creates a score
                      # matrix of size [chunk_size, unique_entities_in_slot] or
                      # [chunk_size, unique_relations_in_slot]. All scores
                      # relevant for positive and negative triples are contained in this
                      # score matrix.
                      forward_time -= time.time()
                      unique_targets, column_indexes = torch.unique(
                          torch.cat((triples[:, [slot]], negative_samples[slot]), 1).view(
                              -1
                          ),
                          return_inverse=True,
                      )
  
                      # compute scores for all unique targets for slot
                      if slot == S:
                          all_scores = self.model.score_po(
                              triples[:, P], triples[:, O], unique_targets
                          )
                      elif slot == P:
                          all_scores = self.model.score_so(
                              triples[:, S], triples[:, O], unique_targets
                          )
                      elif slot == O:
                          all_scores = self.model.score_sp(
                              triples[:, S], triples[:, P], unique_targets
                          )
                      else:
                          raise NotImplementedError
                      forward_time += time.time()
  
                      # determine indexes of relevant scores in scoring matrix
                      prepare_time -= time.time()
                      row_indexes = (
                          torch.arange(chunk_size, device=self.device)
                          .unsqueeze(1)
                          .repeat(1, 1 + num_samples)
                          .view(-1)
                      )  # 000 111 222; each 1+num_negative times (here: 3)
                      prepare_time += time.time()
  
                      # now pick the scores we need
                      forward_time -= time.time()
                      scores = all_scores[row_indexes, column_indexes].view(
                          chunk_size, -1
                      )
                      forward_time += time.time()
  
                  # compute chunk loss (concluding the forward pass of the chunk)
                  forward_time -= time.time()
                  loss_value_torch = (
                      self.loss(scores, labels, num_negatives=num_samples) / batch_size
                  )
                  loss_value += loss_value_torch.item()
                  forward_time += time.time()
  
                  # backward pass for this chunk
                  backward_time -= time.time()
                  loss_value_torch.backward()
                  backward_time += time.time()
  
          # all done
          return TrainingJob._ProcessBatchResult(
              loss_value, batch_size, prepare_time, forward_time, backward_time
          )
  
  
  class TrainingJob1vsAll(TrainingJob):
      """Samples SPO pairs and queries sp_ and _po, treating all other entities as negative."""
  
      def __init__(self, config, dataset, parent_job=None):
          super().__init__(config, dataset, parent_job)
          self.is_prepared = False
          config.log("Initializing spo training job...")
          self.type_str = "1vsAll"
  
          if self.__class__ == TrainingJob1vsAll:
              for f in Job.job_created_hooks:
                  f(self)
  
      def _prepare(self):
          """Construct dataloader"""
  
          if self.is_prepared:
              return
  
          self.num_examples = self.dataset.split(self.train_split).size(0)
          self.loader = torch.utils.data.DataLoader(
              range(self.num_examples),
              collate_fn=lambda batch: {
                  "triples": self.dataset.split(self.train_split)[batch, :].long()
              },
              shuffle=True,
              batch_size=self.batch_size,
              num_workers=self.config.get("train.num_workers"),
              pin_memory=self.config.get("train.pin_memory"),
          )
  
          self.is_prepared = True
  
      def _process_batch(self, batch_index, batch) -> TrainingJob._ProcessBatchResult:
          # prepare
          prepare_time = -time.time()
          triples = batch["triples"].to(self.device)
          batch_size = len(triples)
          prepare_time += time.time()
  
          # forward/backward pass (sp)
          forward_time = -time.time()
          scores_sp = self.model.score_sp(triples[:, 0], triples[:, 1])
          loss_value_sp = self.loss(scores_sp, triples[:, 2]) / batch_size
          loss_value = loss_value_sp.item()
          forward_time += time.time()
          backward_time = -time.time()
          loss_value_sp.backward()
          backward_time += time.time()
  
          # forward/backward pass (po)
          forward_time -= time.time()
          scores_po = self.model.score_po(triples[:, 1], triples[:, 2])
          loss_value_po = self.loss(scores_po, triples[:, 0]) / batch_size
          loss_value += loss_value_po.item()
          forward_time += time.time()
          backward_time -= time.time()
          loss_value_po.backward()
          backward_time += time.time()
  
          # all done
          return TrainingJob._ProcessBatchResult(
              loss_value, batch_size, prepare_time, forward_time, backward_time
          )
  
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 of epochs reached.")
                  break
  
              # start a new epoch
              self.epoch += 1
              self.config.log("Starting epoch {}...".format(self.epoch))
              trace_entry = self.run_epoch()
              for f in self.post_epoch_hooks:
                  f(self, trace_entry)
              self.config.log("Finished epoch {}.".format(self.epoch))
  
              # update model metadata
              self.model.meta["train_job_trace_entry"] = self.trace_entry
              self.model.meta["train_epoch"] = self.epoch
              self.model.meta["train_config"] = self.config
              self.model.meta["train_trace_entry"] = trace_entry
  
              # validate and update learning rate
              if (
                  self.config.get("valid.every") > 0
                  and self.epoch % self.config.get("valid.every") == 0
              ):
                  self.valid_job.epoch = self.epoch
                  trace_entry = self.valid_job.run()
                  self.valid_trace.append(trace_entry)
                  for f in self.post_valid_hooks:
                      f(self, trace_entry)
                  self.model.meta["valid_trace_entry"] = trace_entry
  
                  # metric-based scheduler step
                  self.kge_lr_scheduler.step(trace_entry[metric_name])
              else:
                  self.kge_lr_scheduler.step()
  
              # create checkpoint and delete old one, if necessary
              self.save(self.config.checkpoint_file(self.epoch))
              if self.epoch > 1:
                  delete_checkpoint_epoch = -1
                  if checkpoint_every == 0:
                      # do not keep any old checkpoints
                      delete_checkpoint_epoch = self.epoch - 1
                  elif (self.epoch - 1) % checkpoint_every != 0:
                      # delete checkpoints that are not in the checkpoint.every schedule
                      delete_checkpoint_epoch = self.epoch - 1
                  elif checkpoint_keep > 0:
                      # keep a maximum number of checkpoint_keep checkpoints
                      delete_checkpoint_epoch = (
                          self.epoch - 1 - checkpoint_every * checkpoint_keep
                      )
                  if delete_checkpoint_epoch > 0:
                      if os.path.exists(
                          self.config.checkpoint_file(delete_checkpoint_epoch)
                      ):
                          self.config.log(
                              "Removing old checkpoint {}...".format(
                                  self.config.checkpoint_file(delete_checkpoint_epoch)
                              )
                          )
                          os.remove(self.config.checkpoint_file(delete_checkpoint_epoch))
                      else:
                          self.config.log(
                              "Could not delete old checkpoint {}, does not exits.".format(
                                  self.config.checkpoint_file(delete_checkpoint_epoch)
                              )
                          )
  
          for f in self.post_train_hooks:
              f(self, trace_entry)
          self.trace(event="train_completed")
  
      def save(self, filename) -> None:
          """Save current state to specified file"""
          self.config.log("Saving checkpoint to {}...".format(filename))
          torch.save(
              {
                  "type": "train",
                  "config": self.config,
                  "epoch": self.epoch,
                  "valid_trace": self.valid_trace,
                  "model": self.model.save(),
                  "optimizer_state_dict": self.optimizer.state_dict(),
                  "lr_scheduler_state_dict": self.kge_lr_scheduler.state_dict(),
                  "job_id": self.job_id,
              },
              filename,
          )
  
      def load(self, filename: str) -> str:
          """Load job state from specified file.
  
          Returns job id of the job that created the checkpoint."""
          self.config.log("Loading checkpoint from {}...".format(filename))
          checkpoint = torch.load(filename, map_location="cpu")
          if "model" in checkpoint:
              # new format
              self.model.load(checkpoint["model"])
          else:
              # old format (deprecated, will eventually be removed)
              self.model.load_state_dict(checkpoint["model_state_dict"])
          self.optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
          if "lr_scheduler_state_dict" in checkpoint:
              # new format
              self.kge_lr_scheduler.load_state_dict(checkpoint["lr_scheduler_state_dict"])
          self.epoch = checkpoint["epoch"]
          self.valid_trace = checkpoint["valid_trace"]
          self.model.train()
          return checkpoint.get("job_id")
  
      def resume(self, checkpoint_file: str = None) -> None:
          if checkpoint_file is None:
              last_checkpoint = self.config.last_checkpoint()
              if last_checkpoint is not None:
                  checkpoint_file = self.config.checkpoint_file(last_checkpoint)
  
          if checkpoint_file is not None:
              self.resumed_from_job_id = self.load(checkpoint_file)
              self.trace(
                  event="job_resumed", epoch=self.epoch, checkpoint_file=checkpoint_file
              )
              self.config.log(
                  "Resumed from {} of job {}".format(
                      checkpoint_file, self.resumed_from_job_id
                  )
              )
          else:
              self.config.log("No checkpoint found, starting from scratch...")
  
      def run_epoch(self) -> Dict[str, Any]:
          "Runs an epoch and returns a trace entry."
  
          # prepare the job is not done already
          if not self.is_prepared:
              self._prepare()
              self.model.prepare_job(self)  # let the model add some hooks
              self.is_prepared = True
  
          # variables that record various statitics
          sum_loss = 0.0
          sum_penalty = 0.0
          sum_penalties = defaultdict(lambda: 0.0)
          epoch_time = -time.time()
          prepare_time = 0.0
          forward_time = 0.0
          backward_time = 0.0
          optimizer_time = 0.0
  
          # process each batch
          for batch_index, batch in enumerate(self.loader):
              for f in self.pre_batch_hooks:
                  f(self)
  
              # process batch (preprocessing + forward pass + backward pass on loss)
              self.optimizer.zero_grad()
              batch_result: TrainingJob._ProcessBatchResult = self._process_batch(
                  batch_index, batch
              )
              sum_loss += batch_result.avg_loss * batch_result.size
  
              # determine penalty terms (forward pass)
              batch_forward_time = batch_result.forward_time - time.time()
              penalties_torch = self.model.penalty(
                  epoch=self.epoch,
                  batch_index=batch_index,
                  num_batches=len(self.loader),
                  batch=batch,
              )
              batch_forward_time += time.time()
  
              # backward pass on penalties
              batch_backward_time = batch_result.backward_time - time.time()
              penalty = 0.0
              for index, (penalty_key, penalty_value_torch) in enumerate(penalties_torch):
                  penalty_value_torch.backward()
                  penalty += penalty_value_torch.item()
                  sum_penalties[penalty_key] += penalty_value_torch.item()
              sum_penalty += penalty
              batch_backward_time += time.time()
  
              # determine full cost
              cost_value = batch_result.avg_loss + penalty
  
              # abort on nan
              if self.abort_on_nan and math.isnan(cost_value):
                  raise FloatingPointError("Cost became nan, aborting training job")
  
              # TODO # visualize graph
              # if (
              #     self.epoch == 1
              #     and batch_index == 0
              #     and self.config.get("train.visualize_graph")
              # ):
              #     from torchviz import make_dot
  
              #     f = os.path.join(self.config.folder, "cost_value")
              #     graph = make_dot(cost_value, params=dict(self.model.named_parameters()))
              #     graph.save(f"{f}.gv")
              #     graph.render(f)  # needs graphviz installed
              #     self.config.log("Exported compute graph to " + f + ".{gv,pdf}")
  
              # print memory stats
              if self.epoch == 1 and batch_index == 0:
                  if self.device.startswith("cuda"):
                      self.config.log(
                          "CUDA memory after first batch: allocated={:14,} "
                          "cached={:14,} max_allocated={:14,}".format(
                              torch.cuda.memory_allocated(self.device),
                              torch.cuda.memory_cached(self.device),
                              torch.cuda.max_memory_allocated(self.device),
                          )
                      )
  
              # update parameters
              batch_optimizer_time = -time.time()
              self.optimizer.step()
              batch_optimizer_time += time.time()
  
              # tracing/logging
              if self.trace_batch:
                  batch_trace = {
                      "type": self.type_str,
                      "scope": "batch",
                      "epoch": self.epoch,
                      "split": self.train_split,
                      "batch": batch_index,
                      "size": batch_result.size,
                      "batches": len(self.loader),
                      "lr": [group["lr"] for group in self.optimizer.param_groups],
                      "avg_loss": batch_result.avg_loss,
                      "penalties": [p.item() for k, p in penalties_torch],
                      "penalty": penalty,
                      "cost": cost_value,
                      "prepare_time": batch_result.prepare_time,
                      "forward_time": batch_forward_time,
                      "backward_time": batch_backward_time,
                      "optimizer_time": batch_optimizer_time,
                  }
                  for f in self.post_batch_trace_hooks:
                      f(self, batch_trace)
                  self.trace(**batch_trace, event="batch_completed")
              print(
                  (
                      "\r"  # go back
                      + "{}  batch{: "
                      + str(1 + int(math.ceil(math.log10(len(self.loader)))))
                      + "d}/{}"
                      + ", avg_loss {:.4E}, penalty {:.4E}, cost {:.4E}, time {:6.2f}s"
                      + "\033[K"  # clear to right
                  ).format(
                      self.config.log_prefix,
                      batch_index,
                      len(self.loader) - 1,
                      batch_result.avg_loss,
                      penalty,
                      cost_value,
                      batch_result.prepare_time
                      + batch_forward_time
                      + batch_backward_time
                      + batch_optimizer_time,
                  ),
                  end="",
                  flush=True,
              )
  
              # update times
              prepare_time += batch_result.prepare_time
              forward_time += batch_forward_time
              backward_time += batch_backward_time
              optimizer_time += batch_optimizer_time
  
          # all done; now trace and log
          epoch_time += time.time()
          print("\033[2K\r", end="", flush=True)  # clear line and go back
  
          other_time = (
              epoch_time - prepare_time - forward_time - backward_time - optimizer_time
          )
          trace_entry = dict(
              type=self.type_str,
              scope="epoch",
              epoch=self.epoch,
              split=self.train_split,
              batches=len(self.loader),
              size=self.num_examples,
              lr = [group["lr"] for group in self.optimizer.param_groups],
              avg_loss=sum_loss / self.num_examples,
              avg_penalty=sum_penalty / len(self.loader),
              avg_penalties={k: p / len(self.loader) for k, p in sum_penalties.items()},
              avg_cost=sum_loss / self.num_examples + sum_penalty / len(self.loader),
              epoch_time=epoch_time,
              prepare_time=prepare_time,
              forward_time=forward_time,
              backward_time=backward_time,
              optimizer_time=optimizer_time,
              other_time=other_time,
              event="epoch_completed",
          )
          for f in self.post_epoch_trace_hooks:
              f(self, trace_entry)
          trace_entry = self.trace(**trace_entry, echo=True, echo_prefix="  ", log=True)
          return trace_entry
  
      def _prepare(self):
          """Prepare this job for running.
  
          Sets (at least) the `loader`, `num_examples`, and `type_str` attributes of this
          job to a data loader, number of examples per epoch, and a name for the trainer,
          repectively.
  
          Guaranteed to be called exactly once before running the first epoch.
  
          """
          raise NotImplementedError
  
      @dataclass
      class _ProcessBatchResult:
          """Result of running forward+backward pass on a batch."""
  
          avg_loss: float
          size: int
          prepare_time: float
          forward_time: float
          backward_time: float
  
      def _process_batch(
          self, batch_index: int, batch
      ) -> "TrainingJob._ProcessBatchResult":
          "Run forward and backward pass on batch and return results."
          raise NotImplementedError
  
  
  class TrainingJobKvsAll(TrainingJob):
      def __init__(self, config, dataset, parent_job=None):
          super().__init__(config, dataset, parent_job)
          self.label_smoothing = config.check_range(
              "KvsAll.label_smoothing", float("-inf"), 1.0, max_inclusive=False
          )
          if self.label_smoothing < 0:
              if config.get("train.auto_correct"):
                  config.log(
                      "Setting label_smoothing to 0, "
                      "was set to {}.".format(self.label_smoothing)
                  )
                  self.label_smoothing = 0
              else:
                  raise Exception(
                      "Label_smoothing was set to {}, "
                      "should be at least 0.".format(self.label_smoothing)
                  )
          elif self.label_smoothing > 0 and self.label_smoothing <= (
              1.0 / dataset.num_entities()
          ):
              if config.get("train.auto_correct"):
                  # just to be sure it's used correctly
                  config.log(
                      "Setting label_smoothing to 1/num_entities = {}, "
                      "was set to {}.".format(
                          1.0 / dataset.num_entities(), self.label_smoothing
                      )
                  )
                  self.label_smoothing = 1.0 / dataset.num_entities()
              else:
                  raise Exception(
                      "Label_smoothing was set to {}, "
                      "should be at least {}.".format(
                          self.label_smoothing, 1.0 / dataset.num_entities()
                      )
                  )
  
          config.log("Initializing 1-to-N training job...")
          self.type_str = "KvsAll"
  
          if self.__class__ == TrainingJobKvsAll:
              for f in Job.job_created_hooks:
                  f(self)
  
      def _prepare(self):
          # create sp and po label_coords (if not done before)
          train_sp = self.dataset.index(f"{self.train_split}_sp_to_o")
          train_po = self.dataset.index(f"{self.train_split}_po_to_s")
  
          # convert indexes to pytoch tensors: a nx2 keys tensor (rows = keys),
          # an offset vector (row = starting offset in values for corresponding
          # key), a values vector (entries correspond to values of original
          # index)
          #
          # Afterwards, it holds:
          # index[keys[i]] = values[offsets[i]:offsets[i+1]]
  
          (
              self.train_sp_keys,
              self.train_sp_values,
              self.train_sp_offsets,
          ) = kge.indexing.prepare_index(train_sp)
          (
              self.train_po_keys,
              self.train_po_values,
              self.train_po_offsets,
          ) = kge.indexing.prepare_index(train_po)
  
          # create dataloader
          self.loader = torch.utils.data.DataLoader(
              range(len(train_sp) + len(train_po)),
              collate_fn=self._get_collate_fun(),
              shuffle=True,
              batch_size=self.batch_size,
              num_workers=self.config.get("train.num_workers"),
              pin_memory=self.config.get("train.pin_memory"),
          )
          self.num_examples = len(train_sp) + len(train_po)
  
      def _get_collate_fun(self):
          num_sp = len(self.train_sp_keys)
  
          # create the collate function
          def collate(batch):
              """For a batch of size n, returns a triple of:
  
              - pairs (nx2 tensor, row = sp or po indexes),
              - label coordinates (position of ones in a batch_size x num_entities tensor)
              - is_sp (vector of size n, 1 if corresponding example_index is sp, 0 if po)
              - triples (all true triples in the batch: needed for weighted penalties only)
  
              """
              # count how many labels we have
              num_ones = 0
              for example_index in batch:
                  if example_index < num_sp:
                      num_ones += self.train_sp_offsets[example_index + 1]
                      num_ones -= self.train_sp_offsets[example_index]
                  else:
                      example_index -= num_sp
                      num_ones += self.train_po_offsets[example_index + 1]
                      num_ones -= self.train_po_offsets[example_index]
  
              # now create the results
              sp_po_batch = torch.zeros([len(batch), 2], dtype=torch.long)
              is_sp = torch.zeros([len(batch)], dtype=torch.long)
              label_coords = torch.zeros([num_ones, 2], dtype=torch.int)
              current_index = 0
              triples = torch.zeros([num_ones, 3], dtype=torch.long)
              for batch_index, example_index in enumerate(batch):
                  is_sp[batch_index] = 1 if example_index < num_sp else 0
                  if is_sp[batch_index]:
                      keys = self.train_sp_keys
                      offsets = self.train_sp_offsets
                      values = self.train_sp_values
                      sp_po_col_1, sp_po_col_2, o_s_col = S, P, O
                  else:
                      example_index -= num_sp
                      keys = self.train_po_keys
                      offsets = self.train_po_offsets
                      values = self.train_po_values
                      o_s_col, sp_po_col_1, sp_po_col_2 = S, P, O
  
                  sp_po_batch[batch_index,] = keys[example_index]
                  start = offsets[example_index]
                  end = offsets[example_index + 1]
                  size = end - start
                  label_coords[current_index : (current_index + size), 0] = batch_index
                  label_coords[current_index : (current_index + size), 1] = values[
                      start:end
                  ]
                  triples[current_index : (current_index + size), sp_po_col_1] = keys[
                      example_index
                  ][0]
                  triples[current_index : (current_index + size), sp_po_col_2] = keys[
                      example_index
                  ][1]
                  triples[current_index : (current_index + size), o_s_col] = values[
                      start:end
                  ]
                  current_index += size
  
              # all done
              return {
                  "sp_po_batch": sp_po_batch,
                  "label_coords": label_coords,
                  "is_sp": is_sp,
                  "triples": triples,
              }
  
          return collate
  
      def _process_batch(self, batch_index, batch) -> TrainingJob._ProcessBatchResult:
          # prepare
          prepare_time = -time.time()
          sp_po_batch = batch["sp_po_batch"].to(self.device)
          batch_size = len(sp_po_batch)
          label_coords = batch["label_coords"].to(self.device)
          is_sp = batch["is_sp"]
          sp_indexes = is_sp.nonzero().to(self.device).view(-1)
          po_indexes = (is_sp == 0).nonzero().to(self.device).view(-1)
          labels = kge.job.util.coord_to_sparse_tensor(
              batch_size, self.dataset.num_entities(), label_coords, self.device
          ).to_dense()
          if self.label_smoothing > 0.0:
              # as in ConvE: https://github.com/TimDettmers/ConvE
              labels = (1.0 - self.label_smoothing) * labels + 1.0 / labels.size(1)
          prepare_time += time.time()
  
          # forward/backward pass (sp)
          loss_value = 0.0
          if len(sp_indexes) > 0:
              forward_time = -time.time()
              scores_sp = self.model.score_sp(
                  sp_po_batch[sp_indexes, 0], sp_po_batch[sp_indexes, 1]
              )
              loss_value_sp = self.loss(scores_sp, labels[sp_indexes,]) / batch_size
              loss_value = loss_value_sp.item()
              forward_time += time.time()
              backward_time = -time.time()
              loss_value_sp.backward()
              backward_time += time.time()
  
          # forward/backward pass (po)
          if len(po_indexes) > 0:
              forward_time = -time.time()
              scores_po = self.model.score_po(
                  sp_po_batch[po_indexes, 0], sp_po_batch[po_indexes, 1]
              )
              loss_value_po = self.loss(scores_po, labels[po_indexes,]) / batch_size
              loss_value += loss_value_po.item()
              forward_time += time.time()
              backward_time = -time.time()
              loss_value_po.backward()
              backward_time += time.time()
  
          # all done
          return TrainingJob._ProcessBatchResult(
              loss_value, batch_size, prepare_time, forward_time, backward_time
          )
  
  
  class TrainingJobNegativeSampling(TrainingJob):
      def __init__(self, config, dataset, parent_job=None):
          super().__init__(config, dataset, parent_job)
          self._sampler = KgeSampler.create(config, "negative_sampling", dataset)
          self.is_prepared = False
          self._implementation = self.config.check(
              "negative_sampling.implementation", ["triple", "all", "batch", "auto"],
          )
          if self._implementation == "auto":
              max_nr_of_negs = max(self._sampler.num_samples)
              if self._sampler.shared:
                  self._implementation = "batch"
              elif max_nr_of_negs <= 30:
                  self._implementation = "triple"
              elif max_nr_of_negs > 30:
                  self._implementation = "batch"
          self._max_chunk_size = self.config.get("negative_sampling.chunk_size")
  
          config.log(
              "Initializing negative sampling training job with "
              "'{}' scoring function ...".format(self._implementation)
          )
          self.type_str = "negative_sampling"
  
          if self.__class__ == TrainingJobNegativeSampling:
              for f in Job.job_created_hooks:
                  f(self)
  
      def _prepare(self):
          """Construct dataloader"""
  
          if self.is_prepared:
              return
  
          self.num_examples = self.dataset.split(self.train_split).size(0)
          self.loader = torch.utils.data.DataLoader(
              range(self.num_examples),
              collate_fn=self._get_collate_fun(),
              shuffle=True,
              batch_size=self.batch_size,
              num_workers=self.config.get("train.num_workers"),
              pin_memory=self.config.get("train.pin_memory"),
          )
  
          self.is_prepared = True
  
      def _get_collate_fun(self):
          # create the collate function
          def collate(batch):
              """For a batch of size n, returns a tuple of:
  
              - triples (tensor of shape [n,3], ),
              - negative_samples (list of tensors of shape [n,num_samples]; 3 elements
                in order S,P,O)
              """
  
              triples = self.dataset.split(self.train_split)[batch, :].long()
              # labels = torch.zeros((len(batch), self._sampler.num_negatives_total + 1))
              # labels[:, 0] = 1
              # labels = labels.view(-1)
  
              negative_samples = list()
              for slot in [S, P, O]:
                  negative_samples.append(self._sampler.sample(triples, slot))
              return {"triples": triples, "negative_samples": negative_samples}
  
          return collate
  
      def _process_batch(self, batch_index, batch) -> TrainingJob._ProcessBatchResult:
          # prepare
          prepare_time = -time.time()
          batch_triples = batch["triples"].to(self.device)
          batch_negative_samples = [
              ns.to(self.device) for ns in batch["negative_samples"]
          ]
          batch_size = len(batch_triples)
          prepare_time += time.time()
  
          loss_value = 0.0
          forward_time = 0.0
          backward_time = 0.0
          labels = None
  
          # perform processing of batch in smaller chunks to save memory
          max_chunk_size = (
              self._max_chunk_size if self._max_chunk_size > 0 else batch_size
          )
          for chunk_number in range(math.ceil(batch_size / max_chunk_size)):
              # determine data used for this chunk
              chunk_start = max_chunk_size * chunk_number
              chunk_end = min(max_chunk_size * (chunk_number + 1), batch_size)
              negative_samples = [
                  ns[chunk_start:chunk_end, :] for ns in batch_negative_samples
              ]
              triples = batch_triples[chunk_start:chunk_end, :]
              chunk_size = chunk_end - chunk_start
  
              # process the chunk
              for slot in [S, P, O]:
                  num_samples = self._sampler.num_samples[slot]
                  if num_samples <= 0:
                      continue
  
                  # construct gold labels: first column corresponds to positives,
                  # remaining columns to negatives
                  if labels is None or labels.shape != torch.Size(
                      [chunk_size, 1 + num_samples]
                  ):
                      prepare_time -= time.time()
                      labels = torch.zeros(
                          (chunk_size, 1 + num_samples), device=self.device
                      )
                      labels[:, 0] = 1
                      prepare_time += time.time()
  
                  # compute corresponding scores
                  scores = None
                  if self._implementation == "triple":
                      # construct triples
                      prepare_time -= time.time()
                      triples_to_score = triples.repeat(1, 1 + num_samples).view(-1, 3)
                      triples_to_score[:, slot] = torch.cat(
                          (
                              triples[:, [slot]],  # positives
                              negative_samples[slot],  # negatives
                          ),
                          1,
                      ).view(-1)
                      prepare_time += time.time()
  
                      # and score them
                      forward_time -= time.time()
                      scores = self.model.score_spo(
                          triples_to_score[:, 0],
                          triples_to_score[:, 1],
                          triples_to_score[:, 2],
                          direction="s" if slot == S else ("o" if slot == O else "p"),
                      ).view(chunk_size, -1)
                      forward_time += time.time()
                  elif self._implementation == "all":
                      # Score against all possible targets. Creates a score matrix of size
                      # [chunk_size, num_entities] or [chunk_size, num_relations]. All
                      # scores relevant for positive and negative triples are contained in
                      # this score matrix.
  
                      # compute all scores for slot
                      forward_time -= time.time()
                      if slot == S:
                          all_scores = self.model.score_po(triples[:, P], triples[:, O])
                      elif slot == P:
                          all_scores = self.model.score_so(triples[:, S], triples[:, O])
                      elif slot == O:
                          all_scores = self.model.score_sp(triples[:, S], triples[:, P])
                      else:
                          raise NotImplementedError
                      forward_time += time.time()
  
                      # determine indexes of relevant scores in scoring matrix
                      prepare_time -= time.time()
                      row_indexes = (
                          torch.arange(chunk_size, device=self.device)
                          .unsqueeze(1)
                          .repeat(1, 1 + num_samples)
                          .view(-1)
                      )  # 000 111 222; each 1+num_negative times (here: 3)
                      column_indexes = torch.cat(
                          (
                              triples[:, [slot]],  # positives
                              negative_samples[slot],  # negatives
                          ),
                          1,
                      ).view(-1)
                      prepare_time += time.time()
  
                      # now pick the scores we need
                      forward_time -= time.time()
                      scores = all_scores[row_indexes, column_indexes].view(
                          chunk_size, -1
                      )
                      forward_time += time.time()
                  elif self._implementation == "batch":
                      # Score against all targets contained in the chunk. Creates a score
                      # matrix of size [chunk_size, unique_entities_in_slot] or
                      # [chunk_size, unique_relations_in_slot]. All scores
                      # relevant for positive and negative triples are contained in this
                      # score matrix.
                      forward_time -= time.time()
                      unique_targets, column_indexes = torch.unique(
                          torch.cat((triples[:, [slot]], negative_samples[slot]), 1).view(
                              -1
                          ),
                          return_inverse=True,
                      )
  
                      # compute scores for all unique targets for slot
                      if slot == S:
                          all_scores = self.model.score_po(
                              triples[:, P], triples[:, O], unique_targets
                          )
                      elif slot == P:
                          all_scores = self.model.score_so(
                              triples[:, S], triples[:, O], unique_targets
                          )
                      elif slot == O:
                          all_scores = self.model.score_sp(
                              triples[:, S], triples[:, P], unique_targets
                          )
                      else:
                          raise NotImplementedError
                      forward_time += time.time()
  
                      # determine indexes of relevant scores in scoring matrix
                      prepare_time -= time.time()
                      row_indexes = (
                          torch.arange(chunk_size, device=self.device)
                          .unsqueeze(1)
                          .repeat(1, 1 + num_samples)
                          .view(-1)
                      )  # 000 111 222; each 1+num_negative times (here: 3)
                      prepare_time += time.time()
  
                      # now pick the scores we need
                      forward_time -= time.time()
                      scores = all_scores[row_indexes, column_indexes].view(
                          chunk_size, -1
                      )
                      forward_time += time.time()
  
                  # compute chunk loss (concluding the forward pass of the chunk)
                  forward_time -= time.time()
                  loss_value_torch = (
                      self.loss(scores, labels, num_negatives=num_samples) / batch_size
                  )
                  loss_value += loss_value_torch.item()
                  forward_time += time.time()
  
                  # backward pass for this chunk
                  backward_time -= time.time()
                  loss_value_torch.backward()
                  backward_time += time.time()
  
          # all done
          return TrainingJob._ProcessBatchResult(
              loss_value, batch_size, prepare_time, forward_time, backward_time
          )
  
  
  class TrainingJob1vsAll(TrainingJob):
      """Samples SPO pairs and queries sp* and *po, treating all other entities as negative."""
  
      def __init__(self, config, dataset, parent_job=None):
          super().__init__(config, dataset, parent_job)
          self.is_prepared = False
          config.log("Initializing spo training job...")
          self.type_str = "1vsAll"
  
          if self.__class__ == TrainingJob1vsAll:
              for f in Job.job_created_hooks:
                  f(self)
  
      def _prepare(self):
          """Construct dataloader"""
  
          if self.is_prepared:
              return
  
          self.num_examples = self.dataset.split(self.train_split).size(0)
          self.loader = torch.utils.data.DataLoader(
              range(self.num_examples),
              collate_fn=lambda batch: {
                  "triples": self.dataset.split(self.train_split)[batch, :].long()
              },
              shuffle=True,
              batch_size=self.batch_size,
              num_workers=self.config.get("train.num_workers"),
              pin_memory=self.config.get("train.pin_memory"),
          )
  
          self.is_prepared = True
  
      def _process_batch(self, batch_index, batch) -> TrainingJob._ProcessBatchResult:
          # prepare
          prepare_time = -time.time()
          triples = batch["triples"].to(self.device)
          batch_size = len(triples)
          prepare_time += time.time()
  
          # forward/backward pass (sp)
          forward_time = -time.time()
          scores_sp = self.model.score_sp(triples[:, 0], triples[:, 1])
          loss_value_sp = self.loss(scores_sp, triples[:, 2]) / batch_size
          loss_value = loss_value_sp.item()
          forward_time += time.time()
          backward_time = -time.time()
          loss_value_sp.backward()
          backward_time += time.time()
  
          # forward/backward pass (po)
          forward_time -= time.time()
          scores_po = self.model.score_po(triples[:, 1], triples[:, 2])
          loss_value_po = self.loss(scores_po, triples[:, 0]) / batch_size
          loss_value += loss_value_po.item()
          forward_time += time.time()
          backward_time -= time.time()
          loss_value_po.backward()
          backward_time += time.time()
  
          # all done
          return TrainingJob._ProcessBatchResult(
              loss_value, batch_size, prepare_time, forward_time, backward_time
          )
  
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
          )
  
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
          )
  
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
          )
  
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
          )
  
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
          )
  
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
          )
  
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
          )
  
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
          )
  
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.repeat_interleave(n_o, 0)
              o_embs = o_emb.repeat((n, 1))
              out = self.score_emb_spo(s_embs, p_embs, o_embs)
          elif combine == "_po":
              assert o_emb.size(0) == n
              n_s = s_emb.size(0)
              s_embs = s_emb.repeat((n, 1))
              p_embs = p_emb.repeat_interleave(n_s, 0)
              o_embs = o_emb.repeat_interleave(n_s, 0)
              out = self.score_emb_spo(s_embs, p_embs, o_embs)
          elif combine == "s_o":
              n = s_emb.size(0)
              assert o_emb.size(0) == n
              n_p = p_emb.size(0)
              s_embs = s_emb.repeat_interleave(n_p, 0)
              p_embs = p_emb.repeat((n, 1))
              o_embs = o_emb.repeat_interleave(n_p, 0)
              out = self.score_emb_spo(s_embs, p_embs, o_embs)
          else:
              raise ValueError('cannot handle combine="{}".format(combine)')
  
          return out.view(n, -1)
  
  
  class KgeEmbedder(KgeBase):
      r"""Base class for all embedders of a fixed number of objects.
  
      Objects can be entities, relations, mentions, and so on.
  
      """
  
      def __init__(self, config: Config, dataset: Dataset, configuration_key: str):
          super().__init__(config, dataset, configuration_key)
  
          #: location of the configuration options of this embedder
          self.embedder_type: str = self.get_option("type")
  
          # verify all custom options by trying to set them in a copy of this
          # configuration (quick and dirty, but works)
          try:
              custom_options = Config.flatten(config.get(self.configuration_key))
          except KeyError:
              # there are no custom options
              custom_options = {}
          if "type" in custom_options:
              del custom_options["type"]
          dummy_config = self.config.clone()
          for key, value in custom_options.items():
              try:
                  dummy_config.set(self.embedder_type + "." + key, value)
              except ValueError as ve:
                  raise ValueError(
                      "key {}.{} invalid or of incorrect type, message was {}".format(
                          self.configuration_key, key, ve
                      )
                  )
  
          self.dim: int = self.get_option("dim")
  
      @staticmethod
      def create(
          config: Config, dataset: Dataset, configuration_key: str, vocab_size: int
      ) -> "KgeEmbedder":
          """Factory method for embedder creation."""
  
          try:
              embedder_type = config.get_default(configuration_key + ".type")
              class_name = config.get(embedder_type + ".class_name")
              module = importlib.import_module("kge.model")
          except:
              raise Exception("Can't find {}.type in config".format(configuration_key))
  
          try:
              embedder = getattr(module, class_name)(
                  config, dataset, configuration_key, vocab_size
              )
              return embedder
          except ImportError:
              # perhaps TODO: try class with specified name -> extensibility
              raise ValueError(
                  "Can't find class {} in 'kge.model' for embedder {}".format(
                      class_name, embedder_type
                  )
              )
  
      def forward(self, indexes: Tensor) -> Tensor:
          return self.embed(indexes)
  
      def embed(self, indexes: Tensor) -> Tensor:
          """Computes the embedding."""
          raise NotImplementedError
  
      def embed_all(self) -> Tensor:
          """Returns all embeddings."""
          raise NotImplementedError
  
  
  class KgeModel(KgeBase):
      r"""Generic KGE model for KBs with a fixed set of entities and relations.
  
      This class uses :class:`KgeEmbedder` to associate each subject, relation, and object
      with an embedding, and a :class:`RelationalScorer` to score (subject, predicate,
      object) triples.
  
      """
  
      def __init__(
          self,
          config: Config,
          dataset: Dataset,
          scorer: Union[RelationalScorer, type],
          initialize_embedders=True,
          configuration_key=None,
      ):
          super().__init__(config, dataset, configuration_key)
  
          # TODO support different embedders for subjects and objects
  
          #: Embedder used for entities (both subject and objects)
          self._entity_embedder: KgeEmbedder
  
          #: Embedder used for relations
          self._relation_embedder: KgeEmbedder
  
          if initialize_embedders:
              self._entity_embedder = KgeEmbedder.create(
                  config,
                  dataset,
                  self.configuration_key + ".entity_embedder",
                  dataset.num_entities(),
              )
  
              #: Embedder used for relations
              num_relations = dataset.num_relations()
              self._relation_embedder = KgeEmbedder.create(
                  config,
                  dataset,
                  self.configuration_key + ".relation_embedder",
                  num_relations,
              )
  
          #: Scorer
          self._scorer: RelationalScorer
          if type(scorer) == type:
              # scorer is type of the scorer to use; call its constructor
              self._scorer = scorer(
                  config=config, dataset=dataset, configuration_key=self.configuration_key
              )
          else:
              self._scorer = scorer
  
      # overridden to also set self.model
      def _init_configuration(self, config: Config, configuration_key: Optional[str]):
          Configurable._init_configuration(self, config, configuration_key)
          if not hasattr(self, "model") or not self.model:
              if self.configuration_key:
                  self.model: str = config.get(self.configuration_key + ".type")
              else:
                  self.model: str = config.get("model")
                  self.configuration_key = self.model
  
      @staticmethod
      def create(
          config: Config, dataset: Dataset, configuration_key: str = None
      ) -> "KgeModel":
          """Factory method for model creation."""
  
          try:
              if configuration_key:
                  model_name = config.get(configuration_key + ".type")
              else:
                  model_name = config.get("model")
              class_name = config.get(model_name + ".class_name")
              module = importlib.import_module("kge.model")
          except:
              raise Exception("Can't find {}.type in config".format(configuration_key))
  
          try:
              model = getattr(module, class_name)(config, dataset, configuration_key)
              model.to(config.get("job.device"))
              return model
          except ImportError:
              # perhaps TODO: try class with specified name -> extensibility
              raise ValueError(
                  "Can't find class {} in 'kge.model' for model {}".format(
                      class_name, model_name
                  )
              )
  
      @staticmethod
      def create_default(
          model: Optional[str] = None,
          dataset: Optional[Union[Dataset,str]] = None,
          options: Dict[str, Any] = {},
          folder: Optional[str] = None,
      ) -> "KgeModel":
          """Utility method to create a model, including configuration and dataset.
  
          `model` is the name of the model (takes precedence over
          ``options["model"]``), `dataset` a dataset name or `Dataset` instance (takes
          precedence over ``options["dataset.name"]``), and options arbitrary other
          configuration options.
  
          If `folder` is ``None``, creates a temporary folder. Otherwise uses the
          specified folder.
  
          """
          # load default model config
          if model is None:
              model = options["model"]
          default_config_file = filename_in_module(kge.model, "{}.yaml".format(model))
          config = Config()
          config.load(default_config_file, create=True)
  
          # apply specified options
          config.set("model", model)
          if isinstance(dataset, Dataset):
              config.set("dataset.name", dataset.config.get("dataset.name"))
          elif isinstance(dataset, str):
              config.set("dataset.name", dataset)
          config.set_all(new_options=options)
  
          # create output folder
          if folder is None:
              config.folder = tempfile.mkdtemp(
                  "{}-{}-".format(config.get("dataset.name"), config.get("model"))
              )
          else:
              config.folder = folder
  
          # create dataset and model
          if not isinstance(dataset, Dataset):
              dataset = Dataset.load(config)
          model = KgeModel.create(config, dataset)
          return model
  
      @staticmethod
      def load_from_checkpoint(
          filename: str, dataset=None, use_tmp_log_folder=True, device="cpu"
      ) -> "KgeModel":
          """Loads a model from a checkpoint file of a training job.
  
          If dataset is specified, associates this dataset with the model. Otherwise uses
          the dataset used to train the model.
  
          If `use_tmp_log_folder` is set, the logs and traces are written to a temporary
          file. Otherwise, the files `kge.log` and `trace.yaml` will be created (or
          appended to) in the checkpoint's folder.
  
          """
  
          checkpoint = torch.load(filename, map_location=device)
  
          original_config = checkpoint["config"]
          config = Config()  # round trip to handle deprecated configs
          config.load_options(original_config.options)
          config.set("job.device", device)
          if use_tmp_log_folder:
              import tempfile
  
              config.log_folder = tempfile.mkdtemp(prefix="kge-")
          else:
              config.log_folder = os.path.dirname(filename)
              if not config.log_folder:
                  config.log_folder = "."
          if dataset is None:
              dataset = Dataset.load(config, preload_data=False)
          model = KgeModel.create(config, dataset)
          model.load(checkpoint["model"])
          model.eval()
          return model
  
      def prepare_job(self, job: "Job", **kwargs):
          super().prepare_job(job, **kwargs)
          self._entity_embedder.prepare_job(job, **kwargs)
          self._relation_embedder.prepare_job(job, **kwargs)
  
          def append_num_parameter(job, trace):
              trace["num_parameters"] = sum(map(lambda p: p.numel(), self.parameters()))
  
          job.post_epoch_trace_hooks.append(append_num_parameter)
  
      def penalty(self, **kwargs) -> List[Tensor]:
          if "batch" in kwargs and "triples" in kwargs["batch"]:
              kwargs["batch"]["triples"] = kwargs["batch"]["triples"].to(
                  self.config.get("job.device")
              )
          return (
              # Note: If the subject and object embedder are identical, embeddings may be
              # penalized twice. This is intended (and necessary, e.g., if the penalty is
              # weighted).
              super().penalty(**kwargs)
              + self.get_s_embedder().penalty(slot=S, **kwargs)
              + self.get_p_embedder().penalty(slot=P, **kwargs)
              + self.get_o_embedder().penalty(slot=O, **kwargs)
          )
  
      def get_s_embedder(self) -> KgeEmbedder:
          return self._entity_embedder
  
      def get_o_embedder(self) -> KgeEmbedder:
          return self._entity_embedder
  
      def get_p_embedder(self) -> KgeEmbedder:
          return self._relation_embedder
  
      def get_scorer(self) -> RelationalScorer:
          return self._scorer
  
      def score_spo(self, s: Tensor, p: Tensor, o: Tensor, direction=None) -> Tensor:
          r"""Compute scores for a set of triples.
  
          `s`, `p`, and `o` are vectors of common size :math:`n`, holding the indexes of
          the subjects, relations, and objects to score.
  
          `direction` may influence how scores are computed. For most models, this setting
          has no meaning. For reciprocal relations, direction must be either `"s"` or
          `"o"` (depending on what is predicted).
  
          Returns a vector of size :math:`n`, in which the :math:`i`-th entry holds the
          score of triple :math:`(s_i, p_i, o_i)`.
  
          """
          s = self.get_s_embedder().embed(s)
          p = self.get_p_embedder().embed(p)
          o = self.get_o_embedder().embed(o)
          return self._scorer.score_emb(s, p, o, combine="spo").view(-1)
  
      def score_sp(self, s: Tensor, p: Tensor, o: Tensor = None) -> Tensor:
          r"""Compute scores for triples formed from a set of sp-pairs and all (or a subset of the) objects.
  
          `s` and `p` are vectors of common size :math:`n`, holding the indexes of the
          subjects and relations to score.
  
          Returns an :math:`n\times E` tensor, where :math:`E` is the total number of
          known entities. The :math:`(i,j)`-entry holds the score for triple :math:`(s_i,
          p_i, j)`.
  
          If `o` is not None, it is a vector holding the indexes of the objects to score.
  
          """
          s = self.get_s_embedder().embed(s)
          p = self.get_p_embedder().embed(p)
          if o is None:
              o = self.get_o_embedder().embed_all()
          else:
              o = self.get_o_embedder().embed(o)
  
          return self._scorer.score_emb(s, p, o, combine="sp_")
  
      def score_po(self, p: Tensor, o: Tensor, s: Tensor = None) -> Tensor:
          r"""Compute scores for triples formed from a set of po-pairs and (or a subset of the) subjects.
  
          `p` and `o` are vectors of common size :math:`n`, holding the indexes of the
          relations and objects to score.
  
          Returns an :math:`n\times E` tensor, where :math:`E` is the total number of
          known entities. The :math:`(i,j)`-entry holds the score for triple :math:`(j,
          p_i, o_i)`.
  
          If `s` is not None, it is a vector holding the indexes of the objects to score.
  
          """
  
          if s is None:
              s = self.get_s_embedder().embed_all()
          else:
              s = self.get_s_embedder().embed(s)
          o = self.get_o_embedder().embed(o)
          p = self.get_p_embedder().embed(p)
  
          return self._scorer.score_emb(s, p, o, combine="_po")
  
      def score_so(self, s: Tensor, o: Tensor, p: Tensor = None) -> Tensor:
          r"""Compute scores for triples formed from a set of so-pairs and all (or a subset of the) relations.
  
          `s` and `o` are vectors of common size :math:`n`, holding the indexes of the
          subjects and objects to score.
  
          Returns an :math:`n\times R` tensor, where :math:`R` is the total number of
          known relations. The :math:`(i,j)`-entry holds the score for triple :math:`(s_i,
          j, o_i)`.
  
          If `p` is not None, it is a vector holding the indexes of the relations to score.
  
          """
          s = self.get_s_embedder().embed(s)
          o = self.get_o_embedder().embed(o)
          if p is None:
              p = self.get_p_embedder().embed_all()
          else:
              p = self.get_p_embedder().embed(p)
  
          return self._scorer.score_emb(s, p, o, combine="s_o")
  
      def score_sp_po(
          self, s: Tensor, p: Tensor, o: Tensor, entity_subset: Tensor = None
      ) -> Tensor:
          r"""Combine `score_sp` and `score_po`.
  
          `s`, `p` and `o` are vectors of common size :math:`n`, holding the indexes of
          the subjects, relations, and objects to score.
  
          Each sp-pair and each po-pair is scored against the entities in `entity_subset`
          (also holds indexes). If set to `entity_subset` is `None`, scores against all
          entities.
  
          The result is the horizontal concatenation of the outputs of
          :code:`score_sp(s,p,entity_subset)` and :code:`score_po(p,o,entity_subset)`.
          I.e., returns an :math:`n\times 2E` tensor, where :math:`E` is the size of
          `entity_subset`. For :math:`j<E`, the :math:`(i,j)`-entry holds the score for
          triple :math:`(s_i, p_i, e_j)`. For :math:`j\ge E`, the :math:`(i,j)`-entry
          holds the score for triple :math:`(e_{j-E}, p_i, o_i)`.
  
          """
  
          s = self.get_s_embedder().embed(s)
          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(all_entities, p, o, combine="_po")
          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(all_subjects, p, o, combine="_po")
          return torch.cat((sp_scores, po_scores), dim=1)
  
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 = p_emb.repeat_interleave(n_o, 0)
              o_embs = o_emb.repeat((n, 1))
              out = self.score_emb_spo(s_embs, p_embs, o_embs)
          elif combine == "*po":
              assert o_emb.size(0) == n
              n_s = s_emb.size(0)
              s_embs = s_emb.repeat((n, 1))
              p_embs = p_emb.repeat_interleave(n_s, 0)
              o_embs = o_emb.repeat_interleave(n_s, 0)
              out = self.score_emb_spo(s_embs, p_embs, o_embs)
          elif combine == "s*o":
              n = s_emb.size(0)
              assert o_emb.size(0) == n
              n_p = p_emb.size(0)
              s_embs = s_emb.repeat_interleave(n_p, 0)
              p_embs = p_emb.repeat((n, 1))
              o_embs = o_emb.repeat_interleave(n_p, 0)
              out = self.score_emb_spo(s_embs, p_embs, o_embs)
          else:
              raise ValueError('cannot handle combine="{}".format(combine)')
  
          return out.view(n, -1)
  
  
  class KgeEmbedder(KgeBase):
      r"""Base class for all embedders of a fixed number of objects.
  
      Objects can be entities, relations, mentions, and so on.
  
      """
  
      def __init__(self, config: Config, dataset: Dataset, configuration_key: str):
          super().__init__(config, dataset, configuration_key)
  
          #: location of the configuration options of this embedder
          self.embedder_type: str = self.get_option("type")
  
          # verify all custom options by trying to set them in a copy of this
          # configuration (quick and dirty, but works)
          try:
              custom_options = Config.flatten(config.get(self.configuration_key))
          except KeyError:
              # there are no custom options
              custom_options = {}
          if "type" in custom_options:
              del custom_options["type"]
          dummy_config = self.config.clone()
          for key, value in custom_options.items():
              try:
                  dummy_config.set(self.embedder_type + "." + key, value)
              except ValueError as ve:
                  raise ValueError(
                      "key {}.{} invalid or of incorrect type, message was {}".format(
                          self.configuration_key, key, ve
                      )
                  )
  
          self.dim: int = self.get_option("dim")
  
      @staticmethod
      def create(
          config: Config, dataset: Dataset, configuration_key: str, vocab_size: int
      ) -> "KgeEmbedder":
          """Factory method for embedder creation."""
  
          try:
              embedder_type = config.get_default(configuration_key + ".type")
              class_name = config.get(embedder_type + ".class_name")
              module = importlib.import_module("kge.model")
          except:
              raise Exception("Can't find {}.type in config".format(configuration_key))
  
          try:
              embedder = getattr(module, class_name)(
                  config, dataset, configuration_key, vocab_size
              )
              return embedder
          except ImportError:
              # perhaps TODO: try class with specified name -> extensibility
              raise ValueError(
                  "Can't find class {} in 'kge.model' for embedder {}".format(
                      class_name, embedder_type
                  )
              )
  
      def forward(self, indexes: Tensor) -> Tensor:
          return self.embed(indexes)
  
      def embed(self, indexes: Tensor) -> Tensor:
          """Computes the embedding."""
          raise NotImplementedError
  
      def embed_all(self) -> Tensor:
          """Returns all embeddings."""
          raise NotImplementedError
  
  
  class KgeModel(KgeBase):
      r"""Generic KGE model for KBs with a fixed set of entities and relations.
  
      This class uses :class:`KgeEmbedder` to associate each subject, relation, and object
      with an embedding, and a :class:`RelationalScorer` to score (subject, predicate,
      object) triples.
  
      """
  
      def __init__(
          self,
          config: Config,
          dataset: Dataset,
          scorer: Union[RelationalScorer, type],
          initialize_embedders=True,
          configuration_key=None,
      ):
          super().__init__(config, dataset, configuration_key)
  
          # TODO support different embedders for subjects and objects
  
          #: Embedder used for entities (both subject and objects)
          self._entity_embedder: KgeEmbedder
  
          #: Embedder used for relations
          self._relation_embedder: KgeEmbedder
  
          if initialize_embedders:
              self._entity_embedder = KgeEmbedder.create(
                  config,
                  dataset,
                  self.configuration_key + ".entity_embedder",
                  dataset.num_entities(),
              )
  
              #: Embedder used for relations
              num_relations = dataset.num_relations()
              self._relation_embedder = KgeEmbedder.create(
                  config,
                  dataset,
                  self.configuration_key + ".relation_embedder",
                  num_relations,
              )
  
          #: Scorer
          self._scorer: RelationalScorer
          if type(scorer) == type:
              # scorer is type of the scorer to use; call its constructor
              self._scorer = scorer(
                  config=config, dataset=dataset, configuration_key=self.configuration_key
              )
          else:
              self._scorer = scorer
  
      # overridden to also set self.model
      def _init_configuration(self, config: Config, configuration_key: Optional[str]):
          Configurable._init_configuration(self, config, configuration_key)
          if not hasattr(self, "model") or not self.model:
              if self.configuration_key:
                  self.model: str = config.get(self.configuration_key + ".type")
              else:
                  self.model: str = config.get("model")
                  self.configuration_key = self.model
  
      @staticmethod
      def create(
          config: Config, dataset: Dataset, configuration_key: str = None
      ) -> "KgeModel":
          """Factory method for model creation."""
  
          try:
              if configuration_key:
                  model_name = config.get(configuration_key + ".type")
              else:
                  model_name = config.get("model")
              class_name = config.get(model_name + ".class_name")
              module = importlib.import_module("kge.model")
          except:
              raise Exception("Can't find {}.type in config".format(configuration_key))
  
          try:
              model = getattr(module, class_name)(config, dataset, configuration_key)
              model.to(config.get("job.device"))
              return model
          except ImportError:
              # perhaps TODO: try class with specified name -> extensibility
              raise ValueError(
                  "Can't find class {} in 'kge.model' for model {}".format(
                      class_name, model_name
                  )
              )
  
      @staticmethod
      def create_default(
          model: Optional[str] = None,
          dataset: Optional[Union[Dataset,str]] = None,
          options: Dict[str, Any] = {},
          folder: Optional[str] = None,
      ) -> "KgeModel":
          """Utility method to create a model, including configuration and dataset.
  
          `model` is the name of the model (takes precedence over
          ``options["model"]``), `dataset` a dataset name or `Dataset` instance (takes
          precedence over ``options["dataset.name"]``), and options arbitrary other
          configuration options.
  
          If `folder` is ``None``, creates a temporary folder. Otherwise uses the
          specified folder.
  
          """
          # load default model config
          if model is None:
              model = options["model"]
          default_config_file = filename_in_module(kge.model, "{}.yaml".format(model))
          config = Config()
          config.load(default_config_file, create=True)
  
          # apply specified options
          config.set("model", model)
          if isinstance(dataset, Dataset):
              config.set("dataset.name", dataset.config.get("dataset.name"))
          elif isinstance(dataset, str):
              config.set("dataset.name", dataset)
          config.set_all(new_options=options)
  
          # create output folder
          if folder is None:
              config.folder = tempfile.mkdtemp(
                  "{}-{}-".format(config.get("dataset.name"), config.get("model"))
              )
          else:
              config.folder = folder
  
          # create dataset and model
          if not isinstance(dataset, Dataset):
              dataset = Dataset.load(config)
          model = KgeModel.create(config, dataset)
          return model
  
      @staticmethod
      def load_from_checkpoint(
          filename: str, dataset=None, use_tmp_log_folder=True, device="cpu"
      ) -> "KgeModel":
          """Loads a model from a checkpoint file of a training job.
  
          If dataset is specified, associates this dataset with the model. Otherwise uses
          the dataset used to train the model.
  
          If `use_tmp_log_folder` is set, the logs and traces are written to a temporary
          file. Otherwise, the files `kge.log` and `trace.yaml` will be created (or
          appended to) in the checkpoint's folder.
  
          """
  
          checkpoint = torch.load(filename, map_location=device)
  
          original_config = checkpoint["config"]
          config = Config()  # round trip to handle deprecated configs
          config.load_options(original_config.options)
          config.set("job.device", device)
          if use_tmp_log_folder:
              import tempfile
  
              config.log_folder = tempfile.mkdtemp(prefix="kge-")
          else:
              config.log_folder = os.path.dirname(filename)
              if not config.log_folder:
                  config.log_folder = "."
          if dataset is None:
              dataset = Dataset.load(config, preload_data=False)
          model = KgeModel.create(config, dataset)
          model.load(checkpoint["model"])
          model.eval()
          return model
  
      def prepare_job(self, job: "Job", **kwargs):
          super().prepare_job(job, **kwargs)
          self._entity_embedder.prepare_job(job, **kwargs)
          self._relation_embedder.prepare_job(job, **kwargs)
  
          def append_num_parameter(job, trace):
              trace["num_parameters"] = sum(map(lambda p: p.numel(), self.parameters()))
  
          job.post_epoch_trace_hooks.append(append_num_parameter)
  
      def penalty(self, **kwargs) -> List[Tensor]:
          if "batch" in kwargs and "triples" in kwargs["batch"]:
              kwargs["batch"]["triples"] = kwargs["batch"]["triples"].to(
                  self.config.get("job.device")
              )
          return (
              # Note: If the subject and object embedder are identical, embeddings may be
              # penalized twice. This is intended (and necessary, e.g., if the penalty is
              # weighted).
              super().penalty(**kwargs)
              + self.get_s_embedder().penalty(slot=S, **kwargs)
              + self.get_p_embedder().penalty(slot=P, **kwargs)
              + self.get_o_embedder().penalty(slot=O, **kwargs)
          )
  
      def get_s_embedder(self) -> KgeEmbedder:
          return self._entity_embedder
  
      def get_o_embedder(self) -> KgeEmbedder:
          return self._entity_embedder
  
      def get_p_embedder(self) -> KgeEmbedder:
          return self._relation_embedder
  
      def get_scorer(self) -> RelationalScorer:
          return self._scorer
  
      def score_spo(self, s: Tensor, p: Tensor, o: Tensor, direction=None) -> Tensor:
          r"""Compute scores for a set of triples.
  
          `s`, `p`, and `o` are vectors of common size :math:`n`, holding the indexes of
          the subjects, relations, and objects to score.
  
          `direction` may influence how scores are computed. For most models, this setting
          has no meaning. For reciprocal relations, direction must be either `"s"` or
          `"o"` (depending on what is predicted).
  
          Returns a vector of size :math:`n`, in which the :math:`i`-th entry holds the
          score of triple :math:`(s_i, p_i, o_i)`.
  
          """
          s = self.get_s_embedder().embed(s)
          p = self.get_p_embedder().embed(p)
          o = self.get_o_embedder().embed(o)
          return self._scorer.score_emb(s, p, o, combine="spo").view(-1)
  
      def score_sp(self, s: Tensor, p: Tensor, o: Tensor = None) -> Tensor:
          r"""Compute scores for triples formed from a set of sp-pairs and all (or a subset of the) objects.
  
          `s` and `p` are vectors of common size :math:`n`, holding the indexes of the
          subjects and relations to score.
  
          Returns an :math:`n\times E` tensor, where :math:`E` is the total number of
          known entities. The :math:`(i,j)`-entry holds the score for triple :math:`(s_i,
          p_i, j)`.
  
          If `o` is not None, it is a vector holding the indexes of the objects to score.
  
          """
          s = self.get_s_embedder().embed(s)
          p = self.get_p_embedder().embed(p)
          if o is None:
              o = self.get_o_embedder().embed_all()
          else:
              o = self.get_o_embedder().embed(o)
  
          return self._scorer.score_emb(s, p, o, combine="sp*")
  
      def score_po(self, p: Tensor, o: Tensor, s: Tensor = None) -> Tensor:
          r"""Compute scores for triples formed from a set of po-pairs and (or a subset of the) subjects.
  
          `p` and `o` are vectors of common size :math:`n`, holding the indexes of the
          relations and objects to score.
  
          Returns an :math:`n\times E` tensor, where :math:`E` is the total number of
          known entities. The :math:`(i,j)`-entry holds the score for triple :math:`(j,
          p_i, o_i)`.
  
          If `s` is not None, it is a vector holding the indexes of the objects to score.
  
          """
  
          if s is None:
              s = self.get_s_embedder().embed_all()
          else:
              s = self.get_s_embedder().embed(s)
          o = self.get_o_embedder().embed(o)
          p = self.get_p_embedder().embed(p)
  
          return self._scorer.score_emb(s, p, o, combine="*po")
  
      def score_so(self, s: Tensor, o: Tensor, p: Tensor = None) -> Tensor:
          r"""Compute scores for triples formed from a set of so-pairs and all (or a subset of the) relations.
  
          `s` and `o` are vectors of common size :math:`n`, holding the indexes of the
          subjects and objects to score.
  
          Returns an :math:`n\times R` tensor, where :math:`R` is the total number of
          known relations. The :math:`(i,j)`-entry holds the score for triple :math:`(s_i,
          j, o_i)`.
  
          If `p` is not None, it is a vector holding the indexes of the relations to score.
  
          """
          s = self.get_s_embedder().embed(s)
          o = self.get_o_embedder().embed(o)
          if p is None:
              p = self.get_p_embedder().embed_all()
          else:
              p = self.get_p_embedder().embed(p)
  
          return self._scorer.score_emb(s, p, o, combine="s*o")
  
      def score_sp_po(
          self, s: Tensor, p: Tensor, o: Tensor, entity_subset: Tensor = None
      ) -> Tensor:
          r"""Combine `score_sp` and `score_po`.
  
          `s`, `p` and `o` are vectors of common size :math:`n`, holding the indexes of
          the subjects, relations, and objects to score.
  
          Each sp-pair and each po-pair is scored against the entities in `entity_subset`
          (also holds indexes). If set to `entity_subset` is `None`, scores against all
          entities.
  
          The result is the horizontal concatenation of the outputs of
          :code:`score_sp(s,p,entity_subset)` and :code:`score_po(p,o,entity_subset)`.
          I.e., returns an :math:`n\times 2E` tensor, where :math:`E` is the size of
          `entity_subset`. For :math:`j<E`, the :math:`(i,j)`-entry holds the score for
          triple :math:`(s_i, p_i, e_j)`. For :math:`j\ge E`, the :math:`(i,j)`-entry
          holds the score for triple :math:`(e_{j-E}, p_i, o_i)`.
  
          """
  
          s = self.get_s_embedder().embed(s)
          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(all_entities, p, o, combine="*po")
          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(all_subjects, p, o, combine="*po")
          return torch.cat((sp_scores, po_scores), dim=1)
  
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)
  
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)
  
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)
  
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)
  
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_im = torch.stack((x_re, x_im), dim=0)  # dim0: real, imaginary
      return torch.norm(x_re_im, dim=0)  # sqrt(real^2+imaginary^2)
  
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_im = torch.stack((x_re, x_im), dim=0)  # dim0: real, imaginary
      return torch.norm(x_re_im, dim=0)  # sqrt(real^2+imaginary^2)
  
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
          )
  
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
          )
  
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
          )
  
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
          )
  
files: kge/model/transe.py