item tracking more carefully

data/local_completions_preference_movies.json

@@ -0,0 +1 @@
+{"0": 1, "1": 1, "2": 1, "3": 1, "4": 1, "5": 1, "6": 1, "7": 1, "8": 1, "9": 1, "10": 1, "11": 1, "12": 1, "13": 1, "14": 1, "15": 1, "16": 1, "17": 1, "18": 1, "19": 1, "20": 1, "21": 1, "22": 1, "23": 1, "24": 1, "25": 1, "26": 1, "27": 1, "28": 1, "29": 1, "30": 1, "31": 1, "32": 1, "33": 1, "34": 1, "35": 1, "36": 1, "37": 1, "38": 1, "39": 1, "40": 1, "41": 1, "42": 1, "43": 1, "44": 1, "45": 1, "46": 1, "47": 1, "48": 1, "49": 1}

data/reservations.json

@@ -0,0 +1 @@
+{}

reject_submissions.py

@@ -0,0 +1,57 @@
+"""
+Helper script to reject a submission.
+
+Usage:
+    python scripts/reject_submission.py \
+        --repo lms-shape-preferences/results-preference-base \
+        --token hf_... \
+        --path json/worker123/submission456.json
+
+This moves the file from json/ to rejected/ so the item gets re-assigned
+to the next available participant.
+"""
+import argparse
+import os
+import tempfile
+from huggingface_hub import HfApi
+
+def reject_submission(repo_id: str, file_path: str, token: str) -> None:
+    assert file_path.startswith("json/"), \
+        f"Expected path starting with json/, got: {file_path}"
+
+    api          = HfApi(token=token)
+    rejected_path = file_path.replace("json/", "rejected/", 1)
+
+    print(f"Downloading {file_path} ...")
+    local = api.hf_hub_download(
+        repo_id=repo_id,
+        filename=file_path,
+        repo_type="dataset",
+        token=token,
+    )
+
+    print(f"Re-uploading to {rejected_path} ...")
+    api.upload_file(
+        path_or_fileobj=local,
+        path_in_repo=rejected_path,
+        repo_id=repo_id,
+        repo_type="dataset",
+    )
+
+    print(f"Deleting {file_path} ...")
+    api.delete_file(
+        path_in_repo=file_path,
+        repo_id=repo_id,
+        repo_type="dataset",
+    )
+
+    print(f"Done. {file_path} → {rejected_path}")
+    print("The item will be re-assigned within 5 minutes (cache TTL).")
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--repo",  required=True)
+    parser.add_argument("--token", default=os.getenv("HF_TOKEN"))
+    parser.add_argument("--path",  required=True, help="e.g. json/worker123/submission456.json")
+    args = parser.parse_args()
+    reject_submission(args.repo, args.path, args.token)

scripts/test_coverage.py

