files.new_contents: 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
     
  
  [... the rest is omitted ...]
files.new_path: kge/job/train.py
files.old_contents: 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
     
  
  [... the rest is omitted ...]
files.old_path: kge/job/train.py
files.new_contents: 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(
  
  [... the rest is omitted ...]
files.new_path: kge/model/complex.py
files.old_contents: 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(
  
  [... the rest is omitted ...]
files.old_path: kge/model/complex.py
files.new_contents: 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
  
  [... the rest is omitted ...]
files.new_path: kge/model/conve.py
files.old_contents: 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
  
  [... the rest is omitted ...]
files.old_path: kge/model/conve.py
files.new_contents: 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
  
  [... the rest is omitted ...]
files.new_path: kge/model/cp.py
files.old_contents: 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
  
  [... the rest is omitted ...]
files.old_path: kge/model/cp.py
files.new_contents: 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))
  
  
  [... the rest is omitted ...]
files.new_path: kge/model/distmult.py
files.old_contents: 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))
  
  
  [... the rest is omitted ...]
files.old_path: kge/model/distmult.py
files.new_contents: 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
  
  [... the rest is omitted ...]
files.new_path: kge/model/kge_model.py
files.old_contents: 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
  
  [... the rest is omitted ...]
files.old_path: kge/model/kge_model.py
files.new_contents: 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 = K
  
  [... the rest is omitted ...]
files.new_path: kge/model/reciprocal_relations_model.py
files.old_contents: 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 = K
  
  [... the rest is omitted ...]
files.old_path: kge/model/reciprocal_relations_model.py
files.new_contents: 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(-
  
  [... the rest is omitted ...]
files.new_path: kge/model/rescal.py
files.old_contents: 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(-
  
  [... the rest is omitted ...]
files.old_path: kge/model/rescal.py
files.new_contents: 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")
  
     
  
  [... the rest is omitted ...]
files.new_path: kge/model/rotate.py
files.old_contents: 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")
  
     
  
  [... the rest is omitted ...]
files.old_path: kge/model/rotate.py
files.new_contents: 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
  
  [... the rest is omitted ...]
files.new_path: kge/model/simple.py
files.old_contents: 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
  
  [... the rest is omitted ...]
files.old_path: kge/model/simple.py
files.new_contents: 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
  
  [... the rest is omitted ...]
files.new_path: kge/model/transe.py
files.old_contents: 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
  
  [... the rest is omitted ...]
files.old_path: kge/model/transe.py