files: #!/bin/sh
  
  BASEDIR=`pwd`
  
  if [ ! -d "$BASEDIR" ]; then
      mkdir "$BASEDIR"
  fi
  
  
  # toy (testing dataset)
  if [ ! -d "$BASEDIR/toy" ]; then
      echo Downloading toy
      cd $BASEDIR
      curl -O https://www.uni-mannheim.de/media/Einrichtungen/dws/pi1/kge_datasets/toy.tar.gz
      tar xvf toy.tar.gz
      python preprocess.py toy
  else
      echo toy already present
  fi
  
  
  # fb15k
  if [ ! -d "$BASEDIR/fb15k" ]; then
    echo Downloading fb15k
    cd $BASEDIR
    curl -O https://www.uni-mannheim.de/media/Einrichtungen/dws/pi1/kge_datasets/fb15k.tar.gz
    tar xvf fb15k.tar.gz
    cd fb15k
    case "$(uname -s)" in
        CYGWIN*|MINGW32*|MSYS*)
            cmd.exe /c mklink train.txt freebase_mtr100_mte100-train.txt
            cmd.exe /c mklink valid.txt freebase_mtr100_mte100-valid.txt
            cmd.exe /c mklink test.txt freebase_mtr100_mte100-test.txt
            ;;
        *)
            ln -s freebase_mtr100_mte100-train.txt train.txt
            ln -s freebase_mtr100_mte100-valid.txt valid.txt
            ln -s freebase_mtr100_mte100-test.txt test.txt
            ;;
    esac
    cd ..
    python preprocess.py fb15k
  else
      echo fb15k already present
  fi
  
  # fb15k-237
  if [ ! -d "$BASEDIR/fb15k-237" ]; then
    echo Downloading fb15k-237
    cd $BASEDIR
    curl -O https://www.uni-mannheim.de/media/Einrichtungen/dws/pi1/kge_datasets/fb15k-237.tar.gz
    tar xvf fb15k-237.tar.gz
    python preprocess.py fb15k-237
  else
      echo fb15k-237 already present
  fi
  
  # wn18
  if [ ! -d "$BASEDIR/wn18" ]; then
    echo Downloading wn18
    cd $BASEDIR
    curl -O https://www.uni-mannheim.de/media/Einrichtungen/dws/pi1/kge_datasets/wn18.tar.gz
    tar xvf wn18.tar.gz
    cd wn18
    case "$(uname -s)" in
        CYGWIN*|MINGW32*|MSYS*)
            cmd.exe /c mklink train.txt wordnet-mlj12-train.txt
            cmd.exe /c mklink valid.txt wordnet-mlj12-valid.txt
            cmd.exe /c mklink test.txt wordnet-mlj12-test.txt
            ;;
        *)
            ln -s wordnet-mlj12-train.txt train.txt
            ln -s wordnet-mlj12-valid.txt valid.txt
            ln -s wordnet-mlj12-test.txt test.txt
            ;;
    esac
    cd ..
    python preprocess.py wn18
  else
      echo wn18 already present
  fi
  
  # wnrr
  if [ ! -d "$BASEDIR/wnrr" ]; then
    echo Downloading wnrr
    cd $BASEDIR
    curl -O https://www.uni-mannheim.de/media/Einrichtungen/dws/pi1/kge_datasets/wnrr.tar.gz
    tar xvf wnrr.tar.gz
    python preprocess.py wnrr
  else
      echo wnrr already present
  fi
  
  
  # dbpedia50
  if [ ! -d "$BASEDIR/dbpedia50" ]; then
    echo Downloading dbpedia50
    cd $BASEDIR
    curl -O https://www.uni-mannheim.de/media/Einrichtungen/dws/pi1/kge_datasets/dbpedia50.tar.gz
    tar xvf dbpedia50.tar.gz
    python preprocess.py dbpedia50
  else
      echo dbpedia50 already present
  fi
  
  # dbpedia500
  if [ ! -d "$BASEDIR/dbpedia500" ]; then
    echo Downloading dbpedia500
    cd $BASEDIR
    curl -O https://www.uni-mann
  
  [... the rest is omitted ...]
