refactor

pgx/bridge_bidding.py

@@ -43,16 +43,14 @@
 
 def download_dds_results(download_dir="dds_results"):
  """Download and split the results into 100K chunks."""
- os.makedirs(download_dir, exist_ok=True)
- train_fname = os.path.join(download_dir, "dds_results_2.5M.npy")
- if not os.path.exists(train_fname):
- _download(DDS_RESULTS_TRAIN_URL, train_fname)
- with open(train_fname, "rb") as f:
+
+ def split_data(data, prefix, base_i=0):
+ with open(data, "rb") as f:
  keys, values = jnp.load(f)
  n = 100_000
  m = keys.shape[0] // n
  for i in range(m):
- fname = os.path.join(download_dir, f"train_{i:03d}.npy")
+ fname = os.path.join(download_dir, f"{prefix}_{i:03d}.npy")
  with open(fname, "wb") as f:
  print(
  f"saving {fname} ... [{i * n}, {(i + 1) * n})",
@@ -66,27 +64,17 @@ def download_dds_results(download_dir="dds_results"):
  ),
  )
 
+ os.makedirs(download_dir, exist_ok=True)
+
+ train_fname = os.path.join(download_dir, "dds_results_2.5M.npy")
+ if not os.path.exists(train_fname):
+ _download(DDS_RESULTS_TRAIN_URL, train_fname)
+ split_data(train_fname, "train") 
+
  test_fname = os.path.join(download_dir, "dds_results_500K.npy")
  if not os.path.exists(test_fname):
  _download(DDS_RESULTS_TEST_URL, test_fname)
- with open(test_fname, "rb") as f:
- keys, values = jnp.load(f)
- n = 100_000
- m = keys.shape[0] // n
- for i in range(m):
- fname = os.path.join(download_dir, f"test_{i:03d}.npy")
- with open(fname, "wb") as f:
- print(
- f"saving {fname} ... [{i * n}, {(i + 1) * n})",
- file=sys.stderr,
- )
- jnp.save(
- f,
- (
- keys[i * n : (i + 1) * n],
- values[i * n : (i + 1) * n],
- ),
- )
+ split_data(test_fname, "test")
 
 
 @dataclass