app.py

"""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 = """
<div style="text-align:center; padding: 12px 16px 4px;">
  <h2 style="margin-bottom: 8px;">Caption Preference Study</h2>
  <p style="font-size: 1.05em; margin: 0;">
    You will see images with two captions. Click the caption that better
    describes the image.
  </p>
</div>
"""

DONE_NEW_HTML = """
<div style="text-align:center; padding: 32px;">
  <h2>All done — thank you for participating!</h2>
  <p>You can close this tab now.</p>
</div>
"""

DONE_ALREADY_HTML_TMPL = """
<div style="text-align:center; padding: 32px;">
  <h2>You've already completed this study.</h2>
  <p>Thank you, {name}! Our records show you finished all
  {total} trials. There's nothing more to do — feel free to close this tab.</p>
</div>
"""


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)])