files: data/download_all.sh
files: #!/bin/sh
  
  BASEDIR=`pwd`
  
  if [ ! -d "$BASEDIR" ]; then
      mkdir "$BASEDIR"
  fi
  
  
  # toy (testing dataset)
  if [ ! -d "$BASEDIR/toy" ]; then
      echo Downloading toy
      cd $BASEDIR
      curl -O https://www.uni-mannheim.de/media/Einrichtungen/dws/pi1/kge_datasets/toy.tar.gz
      tar xvf toy.tar.gz
      python preprocess.py --folder toy
  else
      echo toy already present
  fi
  
  
  # fb15k
  if [ ! -d "$BASEDIR/fb15k" ]; then
    echo Downloading fb15k
    cd $BASEDIR
    curl -O https://www.uni-mannheim.de/media/Einrichtungen/dws/pi1/kge_datasets/fb15k.tar.gz
    tar xvf fb15k.tar.gz
    cd fb15k
    case "$(uname -s)" in
        CYGWIN*|MINGW32*|MSYS*)
            cmd.exe /c mklink train.txt freebase_mtr100_mte100-train.txt
            cmd.exe /c mklink valid.txt freebase_mtr100_mte100-valid.txt
            cmd.exe /c mklink test.txt freebase_mtr100_mte100-test.txt
            ;;
        *)
            ln -s freebase_mtr100_mte100-train.txt train.txt
            ln -s freebase_mtr100_mte100-valid.txt valid.txt
            ln -s freebase_mtr100_mte100-test.txt test.txt
            ;;
    esac
    cd ..
    python preprocess.py --folder fb15k
  else
      echo fb15k already present
  fi
  
  # fb15k-237
  if [ ! -d "$BASEDIR/fb15k-237" ]; then
    echo Downloading fb15k-237
    cd $BASEDIR
    curl -O https://www.uni-mannheim.de/media/Einrichtungen/dws/pi1/kge_datasets/fb15k-237.tar.gz
    tar xvf fb15k-237.tar.gz
    python preprocess.py --folder fb15k-237
  else
      echo fb15k-237 already present
  fi
  
  # wn18
  if [ ! -d "$BASEDIR/wn18" ]; then
    echo Downloading wn18
    cd $BASEDIR
    curl -O https://www.uni-mannheim.de/media/Einrichtungen/dws/pi1/kge_datasets/wn18.tar.gz
    tar xvf wn18.tar.gz
    cd wn18
    case "$(uname -s)" in
        CYGWIN*|MINGW32*|MSYS*)
            cmd.exe /c mklink train.txt wordnet-mlj12-train.txt
            cmd.exe /c mklink valid.txt wordnet-mlj12-valid.txt
            cmd.exe /c mklink test.txt wordnet-mlj12-test.txt
            ;;
        *)
            ln -s wordnet-mlj12-train.txt train.txt
            ln -s wordnet-mlj12-valid.txt valid.txt
            ln -s wordnet-mlj12-test.txt test.txt
            ;;
    esac
    cd ..
    python preprocess.py --folder wn18
  else
      echo wn18 already present
  fi
  
  # wnrr
  if [ ! -d "$BASEDIR/wnrr" ]; then
    echo Downloading wnrr
    cd $BASEDIR
    curl -O https://www.uni-mannheim.de/media/Einrichtungen/dws/pi1/kge_datasets/wnrr.tar.gz
    tar xvf wnrr.tar.gz
    python preprocess.py --folder wnrr
  else
      echo wnrr already present
  fi
  
  
  # dbpedia50
  if [ ! -d "$BASEDIR/dbpedia50" ]; then
    echo Downloading dbpedia50
    cd $BASEDIR
    curl -O https://www.uni-mannheim.de/media/Einrichtungen/dws/pi1/kge_datasets/dbpedia50.tar.gz
    tar xvf dbpedia50.tar.gz
    python preprocess.py --folder dbpedia50
  else
      echo dbpedia50 already present
  fi
  
  # dbpedia500
  if [ ! -d "$BASEDIR/dbpedia500" ]; then
    echo Downloading dbpedia500
    cd $BASEDIR
  
  [... the rest is omitted ...]
