Update scripts/cache_dataset.py: speed + balance + correctness fixes

scripts/cache_dataset.py

@@ -146,6 +146,8 @@ def main():
     parser.add_argument("--context_length", type=int, default=8192)
     parser.add_argument("--min_trades", type=int, default=10)
     parser.add_argument("--samples_per_token", type=int, default=1)
+    parser.add_argument("--horizons_seconds", type=int, nargs="+", default=[30, 60, 120, 240, 420])
+    parser.add_argument("--quantiles", type=float, nargs="+", default=[0.1, 0.5, 0.9])
     parser.add_argument("--num_workers", type=int, default=1)
     parser.add_argument("--clickhouse_host", type=str, default=os.getenv("CLICKHOUSE_HOST", "localhost"))
     parser.add_argument("--clickhouse_port", type=int, default=int(os.getenv("CLICKHOUSE_PORT", 9000)))
@@ -180,7 +182,7 @@ def main():
         quality_scores_map = get_token_quality_scores(clickhouse_client)
         print(f"INFO: Loaded {len(quality_scores_map)} quality scores.")
 
-        dataset = OracleDataset(data_fetcher=data_fetcher, max_samples=args.max_samples, start_date=start_date_dt, ohlc_stats_path=args.ohlc_stats_path, horizons_seconds=[60, 180, 300, 600, 1800, 3600, 7200], quantiles=[0.5], min_trade_usd=args.min_trade_usd, max_seq_len=args.context_length)
+        dataset = OracleDataset(data_fetcher=data_fetcher, max_samples=args.max_samples, start_date=start_date_dt, ohlc_stats_path=args.ohlc_stats_path, horizons_seconds=args.horizons_seconds, quantiles=args.quantiles, min_trade_usd=args.min_trade_usd, max_seq_len=args.context_length)
 
         if len(dataset) == 0:
             print("WARNING: No samples. Exiting.")
@@ -189,7 +191,25 @@ def main():
         # Filter mints by return_class_map
         original_size = len(dataset.sampled_mints)
         filtered_mints = [m for m in dataset.sampled_mints if m['mint_address'] in return_class_map]
-        print(f"INFO: Filtered {original_size} -> {len(filtered_mints)} tokens")
+        print(f"INFO: Filtered by class map: {original_size} -> {len(filtered_mints)} tokens")
+
+        # Pre-filter: only keep tokens with >= min_trades trades (fast ClickHouse count query)
+        print(f"INFO: Pre-filtering tokens by trade count (>= {args.min_trades} trades)...")
+        trade_counts = clickhouse_client.execute("""
+            SELECT base_address, count() as cnt
+            FROM trades
+            GROUP BY base_address
+            HAVING cnt >= %(min_trades)s
+        """, {'min_trades': args.min_trades})
+        valid_tokens = {row[0] for row in trade_counts}
+        pre_filter_size = len(filtered_mints)
+        filtered_mints = [m for m in filtered_mints if m['mint_address'] in valid_tokens]
+        print(f"INFO: Pre-filtered by trade count: {pre_filter_size} -> {len(filtered_mints)} tokens (removed {pre_filter_size - len(filtered_mints)} with < {args.min_trades} trades)")
+
+        # Also filter by quality score availability
+        pre_quality_size = len(filtered_mints)
+        filtered_mints = [m for m in filtered_mints if m['mint_address'] in quality_scores_map]
+        print(f"INFO: Filtered by quality score: {pre_quality_size} -> {len(filtered_mints)} tokens")
 
         if len(filtered_mints) == 0:
             print("WARNING: No tokens after filtering.")
@@ -198,16 +218,69 @@ def main():
         print(f"INFO: Cache mode: {args.cache_mode}, Workers: {args.num_workers}")
 
         db_config = {'clickhouse_host': args.clickhouse_host, 'clickhouse_port': args.clickhouse_port, 'neo4j_uri': args.neo4j_uri, 'neo4j_user': args.neo4j_user, 'neo4j_password': args.neo4j_password}
