text: import time
  import torch
  import torch.utils.data
  from typing import List
  
  from kge.job import Job
  import kge.job.util
  from kge.job.train import TrainingJob, _generate_worker_init_fn
  
  SLOTS = [0, 1, 2]
  S, P, O = SLOTS
  SLOT_STR = ["s", "p", "o"]
  
  
  class TrainingJobKvsAll(TrainingJob):
      """Train with examples consisting of a query and its answers.
  
      Terminology:
      - Query type: which queries to ask (sp_, s_o, and/or _po), can be configured via
        configuration key `KvsAll.query_type` (which see)
      - Query: a particular query, e.g., (John,marriedTo) of type sp_
      - Labels: list of true answers of a query (e.g., [Jane])
      - Example: a query + its labels, e.g., (John,marriedTo), [Jane]
      """
  
      from kge.indexing import KvsAllIndex
  
      def __init__(
          self, config, dataset, parent_job=None, model=None, forward_only=False
      ):
          super().__init__(
              config, dataset, parent_job, model=model, forward_only=forward_only
          )
          self.label_smoothing = config.check_range(
              "KvsAll.label_smoothing", float("-inf"), 1.0, max_inclusive=False
          )
          if self.label_smoothing < 0:
              if config.get("train.auto_correct"):
                  config.log(
                      "Setting label_smoothing to 0, "
                      "was set to {}.".format(self.label_smoothing)
                  )
                  self.label_smoothing = 0
              else:
                  raise Exception(
                      "Label_smoothing was set to {}, "
                      "should be at least 0.".format(self.label_smoothing)
                  )
          elif self.label_smoothing > 0 and self.label_smoothing <= (
              1.0 / dataset.num_entities()
          ):
              if config.get("train.auto_correct"):
                  # just to be sure it's used correctly
                  config.log(
                      "Setting label_smoothing to 1/num_entities = {}, "
                      "was set to {}.".format(
                          1.0 / dataset.num_entities(), self.label_smoothing
                      )
                  )
                  self.label_smoothing = 1.0 / dataset.num_entities()
              else:
                  raise Exception(
                      "Label_smoothing was set to {}, "
                      "should be at least {}.".format(
                          self.label_smoothing, 1.0 / dataset.num_entities()
                      )
                  )
  
          config.log("Initializing 1-to-N training job...")
          self.type_str = "KvsAll"
  
          if self.__class__ == TrainingJobKvsAll:
              for f in Job.job_created_hooks:
                  f(self)
  
      def _prepare(self):
          super()._prepare()
          # determine enabled query types
          self.query_types = [
              key
              for key, enabled in self.config.get("KvsAll.query_types").items()
              if enabled
          ]
  
          # corresponding indexes
          self.query_indexes: List[KvsAllIndex] = []
  
          #' for each query type (ordered as in self.query_types), index right after last
          #' example of that type in the list of all examples (over all query types)
          self.query_last_example = []
  
          # construct relevant data structures
          self.num_examples = 0
          for query_type in self.query_types:
              index_type = (
                  "sp_to_o"
                  if query_type == "sp_"
                  else ("so_to_p" if query_type == "s_o" else "po_to_s")
              )
              index = self.dataset.index(f"{self.train_split}_{index_type}")
              self.query_indexes.append(index)
              self.num_examples += len(index)
              self.query_last_example.append(self.num_examples)
  
          # create dataloader
          self.loader = torch.utils.data.DataLoader(
              range(self.num_examples),
              collate_fn=self._get_collate_fun(),
              shuffle=True,
              batch_size=self.batch_size,
              num_workers=self.config.get("train.num_workers"),
              worker_init_fn=_generate_worker_init_fn(self.config),
              pin_memory=self.config.get("train.pin_memory"),
          )
  
      def _get_collate_fun(self):
          # create the collate function
          def collate(batch):
              """For a batch of size n, returns a dictionary of:
  
              - queries: nx2 tensor, row = query (sp, po, or so indexes)
              - label_coords: for each query, position of true answers (an Nx2 tensor,
                first columns holds query index, second colum holds index of label)
              - query_type_indexes (vector of size n holding the query type of each query)
              - triples (all true triples in the batch; e.g., needed for weighted
                penalties)
  
              """
  
              # count how many labels we have across the entire batch
              num_ones = 0
              for example_index in batch:
                  start = 0
                  for query_type_index in range(len(self.query_types)):
                      end = self.query_last_example[query_type_index]
                      if example_index < end:
                          example_index -= start
                          num_ones += self.query_indexes[query_type_index]._values_offset[
                              example_index + 1
                          ]
                          num_ones -= self.query_indexes[query_type_index]._values_offset[
                              example_index
                          ]
                          break
                      start = end
  
              # now create the batch elements
              queries_batch = torch.zeros([len(batch), 2], dtype=torch.long)
              query_type_indexes_batch = torch.zeros([len(batch)], dtype=torch.long)
              label_coords_batch = torch.zeros([num_ones, 2], dtype=torch.int)
              triples_batch = torch.zeros([num_ones, 3], dtype=torch.long)
              current_index = 0
              for batch_index, example_index in enumerate(batch):
                  start = 0
                  for query_type_index, query_type in enumerate(self.query_types):
                      end = self.query_last_example[query_type_index]
                      if example_index < end:
                          example_index -= start
                          query_type_indexes_batch[batch_index] = query_type_index
                          queries = self.query_indexes[query_type_index]._keys
                          label_offsets = self.query_indexes[
                              query_type_index
                          ]._values_offset
                          labels = self.query_indexes[query_type_index]._values
                          if query_type == "sp_":
                              query_col_1, query_col_2, target_col = S, P, O
                          elif query_type == "s_o":
                              query_col_1, target_col, query_col_2 = S, P, O
                          else:
                              target_col, query_col_1, query_col_2 = S, P, O
                          break
                      start = end
  
                  queries_batch[batch_index,] = queries[example_index]
                  start = label_offsets[example_index]
                  end = label_offsets[example_index + 1]
                  size = end - start
                  label_coords_batch[
                      current_index : (current_index + size), 0
                  ] = batch_index
                  label_coords_batch[current_index : (current_index + size), 1] = labels[
                      start:end
                  ]
                  triples_batch[
                      current_index : (current_index + size), query_col_1
                  ] = queries[example_index][0]
                  triples_batch[
                      current_index : (current_index + size), query_col_2
                  ] = queries[example_index][1]
                  triples_batch[
                      current_index : (current_index + size), target_col
                  ] = labels[start:end]
                  current_index += size
  
              # all done
              return {
                  "queries": queries_batch,
                  "label_coords": label_coords_batch,
                  "query_type_indexes": query_type_indexes_batch,
                  "triples": triples_batch,
              }
  
          return collate
  
      def _prepare_batch(
          self, batch_index, batch, result: TrainingJob._ProcessBatchResult
      ):
          # move labels to GPU for entire batch (else somewhat costly, but this should be
          # reasonably small)
          result.prepare_time -= time.time()
          batch["label_coords"] = batch["label_coords"].to(self.device)
          result.size = len(batch["queries"])
          result.prepare_time += time.time()
  
      def _process_subbatch(
          self,
          batch_index,
          batch,
          subbatch_slice,
          result: TrainingJob._ProcessBatchResult
      ):
          batch_size = result.size
  
          # prepare
          result.prepare_time -= time.time()
          queries_subbatch = batch["queries"][subbatch_slice].to(self.device)
          subbatch_size = len(queries_subbatch)
          label_coords_batch = batch["label_coords"]
          query_type_indexes_subbatch = batch["query_type_indexes"][subbatch_slice]
  
          # in this method, example refers to the index of an example in the batch, i.e.,
          # it takes values in 0,1,...,batch_size-1
          examples_for_query_type = {}
          for query_type_index, query_type in enumerate(self.query_types):
              examples_for_query_type[query_type] = (
                  (query_type_indexes_subbatch == query_type_index)
                  .nonzero(as_tuple=False)
                  .to(self.device)
                  .view(-1)
              )
  
          labels_subbatch = kge.job.util.coord_to_sparse_tensor(
              subbatch_size,
              max(self.dataset.num_entities(), self.dataset.num_relations()),
              label_coords_batch,
              self.device,
              row_slice=subbatch_slice,
          ).to_dense()
          labels_for_query_type = {}
          for query_type, examples in examples_for_query_type.items():
              if query_type == "s_o":
                  labels_for_query_type[query_type] = labels_subbatch[
                      examples, : self.dataset.num_relations()
                  ]
              else:
                  labels_for_query_type[query_type] = labels_subbatch[
                      examples, : self.dataset.num_entities()
                  ]
  
          if self.label_smoothing > 0.0:
              # as in ConvE: https://github.com/TimDettmers/ConvE
              for query_type, labels in labels_for_query_type.items():
                  if query_type != "s_o":  # entity targets only for now
                      labels_for_query_type[query_type] = (
                          1.0 - self.label_smoothing
                      ) * labels + 1.0 / labels.size(1)
  
          result.prepare_time += time.time()
  
          # forward/backward pass (sp)
          for query_type, examples in examples_for_query_type.items():
              if len(examples) > 0:
                  result.forward_time -= time.time()
                  if query_type == "sp_":
                      scores = self.model.score_sp(
                          queries_subbatch[examples, 0], queries_subbatch[examples, 1]
                      )
                  elif query_type == "s_o":
                      scores = self.model.score_so(
                          queries_subbatch[examples, 0], queries_subbatch[examples, 1]
                      )
                  else:
                      scores = self.model.score_po(
                          queries_subbatch[examples, 0], queries_subbatch[examples, 1]
                      )
                  # note: average on batch_size, not on subbatch_size
                  loss_value = self._loss_weight * (
                      self.loss(scores, labels_for_query_type[query_type]) / batch_size
                  )
                  result.avg_loss += loss_value.item()
                  result.forward_time += time.time()
                  result.backward_time -= time.time()
                  if not self.is_forward_only:
                      loss_value.backward()
                  result.backward_time += time.time()
  