Upload scripts/cache_dataset.py with huggingface_hub

scripts/cache_dataset.py

@@ -288,11 +288,37 @@ def main():
         class_distribution = {}
         process_fn = _process_single_token_context if args.cache_mode == "context" else _process_single_token_raw
 
+        import time as _time
+
+        def _log_progress(task_num, total, start_time, recent_times, success_count, skipped_count, error_count):
+            """Print progress with rolling ETA every 10 tokens."""
+            if (task_num + 1) % 10 == 0 and recent_times:
+                avg_time = sum(recent_times) / len(recent_times)
+                remaining = total - (task_num + 1)
+                eta_seconds = avg_time * remaining
+                eta_hours = eta_seconds / 3600
+                wall_elapsed = _time.perf_counter() - start_time
+                speed = (task_num + 1) / wall_elapsed
+                tqdm.write(
+                    f"  [PROGRESS] {task_num+1}/{total} | "
+                    f"Speed: {speed:.1f} tok/s ({speed*60:.0f} tok/min) | "
+                    f"Avg: {avg_time:.1f}s/tok | "
+                    f"ETA: {eta_hours:.1f}h | "
+                    f"OK: {success_count} Skip: {skipped_count} Err: {error_count}"
+                )
+
         if args.num_workers == 1:
             print("INFO: Single-threaded mode...")
             _init_worker(db_config, dataset_config, return_class_map, quality_scores_map)
-            for task in tqdm(tasks, desc="Caching"):
+            start_time = _time.perf_counter()
+            recent_times = []
+            for task_num, task in enumerate(tqdm(tasks, desc="Caching", unit="tok")):
+                t0 = _time.perf_counter()
                 result = process_fn(task)
+                elapsed = _time.perf_counter() - t0
+                recent_times.append(elapsed)
+                if len(recent_times) > 50:
+                    recent_times.pop(0)
                 if result['status'] == 'success':
                     success_count += 1
                     class_distribution[result['class_id']] = class_distribution.get(result['class_id'], 0) + 1
@@ -301,13 +327,21 @@ def main():
                 else:
                     error_count += 1
                     tqdm.write(f"ERROR: {result['mint'][:16]} - {result['error']}")
+                _log_progress(task_num, len(tasks), start_time, recent_times, success_count, skipped_count, error_count)
         else:
             print(f"INFO: Running with {args.num_workers} workers...")
+            start_time = _time.perf_counter()
+            recent_times = []
             with ProcessPoolExecutor(max_workers=args.num_workers, initializer=_init_worker, initargs=(db_config, dataset_config, return_class_map, quality_scores_map)) as executor:
                 futures = {executor.submit(process_fn, task): task for task in tasks}
-                for future in tqdm(as_completed(futures), total=len(futures), desc="Caching"):
+                for task_num, future in enumerate(tqdm(as_completed(futures), total=len(futures), desc="Caching", unit="tok")):
+                    t0 = _time.perf_counter()
                     try:
                         result = future.result(timeout=300)
+                        elapsed = _time.perf_counter() - t0
+                        recent_times.append(elapsed)
+                        if len(recent_times) > 50:
+                            recent_times.pop(0)
                         if result['status'] == 'success':
                             success_count += 1
                             class_distribution[result['class_id']] = class_distribution.get(result['class_id'], 0) + 1
@@ -318,6 +352,7 @@ def main():
                     except Exception as e:
                         error_count += 1
                         tqdm.write(f"WORKER ERROR: {e}")
+                    _log_progress(task_num, len(tasks), start_time, recent_times, success_count, skipped_count, error_count)
 
         print("INFO: Building metadata...")
         file_class_map = {}