-        dataset_config = {'max_samples': args.max_samples, 'start_date': start_date_dt, 'ohlc_stats_path': args.ohlc_stats_path, 'horizons_seconds': [60, 180, 300, 600, 1800, 3600, 7200], 'quantiles': [0.5], 'min_trade_usd': args.min_trade_usd, 'max_seq_len': args.context_length, 'sampled_mints': filtered_mints}
+        dataset_config = {'max_samples': args.max_samples, 'start_date': start_date_dt, 'ohlc_stats_path': args.ohlc_stats_path, 'horizons_seconds': args.horizons_seconds, 'quantiles': args.quantiles, 'min_trade_usd': args.min_trade_usd, 'max_seq_len': args.context_length, 'sampled_mints': filtered_mints}
+
+        # Build tasks with class-aware multi-sampling for balanced cache
+        import random
+        from collections import Counter, defaultdict
+
+        # Count eligible tokens per class
+        eligible_class_counts = Counter()
+        mints_by_class = defaultdict(list)
+        for i, m in enumerate(filtered_mints):
+            cid = return_class_map.get(m['mint_address'])
+            if cid is not None:
+                eligible_class_counts[cid] += 1
+                mints_by_class[cid].append((i, m))
+
+        print(f"INFO: Eligible tokens per class: {dict(sorted(eligible_class_counts.items()))}")
+
+        # Compute balanced samples_per_token for each class
+        num_classes = len(eligible_class_counts)
+        if args.max_samples:
+            target_total = args.max_samples
+        else:
+            target_total = 15000  # Default target: 15k balanced files
+        target_per_class = target_total // max(num_classes, 1)
+
+        class_multipliers = {}
+        class_token_caps = {}
+        for cid, count in eligible_class_counts.items():
+            if count >= target_per_class:
+                # Enough tokens — 1 sample each, cap token count
+                class_multipliers[cid] = 1
+                class_token_caps[cid] = target_per_class
+            else:
+                # Not enough tokens — multi-sample, use all tokens
+                class_multipliers[cid] = min(10, max(1, math.ceil(target_per_class / max(count, 1))))
+                class_token_caps[cid] = count
+
+        print(f"INFO: Target total: {target_total}, Target per class: {target_per_class}")
+        print(f"INFO: Class multipliers: {dict(sorted(class_multipliers.items()))}")
+        print(f"INFO: Class token caps: {dict(sorted(class_token_caps.items()))}")
 
-        # Build tasks from filtered_mints directly
+        # Build balanced task list
         tasks = []
-        for i, mint_record in enumerate(filtered_mints):
-            mint_addr = mint_record['mint_address']
-            if args.cache_mode == "context":
-                tasks.append((i, mint_addr, args.samples_per_token, str(output_dir)))
-            else:
-                tasks.append((i, mint_addr, str(output_dir)))
+        for cid, mint_list in mints_by_class.items():
+            random.shuffle(mint_list)
+            cap = class_token_caps.get(cid, len(mint_list))
+            spt = class_multipliers.get(cid, 1)
+            # Override with CLI --samples_per_token if explicitly set > 1
+            if args.samples_per_token > 1:
+                spt = args.samples_per_token
+            for i, m in mint_list[:cap]:
+                mint_addr = m['mint_address']
+                if args.cache_mode == "context":
+                    tasks.append((i, mint_addr, spt, str(output_dir)))
+                else:
+                    tasks.append((i, mint_addr, str(output_dir)))
+
+        random.shuffle(tasks)  # Shuffle tasks for even load distribution across workers
+        expected_files = sum(
+            class_multipliers.get(cid, 1) * min(class_token_caps.get(cid, len(ml)), len(ml))
+            for cid, ml in mints_by_class.items()
+        )
+        print(f"INFO: Total tasks: {len(tasks)} (expected ~{expected_files} output files, target ~{target_total})")
 
         print(f"INFO: Starting to cache {len(tasks)} tokens...")
 
@@ -255,7 +328,18 @@ def main():
                 pass
 
         with open(output_dir / "class_metadata.json", 'w') as f:
-            json.dump({'file_class_map': file_class_map, 'class_distribution': {str(k): v for k, v in class_distribution.items()}, 'cache_mode': args.cache_mode, 'num_workers': args.num_workers}, f, indent=2)
+            json.dump({
+                'file_class_map': file_class_map,
+                'class_distribution': {str(k): v for k, v in class_distribution.items()},
+                'cache_mode': args.cache_mode,
+                'num_workers': args.num_workers,
+                'horizons_seconds': args.horizons_seconds,
+                'quantiles': args.quantiles,
+                'class_multipliers': {str(k): v for k, v in class_multipliers.items()},
+                'class_token_caps': {str(k): v for k, v in class_token_caps.items()},
+                'target_total': target_total,
+                'target_per_class': target_per_class,
+            }, f, indent=2)
 
         print(f"\n--- Done ---\nSuccess: {success_count}, Skipped: {skipped_count}, Errors: {error_count}\nFiles: {len(file_class_map)}\nLocation: {output_dir.resolve()}")