Upload ai/data_generation/consolidate_data.py with huggingface_hub

ai/data_generation/consolidate_data.py

@@ -0,0 +1,40 @@
+import os
+
+import numpy as np
+
+
+def consolidate_data(files, output_file):
+    all_states = []
+    all_policies = []
+    all_winners = []
+
+    for f in files:
+        if not os.path.exists(f):
+            print(f"Skipping {f}, not found.")
+            continue
+        print(f"Loading {f}...")
+        data = np.load(f)
+        all_states.append(data["states"])
+        all_policies.append(data["policies"])
+        all_winners.append(data["winners"])
+
+    if not all_states:
+        print("No data to consolidate.")
+        return
+
+    np_states = np.concatenate(all_states, axis=0)
+    np_policies = np.concatenate(all_policies, axis=0)
+    np_winners = np.concatenate(all_winners, axis=0)
+
+    np.savez_compressed(output_file, states=np_states, policies=np_policies, winners=np_winners)
+    print(f"Consolidated {len(np_states)} samples to {output_file}")
+
+
+if __name__ == "__main__":
+    files = [
+        "ai/data/data_poc_800.npz",
+        "ai/data/data_batch_strat_1.npz",
+        "ai/data/data_batch_0.npz",
+        "ai/data/data_batch_strat_0.npz",
+    ]
+    consolidate_data(files, "ai/data/data_consolidated.npz")