"""Caption Preference Study — Gradio Space.
Participants register with their full name + email, then see an image and two
captions (human vs. model) and pick a preference. Per-participant results are
stored as ``firstname-lastname.csv`` in a private HF dataset. If a participant
returns later their session resumes from wherever they left off, and if they
have already completed the study they are told so.
"""
from __future__ import annotations
import io
import json
import os
import random
import re
import threading
import time
from datetime import datetime, timezone
from pathlib import Path
from typing import Any
import gradio as gr
import pandas as pd
from huggingface_hub import HfApi, hf_hub_download, snapshot_download
from huggingface_hub.utils import EntryNotFoundError, RepositoryNotFoundError
HF_USER = "pmadinei"
IMAGES_REPO = f"{HF_USER}/caption-preference-images"
RESULTS_REPO = f"{HF_USER}/caption-preference-results"
HF_TOKEN = os.environ.get("HF_TOKEN")
RESPONSE_TIME_CAP = 100.0
CSV_PATH = Path(__file__).parent / "Qwen3-VL-8B-Instruct.csv"
IMAGE_DIR = Path(os.environ.get("IMAGE_DIR", "/tmp/caption_experiment_images"))
IMAGE_DIR.mkdir(parents=True, exist_ok=True)
RESULTS_COLUMNS = [
"id",
"image_id",
"filename",
"type",
"human_caption",
"model_caption",
"preference",
"response_time",
]
EMAIL_RE = re.compile(r"^[^@\s]+@[^@\s]+\.[^@\s]+$")
SLUG_RE = re.compile(r"[^a-z0-9]+")
api = HfApi(token=HF_TOKEN)
# ---------------------------------------------------------------------------
# Data loading
# ---------------------------------------------------------------------------
def _clean_caption(value: Any) -> str:
if value is None:
return ""
text = str(value)
if len(text) >= 2 and text[0] == text[-1] and text[0] in ('"', "'"):
text = text[1:-1]
return text
print(f"[startup] Loading CSV from {CSV_PATH}")
df = pd.read_csv(CSV_PATH)
df["human_caption"] = df["human_caption"].map(_clean_caption)
df["model_caption"] = df["model_caption"].map(_clean_caption)
_test_mask = df["image_id"].astype(str).str.contains("test", case=False, na=False)
TEST_DF = df[_test_mask].reset_index(drop=True)
NONTEST_DF = df[~_test_mask].reset_index(drop=True)
NONTEST_IMAGE_IDS: list = list(NONTEST_DF["image_id"].unique())
NONTEST_IMAGE_ID_SET = set(NONTEST_IMAGE_IDS)
IMAGE_ID_TO_FILENAMES: dict = {
img_id: list(NONTEST_DF[NONTEST_DF["image_id"] == img_id]["filename"].unique())
for img_id in NONTEST_IMAGE_IDS
}
TEST_ROW_IDS = set(int(x) for x in TEST_DF["id"]) if len(TEST_DF) else set()
TOTAL_TRIALS_PER_PARTICIPANT = len(NONTEST_IMAGE_IDS) + len(TEST_DF)
print(
f"[startup] {len(df)} rows | {len(NONTEST_IMAGE_IDS)} non-test image_ids | "
f"{len(TEST_DF)} test rows | {TOTAL_TRIALS_PER_PARTICIPANT} trials per participant"
)
# ---------------------------------------------------------------------------
# Image download
# ---------------------------------------------------------------------------
def _ensure_images_downloaded() -> None:
if not HF_TOKEN:
print("[startup] WARNING: HF_TOKEN is not set; cannot download images.")
return
print(f"[startup] Downloading images from {IMAGES_REPO} to {IMAGE_DIR}...")
snapshot_download(
repo_id=IMAGES_REPO,
repo_type="dataset",
local_dir=str(IMAGE_DIR),
token=HF_TOKEN,
max_workers=16,
)
print("[startup] Image download complete.")
_ensure_images_downloaded()
# ---------------------------------------------------------------------------
# Round-robin state (persisted to RESULTS_REPO/state.json)
# ---------------------------------------------------------------------------
_STATE_LOCK = threading.Lock()
_STATE: dict = {"image_id_used": {}}
def _state_key(image_id: Any) -> str:
return str(image_id)
def _load_state() -> None:
global _STATE
if not HF_TOKEN:
return
try:
path = hf_hub_download(
repo_id=RESULTS_REPO,
repo_type="dataset",
filename="state.json",
token=HF_TOKEN,
force_download=True,
)
with open(path) as f:
loaded = json.load(f)
_STATE = {"image_id_used": loaded.get("image_id_used", {})}
print(
f"[state] Loaded round-robin state with "
f"{len(_STATE['image_id_used'])} image_ids tracked."
)
except (EntryNotFoundError, RepositoryNotFoundError, FileNotFoundError):
print("[state] No existing state.json found, starting fresh.")
_STATE = {"image_id_used": {}}
except Exception as exc: # noqa: BLE001
print(f"[state] Could not load state.json ({exc}); starting fresh.")
_STATE = {"image_id_used": {}}
def _save_state() -> None:
if not HF_TOKEN:
return
payload = json.dumps(_STATE, indent=2).encode()
api.upload_file(
path_or_fileobj=io.BytesIO(payload),
path_in_repo="state.json",
repo_id=RESULTS_REPO,
repo_type="dataset",
commit_message="Update round-robin state",
)
_load_state()
def _assign_filenames(image_ids_to_assign: list) -> dict:
"""Round-robin filename pick for a given set of image_ids.
For each image_id, choose uniformly from filenames not yet used since the
last reset. When all filenames have been used, reset and start a fresh
cycle. Independent per image_id.
"""
with _STATE_LOCK:
assignments: dict = {}
for img_id in image_ids_to_assign:
all_fns = IMAGE_ID_TO_FILENAMES[img_id]
key = _state_key(img_id)
used = list(_STATE["image_id_used"].get(key, []))
available = [fn for fn in all_fns if fn not in used]
if not available:
used = []
available = list(all_fns)
chosen = random.choice(available)
used.append(chosen)
_STATE["image_id_used"][key] = used
assignments[img_id] = chosen
if assignments:
try:
_save_state()
except Exception as exc: # noqa: BLE001
print(f"[state] WARNING: could not persist state.json ({exc}).")
return assignments
# ---------------------------------------------------------------------------
# Per-participant CSV + registry
# ---------------------------------------------------------------------------
def _slugify(s: str) -> str:
s = (s or "").strip().lower()
s = SLUG_RE.sub("-", s)
return s.strip("-")
def _participant_filename(first: str, last: str) -> str:
return f"results/{_slugify(first)}-{_slugify(last)}.csv"
def _load_participant_results(participant_file: str) -> list[dict]:
if not HF_TOKEN:
return []
try:
path = hf_hub_download(
repo_id=RESULTS_REPO,
repo_type="dataset",
filename=participant_file,
token=HF_TOKEN,
force_download=True,
)
frame = pd.read_csv(path)
return frame.to_dict(orient="records")
except (EntryNotFoundError, RepositoryNotFoundError, FileNotFoundError):
return []
except Exception as exc: # noqa: BLE001
print(f"[participant] Could not load {participant_file} ({exc})")
return []
def _completed_keys(prior_results: list[dict]) -> tuple[set, set]:
"""Return (done_nontest_image_ids, done_test_row_ids) from a CSV-loaded list."""
done_image_ids = set()
done_test_ids = set()
for r in prior_results:
try:
row_id = int(r["id"])
except (KeyError, TypeError, ValueError):
continue
if row_id in TEST_ROW_IDS:
done_test_ids.add(row_id)
continue
img_id_str = str(r.get("image_id"))
if "test" in img_id_str.lower():
done_test_ids.add(row_id)
continue
img_id_val = r.get("image_id")
if img_id_val in NONTEST_IMAGE_ID_SET:
done_image_ids.add(img_id_val)
else:
try:
coerced = int(img_id_val)
if coerced in NONTEST_IMAGE_ID_SET:
done_image_ids.add(coerced)
except (TypeError, ValueError):
pass
return done_image_ids, done_test_ids
def _is_complete(prior_results: list[dict]) -> bool:
done_image_ids, done_test_ids = _completed_keys(prior_results)
return done_image_ids >= NONTEST_IMAGE_ID_SET and done_test_ids >= TEST_ROW_IDS
def _build_remaining_trials(prior_results: list[dict]) -> list[dict]:
done_image_ids, done_test_ids = _completed_keys(prior_results)
remaining_image_ids = [
iid for iid in NONTEST_IMAGE_IDS if iid not in done_image_ids
]
assignments = _assign_filenames(remaining_image_ids)
trials: list[dict] = []
for img_id in remaining_image_ids:
fn = assignments[img_id]
match = NONTEST_DF[
(NONTEST_DF["image_id"] == img_id) & (NONTEST_DF["filename"] == fn)
]
if match.empty:
continue
trials.append(_row_to_trial(match.iloc[0]))
for _, row in TEST_DF.iterrows():
if int(row["id"]) in done_test_ids:
continue
trials.append(_row_to_trial(row))
random.shuffle(trials)
return trials
def _row_to_trial(row: pd.Series) -> dict:
raw_image_id = row["image_id"]
if isinstance(raw_image_id, (int,)) or (
isinstance(raw_image_id, str) and raw_image_id.lstrip("-").isdigit()
):
image_id_out: Any = int(raw_image_id)
else:
image_id_out = str(raw_image_id)
return {
"id": int(row["id"]),
"image_id": image_id_out,
"filename": str(row["filename"]),
"type": str(row["type"]),
"human_caption": str(row["human_caption"]),
"model_caption": str(row["model_caption"]),
"human_on_left": random.choice([True, False]),
}
def _save_results(participant_file: str, results: list[dict]) -> None:
if not HF_TOKEN or not results:
return
frame = pd.DataFrame(results, columns=RESULTS_COLUMNS)
buf = io.BytesIO()
frame.to_csv(buf, index=False)
buf.seek(0)
api.upload_file(
path_or_fileobj=buf,
path_in_repo=participant_file,
repo_id=RESULTS_REPO,
repo_type="dataset",
commit_message=f"Update {participant_file} (n={len(results)})",
)
def _load_participants_registry() -> dict:
if not HF_TOKEN:
return {}
try:
path = hf_hub_download(
repo_id=RESULTS_REPO,
repo_type="dataset",
filename="participants.json",
token=HF_TOKEN,
force_download=True,
)
with open(path) as f:
return json.load(f)
except (EntryNotFoundError, RepositoryNotFoundError, FileNotFoundError):
return {}
except Exception as exc: # noqa: BLE001
print(f"[participants] Could not load registry ({exc})")
return {}
_REGISTRY_LOCK = threading.Lock()
def _register_participant(slug: str, first: str, last: str, email: str) -> None:
if not HF_TOKEN:
return
with _REGISTRY_LOCK:
registry = _load_participants_registry()
entry = registry.get(slug, {})
now_iso = datetime.now(timezone.utc).isoformat()
if not entry:
entry = {
"full_name": f"{first} {last}".strip(),
"first_name": first,
"last_name": last,
"email": email,
"registered_at": now_iso,
"last_session_at": now_iso,
}
else:
entry.setdefault("first_name", first)
entry.setdefault("last_name", last)
entry.setdefault("registered_at", now_iso)
entry["full_name"] = f"{first} {last}".strip()
entry["email"] = email
entry["last_session_at"] = now_iso
registry[slug] = entry
payload = json.dumps(registry, indent=2).encode()
try:
api.upload_file(
path_or_fileobj=io.BytesIO(payload),
path_in_repo="participants.json",
repo_id=RESULTS_REPO,
repo_type="dataset",
commit_message=f"Register/update participant {slug}",
)
except Exception as exc: # noqa: BLE001
print(f"[participants] WARNING: could not save registry ({exc}).")
# ---------------------------------------------------------------------------
# Gradio handlers
# ---------------------------------------------------------------------------
WELCOME_HTML = """
Caption Preference Study
You will see images with two captions. Click the caption that better
describes the image.
"""
DONE_NEW_HTML = """
All done — thank you for participating!
You can close this tab now.
"""
DONE_ALREADY_HTML_TMPL = """
You've already completed this study.
Thank you, {name}! Our records show you finished all
{total} trials. There's nothing more to do — feel free to close this tab.
"""
def _validation_error(message: str):
return (
None, # state
gr.update(visible=True), # intro
gr.update(visible=False), # trial group
gr.update(visible=False, value=""), # done panel
None, # image
gr.update(value=""), # left button
gr.update(value=""), # right button
"", # progress
gr.update(value=message, visible=True), # error markdown
)
def start_session(first_name: str, last_name: str, email: str):
first = (first_name or "").strip()
last = (last_name or "").strip()
email_v = (email or "").strip()
if not first:
return _validation_error("Please enter your **first name**.")
if not last:
return _validation_error("Please enter your **last name**.")
if not EMAIL_RE.match(email_v):
return _validation_error("Please enter a valid **email address**.")
slug_first = _slugify(first)
slug_last = _slugify(last)
if not slug_first or not slug_last:
return _validation_error(
"Your name must include at least one letter or digit."
)
participant_file = _participant_filename(first, last)
prior = _load_participant_results(participant_file)
if _is_complete(prior):
msg = DONE_ALREADY_HTML_TMPL.format(
name=f"{first} {last}",
total=TOTAL_TRIALS_PER_PARTICIPANT,
)
# Still log that they came back (no overwrite of prior CSV).
threading.Thread(
target=_register_participant,
args=(f"{slug_first}-{slug_last}", first, last, email_v),
daemon=True,
).start()
return (
None,
gr.update(visible=False),
gr.update(visible=False),
gr.update(value=msg, visible=True),
None,
gr.update(value=""),
gr.update(value=""),
"",
gr.update(value="", visible=False),
)
trials = _build_remaining_trials(prior)
if not trials:
# Defensive: no trials remaining but not "complete" by the strict
# check — treat as done so the participant isn't stuck.
msg = DONE_ALREADY_HTML_TMPL.format(
name=f"{first} {last}",
total=TOTAL_TRIALS_PER_PARTICIPANT,
)
return (
None,
gr.update(visible=False),
gr.update(visible=False),
gr.update(value=msg, visible=True),
None,
gr.update(value=""),
gr.update(value=""),
"",
gr.update(value="", visible=False),
)
_register_participant(f"{slug_first}-{slug_last}", first, last, email_v)
state = {
"participant_file": participant_file,
"trials": trials,
"current_idx": 0,
"trial_start_time": time.time(),
"results": list(prior),
"prior_count": len(prior),
"total_trials": TOTAL_TRIALS_PER_PARTICIPANT,
}
img_path, left, right, progress = _current_display(state)
return (
state,
gr.update(visible=False), # intro
gr.update(visible=True), # trial group
gr.update(value="", visible=False), # done panel
img_path, # image
gr.update(value=left), # left button
gr.update(value=right), # right button
progress, # progress
gr.update(value="", visible=False), # error
)
def _current_display(state: dict) -> tuple:
if state is None or state["current_idx"] >= len(state["trials"]):
return None, "", "", ""
trial = state["trials"][state["current_idx"]]
img_path = str(IMAGE_DIR / trial["filename"])
if trial["human_on_left"]:
left, right = trial["human_caption"], trial["model_caption"]
else:
left, right = trial["model_caption"], trial["human_caption"]
completed = state["prior_count"] + state["current_idx"]
total = state["total_trials"]
progress = f"Trial {completed + 1} of {total}"
return img_path, left, right, progress
def _make_choice(state: dict, side: str):
if state is None:
return (
state,
gr.update(visible=False),
gr.update(visible=False),
None,
gr.update(value=""),
gr.update(value=""),
"",
)
elapsed = min(time.time() - state["trial_start_time"], RESPONSE_TIME_CAP)
trial = state["trials"][state["current_idx"]]
chose_human = trial["human_on_left"] if side == "left" else not trial["human_on_left"]
state["results"].append(
{
"id": trial["id"],
"image_id": trial["image_id"],
"filename": trial["filename"],
"type": trial["type"],
"human_caption": trial["human_caption"],
"model_caption": trial["model_caption"],
"preference": "H" if chose_human else "M",
"response_time": round(elapsed, 3),
}
)
threading.Thread(
target=_save_results,
args=(state["participant_file"], list(state["results"])),
daemon=True,
).start()
state["current_idx"] += 1
if state["current_idx"] >= len(state["trials"]):
total = state["total_trials"]
return (
state,
gr.update(visible=False),
gr.update(value=DONE_NEW_HTML, visible=True),
None,
gr.update(value=""),
gr.update(value=""),
f"Done — {total} / {total}",
)
state["trial_start_time"] = time.time()
img_path, left, right, progress = _current_display(state)
return (
state,
gr.update(visible=True),
gr.update(visible=False),
img_path,
gr.update(value=left),
gr.update(value=right),
progress,
)
# ---------------------------------------------------------------------------
# UI
# ---------------------------------------------------------------------------
custom_css = """
.caption-btn {
min-height: 140px !important;
font-size: 1.05em !important;
white-space: normal !important;
line-height: 1.4 !important;
padding: 16px !important;
text-align: left !important;
}
.center-img img { max-height: 60vh !important; object-fit: contain !important; }
.form-error { color: #b91c1c !important; }
"""
with gr.Blocks(title="Caption Preference Study", css=custom_css) as demo:
state = gr.State()
intro = gr.Group(visible=True)
with intro:
gr.HTML(WELCOME_HTML)
with gr.Row():
with gr.Column(scale=1):
pass
with gr.Column(scale=2):
first_input = gr.Textbox(
label="First name", placeholder="e.g. Jane", max_lines=1
)
last_input = gr.Textbox(
label="Last name", placeholder="e.g. Smith", max_lines=1
)
email_input = gr.Textbox(
label="Email address",
placeholder="you@example.com",
max_lines=1,
)
start_btn = gr.Button("Start", variant="primary", size="lg")
error_md = gr.Markdown("", visible=False, elem_classes=["form-error"])
with gr.Column(scale=1):
pass
trial_group = gr.Group(visible=False)
with trial_group:
progress = gr.Markdown("")
image = gr.Image(
label=None,
show_label=False,
interactive=False,
elem_classes=["center-img"],
)
with gr.Row():
left_btn = gr.Button("", elem_classes=["caption-btn"])
right_btn = gr.Button("", elem_classes=["caption-btn"])
done_panel = gr.HTML(visible=False)
start_btn.click(
start_session,
inputs=[first_input, last_input, email_input],
outputs=[
state,
intro,
trial_group,
done_panel,
image,
left_btn,
right_btn,
progress,
error_md,
],
)
left_btn.click(
lambda s: _make_choice(s, "left"),
inputs=[state],
outputs=[state, trial_group, done_panel, image, left_btn, right_btn, progress],
)
right_btn.click(
lambda s: _make_choice(s, "right"),
inputs=[state],
outputs=[state, trial_group, done_panel, image, left_btn, right_btn, progress],
)
if __name__ == "__main__":
demo.queue(default_concurrency_limit=8).launch(allowed_paths=[str(IMAGE_DIR)])