feat: add tqdm progress bar to VLM conversion and download benchmark test
This commit is contained in:
parent
50885a7aa3
commit
f59eaad212
2 changed files with 62 additions and 1 deletions
|
|
@ -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
|
||||
|
|
|
|||
54
studio/tests/test_url_download_benchmark.py
Normal file
54
studio/tests/test_url_download_benchmark.py
Normal 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.")
|
||||
Loading…
Add table
Add a link
Reference in a new issue