files: data/download_all.sh
files: #!/usr/bin/env python
  """Preprocess a KGE dataset into a the format expected by libkge.
  
  Call as `preprocess.py --folder <name>`. The original dataset should be stored in
  subfolder `name` and have files "train.txt", "valid.txt", and "test.txt". Each file
  contains one SPO triple per line, separated by tabs.
  
  During preprocessing, each distinct entity name and each distinct distinct relation name
  is assigned an index (dense). The index-to-object mapping is stored in files
  "entity_map.del" and "relation_map.del", resp. The triples (as indexes) are stored in
  files "train.del", "valid.del", and "test.del". Metadata information is stored in a file
  "dataset.yaml".
  
  """
  
  import argparse
  import yaml
  import os.path
  from collections import OrderedDict
  
  
  def store_map(symbol_map, filename):
      with open(filename, "w") as f:
          for symbol, index in symbol_map.items():
              f.write(f"{index}\t{symbol}\n")
  
  
  if __name__ == "__main__":
      parser = argparse.ArgumentParser()
      parser.add_argument("folder", type=str)
      parser.add_argument("--order_sop", action="store_true")
      args = parser.parse_args()
  
      print(f"Preprocessing {args.folder}...")
      raw_split_files = {"train": "train.txt", "valid": "valid.txt", "test": "test.txt"}
      split_files = {"train": "train.del", "valid": "valid.del", "test": "test.del"}
      split_sizes = {}
  
      if args.order_sop:
          S, P, O = 0, 2, 1
      else:
          S, P, O = 0, 1, 2
  
      # read data and collect entities and relations
      raw = {}
      entities = {}
      relations = {}
      ent_id = 0
      rel_id = 0
      for split, filename in raw_split_files.items():
          with open(args.folder + "/" + filename, "r") as f:
              raw[split] = list(map(lambda s: s.strip().split("\t"), f.readlines()))
              for t in raw[split]:
                  if t[S] not in entities:
                      entities[t[S]] = ent_id
                      ent_id += 1
                  if t[P] not in relations:
                      relations[t[P]] = rel_id
                      rel_id += 1
                  if t[O] not in entities:
                      entities[t[O]] = ent_id
                      ent_id += 1
              print(
                  f"Found {len(raw[split])} triples in {split} split "
                  f"(file: {filename})."
              )
              split_sizes[split + "_size"] = len(raw[split])
  
      print(f"{len(relations)} distinct relations")
      print(f"{len(entities)} distinct entities")
      print("Writing relation and entity map...")
      store_map(relations, os.path.join(args.folder, "relation_map.del"))
      store_map(entities, os.path.join(args.folder, "entity_map.del"))
      print("Done.")
  
      # write config
      print("Writing dataset.yaml...")
      dataset_config = dict(
          name=args.folder,
          entity_map="entity_map.del",
          relation_map="relation_map.del",
          num_entities=len(entities),
          num_relations=len(relations),
          **split_files,
          **split_sizes,
      )
      print(yaml.dump(dict(dataset=dataset_config)))
      with open(os.path.join(args.folder, "dataset.yaml"), "w+") as filename:
          filename.write(yaml.dump(dict(dataset=dataset_config)))
  
      # write out triples using indexes
      print("Writing triples...")
      for split, filename in split_files.items():
          with open(os.path.join(args.folder, filename), "w") as f:
              for t in raw[split]:
                  f.write(
                      str(entities[t[S]])
                      + "\t"
                      + str(relations[t[P]])
                      + "\t"
                      + str(entities[t[O]])
                      + "\n"
                  )
  
  
  [... the rest is omitted ...]