@@ -0,0 +1,120 @@
+"""
+Simulate N users going through the study and verify all 50 items get covered.
+
+Usage:
+    cd /dfs/scratch1/echoi1/prolific_preferences
+    HF_TOKEN=hf_... python scripts/test_coverage.py
+"""
+import sys
+import uuid
+from pathlib import Path
+
+sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
+
+from src.config import load_config
+from src.data import (
+    ensure_datasets,
+    assign_items,
+    release_reservation,
+    record_completion,
+    _load_pool,
+    _pool_path,
+    _data_dir,
+)
+
+
+def simulate_user(cfg: dict, complete: bool = True) -> dict:
+    user_id = str(uuid.uuid4())
+    items   = assign_items(cfg, user_id)
+    if complete:
+        release_reservation(user_id, cfg)
+        record_completion(user_id, items, cfg)
+    item_ids = [(item.get("pair_id") or item.get("item_id", ""), item.get("category", ""))
+                for item in items]
+    return {"user_id": user_id, "items": item_ids, "raw_items": items, "completed": complete}
+
+
+def clear_local_state(cfg: dict):
+    data_dir = _data_dir(cfg)
+    for pattern in ["reservations*", "completion_cache*", "local_completions*",
+                    "variant_counter*", "alternation_counter*"]:
+        for f in data_dir.glob(pattern):
+            f.unlink()
+
+
+def analyse_coverage(results: list, cfg: dict) -> bool:
+    cats      = [c["name"] for c in cfg["categories"]]
+    all_passed = True
+
+    print()
+    print("=" * 60)
+    print("COVERAGE ANALYSIS")
+    print("=" * 60)
+
+    for cat in cats:
+        pool     = _load_pool(str(_pool_path(cat, cfg)))
+        pool_ids = [p.get("pair_id") or p.get("item_id", "") for p in pool]
+        covered  = {pid: 0 for pid in pool_ids}
+
+        for result in results:
+            if not result["completed"]:
+                continue
+            for item_id, item_cat in result["items"]:
+                if item_cat == cat and item_id in covered:
+                    covered[item_id] += 1
+
+        covered_once  = sum(1 for c in covered.values() if c >= 1)
+        never_covered = [pid[:8] for pid, c in covered.items() if c == 0]
+        over_covered  = [pid[:8] for pid, c in covered.items() if c > 1]
+
+        print(f"\nCategory: {cat}")
+        print(f"  Pool size:     {len(pool)}")
+        print(f"  Covered >= 1x: {covered_once} / {len(pool)}")
+        print(f"  Never covered: {len(never_covered)} {never_covered[:5]}")
+        print(f"  Over-covered:  {len(over_covered)} {over_covered[:5]}")
+
+        if covered_once == len(pool):
+            print(f"  ✅ PASS — all {len(pool)} items covered")
+        else:
+            print(f"  ❌ FAIL — {len(pool) - covered_once} items not covered")
+            all_passed = False
+
+    print()
+    print("=" * 60)
+    print("OVERALL:", "✅ PASS" if all_passed else "❌ FAIL")
+    print("=" * 60)
+    return all_passed
+
+
+def run_simulation(label: str, n_users: int, dropout_indices: list = None):
+    dropout_indices = dropout_indices or []
+    cfg = load_config()
+    ensure_datasets(cfg)
+    clear_local_state(cfg)
+
+    print(f"\n── {label} ──")
+    print(f"[TEST] {n_users} users, dropouts at: {dropout_indices}")
+
+    results = []
+    for i in range(n_users):
+        complete = i not in dropout_indices
+        result   = simulate_user(cfg, complete=complete)
+        results.append(result)
+        status = "✅ completed" if complete else "❌ abandoned"
+        print(f"  User {i+1:2d} ({status}): "
+              f"indices = {[r[0][:8] for r in result['items']]}")
+
+    return analyse_coverage(results, cfg)
+
+
+if __name__ == "__main__":
+    # Test 1: perfect run — all 10 users complete, all 50 items covered exactly once
+    run_simulation("Test 1: Perfect run", n_users=10)
+
+    # Test 2: 2 dropouts — abandoned items should be picked up by extra users
+    # The new sort_key means uncovered+reserved items are preferred over covered+unreserved
+    # so items 35-39 (abandoned) get picked up by users 11-12 instead of re-covering 0-9
+    run_simulation("Test 2: 2 dropouts, 12 users", n_users=12, dropout_indices=[7, 3])
+
+    # Test 3: first user drops out — 11 users needed to cover all 50
+    run_simulation("Test 3: First user drops out", n_users=11, dropout_indices=[0])

src/data.py

@@ -1,18 +1,32 @@
 """
-Dataset download, item-pool caching, round-robin assignment, and session-state init.
-
-Pool selection
---------------
-50 items per (study_type, category) are selected from the HF test split using the
-configured seed, written to data/ on first startup, and reloaded from disk after that.
-
-Assignment
-----------
-Items are assigned round-robin via a file-locked counter per category so that
-concurrent users get different items.  For two-category studies the per-user split
-(e.g. 3 movies / 2 groceries) alternates on successive calls via a separate
-alternation counter that increments by exactly 1 per user — fixing the original
-bug where the movies product counter was reused.
+Dataset download, item-pool caching, completion-aware assignment, and session-state init.
+
+Assignment strategy
+-------------------
+Items are assigned based on how many *accepted* completions they already have,
+ensuring the least-covered items are always prioritised.
+
+Each assigned item is stamped with _pool_index and _pool_category at assignment
+time so record_completion never needs to do a fuzzy pair_id match — it reads
+the index directly.
+
+Accepted completions = JSON files under json/ in the output repo.
+Rejected completions = JSON files moved to rejected/ by the admin.
+  → moving a file to rejected/ automatically makes that item available again.
+
+Reservations
+------------
+When a user starts, their items are "reserved" in a local file for 80 min.
+Concurrent users (up to 5) each get a FileLock on the reservation file so they
+never receive the same items. Reservations expire automatically so abandoned
+sessions don't permanently block items.
+
+Dropout / rejection recovery
+-----------------------------
+- Dropout: reservation expires after 80 min → item re-enters the pool.
+- Rejection: admin moves json/{worker}/{id}.json → rejected/{worker}/{id}.json
+  in the HF dataset repo. On next Space restart (or cache expiry) the item's
+  accepted count drops to 0 and it gets re-assigned.
 """
 import json
 import random
