Enhance batch processing with progress tracking and update AIAgent for OpenRouter detection

- Integrated tqdm for progress tracking in batch processing, replacing map with imap_unordered for improved performance.
- Added base_url attribute in AIAgent to facilitate OpenRouter detection.
This commit is contained in:
teknium
2026-01-14 13:46:16 +00:00
parent b66c093316
commit 6e3dbb8d8b
2 changed files with 10 additions and 2 deletions

View File

@@ -30,6 +30,7 @@ from datetime import datetime
from multiprocessing import Pool, Manager, Lock
import traceback
from tqdm import tqdm
import fire
from run_agent import AIAgent
@@ -642,8 +643,14 @@ class BatchRunner:
print(f"✅ Created {len(tasks)} batch tasks")
print(f"🚀 Starting parallel batch processing...\n")
# Use map to process batches in parallel
results = pool.map(_process_batch_worker, tasks)
# Use imap_unordered with tqdm for progress tracking
results = list(tqdm(
pool.imap_unordered(_process_batch_worker, tasks),
total=len(tasks),
desc="📦 Batches",
unit="batch",
ncols=80
))
# Aggregate all batch statistics and update checkpoint
all_completed_prompts = list(completed_prompts_set)