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-mannheim.de/media/Einrichtungen/dws/pi1/kge_datasets/dbpedia500.tar.gz
    tar xvf dbpedia500.tar.gz
    python preprocess.py dbpedia500 --order_sop
  else
      echo dbpedia500 already present
  fi
  
  # db100k
  if [ ! -d "$BASEDIR/db100k" ]; then
    echo Downloading db100k
    cd $BASEDIR
    curl -O https://www.uni-mannheim.de/media/Einrichtungen/dws/pi1/kge_datasets/db100k.tar.gz
    tar xvf db100k.tar.gz
    cd db100k
    case "$(uname -s)" in
        CYGWIN*|MINGW32*|MSYS*)
            cmd.exe /c mklink train.txt _train.txt
            cmd.exe /c mklink valid.txt _valid.txt
            cmd.exe /c mklink test.txt _test.txt
            ;;
        *)
            ln -s _train.txt train.txt
            ln -s _valid.txt valid.txt
            ln -s _test.txt test.txt
            ;;
    esac
    cd ..
    python preprocess.py db100k
  else
      echo db100k already present
  fi
  
  # yago3-10
  if [ ! -d "$BASEDIR/yago3-10" ]; then
    echo Downloading yago3-10
    cd $BASEDIR
    curl -O https://www.uni-mannheim.de/media/Einrichtungen/dws/pi1/kge_datasets/yago3-10.tar.gz
    tar xvf yago3-10.tar.gz
    python preprocess.py yago3-10
  else
      echo yago3-10 already present
  fi
  
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
    curl -O https://www.uni-mannheim.de/media/Einrichtungen/dws/pi1/kge_datasets/dbpedia500.tar.gz
    tar xvf dbpedia500.tar.gz
    python preprocess.py --folder dbpedia500 --order_sop
  else
      echo dbpedia500 already present
  fi
  
  # db100k
  if [ ! -d "$BASEDIR/db100k" ]; then
    echo Downloading db100k
    cd $BASEDIR
    curl -O https://www.uni-mannheim.de/media/Einrichtungen/dws/pi1/kge_datasets/db100k.tar.gz
    tar xvf db100k.tar.gz
    cd db100k
    case "$(uname -s)" in
        CYGWIN*|MINGW32*|MSYS*)
            cmd.exe /c mklink train.txt _train.txt
            cmd.exe /c mklink valid.txt _valid.txt
            cmd.exe /c mklink test.txt _test.txt
            ;;
        *)
            ln -s _train.txt train.txt
            ln -s _valid.txt valid.txt
            ln -s _test.txt test.txt
            ;;
    esac
    cd ..
    python preprocess.py --folder db100k
  else
      echo db100k already present
  fi
  
  # yago3-10
  if [ ! -d "$BASEDIR/yago3-10" ]; then
    echo Downloading yago3-10
    cd $BASEDIR
    curl -O https://www.uni-mannheim.de/media/Einrichtungen/dws/pi1/kge_datasets/yago3-10.tar.gz
    tar xvf yago3-10.tar.gz
    python preprocess.py --folder yago3-10
  else
      echo yago3-10 already present
  fi
  
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"
                  )
  
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")
  
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         # 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)
  
  
  # Options for KvsAll training (train.type=="KvsAll)
  KvsAll:
    # Amount of label smoothing (disabled when 0). Disencourages models to
    # perform extreme predictions (0 or 1).
    #
    # Technically, reduces all labels by fraction given by this value and
    # subsequently increases them by 1.0/num_entities. For example, 0s become
    # 1.0/num_entities and 1s become (1.0-label_smoothing)+1.0/num_entities.
    #
    # This form of label smoothing was used by ConvE
    # ( https://github.com/TimDettmers/ConvE/blob/853d7010a4b569c3d24c9e6eeaf9dc3535d78832/main.py#L156) with a default value of 0.1.
    label_smoothing: 0.0
  
  # Options for negative sampling training (train.type=="negative_sampling")
  negative_sampling:
    sampling_type: uniform
    num_negatives_s: 3          # this is default value if others are -1, must not be -1
    num_negatives_p: -1
    num_negatives_o: -1
    filter_true_s: False        # default is not filtering
    filter_true_p: False
    filter_true_o: False
  
    # set to spo (for small negative samples), sp_po (for large negative samples),
    # or set to auto which sets spo or sp_po automatically if the
    # max(num_negatives_s, num_negatives_p, num_negatives_o) > 30. Recommended for
    # models which compute sp_po with an inner product, but not for all other
    # models (e.g., not for TransE in the current implementation).
    implementation: spo
  
  
  ## VALIDATION ##################################################################
  
  # Configuration options for model validation/selection during training.
  valid:
    # Validation is run every this many epochs during training (disable validation
    # with 0).
    every: 5
  
    # Name of the trace entry that holds the validation metric (higher value is
    # better)
    metric: mean_reciprocal_rank_filtered
  
    # If the above metric is not present in trace (e.g., because a custom metric
    # should be used), a Python expression to compute the metric. Can refer to
    # trace entries directly and to configuration options via config.
    # Example: 'math.sqrt(mean_reciprocal_rank) + config.get("user.par")'
    metric_expr: 'float("nan")'
  
    early_stopping:
      # Grace period of validation runs before a training run is stopped early
      # (disable early stopping with 0). If the value is set to n, then training is
      # stopped when there has been no improvement (compared to the best overall
      # result so far) in the validation metric during the last n validation runs.
      patience: 5
  
      # A minimum validation metric value threshold that should be reached after
      # n epochs, set to 0 epoch to turn off. Should be set very very conservatively
      # and the main purpose is for pruning completely useless hyperparameter
      # settings during hyper-parameter optimization.
      min_threshold:
        epochs: 0
        metric_value: 0.0
  
    # When to write a trace entry. Possible values: example, batch, epoch
    trace_level: epoch
  
    # Whether test data should used during validation. Filtered
    # metrics by default only use train and validation data. When this is set to
    # True, additionally produces "filtered_with_test" validation metrics (such as
    # MRR or HITS@k). Apparently, many existing models have been trained with this
    # set to True during model selection and using a metric such as
    # mean_reciprocal_rank_filtered_with_test.
    filter_with_test: False
  
  
  ## EVALUATION ##################################################################
  
  # Options of evaluation jobs (job.type=="eval"). Also used during validation
  # when training. Right now, evaluation jobs compute a fixed set of metrics: the
  # MRR and the specified HITS@k's.
  eval:
    data: valid
    type: entity_ranking
    hits_at_k_s: [1, 3, 10, 50, 100, 200, 300, 400, 500, 1000]     # ks for HITS@k
    batch_size: 100
    num_workers: 0
    pin_memory: False
    trace_level: example            # example or batch or epoch
    chunk_size: -1                  # default no chunking
  
    # Metrics are always computed over the entire evaluation data. Optionally,
    # certain more specific metric can be computed in addition.
    metrics_per:
      relation_type: False          # 1-to-1, 1-to-N, N-to-1, N-to-N
      head_and_tail: False          # head, tail
      argument_frequency: False     # 25%, 50%, 75%, top quantiles per argument
  
  
  ## HYPERPARAMETER SEARCH #######################################################
  
  # Options of hyperparameter search jobs (job.type=="search").
  search:
    # The type of search to run (see descriptions below). Possible values: manual,
    # grid, ax
    type: ax
  
    # Maximum number of parallel training jobs to run during a search.
    num_workers: 1
  
    # Device pool to use for training jobs. If this list is empty, `job.device` is
    # used for all parallel searches. Otherwise, the first `search.num_workers`
    # devices from this list are used. If the number of devices specified here is
    # less than `search.num_workers`, the list wraps around so that devices are
    # used by multiple jobs in parallel.
    device_pool: [ ]
  
    # What to do when an error occurs during a training job. Possible values:
    # continue, abort
    on_error: abort
  
  # Manually specify all configurations to try
  manual_search:
    # If false, only creates training job folders but does not run the jobs.
    run: True
  
    # List of configurations to search. Each entry is a record with a field
    # 'folder' (where the training job is stored) and an arbitrary number of other
    # fields that define the search configuration (e.g.
    # 'train.optimizer_args.lr').
    configurations: []
  
  
  # Metajob for a grid search. Creates a manual search job with all points on the
  # grid.
  grid_search:
    # If false, only creates manual search job configuration file but does not run
    # it. This may be useful to edit the configuration (e.g., change folder names)
    # or to copy the created configurations to an existing manual search job.
    run: True
  
    # Define the grid. This is a dict where the key is a (flattened or nested)
    # configuration option (e.g., "train.optimizer_args.lr") and the value is an
    # array of grid-search values (e.g., [ 0.1, 0.01 ]). No default values
    # specified.
    parameters:
      +++: +++
  
  # Dynamic search job that picks configurations using ax
  ax_search:
    # Total number of trials to run. Can be increased when a search job is
    # resumed.
    num_trials: 10
  
    # Number of sobol trials to run (-1: automatic); remaining trials are GP+EI.
    # If equal or larger than num_trials, only Sobal trials will be run.
    num_sobol_trials: -1
  
    # Search space definition passed to ax. See create_experiment in
    # https://ax.dev/api/service.html#module-ax.service.ax_client
    parameters: []
    parameter_constraints: []
  
    # 'fixed_parameters' is a dictionary of parameters and the values these
    # parameters should be fixed to for *generating* trials during BO. The purpose
    # of 'fixed_parameters' is to initialize a search (f.ex. with Sobol) and then
    # fix a range or choice parameter when generating new trials with BO to a
    # certain value. The names and values should correspond with the definitions
    # in 'ax_search.parameters'.
    # Attention: there is a name clash with regards to the meaning of 'fixed' in Ax
    # and Botorch! The difference to 'ax_search.parameters#type = fixed' is, that
    # those are filtered out for fitting a BO model, while the parameters for
    # 'fixed_parameters' are not. This setting is used in /botorch/gen.py
    #
    # Example:
    #
    #  parameters:
    #    - name: rescal.entity_embedder.dim
    #      type: choice
    #      values: [100, 200]
    #
    #  fixed_parameters:
    #    - name: rescal.entity_embedder.dim
    #      value: 100
  
    fixed_parameters: []
  
  
  ## USER PARAMETERS #####################################################################
  
  # These parameters are not used by the kge framework itself. It can be used to
  # add additional configuration options or information to this config.
  user:
    +++: +++
  
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)
  
  
  # Options for KvsAll training (train.type=="KvsAll)
  KvsAll:
    # Amount of label smoothing (disabled when 0). Disencourages models to
    # perform extreme predictions (0 or 1).
    #
    # Technically, reduces all labels by fraction given by this value and
    # subsequently increases them by 1.0/num_entities. For example, 0s become
    # 1.0/num_entities and 1s become (1.0-label_smoothing)+1.0/num_entities.
    #
    # This form of label smoothing was used by ConvE
    # ( https://github.com/TimDettmers/ConvE/blob/853d7010a4b569c3d24c9e6eeaf9dc3535d78832/main.py#L156) with a default value of 0.1.
    label_smoothing: 0.0
  
  # Options for negative sampling training (train.type=="negative_sampling")
  negative_sampling:
    sampling_type: uniform
    num_negatives_s: 3          # this is default value if others are -1, must not be -1
    num_negatives_p: -1
    num_negatives_o: -1
    filter_true_s: False        # default is not filtering
    filter_true_p: False
    filter_true_o: False
  
    # set to spo (for small negative samples), sp_po (for large negative samples),
    # or set to auto which sets spo or sp_po automatically if the
    # max(num_negatives_s, num_negatives_p, num_negatives_o) > 30. Recommended for
    # models which compute sp_po with an inner product, but not for all other
    # models (e.g., not for TransE in the current implementation).
    implementation: spo
  
  
  ## VALIDATION ##################################################################
  
  # Configuration options for model validation/selection during training.
  valid:
    # Validation is run every this many epochs during training (disable validation
    # with 0).
    every: 5
  
    # Name of the trace entry that holds the validation metric (higher value is
    # better)
    metric: mean_reciprocal_rank_filtered
  
    # If the above metric is not present in trace (e.g., because a custom metric
    # should be used), a Python expression to compute the metric. Can refer to
    # trace entries directly and to configuration options via config.
    # Example: 'math.sqrt(mean_reciprocal_rank) + config.get("user.par")'
    metric_expr: 'float("nan")'
  
    early_stopping:
      # Grace period of validation runs before a training run is stopped early
      # (disable early stopping with 0). If the value is set to n, then training is
      # stopped when there has been no improvement (compared to the best overall
      # result so far) in the validation metric during the last n validation runs.
      patience: 5
  
      # A minimum validation metric value threshold that should be reached after
      # n epochs, set to 0 epoch to turn off. Should be set very very conservatively
      # and the main purpose is for pruning completely useless hyperparameter
      # settings during hyper-parameter optimization.
      min_threshold:
        epochs: 0
        metric_value: 0.0
  
    # When to write a trace entry. Possible values: example, batch, epoch
    trace_level: epoch
  
    # Whether test data should used during validation. Filtered
    # metrics by default only use train and validation data. When this is set to
    # True, additionally produces "filtered_with_test" validation metrics (such as
    # MRR or HITS@k). Apparently, many existing models have been trained with this
    # set to True during model selection and using a metric such as
    # mean_reciprocal_rank_filtered_with_test.
    filter_with_test: False
  
  
  ## EVALUATION ##################################################################
  
  # Options of evaluation jobs (job.type=="eval"). Also used during validation
  # when training. Right now, evaluation jobs compute a fixed set of metrics: the
  # MRR and the specified HITS@k's.
  eval:
    data: valid
    type: entity_ranking
    hits_at_k_s: [1, 3, 10, 50, 100, 200, 300, 400, 500, 1000]     # ks for HITS@k
    batch_size: 100
    num_workers: 0
    pin_memory: False
    trace_level: example            # example or batch or epoch
    chunk_size: -1                  # default no chunking
  
    # Metrics are always computed over the entire evaluation data. Optionally,
    # certain more specific metric can be computed in addition.
    metrics_per:
      relation_type: False          # 1-to-1, 1-to-N, N-to-1, N-to-N
      head_and_tail: False          # head, tail
      argument_frequency: False     # 25%, 50%, 75%, top quantiles per argument
  
  
  ## HYPERPARAMETER SEARCH #######################################################
  
  # Options of hyperparameter search jobs (job.type=="search").
  search:
    # The type of search to run (see descriptions below). Possible values: manual,
    # grid, ax
    type: ax
  
    # Maximum number of parallel training jobs to run during a search.
    num_workers: 1
  
    # Device pool to use for training jobs. If this list is empty, `job.device` is
    # used for all parallel searches. Otherwise, the first `search.num_workers`
    # devices from this list are used. If the number of devices specified here is
    # less than `search.num_workers`, the list wraps around so that devices are
    # used by multiple jobs in parallel.
    device_pool: [ ]
  
    # What to do when an error occurs during a training job. Possible values:
    # continue, abort
    on_error: abort
  
  # Manually specify all configurations to try
  manual_search:
    # If false, only creates training job folders but does not run the jobs.
    run: True
  
    # List of configurations to search. Each entry is a record with a field
    # 'folder' (where the training job is stored) and an arbitrary number of other
    # fields that define the search configuration (e.g.
    # 'train.optimizer_args.lr').
    configurations: []
  
  
  # Metajob for a grid search. Creates a manual search job with all points on the
  # grid.
  grid_search:
    # If false, only creates manual search job configuration file but does not run
    # it. This may be useful to edit the configuration (e.g., change folder names)
    # or to copy the created configurations to an existing manual search job.
    run: True
  
    # Define the grid. This is a dict where the key is a (flattened or nested)
    # configuration option (e.g., "train.optimizer_args.lr") and the value is an
    # array of grid-search values (e.g., [ 0.1, 0.01 ]). No default values
    # specified.
    parameters:
      +++: +++
  
  # Dynamic search job that picks configurations using ax
  ax_search:
    # Total number of trials to run. Can be increased when a search job is
    # resumed.
    num_trials: 10
  
    # Number of sobol trials to run (-1: automatic); remaining trials are GP+EI.
    # If equal or larger than num_trials, only Sobal trials will be run.
    num_sobol_trials: -1
  
    # Search space definition passed to ax. See create_experiment in
    # https://ax.dev/api/service.html#module-ax.service.ax_client
    parameters: []
    parameter_constraints: []
  
    # 'fixed_parameters' is a dictionary of parameters and the values these
    # parameters should be fixed to for *generating* trials during BO. The purpose
    # of 'fixed_parameters' is to initialize a search (f.ex. with Sobol) and then
    # fix a range or choice parameter when generating new trials with BO to a
    # certain value. The names and values should correspond with the definitions
    # in 'ax_search.parameters'.
    # Attention: there is a name clash with regards to the meaning of 'fixed' in Ax
    # and Botorch! The difference to 'ax_search.parameters#type = fixed' is, that
    # those are filtered out for fitting a BO model, while the parameters for
    # 'fixed_parameters' are not. This setting is used in /botorch/gen.py
    #
    # Example:
    #
    #  parameters:
    #    - name: rescal.entity_embedder.dim
    #      type: choice
    #      values: [100, 200]
    #
    #  fixed_parameters:
    #    - name: rescal.entity_embedder.dim
    #      value: 100
  
    fixed_parameters: []
  
  
  ## USER PARAMETERS #####################################################################
  
  # These parameters are not used by the kge framework itself. It can be used to
  # add additional configuration options or information to this config.
  user:
    +++: +++
  
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'')
          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_sp_po''. If
          this index is already present, does not recompute it.
  
          """
          if split == "train":
              triples = self.train
          elif split == "valid":
              triples = self.valid
          elif split == "test":
              triples = self.test
          else:
              raise ValueError()
  
          if sp_po == "sp":
              sp_po_cols = [0, 1]
              value_column = 2
          elif sp_po == "po":
              sp_po_cols = [1, 2]
              value_column = 0
          else:
              raise ValueError()
  
          name = split + "_" + sp_po
          if not self.indexes.get(name):
              self.indexes[name] = Dataset.group_by_sp_po(
                  triples[:, sp_po_cols], triples[:, value_column]
              )
              self.config.log(
                  "{} distinct {} pairs in {}".format(
                      len(self.indexes[name]), sp_po, split
                  ),
                  prefix="  ",
              )
  
          return self.indexes.get(name)
  
      @staticmethod
      def group_by_sp_po(sp_po_list, o_s_list) -> dict:
          result = defaultdict(list)
          for sp_po, o_s in zip(sp_po_list.tolist(), o_s_list.tolist()):
              result[tuple(sp_po)].append(o_s)
          for sp_po, o_s in result.items():
              result[sp_po] = torch.IntTensor(sorted(o_s))
          return OrderedDict(result)
  
      @staticmethod
      def prepare_index(index):
          sp_po = torch.tensor(list(index.keys()), dtype=torch.int)
          o_s = torch.cat(list(index.values()))
          offsets = torch.cumsum(
              torch.tensor([0] + list(map(len, index.values())), dtype=torch.int), 0
          )
          return sp_po, o_s, offsets
  
      def index_relation_types(self):
          """
          create dictionary mapping from {1-N, M-1, 1-1, M-N} -> set of relations
          """
          if "relation_types" in self.indexes:
              return
          relation_types = self._get_relation_types()
          relations_per_type = {}
          for k, v in relation_types.items():
              relations_per_type.setdefault(v, set()).add(k)
          for k,v in relations_per_type.items():
              self.config.log("{} relations of type {}".format(len(v), k), prefix="  ")
          self.indexes["relation_types"] = relation_types
          self.indexes["relations_per_type"] = relations_per_type
  
      def _get_relation_types(self,):
          """
          Classify relations into 1-N, M-1, 1-1, M-N
  
          Bordes, Antoine, et al.
          "Translating embeddings for modeling multi-relational data."
          Advances in neural information processing systems. 2013.
  
          :return: dictionary mapping from int -> {1-N, M-1, 1-1, M-N}
          """
          relation_stats = torch.zeros((self.num_relations, 6))
          for index, p in [
              (self.index_KvsAll('train', 'sp'), 1),
              (self.index_KvsAll('train', 'po'), 0),
          ]:
              for prefix, labels in index.items():
                  relation_stats[prefix[p], 0+p*2] = labels.float().sum()
                  relation_stats[prefix[p], 1+p*2] = relation_stats[prefix[p], 1+p*2] + 1.
          relation_stats[:,4] = (relation_stats[:,0]/relation_stats[:,1]) > 1.5
          relation_stats[:,5] = (relation_stats[:,2]/relation_stats[:,3]) > 1.5
          result = dict()
          for i, relation in enumerate(self.relations):
              result[i] = '{}-{}'.format(
                  '1' if relation_stats[i,4].item() == 0 else 'M',
                  '1' if relation_stats[i,5].item() == 0 else 'N', )
          return result
  
      # TODO this is metadata; refine API
      def index_frequency_percentiles(self, recompute=False):
          """
          :return: dictionary mapping from
          {
           'subject':
              {25%, 50%, 75%, top} -> set of entities
           'relations':
              {25%, 50%, 75%, top} -> set of relations
           'object':
              {25%, 50%, 75%, top} -> set of entities
          }
          """
          if "frequency_percentiles" in self.indexes:
              return
          subject_stats = torch.zeros((self.num_entities, 1))
          relation_stats = torch.zeros((self.num_relations, 1))
          object_stats = torch.zeros((self.num_entities, 1))
          for (s,p,o) in self.train:
              subject_stats[s] += 1
              relation_stats[p] += 1
              object_stats[o] += 1
          result = dict()
          for arg, stats, num in [
              ('subject', [i for i,j in list(sorted(enumerate(subject_stats.tolist()), key=lambda x:x[1]))], self.num_entities),
              ('relation', [i for i,j in list(sorted(enumerate(relation_stats.tolist()), key=lambda x:x[1]))], self.num_relations),
              ('object', [i for i,j in list(sorted(enumerate(object_stats.tolist()), key=lambda x:x[1]))], self.num_entities),
          ]:
              for percentile, (begin, end) in [('25%', (0., 0.25)), ('50%', (0.25, 0.5)), ('75%', (0.5, 0.75)), ('top', (0.75, 1.))]:
                  if arg not in result:
                      result[arg] = dict()
                  result[arg][percentile] = set(stats[int(begin*num):int(end*num)])
          self.indexes["frequency_percentiles"] = result
  
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_sp_po''. If
          this index is already present, does not recompute it.
  
          """
          if split == "train":
              triples = self.train
          elif split == "valid":
              triples = self.valid
          elif split == "test":
              triples = self.test
          else:
              raise ValueError()
  
          if sp_po == "sp":
              sp_po_cols = [0, 1]
              value_column = 2
          elif sp_po == "po":
              sp_po_cols = [1, 2]
              value_column = 0
          else:
              raise ValueError()
  
          name = split + "_" + sp_po
          if not self.indexes.get(name):
              self.indexes[name] = Dataset.group_by_sp_po(
                  triples[:, sp_po_cols], triples[:, value_column]
              )
              self.config.log(
                  "{} distinct {} pairs in {}".format(
                      len(self.indexes[name]), sp_po, split
                  ),
                  prefix="  ",
              )
  
          return self.indexes.get(name)
  
      @staticmethod
      def group_by_sp_po(sp_po_list, o_s_list) -> dict:
          result = defaultdict(list)
          for sp_po, o_s in zip(sp_po_list.tolist(), o_s_list.tolist()):
              result[tuple(sp_po)].append(o_s)
          for sp_po, o_s in result.items():
              result[sp_po] = torch.IntTensor(sorted(o_s))
          return OrderedDict(result)
  
      @staticmethod
      def prepare_index(index):
          sp_po = torch.tensor(list(index.keys()), dtype=torch.int)
          o_s = torch.cat(list(index.values()))
          offsets = torch.cumsum(
              torch.tensor([0] + list(map(len, index.values())), dtype=torch.int), 0
          )
          return sp_po, o_s, offsets
  
      def index_relation_types(self):
          """
          create dictionary mapping from {1-N, M-1, 1-1, M-N} -> set of relations
          """
          if "relation_types" in self.indexes:
              return
          relation_types = self._get_relation_types()
          relations_per_type = {}
          for k, v in relation_types.items():
              relations_per_type.setdefault(v, set()).add(k)
          for k,v in relations_per_type.items():
              self.config.log("{} relations of type {}".format(len(v), k), prefix="  ")
          self.indexes["relation_types"] = relation_types
          self.indexes["relations_per_type"] = relations_per_type
  
      def _get_relation_types(self,):
          """
          Classify relations into 1-N, M-1, 1-1, M-N
  
          Bordes, Antoine, et al.
          "Translating embeddings for modeling multi-relational data."
          Advances in neural information processing systems. 2013.
  
          :return: dictionary mapping from int -> {1-N, M-1, 1-1, M-N}
          """
          relation_stats = torch.zeros((self.num_relations, 6))
          for index, p in [
              (self.index_KvsAll('train', 'sp'), 1),
              (self.index_KvsAll('train', 'po'), 0),
          ]:
              for prefix, labels in index.items():
                  relation_stats[prefix[p], 0+p*2] = labels.float().sum()
                  relation_stats[prefix[p], 1+p*2] = relation_stats[prefix[p], 1+p*2] + 1.
          relation_stats[:,4] = (relation_stats[:,0]/relation_stats[:,1]) > 1.5
          relation_stats[:,5] = (relation_stats[:,2]/relation_stats[:,3]) > 1.5
          result = dict()
          for i, relation in enumerate(self.relations):
              result[i] = '{}-{}'.format(
                  '1' if relation_stats[i,4].item() == 0 else 'M',
                  '1' if relation_stats[i,5].item() == 0 else 'N', )
          return result
  
      # TODO this is metadata; refine API
      def index_frequency_percentiles(self, recompute=False):
          """
          :return: dictionary mapping from
          {
           'subject':
              {25%, 50%, 75%, top} -> set of entities
           'relations':
              {25%, 50%, 75%, top} -> set of relations
           'object':
              {25%, 50%, 75%, top} -> set of entities
          }
          """
          if "frequency_percentiles" in self.indexes:
              return
          subject_stats = torch.zeros((self.num_entities, 1))
          relation_stats = torch.zeros((self.num_relations, 1))
          object_stats = torch.zeros((self.num_entities, 1))
          for (s,p,o) in self.train:
              subject_stats[s] += 1
              relation_stats[p] += 1
              object_stats[o] += 1
          result = dict()
          for arg, stats, num in [
              ('subject', [i for i,j in list(sorted(enumerate(subject_stats.tolist()), key=lambda x:x[1]))], self.num_entities),
              ('relation', [i for i,j in list(sorted(enumerate(relation_stats.tolist()), key=lambda x:x[1]))], self.num_relations),
              ('object', [i for i,j in list(sorted(enumerate(object_stats.tolist()), key=lambda x:x[1]))], self.num_entities),
          ]:
              for percentile, (begin, end) in [('25%', (0., 0.25)), ('50%', (0.25, 0.5)), ('75%', (0.5, 0.75)), ('top', (0.75, 1.))]:
                  if arg not in result:
                      result[arg] = dict()
                  result[arg][percentile] = set(stats[int(begin*num):int(end*num)])
          self.indexes["frequency_percentiles"] = result
  
files: kge/dataset.py