@@ -25,7 +39,9 @@ from filelock import FileLock
 
 from src.config import CATEGORY_TO_REPO
 
-POOL_SIZE = 50  # items selected per (study_type, category)
+POOL_SIZE            = 50    # items selected per (study_type, category)
+RESERVATION_TTL      = 60 * 80   # 80 min: 30 min expected + ~2.5x buffer
+COMPLETION_CACHE_TTL = 300        # re-scan HF repo every 5 minutes
 
 
 # ── Path helpers ──────────────────────────────────────────────────────────────
@@ -40,32 +56,22 @@ def _pool_path(category: str, cfg: dict) -> Path:
     return _data_dir(cfg) / f"pool_{cfg['study_type']}_{category}.json"
 
 
-def _counter_path(category: str, cfg: dict) -> Path:
-    return _data_dir(cfg) / f"counter_{cfg['study_type']}_{category}.txt"
+def _reservation_path(cfg: dict) -> Path:
+    return _data_dir(cfg) / "reservations.json"
 
 
-def _counter_lock_path(category: str, cfg: dict) -> Path:
-    return _data_dir(cfg) / f"counter_{cfg['study_type']}_{category}.lock"
+def _reservation_lock_path(cfg: dict) -> Path:
+    return _data_dir(cfg) / "reservations.lock"
 
 
-def _alternation_path(cfg: dict) -> Path:
-    return _data_dir(cfg) / "alternation_counter.txt"
-
-
-def _alternation_lock_path(cfg: dict) -> Path:
-    return _data_dir(cfg) / "alternation_counter.lock"
-
-
-# ── Counter helpers ───────────────────────────────────────────────────────────
-
-def _read_counter(path: Path) -> int:
-    if not path.exists():
-        return 0
-    return int(path.read_text().strip() or "0")
-
-
-def _write_counter(path: Path, value: int) -> None:
-    path.write_text(str(value))
+def _local_completions_path(category: str, cfg: dict) -> Path:
+    """
+    Local file tracking completed item counts this container session.
+    Updated immediately on each completion so subsequent assignments
+    see accurate counts without waiting for an HF re-scan.
+    Reset on container restart — HF is the durable source of truth.
+    """
+    return _data_dir(cfg) / f"local_completions_{cfg['study_type']}_{category}.json"
 
 
 # ── Dataset download + normalisation ─────────────────────────────────────────
@@ -78,10 +84,6 @@ def _download_and_cache(
     hf_token: str,
     data_dir: str,
 ) -> None:
-    """
-    Download from HuggingFace, select POOL_SIZE items reproducibly, cache to disk.
-    No-op if the pool file already exists.
-    """
     pool_path = Path(data_dir) / f"pool_{study_type}_{category}.json"
     if pool_path.exists():
         print(f"[DATA] Pool already cached: {pool_path}")