files: data/preprocess.py
files: import argparse
  
  
  def index(symbols, file):
      with open(file, "w") as f:
          for i, k in symbols.items():
              f.write(str(k) + "\t" + str(i) + "\n")
  
  
  if __name__ == "__main__":
      parser = argparse.ArgumentParser()
      parser.add_argument("--folder", type=str)
      parser.add_argument("--order_sop", action="store_true")
      args = parser.parse_args()
  
      print("Preprocessing " + args.folder)
      raw_split_files = {"train": "train.txt", "valid": "valid.txt", "test": "test.txt"}
      split_files = {"train": "train.del", "valid": "valid.del", "test": "test.del"}
  
      sub = 0
      rel = 1
      obj = 2
      if args.order_sop:
          obj = 1
          rel = 2
  
      # read data and collect entity and relation names
      raw = {}
      entities = {}
      relations = {}
      ent_id = 0
      rel_id = 0
      for k, file in raw_split_files.items():
          with open(args.folder + "/" + file, "r") as f:
              raw[k] = list(map(lambda s: s.strip().split("\t"), f.readlines()))
              for t in raw[k]:
                  if t[sub] not in entities:
                      entities[t[sub]] = ent_id
                      ent_id += 1
                  if t[rel] not in relations:
                      relations[t[rel]] = rel_id
                      rel_id += 1
                  if t[obj] not in entities:
                      entities[t[obj]] = ent_id
                      ent_id += 1
              print(str(len(raw[k])) + " triples in " + file)
  
      print(str(len(relations)) + " distinct relations")
      print(str(len(entities)) + " distinct entities")
      print("Writing indexes...")
      index(relations, args.folder + "/relation_map.del")
      index(entities, args.folder + "/entity_map.del")
  
      # write out
      print("Writing triples...")
      for k, file in split_files.items():
          with open(args.folder + "/" + file, "w") as f:
              for t in raw[k]:
                  f.write(str(entities[t[sub]]))
                  f.write("\t")
                  f.write(str(relations[t[rel]]))
                  f.write("\t")
                  f.write(str(entities[t[obj]]))
                  f.write("\n")
      print("Done")
  
  
  [... the rest is omitted ...]
files: data/preprocess.py
files: # Job-level configuration options.
  job:
    # Type of job to run. Possible values are "train", "eval", and "search". See
    # the corresponding configuration keys for mode information.
    type: train
  
    # Main device to use for this job (e.g., 'cpu', 'cuda', 'cuda:0')
    device: 'cuda'  
  
  # The seeds of the PRNGs can be set manually for (increased) reproducability.
  # Use -1 to use default seed.
  random_seed:
    python: -1
    torch: -1
    numpy: -1
  
  
  ## DATASET #####################################################################
  
  dataset:
    # Specify a dataset here. There must be a folder of that name under "data/".
    # If this folder contains a dataset.yaml file, it overrides the defaults
    # specified below.
    name: 'toy'
  
    # Names of the training, validation, and test data files. Each file has the
    # following fields, all tab-separated:
    # - 0: subject index
    # - 1: relation index
    # - 2: object index
    # - 3-...: arbitrary metadata fields
    # Indexes are assumed to be dense throughout.
    train: train.del
    valid: valid.del
    test: test.del
  
    # Names of entity and relation map files, which store auxiliary information
    # about each entity/predicate. Fields are tab-separated:
    # - 0: entity/predicate index (as in train/valid/test)
    # - 1...: arbitrary metadata fields
    entity_map: entity_map.del
    relation_map: relation_map.del
  
    # other dataset specific keys
    +++: +++
  
  ## MODEL #######################################################################
  
  # Which KGE model to use. Examples include "rescal", "complex", "distmult",
  # "transe", or "conve". For available models, see the project homepage and/or
  # configuration files under kge/models and kge/models/experimental.
  model: ''
  
  
  ## TRAINING ####################################################################
  
  # Options of training jobs (job.type=="train")
  train:
    # Type of training job.
    # - KvsAll: scores each unique sp/po pair along with all possible objectes/subjects.
    # - negative_sampling: scores each unique spo triple along with sampled corrupted
    #   triples
    # - 1vsAll: scores each spo triples against the complete set of s/p-corrputed triples
    #   (all treated negative)
    type: KvsAll
  
    # Loss used for training. Not all combination of losses and job types are supported.
    # - bce: binary cross entropy after taking logits (all training types)
    # - kl: KL divergence between model and data distribution (all training types)
    # - margin_ranking (negative_sampling only)
    loss: bce
  
    ## Argument of loss function (if any)
    ## - margin_ranking: the margin to use
    loss_arg: 1.0
  
    # Maximum number of epochs used for training
    max_epochs: 20
  
    # Batch size used for training.
    batch_size: 100
  
    # Number of workers used to construct batches. Leave at 0 for default.
    num_workers: 0
  
    # Optimizer used for training.
    optimizer: Adagrad           # sgd, adagrad, adam
  
    # Additional arguments for the optimizer. Arbitrary key-value pairs can be
    # added here and will be passed along to the optimizer. E.g., use entry lr:0.1
    # to set the learning rate to 0.1.
    optimizer_args:
      +++: +++
  
    # Learning rate scheduler to use. Any scheduler from torch.optim.lr_scheduler
    # can be used (e.g., ConstantLRScheduler or ReduceLROnPlateau).
    lr_scheduler: ConstantLRScheduler #
  
    # Additional arguments for the scheduler.
    lr_scheduler_args:
      +++: +++
  
    # When to write entries to the trace file.
    trace_level: epoch           # batch, epoch
  
    # When to create checkpoints
    checkpoint:
      # In addition the the checkpoint of the last epoch (which is transient),
      # create an additional checkpoint every this many epochs. Disable additional
      # checkpoints with 0.
      every: 5
  
      # Keep this many most recent additional checkpoints.
      keep: 3
  
    # Other options
    pin_memory: False
    auto_correct: False         #
  
  [... the rest is omitted ...]
