feat: add tqdm progress bar to VLM conversion and download benchmark test

This commit is contained in:
Roland Tannous 2026-03-04 13:30:27 +00:00
commit f59eaad212
2 changed files with 62 additions and 1 deletions

View file

@ -332,20 +332,26 @@ def convert_to_vlm_format(
has_urls = isinstance(next(iter(dataset))[image_column], str)
probe_needed = has_urls and total > PROBE_SIZE
from tqdm import tqdm
print(f"🔄 Converting {total} samples to VLM format...")
converted_list = []
failed_count = 0
for i, sample in enumerate(dataset):
pbar = tqdm(dataset, total=total, desc="Converting VLM samples", unit="sample")
for i, sample in enumerate(pbar):
try:
converted_list.append(_convert_single_sample(sample))
except Exception as e:
failed_count += 1
pbar.set_postfix(ok=len(converted_list), failed=failed_count, refresh=False)
# Early exit check after probing the first batch
if probe_needed and (i + 1) == PROBE_SIZE:
fail_rate = failed_count / PROBE_SIZE
if fail_rate >= MAX_FAIL_RATE:
pbar.close()
raise ValueError(
f"{fail_rate:.0%} of the first {PROBE_SIZE} images failed to download "
f"({failed_count}/{PROBE_SIZE}). "
@ -353,6 +359,7 @@ def convert_to_vlm_format(
"Consider using a dataset with embedded images instead."
)
print(f"✅ Probe passed: {failed_count}/{PROBE_SIZE} ({fail_rate:.0%}) failures in first batch, continuing...")
pbar.close()
if failed_count > 0:
fail_rate = failed_count / total

View file

@ -0,0 +1,54 @@
"""
Benchmark: fsspec URL image download throughput at different dataset sizes.
Dataset: google-research-datasets/conceptual_captions (subset: labeled)
Tests sizes: 100, 200, 300, 500, 1000, 1500, 2000
Reports: time, success/fail rate, throughput (images/sec)
"""
from datasets import load_dataset, Dataset
from PIL import Image as PILImage
from io import BytesIO
from itertools import islice
import fsspec
import time
DATASET = "google-research-datasets/conceptual_captions"
SUBSET = "labeled"
SPLIT = "train"
SIZES = [100, 200, 300, 500, 1000, 1500, 2000]
# Load the max we need in one go
max_size = max(SIZES)
print(f"Loading {max_size} samples from {DATASET} (streaming)...")
ds = load_dataset(DATASET, name=SUBSET, split=SPLIT, streaming=True)
rows = list(islice(ds, max_size))
full_dataset = Dataset.from_list(rows)
print(f"Loaded {len(full_dataset)} samples")
print(f"Columns: {full_dataset.column_names}")
print()
print(f"{'Size':>6} | {'Time':>8} | {'OK':>6} | {'Fail':>6} | {'Fail%':>6} | {'img/s':>7}")
print("-" * 55)
for size in SIZES:
dataset = full_dataset.select(range(size))
success, fail = 0, 0
t0 = time.time()
for sample in dataset:
url = sample["image_url"]
try:
with fsspec.open(url, "rb", expand=True) as f:
img = PILImage.open(BytesIO(f.read())).convert("RGB")
success += 1
except Exception:
fail += 1
elapsed = time.time() - t0
fail_pct = (fail / size) * 100
throughput = success / elapsed if elapsed > 0 else 0
print(f"{size:>6} | {elapsed:>7.1f}s | {success:>6} | {fail:>6} | {fail_pct:>5.1f}% | {throughput:>6.1f}/s")
print()
print("Done.")