@@ -96,14 +98,11 @@ def _download_and_cache(
     ds = load_dataset(repo_id, token=token_arg, trust_remote_code=True)
 
     if study_type == "preference":
-        # Preference repos: rows have pair_id, category, product_a, product_b, split.
-        # Use the "test" split key when available; otherwise filter by split=="test".
         if "test" in ds:
             rows = [dict(r) for r in ds["test"]]
         else:
             rows = [dict(r) for r in ds["train"] if r.get("split") == "test"]
     else:
-        # Likelihood repos: use test split if present, otherwise first available split.
         split_key = "test" if "test" in ds else list(ds.keys())[0]
         rows = [dict(r) for r in ds[split_key]]
 
@@ -112,7 +111,6 @@ def _download_and_cache(
     selected = rows[:POOL_SIZE]
 
     if study_type == "likelihood":
-        # Normalise: extract the metadata dict, add a stable item_id and category.
         normalised = []
         for i, row in enumerate(selected):
             meta = row["metadata"]
@@ -125,8 +123,6 @@ def _download_and_cache(
             normalised.append(meta)
         selected = normalised
     else:
-        # Preference: deep-copy so nested product dicts are plain Python dicts,
-        # and stamp each product with its category.
         cleaned = []
         for row in selected:
             r = dict(row)
@@ -145,7 +141,6 @@ def _download_and_cache(
 
 
 def ensure_datasets(cfg: dict) -> None:
-    """Download and cache all category pools needed for this study config."""
     for cat_cfg in cfg["categories"]:
         _download_and_cache(
             study_type=cfg["study_type"],
@@ -162,113 +157,320 @@ def _load_pool(pool_path_str: str) -> list:
         return json.load(f)
 
 
-# ── Round-robin assignment ────────────────────────────────────────────────────
+# ── Accepted completion counts ────────────────────────────────────────────────
+
+def _get_accepted_counts(category: str, cfg: dict) -> dict:
+    """
+    Return how many times each pool item has been accepted.
+
+    Sources (merged, highest count wins):
+    1. Local completions file — written immediately on each completion this session.
+    2. HF output repo scan — authoritative after a container restart.
+       Results cached for COMPLETION_CACHE_TTL seconds.
+
+    Rejected submissions live under rejected/ and are NOT counted.
+    """
+    pool   = _load_pool(str(_pool_path(category, cfg)))
+    counts = {str(i): 0 for i in range(len(pool))}
+
+    # ── Source 1: local completions (most up-to-date within this session) ────
+    local_path = _local_completions_path(category, cfg)
+    if local_path.exists():
+        try:
+            with open(local_path) as f:
+                local = json.load(f)
+            for k, v in local.items():
+                counts[k] = max(counts.get(k, 0), v)
+            print(f"[ASSIGN] Local completions for {category}: "
+                  f"{sum(1 for v in local.values() if v > 0)} items completed")
+        except Exception as e:
+            print(f"[ASSIGN] Could not read local completions: {e}")
+
+    # ── Source 2: HF scan (authoritative after restart, with 5-min cache) ───
+    cache_path = _data_dir(cfg) / f"completion_cache_{cfg['study_type']}_{category}.json"
+    now        = time.time()
+    hf_counts  = None
+
+    if cache_path.exists():
+        try:
+            with open(cache_path) as f:
+                cache = json.load(f)
+            if now - cache.get("timestamp", 0) < COMPLETION_CACHE_TTL:
+                hf_counts = cache["counts"]
+        except Exception:
+            pass
+
+    if hf_counts is None:
+        hf_counts   = {str(i): 0 for i in range(len(pool))}
+        hf_token    = cfg.get("hf_token", "")
+        output_repo = cfg.get("output_dataset_repo", "")
+        if hf_token and output_repo:
+            try:
+                from huggingface_hub import HfApi
+                api        = HfApi(token=hf_token)
+                files      = list(api.list_repo_files(repo_id=output_repo, repo_type="dataset"))
+                json_files = [f for f in files if f.startswith("json/") and f.endswith(".json")]
+                for filepath in json_files:
+                    try:
+                        content = api.hf_hub_download(
+                            repo_id=output_repo,
+                            filename=filepath,
+                            repo_type="dataset",
+                            token=hf_token,
+                        )
+                        with open(content) as f:
+                            submission = json.load(f)
+                        for item in submission.get("items", []):
+                            if item.get("category") != category:
+                                continue
+                            idx = item.get("_pool_index")
+                            if idx is not None:
+                                hf_counts[str(idx)] = hf_counts.get(str(idx), 0) + 1
+                    except Exception as e:
+                        print(f"[ASSIGN] Could not parse {filepath}: {e}")
+            except Exception as e:
+                print(f"[ASSIGN] Could not scan HF repo: {e}")
+        try:
+            with open(cache_path, "w") as f:
+                json.dump({"timestamp": now, "counts": hf_counts}, f)
+        except Exception:
+            pass
+
+    for k, v in hf_counts.items():
+        counts[k] = max(counts.get(k, 0), v)
+
+    return counts
+
+
+# ── Reservation management ────────────────────────────────────────────────────
+
+def _load_reservations(cfg: dict) -> dict:
+    path = _reservation_path(cfg)
+    if not path.exists():
+        return {}
+    try:
+        with open(path) as f:
+            return json.load(f)
+    except Exception:
+        return {}
+
+
+def _save_reservations(reservations: dict, cfg: dict) -> None:
+    with open(_reservation_path(cfg), "w") as f:
+        json.dump(reservations, f)
+
+
+def _expire_reservations(reservations: dict) -> dict:
+    now     = time.time()
+    expired = [k for k, v in reservations.items() if v["expiry"] < now]
+    for k in expired:
+        print(f"[ASSIGN] Reservation expired for item index {k}")
+        del reservations[k]
+    return reservations
+
+
+def release_reservation(user_id: str, cfg: dict) -> None:
+    """Release all reservations held by this user immediately after completion."""
+    lock = FileLock(str(_reservation_lock_path(cfg)), timeout=10)
+    with lock:
+        reservations = _load_reservations(cfg)
+        _expire_reservations(reservations)
+        released = [k for k, v in reservations.items() if v["user_id"] == user_id]
+        for k in released:
+            del reservations[k]
+        _save_reservations(reservations, cfg)
+        print(f"[ASSIGN] Released {len(released)} reservations for user {user_id}")
+
+
+def record_completion(user_id: str, items: list, cfg: dict) -> None:
+    """
+    Record completed item indices to the local completions file immediately.
+    Uses _pool_index stamped on each item at assignment time — no fuzzy matching.
+    Called after successful HF upload AND by the simulation script.
+    """
+    # Group by category using the stamped _pool_category and _pool_index
+    by_category: dict = {}
+    for item in items:
+        cat = item.get("_pool_category") or item.get("category", "")
+        idx = item.get("_pool_index")
+        if idx is None:
+            print(f"[ASSIGN] WARNING: item missing _pool_index, skipping: "
+                  f"{item.get('pair_id') or item.get('item_id', '?')}")
+            continue
+        by_category.setdefault(cat, []).append(idx)
+
+    for cat, indices in by_category.items():
+        pool             = _load_pool(str(_pool_path(cat, cfg)))
+        completions_path = _local_completions_path(cat, cfg)
+
+        if completions_path.exists():
+            try:
+                with open(completions_path) as f:
+                    completions = json.load(f)
+            except Exception:
+                completions = {str(i): 0 for i in range(len(pool))}
+        else:
+            completions = {str(i): 0 for i in range(len(pool))}
+
+        for idx in indices:
+            completions[str(idx)] = completions.get(str(idx), 0) + 1
+
+        with open(completions_path, "w") as f:
+            json.dump(completions, f)
+
+        # Invalidate HF cache so next scan re-reads fresh
+        cache_path = _data_dir(cfg) / f"completion_cache_{cfg['study_type']}_{cat}.json"
+        if cache_path.exists():
+            try:
+                cache_path.unlink()
+            except Exception:
+                pass
+
+        print(f"[ASSIGN] Recorded completions for {cat}: indices {indices} "
+              f"(user {user_id[:8]})")
+
+
+# ── Core assignment ───────────────────────────────────────────────────────────
+
+def _assign_from_category(category: str, n: int, user_id: str, cfg: dict) -> list:
+    """
+    Assign n items using least-coverage-first strategy.
+
+    Priority order:
+      1. Uncovered + unreserved         (count=0, not reserved)
+      2. Uncovered + reserved by other  (count=0, reserved — likely abandoned user)
+      3. Covered   + unreserved         (count>0, not reserved)
+      4. Covered   + reserved by other  (count>0, reserved)
 
-def _assign_from_category(category: str, n: int, cfg: dict) -> list:
-    """Atomically assign n items from the pool, wrapping around when exhausted."""
-    pool  = _load_pool(str(_pool_path(category, cfg)))
-    total = len(pool)
-    lock  = FileLock(str(_counter_lock_path(category, cfg)))
+    This ensures abandoned users' items get picked up by subsequent users
+    rather than already-covered items being re-assigned.
+    """
+    pool            = _load_pool(str(_pool_path(category, cfg)))
+    accepted_counts = _get_accepted_counts(category, cfg)
+    lock            = FileLock(str(_reservation_lock_path(cfg)), timeout=10)
 
     with lock:
-        ctr      = _read_counter(_counter_path(category, cfg))
-        assigned = [pool[(ctr + i) % total] for i in range(n)]
-        _write_counter(_counter_path(category, cfg), ctr + n)
+        reservations = _load_reservations(cfg)
+        _expire_reservations(reservations)
+
+        def is_reserved_by_other(i):
+            r = reservations.get(str(i))
+            return r is not None and r["user_id"] != user_id
+
+        def sort_key(i):
+            count    = accepted_counts.get(str(i), 0)
+            reserved = int(is_reserved_by_other(i))
+            return (count, reserved)
+
+        # All indices sorted by (count, is_reserved_by_other)
+        all_indices = sorted(range(len(pool)), key=sort_key)
+        selected_indices = all_indices[:n]
+
+        # Reserve selected items (overrides stale reservations from abandoned users)
+        expiry = time.time() + RESERVATION_TTL
+        for i in selected_indices:
+            reservations[str(i)] = {"user_id": user_id, "expiry": expiry}
+
+        _save_reservations(reservations, cfg)
+
+    selected = []
+    for i in selected_indices:
+        item = dict(pool[i])
+        item["_pool_index"]    = i
+        item["_pool_category"] = category
+        selected.append(item)
 
-    return assigned
+    print(f"[ASSIGN] {category}: assigned indices {selected_indices} "
+          f"(counts: {[accepted_counts.get(str(i), 0) for i in selected_indices]})")
+    return selected
+
+
+# ── Variant assignment ────────────────────────────────────────────────────────
 
 def _assign_variants(cfg: dict, n: int) -> list:
     variants = cfg.get("model_variants")
     if not variants:
-        # Old-style config with top-level model_name/prompt_variant
-        return [{"model_name": cfg["model_name"], "prompt_variant": cfg["prompt_variant"]}] * n
+        return [{"name": "default",
+                 "model_name":     cfg["model_name"],
+                 "prompt_variant": cfg["prompt_variant"]}] * n
 
     if len(variants) == 1:
-        # Single variant — assign it to all items, no alternation needed
         return [variants[0]] * n
 
-    # Two-variant alternation
-    lock = FileLock(str(_data_dir(cfg) / "variant_counter.lock"))
+    lock = FileLock(str(_data_dir(cfg) / "variant_counter.lock"), timeout=10)
     with lock:
-        ctr = _read_counter(_data_dir(cfg) / "variant_counter.txt")
-        _write_counter(_data_dir(cfg) / "variant_counter.txt", ctr + 1)
+        counter_path = _data_dir(cfg) / "variant_counter.txt"
+        ctr = int(counter_path.read_text().strip()) if counter_path.exists() else 0
+        counter_path.write_text(str(ctr + 1))
 
     v0, v1 = variants[0], variants[1]
     if ctr % 2 == 1:
         v0, v1 = v1, v0
 
-    assigned = [v0] * v0["count"] + [v1] * v1["count"]
-    random.shuffle(assigned)
-    print(f"[VARIANTS] user {ctr}: {[v['name'] for v in assigned]}")
-    return assigned
-    
+    from itertools import zip_longest
+    interleaved = []
+    for a, b in zip_longest([v0] * v0["count"], [v1] * v1["count"]):
+        if a: interleaved.append(a)
+        if b: interleaved.append(b)
+
+    print(f"[VARIANTS] user {ctr}: {[v['name'] for v in interleaved]}")
+    return interleaved
+
+
+# ── Category count computation ────────────────────────────────────────────────
 
 def _compute_counts(cfg: dict) -> dict:
-    """
-    Determine how many items to assign from each category for one user.
-
-    Single category → all pairs_per_user go to that category.
-    Two categories  → use the configured counts but swap them on every other
-                      call so the cumulative totals stay balanced:
-                        user 1: movies=3, groceries=2
-                        user 2: movies=2, groceries=3
-                        user 3: movies=3, groceries=2 …
-    """
     cats = cfg["categories"]
     n    = cfg["pairs_per_user"]
 
     if len(cats) == 1:
         return {cats[0]["name"]: n}
 
-    # Alternation counter increments by exactly 1 per user (separate from item counters)
-    lock = FileLock(str(_alternation_lock_path(cfg)))
+    lock = FileLock(str(_data_dir(cfg) / "alternation_counter.lock"), timeout=10)
     with lock:
-        call_count = _read_counter(_alternation_path(cfg))
-        _write_counter(_alternation_path(cfg), call_count + 1)
+        path = _data_dir(cfg) / "alternation_counter.txt"
+        ctr  = int(path.read_text().strip()) if path.exists() else 0
+        path.write_text(str(ctr + 1))
 
     base = {c["name"]: c["count"] for c in cats}
-
-    # Sanity-check: if configured counts don't sum to pairs_per_user, split evenly
     if sum(base.values()) != n:
         base = {}
         for i, c in enumerate(cats):
             base[c["name"]] = n // len(cats) + (1 if i < n % len(cats) else 0)
         return base
 
-    # On odd calls swap the two counts
-    if call_count % 2 == 1:
+    if ctr % 2 == 1:
         names = [c["name"] for c in cats]
         base[names[0]], base[names[1]] = base[names[1]], base[names[0]]
 
     return base
 
 
-def assign_items(cfg: dict) -> list:
-    """Assign a full set of items for one participant, interleaved across categories."""
+def assign_items(cfg: dict, user_id: str) -> list:
     counts = _compute_counts(cfg)
     items  = []
     for cat_name, n in counts.items():
-        items.extend(_assign_from_category(cat_name, n, cfg))
-    random.shuffle(items)   # interleave so participant doesn't see all of one category first
+        items.extend(_assign_from_category(cat_name, n, user_id, cfg))
+    random.shuffle(items)
     return items
 
 
-# ── Session-state construction ────────────────────────────────────────────────
+# ── Item slot construction ────────────────────────────────────────────────────
 
 def _make_item_slot(item: dict, study_type: str) -> dict:
-    """Create a blank result slot for one item/pair, ready to be filled during the study."""
     base = {
-        # conversation stores ALL turns including the two synthetic opening turns.
-        # num_turns counts only real human exchanges (not the synthetic ones).
+        # Preserve pool index and category for record_completion in upload.py
+        "_pool_index":    item.get("_pool_index"),
+        "_pool_category": item.get("_pool_category", item.get("category", "")),
         "conversation": {
             "system_prompt":   "",
-            "closing_message": "",  # vote_final_message equivalent — stored but never shown
+            "closing_message": "",
             "turns":           [],
             "num_turns":       0,
         },
         "reflection":   {},
-        "pre_rating":   None,   # 1–7 int set on the intro screen
-        "post_rating":  None,   # 1–7 int set on the post_rating screen
+        "pre_rating":   None,
+        "post_rating":  None,
         "rating_delta": None,
     }
     if study_type == "preference":
@@ -290,30 +492,34 @@ def _make_item_slot(item: dict, study_type: str) -> dict:
     return base
 
 
+# ── Session-state construction ────────────────────────────────────────────────
+
 def init_state(cfg: dict) -> dict:
     """Build the initial session-state dict for a new participant."""
     n        = cfg["pairs_per_user"]
+    user_id  = str(uuid.uuid4())
     variants = _assign_variants(cfg, n)
-    items    = assign_items(cfg)[:n]
- 
+    items    = assign_items(cfg, user_id)[:n]
+
     slots = [_make_item_slot(it, cfg["study_type"]) for it in items]
     for slot, variant in zip(slots, variants):
         slot["model_name"]     = variant["model_name"]
         slot["prompt_variant"] = variant["prompt_variant"]
- 
+
     for i, slot in enumerate(slots):
         print(f"[ITEM {i}] category={slot.get('category')} "
+              f"pool_index={slot.get('_pool_index')} "
               f"model={slot.get('model_name')} "
-              f"personalization={slot.get('prompt_variant',{}).get('personalization')}")
- 
+              f"personalization={slot.get('prompt_variant', {}).get('personalization')}")
+
     try:
         params = st.query_params
     except Exception:
         params = {}
- 
+
     return {
         "submission_id": str(uuid.uuid4()),
-        "user_id":       str(uuid.uuid4()),
+        "user_id":       user_id,
         "prolific_pid":  params.get("PROLIFIC_PID", ""),
         "study_id":      params.get("STUDY_ID",     ""),
         "session_id":    params.get("SESSION_ID",   ""),

src/upload.py

@@ -9,6 +9,7 @@ from pathlib import Path
 
 import streamlit as st
 from huggingface_hub import HfApi
+from src.data import release_reservation, record_completion
 
 
 @st.cache_resource
@@ -37,6 +38,10 @@ def save_and_upload(state: dict, cfg: dict) -> None:
     submission_id = state.get("submission_id", str(uuid.uuid4()))
     safe_worker   = "".join(c if c.isalnum() else "_" for c in str(worker_id))
 
+    print(f"[SAVE] starting save_and_upload")
+    print(f"[SAVE] output_repo={output_repo}")
+    print(f"[SAVE] hf_token set={bool(hf_token)}")
+
     # ── Write JSON ────────────────────────────────────────────────────────────
     ann_dir   = Path(cfg["annotations_dir"]) / safe_worker
     ann_dir.mkdir(parents=True, exist_ok=True)
@@ -46,6 +51,7 @@ def save_and_upload(state: dict, cfg: dict) -> None:
         json.dump(state, f, indent=2)
     print(f"[SAVE] JSON written: {json_path}")
 
+    uploaded = False
     if hf_token:
         try:
             hf_api.upload_file(
@@ -55,9 +61,17 @@ def save_and_upload(state: dict, cfg: dict) -> None:
                 repo_type="dataset",
             )
             print("[HF] JSON uploaded.")
+            uploaded = True
         except Exception as e:
             print(f"[HF] JSON upload error: {e}")
 
+    if uploaded:
+        # Release reservations so items are immediately available for re-assignment
+        release_reservation(state.get("user_id", ""), cfg)
+        # Record completion locally — updates counts immediately without waiting
+        # for an HF re-scan. Also invalidates the HF cache.
+        record_completion(state.get("user_id", ""), state.get("items", []), cfg)
+
     # ── Write + upload CSV ────────────────────────────────────────────────────
     _save_and_upload_csv(state, cfg, hf_api, safe_worker, submission_id)
 
@@ -135,7 +149,6 @@ def _save_and_upload_csv(
             pv.get("detailed_instruction", True),
             cfg.get("pair_selection_seed", 42),
             item.get("category", ""),
-            # Demographics (14 fields)
             demographics.get("age",                  ""),
             demographics.get("gender",               ""),
             demographics.get("geographic_region",    ""),
@@ -150,19 +163,15 @@ def _save_and_upload_csv(
             demographics.get("political_views",      ""),
             demographics.get("household_size",       ""),
             demographics.get("employment_status",    ""),
-            # Background (6 fixed keys; empty string when category not in study)
             background.get("movies_criteria",    ""),
             background.get("movies_enjoy",        ""),
             background.get("movies_avoid",        ""),
             background.get("groceries_criteria",  ""),
             background.get("groceries_enjoy",     ""),
             background.get("groceries_avoid",     ""),
-            # Ratings
             pre, post, delta,
-            # Conversation — full turn list as JSON string
             conv.get("num_turns", 0),
             json.dumps(conv.get("turns", [])),
-            # Reflection
             refl.get("standout_moment", ""),
             refl.get("thinking_change",  ""),
         ]
@@ -188,7 +197,6 @@ def _save_and_upload_csv(
 
         rows.append(common + extra)
 
-    # Write temp CSV and upload
     timestamp  = datetime.now().strftime("%Y%m%d_%H%M%S")
     unique_tag = uuid.uuid4().hex[:8]
     repo_path  = f"csv/{timestamp}_{safe_worker}_{unique_tag}.csv"