files: kge/config-default.yaml
files: # Job-level configuration options.
  job:
    # Type of job to run. Possible values are "train", "eval", and "search". See
    # the corresponding configuration keys for mode information.
    type: train
  
    # Main device to use for this job (e.g., 'cpu', 'cuda', 'cuda:0')
    device: 'cuda'  
  
  # The seeds of the PRNGs can be set manually for (increased) reproducability.
  # Use -1 to use default seed.
  random_seed:
    python: -1
    torch: -1
    numpy: -1
  
  
  ## DATASET #####################################################################
  
  dataset:
    # Specify a dataset here. There must be a folder of that name under "data/".
    name: 'toy'
  
    # Names of the training, validation, and test data files. Each file has the
    # following fields, all tab-separated:
    # - 0: subject index
    # - 1: relation index
    # - 2: object index
    # - 3-...: arbitrary metadata fields
    # Indexes are assumed to be dense throughout.
    train: train.del
    valid: valid.del
    test: test.del
  
    # Names of entity and relation map files, which store auxiliary information
    # about each entity/predicate. Fields are tab-separated:
    # - 0: entity/predicate index (as in train/valid/test)
    # - 1...: arbitrary metadata fields
    entity_map: entity_map.del
    relation_map: relation_map.del
  
  
  ## MODEL #######################################################################
  
  # Which KGE model to use. Examples include "rescal", "complex", "distmult",
  # "transe", or "conve". For available models, see the project homepage and/or
  # configuration files under kge/models and kge/models/experimental.
  model: ''
  
  
  ## TRAINING ####################################################################
  
  # Options of training jobs (job.type=="train")
  train:
    # Type of training job.
    # - KvsAll: scores each unique sp/po pair along with all possible objectes/subjects.
    # - negative_sampling: scores each unique spo triple along with sampled corrupted
    #   triples
    # - 1vsAll: scores each spo triples against the complete set of s/p-corrputed triples
    #   (all treated negative)
    type: KvsAll
  
    # Loss used for training. Not all combination of losses and job types are supported.
    # - bce: binary cross entropy after taking logits (all training types)
    # - kl: KL divergence between model and data distribution (all training types)
    # - margin_ranking (negative_sampling only)
    loss: bce
  
    ## Argument of loss function (if any)
    ## - margin_ranking: the margin to use
    loss_arg: 1.0
  
    # Maximum number of epochs used for training
    max_epochs: 20
  
    # Batch size used for training.
    batch_size: 100
  
    # Number of workers used to construct batches. Leave at 0 for default.
    num_workers: 0
  
    # Optimizer used for training.
    optimizer: Adagrad           # sgd, adagrad, adam
  
    # Additional arguments for the optimizer. Arbitrary key-value pairs can be
    # added here and will be passed along to the optimizer. E.g., use entry lr:0.1
    # to set the learning rate to 0.1.
    optimizer_args:
      +++: +++
  
    # Learning rate scheduler to use. Any scheduler from torch.optim.lr_scheduler
    # can be used (e.g., ConstantLRScheduler or ReduceLROnPlateau).
    lr_scheduler: ConstantLRScheduler #
  
    # Additional arguments for the scheduler.
    lr_scheduler_args:
      +++: +++
  
    # When to write entries to the trace file.
    trace_level: epoch           # batch, epoch
  
    # When to create checkpoints
    checkpoint:
      # In addition the the checkpoint of the last epoch (which is transient),
      # create an additional checkpoint every this many epochs. Disable additional
      # checkpoints with 0.
      every: 5
  
      # Keep this many most recent additional checkpoints.
      keep: 3
  
    # Other options
    pin_memory: False
    auto_correct: False         # True: allows for trainer to adjust certain settings, False: throw exception
    visualize_graph: False      # create PDF of compute graph (of first batch of first epoch
  
  [... the rest is omitted ...]
