code: def write_output(
      train, test, fn, metric, point_type="float", count=100
  ):
      f = h5py.File(fn, "w")
      f.attrs["distance"] = metric
      f.attrs["point_type"] = point_type
      f.create_dataset(
          "train", (len(train), len(train[0])), dtype=train.dtype
      )[:] = train
      f.create_dataset(
          "test", (len(test), len(test[0])), dtype=test.dtype
      )[:] = test
      neighbors = f.create_dataset(
          "neighbors", (len(test), count), dtype="i"
      )
      distances = f.create_dataset(
          "distances", (len(test), count), dtype="f"
      )
      bf = BruteForceBLAS(metric, precision=train.dtype)
      bf.fit(train)
      for i, x in enumerate(test):
          sys.stderr.write(
              "computing queries %d/%d ...\r" % (i + 1, len(test))
          )
          res = list(bf.query_with_distances(x, count))
          res.sort(key=lambda t: t[-1])
          neighbors[i] = [j for j, _ in res]
          distances[i] = [d for _, d in res]
      sys.stderr.write("\n")
      f.close()
  
file_path: n2/benchmarks/data/create_youtube_dataset.py
imports: import sys
  import h5py
  import numpy
  from sklearn.model_selection import train_test_split
  from metrics import metrics
label_desc: Use image node_id 1 for calling a global function with example usage: write_output(train, test, fn, metric, point_type, count) without return types