code: def generate_dataset(
      config: ConfigSchema, num_entities: int, fractions: List[float]
  ) -> Dataset:
      """Create a randomly-generated dataset compatible with the given config.
  
      Create embeddings for each entity (generating the same given number of
      entities for each type) and produce an edge between them if their dot
      product is positive. The edges are split up into several sets, each one
      containing a fraction of the total, as given in the argument (the fractions
      can sum to less than 1, in which case the leftover edges are discarded).
  
      """
      entity_path = TemporaryDirectory()
      relation_paths = [TemporaryDirectory() for _ in fractions]
  
      embeddings: Dict[str, Tuple[np.ndarray]] = {}
      for entity_name, entity in config.entities.items():
          embeddings[entity_name] = np.split(
              np.random.randn(num_entities, config.dimension),
              np.cumsum(
                  np.random.multinomial(
                      num_entities,
                      [1 / entity.num_partitions]
                      * entity.num_partitions,
                  )[:-1]
              ),
          )
          for partition, embedding in enumerate(
              embeddings[entity_name]
          ):
              with open(
                  os.path.join(
                      entity_path.name,
                      "entity_count_%s_%d.txt"
                      % (entity_name, partition),
                  ),
                  "xt",
              ) as tf:
                  tf.write("%d" % len(embedding))
  
      any_lhs_featurized = any(
          config.entities[relation.lhs].featurized
          for relation in config.relations
      )
      any_rhs_featurized = any(
          config.entities[relation.rhs].featurized
          for relation in config.relations
      )
      num_lhs_partitions = broadcast_nums(
          len(embeddings[relation.lhs]) for relation in config.relations
      )
      num_rhs_partitions = broadcast_nums(
          len(embeddings[relation.rhs]) for relation in config.relations
      )
  
      for lhs_partition in range(num_lhs_partitions):
          for rhs_partition in range(num_rhs_partitions):
              dtype = [
                  ("lhs", np.int64),
                  ("lhs_feat", np.bool_),
                  ("rhs", np.int64),
                  ("rhs_feat", np.bool_),
                  ("rel", np.int64),
              ]
              edges = np.empty((0,), dtype=dtype)
              for rel_idx, relation in enumerate(config.relations):
                  lhs_partitioned = (
                      config.entities[relation.lhs].num_partitions > 1
                  )
                  rhs_partitioned = (
                      config.entities[relation.rhs].num_partitions > 1
                  )
                  lhs_embs = embeddings[relation.lhs][
                      lhs_partition if lhs_partitioned else 0
                  ]
                  rhs_embs = embeddings[relation.rhs][
                      rhs_partition if rhs_partitioned else 0
                  ]
                  scores = np.einsum("ld,rd->lr", lhs_embs, rhs_embs)
                  num_these_edges = np.count_nonzero(scores > 0)
                  these_edges = np.empty(num_these_edges, dtype=dtype)
                  these_edges["lhs"], these_edges["rhs"] = np.nonzero(
                      scores > 0
                  )
                  these_edges["rel"] = rel_idx
                  these_edges["lhs_feat"] = config.entities[
                      relation.lhs
                  ].featurized
                  these_edges["rhs_feat"] = config.entities[
                      relation.rhs
                  ].featurized
                  edges = np.append(edges, these_edges)
              edges = edges[np.random.permutation(len(edges))]
              start_idx = 0
              for fraction, path in zip(fractions, relation_paths):
                  end_idx = start_idx + int(fraction * len(edges))
                  with h5py.File(
                      os.path.join(
                          path.name,
                          "edges_%d_%d.h5"
                          % (lhs_partition, rhs_partition),
                      ),
                      "x",
                  ) as hf:
                      hf.attrs["format_version"] = 1
                      these_edges = edges[start_idx:end_idx]
                      if any_lhs_featurized:
                          hf["lhsd_data"] = these_edges["lhs"][
                              these_edges["lhs_feat"]
                          ]
                          hf["lhsd_offsets"] = np.concatenate(
                              (
                                  np.array([0], dtype=np.int64),
                                  np.cumsum(
                                      these_edges["lhs_feat"],
                                      dtype=np.int64,
                                  ),
                              )
                          )
                          # Poison the non-featurized data.
                          these_edges["lhs"][
                              these_edges["lhs_feat"]
                          ] = -1
                      if any_rhs_featurized:
                          hf["rhsd_data"] = these_edges["rhs"][
                              these_edges["rhs_feat"]
                          ]
                          hf["rhsd_offsets"] = np.concatenate(
                              (
                                  np.array([0], dtype=np.int64),
                                  np.cumsum(
                                      these_edges["rhs_feat"],
                                      dtype=np.int64,
                                  ),
                              )
                          )
                          # Poison the non-featurized data.
                          these_edges["rhs"][
                              these_edges["rhs_feat"]
                          ] = -1
                      hf["lhs"] = these_edges["lhs"]
                      hf["rhs"] = these_edges["rhs"]
                      hf["rel"] = these_edges["rel"]
                  start_idx = end_idx
  
      return Dataset(entity_path, relation_paths)
  
imports: import json
  import logging
  import os.path
  import random
  import time
  import unittest
  from functools import partial
  from tempfile import TemporaryDirectory
  from typing import Dict, Iterable, List, Mapping, NamedTuple, Tuple, Union
  from unittest import main, TestCase
  import attr
  import h5py
  import numpy
  import torch
  import torch.multiprocessing
  from torchbiggraph.config import ConfigSchema, EntitySchema, RelationSchema
  from torchbiggraph.eval import do_eval
  from torchbiggraph.partitionserver import run_partition_server
  from torchbiggraph.stats import SerializedStats
  from torchbiggraph.train import train
  from torchbiggraph.train_gpu import CPP_INSTALLED
  from torchbiggraph.util import call_one_after_the_other, setup_logging, SubprocessInitializer
label_desc: Use image node_id 2 for calling a global function with example usage: generate_dataset(config, num_entities, fractions) and returns: Dataset