files: kge/config-default.yaml
files: import csv
  import os
  from collections import defaultdict, OrderedDict
  
  import torch
  import numpy as np
  
  from kge import Config, Configurable
  from kge.util.misc import kge_base_dir
  
  
  # TODO add support to pickle dataset (and indexes) and reload from there
  class Dataset(Configurable):
      def __init__(
          self,
          config,
          num_entities,
          entities,
          num_relations,
          relations,
          train,
          train_meta,
          valid,
          valid_meta,
          test,
          test_meta,
      ):
          super().__init__(config, "dataset")
          self.num_entities = num_entities
          self.entities = entities  # array: entity index -> metadata array of strings
          self.num_relations = num_relations
          self.relations = relations  # array: relation index -> metadata array of strings
          self.train = train  # (n,3) int32 tensor
          self.train_meta = (
              train_meta
          )  # array: triple row number -> metadata array of strings
          self.valid = valid  # (n,3) int32 tensor
          self.valid_meta = (
              valid_meta
          )  # array: triple row number -> metadata array of strings
          self.test = test  # (n,3) int32 tensor
          self.test_meta = (
              test_meta
          )  # array: triple row number -> metadata array of strings
          self.indexes = {}  # map: name of index -> index (used mainly by training jobs)
  
      @staticmethod
      def load(config: Config):
          name = config.get("dataset.name")
          base_dir = os.path.join(kge_base_dir(), "data", name)
          if os.path.isfile(os.path.join(base_dir, "dataset.yaml")):
              config.log("Loading configuration of dataset " + name + "...")
              config.load(os.path.join(base_dir, "dataset.yaml"))
  
          config.log("Loading dataset " + name + "...")
          num_entities, entities = Dataset._load_map(
              os.path.join(base_dir, config.get("dataset.entity_map"))
          )
          num_relations, relations = Dataset._load_map(
              os.path.join(base_dir, config.get("dataset.relation_map"))
          )
  
          train, train_meta = Dataset._load_triples(
              os.path.join(base_dir, config.get("dataset.train"))
          )
  
          valid, valid_meta = Dataset._load_triples(
              os.path.join(base_dir, config.get("dataset.valid"))
          )
  
          test, test_meta = Dataset._load_triples(
              os.path.join(base_dir, config.get("dataset.test"))
          )
  
          result = Dataset(
              config,
              num_entities,
              entities,
              num_relations,
              relations,
              train,
              train_meta,
              valid,
              valid_meta,
              test,
              test_meta,
          )
  
          config.log(str(num_entities) + " entities", prefix="  ")
          config.log(str(num_relations) + " relations", prefix="  ")
          config.log(str(len(train)) + " training triples", prefix="  ")
          config.log(str(len(valid)) + " validation triples", prefix="  ")
          config.log(str(len(test)) + " test triples", prefix="  ")
  
          return result
  
      @staticmethod
      def _load_map(filename):
          n = 0
          dictionary = {}
          with open(filename, "r") as file:
              reader = csv.reader(file, delimiter="\t")
              for row in reader:
                  index = int(row[0])
                  meta = row[1:]
                  dictionary[index] = meta
                  n = max(n, index + 1)
          array = [[]] * n
          for index, meta in dictionary.items():
              array[index] = meta
          return n, array
  
      @staticmethod
      def _load_triples(filename):
          n = 0
          triples = np.loadtxt(filename, delimiter="\t", usecols=range(0, 3), dtype=int)
          triples = torch.from_numpy(triples)
          num_lines = triples.shape[0]
          meta = [[]] * num_lines
          with open(filename, "r") as file:
              reader = csv.reader(file, delimiter="\t")
              for row in reader:
                  meta[n] = row[3:]
                  n += 1
  
          return triples, meta
  
      def index_KvsAll(self, split: str, sp_po: str):
          """Return an index for the triples in split (''train'', ''valid'', ''test'')
  
  
  [... the rest is omitted ...]
files: kge/dataset.py
files: import csv
  import os
  from collections import defaultdict, OrderedDict
  
  import torch
  import numpy as np
  
  from kge.util.misc import kge_base_dir
  
  
  # TODO add support to pickle dataset (and indexes) and reload from there
  class Dataset:
      def __init__(
          self,
          config,
          num_entities,
          entities,
          num_relations,
          relations,
          train,
          train_meta,
          valid,
          valid_meta,
          test,
          test_meta,
      ):
          self.config = config
          self.num_entities = num_entities
          self.entities = entities  # array: entity index -> metadata array of strings
          self.num_relations = num_relations
          self.relations = relations  # array: relation index -> metadata array of strings
          self.train = train  # (n,3) int32 tensor
          self.train_meta = (
              train_meta
          )  # array: triple row number -> metadata array of strings
          self.valid = valid  # (n,3) int32 tensor
          self.valid_meta = (
              valid_meta
          )  # array: triple row number -> metadata array of strings
          self.test = test  # (n,3) int32 tensor
          self.test_meta = (
              test_meta
          )  # array: triple row number -> metadata array of strings
          self.indexes = {}  # map: name of index -> index (used mainly by training jobs)
  
      @staticmethod
      def load(config):
          name = config.get("dataset.name")
          config.log("Loading dataset " + name + "...")
          base_dir = os.path.join(kge_base_dir(), "data/" + name)
  
          num_entities, entities = Dataset._load_map(
              os.path.join(base_dir, config.get("dataset.entity_map"))
          )
          num_relations, relations = Dataset._load_map(
              os.path.join(base_dir, config.get("dataset.relation_map"))
          )
  
          train, train_meta = Dataset._load_triples(
              os.path.join(base_dir, config.get("dataset.train"))
          )
  
          valid, valid_meta = Dataset._load_triples(
              os.path.join(base_dir, config.get("dataset.valid"))
          )
  
          test, test_meta = Dataset._load_triples(
              os.path.join(base_dir, config.get("dataset.test"))
          )
  
          result = Dataset(
              config,
              num_entities,
              entities,
              num_relations,
              relations,
              train,
              train_meta,
              valid,
              valid_meta,
              test,
              test_meta,
          )
  
          config.log(str(num_entities) + " entities", prefix="  ")
          config.log(str(num_relations) + " relations", prefix="  ")
          config.log(str(len(train)) + " training triples", prefix="  ")
          config.log(str(len(valid)) + " validation triples", prefix="  ")
          config.log(str(len(test)) + " test triples", prefix="  ")
  
          return result
  
      @staticmethod
      def _load_map(filename):
          n = 0
          dictionary = {}
          with open(filename, "r") as file:
              reader = csv.reader(file, delimiter="\t")
              for row in reader:
                  index = int(row[0])
                  meta = row[1:]
                  dictionary[index] = meta
                  n = max(n, index + 1)
          array = [[]] * n
          for index, meta in dictionary.items():
              array[index] = meta
          return n, array
  
      @staticmethod
      def _load_triples(filename):
          n = 0
          triples = np.loadtxt(filename, delimiter="\t", usecols=range(0, 3), dtype=int)
          triples = torch.from_numpy(triples)
          num_lines = triples.shape[0]
          meta = [[]] * num_lines
          with open(filename, "r") as file:
              reader = csv.reader(file, delimiter="\t")
              for row in reader:
                  meta[n] = row[3:]
                  n += 1
  
          return triples, meta
  
      def index_KvsAll(self, split: str, sp_po: str):
          """Return an index for the triples in split (''train'', ''valid'', ''test'')
          from the specified constituents (''sp'' or ''po'') to the indexes of the
          remaining constituent (''o'' or ''s'', respectively.)
  
          The index maps from `tuple' to `torch.LongTensor`.
  
          The index is cached in the provided dataset under name ''split
  
  [... the rest is omitted ...]
files: kge/dataset.py