lspcloud/prolific-preferences-personalized
0
1"""2Dataset download, item-pool caching, completion-aware assignment, and session-state init.3 4Assignment strategy5-------------------6Items are assigned based on how many *accepted* completions they already have,7ensuring the least-covered items are always prioritised.8 9Each assigned item is stamped with _pool_index and _pool_category at assignment10time so record_completion never needs to do a fuzzy pair_id match — it reads11the index directly.12 13Accepted completions = JSON files under json/ in the output repo.14Rejected completions = JSON files moved to rejected/ by the admin.15 → moving a file to rejected/ automatically makes that item available again.16 17Reservations18------------19When a user starts, their items are "reserved" in a local file for 80 min.20Concurrent users each get a FileLock on the reservation file so they21never receive the same items. Reservations expire automatically so abandoned22sessions don't permanently block items.23 24Each reservation stores the user's prolific_pid so we can release their items25immediately when Prolific reports them as RETURNED or TIMED-OUT — no need to26wait for the 80-min TTL.27 28Dropout / rejection recovery29-----------------------------30- Dropout (voluntary return): Prolific marks RETURNED, we query the API and31 release the reservation on the next assignment.32- Dropout (silent): reservation expires after 80 min → item re-enters pool.33- Rejection: admin moves json/{worker}/{id}.json → rejected/{worker}/{id}.json34 in the HF dataset repo. On next Space restart (or cache expiry) the item's35 accepted count drops to 0 and it gets re-assigned.36"""37import json38import random39import time40import uuid41from pathlib import Path42 43import streamlit as st44from filelock import FileLock45 46from src.config import CATEGORY_TO_REPO47 48POOL_SIZE = 50 # items selected per (study_type, category)49RESERVATION_TTL = 60 * 80 # 80 min: 30 min expected + ~2.5x buffer50COMPLETION_CACHE_TTL = 300 # re-scan HF repo every 5 minutes51PROLIFIC_POLL_CACHE_TTL = 120 # re-poll Prolific every 2 minutes52 53 54# ── Path helpers ──────────────────────────────────────────────────────────────55 56def _data_dir(cfg: dict) -> Path:57 p = Path(cfg["data_dir"])58 p.mkdir(parents=True, exist_ok=True)59 return p60 61 62def _pool_path(category: str, cfg: dict) -> Path:63 return _data_dir(cfg) / f"pool_{cfg['study_type']}_{category}.json"64 65 66def _reservation_path(cfg: dict) -> Path:67 return _data_dir(cfg) / "reservations.json"68 69 70def _reservation_lock_path(cfg: dict) -> Path:71 return _data_dir(cfg) / "reservations.lock"72 73 74def _local_completions_path(category: str, cfg: dict) -> Path:75 """76 Local file tracking completed item counts this container session.77 Updated immediately on each completion so subsequent assignments78 see accurate counts without waiting for an HF re-scan.79 Reset on container restart — HF is the durable source of truth.80 """81 return _data_dir(cfg) / f"local_completions_{cfg['study_type']}_{category}.json"82 83 84# ── Dataset download + normalisation ─────────────────────────────────────────85 86@st.cache_resource87def _download_and_cache(88 study_type: str,89 category: str,90 seed: int,91 hf_token: str,92 data_dir: str,93) -> None:94 pool_path = Path(data_dir) / f"pool_{study_type}_{category}.json"95 if pool_path.exists():96 print(f"[DATA] Pool already cached: {pool_path}")97 return98 99 from datasets import load_dataset100 101 repo_id = CATEGORY_TO_REPO[(study_type, category)]102 token_arg = hf_token or None103 print(f"[DATA] Downloading {repo_id} …")104 105 ds = load_dataset(repo_id, token=token_arg, trust_remote_code=True)106 107 if study_type == "preference":108 if "test" in ds:109 rows = [dict(r) for r in ds["test"]]110 else:111 rows = [dict(r) for r in ds["train"] if r.get("split") == "test"]112 else:113 split_key = "test" if "test" in ds else list(ds.keys())[0]114 rows = [dict(r) for r in ds[split_key]]115 116 rng = random.Random(seed)117 rng.shuffle(rows)118 selected = rows[:POOL_SIZE]119 120 if study_type == "likelihood":121 normalised = []122 for i, row in enumerate(selected):123 meta = row["metadata"]124 if isinstance(meta, str):125 meta = json.loads(meta)126 else:127 meta = dict(meta)128 meta["item_id"] = str(uuid.uuid5(uuid.NAMESPACE_DNS, f"{repo_id}_{i}_{seed}"))129 meta["category"] = category130 normalised.append(meta)131 selected = normalised132 else:133 cleaned = []134 for row in selected:135 r = dict(row)136 r["product_a"] = dict(r["product_a"])137 r["product_b"] = dict(r["product_b"])138 r["product_a"].setdefault("category", r.get("category", category))139 r["product_b"].setdefault("category", r.get("category", category))140 cleaned.append(r)141 selected = cleaned142 143 pool_path.parent.mkdir(parents=True, exist_ok=True)144 with open(pool_path, "w") as f:145 json.dump(selected, f, indent=2)146 147 print(f"[DATA] {study_type}/{category}: cached {len(selected)} items (seed={seed}).")148 149 150def ensure_datasets(cfg: dict) -> None:151 for cat_cfg in cfg["categories"]:152 _download_and_cache(153 study_type=cfg["study_type"],154 category=cat_cfg["name"],155 seed=cfg["pair_selection_seed"],156 hf_token=cfg.get("hf_token", ""),157 data_dir=cfg["data_dir"],158 )159 160 161@st.cache_data162def _load_pool(pool_path_str: str) -> list:163 with open(pool_path_str) as f:164 return json.load(f)165 166 167# ── Accepted completion counts ────────────────────────────────────────────────168 169def _get_accepted_counts(category: str, cfg: dict) -> dict:170 """171 Return how many times each pool item has been accepted.172 173 Sources (merged, highest count wins):174 1. Local completions file — written immediately on each completion this session.175 2. HF output repo scan — authoritative after a container restart.176 Results cached for COMPLETION_CACHE_TTL seconds.177 178 Rejected submissions live under rejected/ and are NOT counted.179 """180 pool = _load_pool(str(_pool_path(category, cfg)))181 counts = {str(i): 0 for i in range(len(pool))}182 183 # ── Source 1: local completions (most up-to-date within this session) ────184 local_path = _local_completions_path(category, cfg)185 if local_path.exists():186 try:187 with open(local_path) as f:188 local = json.load(f)189 for k, v in local.items():190 counts[k] = max(counts.get(k, 0), v)191 print(f"[ASSIGN] Local completions for {category}: "192 f"{sum(1 for v in local.values() if v > 0)} items completed")193 except Exception as e:194 print(f"[ASSIGN] Could not read local completions: {e}")195 196 # ── Source 2: HF scan (authoritative after restart, with 5-min cache) ───197 cache_path = _data_dir(cfg) / f"completion_cache_{cfg['study_type']}_{category}.json"198 now = time.time()199 hf_counts = None200 201 if cache_path.exists():202 try:203 with open(cache_path) as f:204 cache = json.load(f)205 if now - cache.get("timestamp", 0) < COMPLETION_CACHE_TTL:206 hf_counts = cache["counts"]207 except Exception:208 pass209 210 if hf_counts is None:211 hf_counts = {str(i): 0 for i in range(len(pool))}212 hf_token = cfg.get("hf_token", "")213 output_repo = cfg.get("output_dataset_repo", "")214 if hf_token and output_repo:215 try:216 from huggingface_hub import HfApi217 api = HfApi(token=hf_token)218 files = list(api.list_repo_files(repo_id=output_repo, repo_type="dataset"))219 json_files = [f for f in files if f.startswith("json/") and f.endswith(".json")]220 221 # Build pair_id → pool_index lookup for fallback matching222 id_to_index = {}223 for i, p in enumerate(pool):224 pid = p.get("pair_id") or p.get("item_id", "")225 if pid:226 id_to_index[pid] = i227 228 for filepath in json_files:229 try:230 content = api.hf_hub_download(231 repo_id=output_repo,232 filename=filepath,233 repo_type="dataset",234 token=hf_token,235 )236 with open(content) as f:237 submission = json.load(f)238 for item in submission.get("items", []):239 if item.get("category") != category:240 continue241 idx = item.get("_pool_index")242 if idx is None:243 pid = item.get("pair_id") or item.get("item_id", "")244 idx = id_to_index.get(pid)245 if idx is not None:246 hf_counts[str(idx)] = hf_counts.get(str(idx), 0) + 1247 except Exception as e:248 print(f"[ASSIGN] Could not parse {filepath}: {e}")249 except Exception as e:250 print(f"[ASSIGN] Could not scan HF repo: {e}")251 try:252 with open(cache_path, "w") as f:253 json.dump({"timestamp": now, "counts": hf_counts}, f)254 except Exception:255 pass256 257 for k, v in hf_counts.items():258 counts[k] = max(counts.get(k, 0), v)259 260 return counts261 262 263# ── Reservation management ────────────────────────────────────────────────────264 265def _load_reservations(cfg: dict) -> dict:266 path = _reservation_path(cfg)267 if not path.exists():268 return {}269 try:270 with open(path) as f:271 return json.load(f)272 except Exception:273 return {}274 275 276def _save_reservations(reservations: dict, cfg: dict) -> None:277 with open(_reservation_path(cfg), "w") as f:278 json.dump(reservations, f)279 280 281def _expire_reservations(reservations: dict) -> dict:282 now = time.time()283 expired = [k for k, v in reservations.items() if v["expiry"] < now]284 for k in expired:285 print(f"[ASSIGN] Reservation expired for item index {k}")286 del reservations[k]287 return reservations288 289 290def release_reservation(user_id: str, cfg: dict) -> None:291 """Release all reservations held by this user immediately after completion."""292 lock = FileLock(str(_reservation_lock_path(cfg)), timeout=10)293 with lock:294 reservations = _load_reservations(cfg)295 _expire_reservations(reservations)296 released = [k for k, v in reservations.items() if v["user_id"] == user_id]297 for k in released:298 del reservations[k]299 _save_reservations(reservations, cfg)300 print(f"[ASSIGN] Released {len(released)} reservations for user {user_id}")301 302 303def record_completion(user_id: str, items: list, cfg: dict) -> None:304 """305 Record completed item indices to the local completions file immediately.306 Uses _pool_index stamped on each item at assignment time — no fuzzy matching.307 Called after successful HF upload AND by the simulation script.308 """309 by_category: dict = {}310 for item in items:311 cat = item.get("_pool_category") or item.get("category", "")312 idx = item.get("_pool_index")313 if idx is None:314 print(f"[ASSIGN] WARNING: item missing _pool_index, skipping: "315 f"{item.get('pair_id') or item.get('item_id', '?')}")316 continue317 by_category.setdefault(cat, []).append(idx)318 319 for cat, indices in by_category.items():320 pool = _load_pool(str(_pool_path(cat, cfg)))321 completions_path = _local_completions_path(cat, cfg)322 323 if completions_path.exists():324 try:325 with open(completions_path) as f:326 completions = json.load(f)327 except Exception:328 completions = {str(i): 0 for i in range(len(pool))}329 else:330 completions = {str(i): 0 for i in range(len(pool))}331 332 for idx in indices:333 completions[str(idx)] = completions.get(str(idx), 0) + 1334 335 with open(completions_path, "w") as f:336 json.dump(completions, f)337 338 # Invalidate HF cache so next scan re-reads fresh339 cache_path = _data_dir(cfg) / f"completion_cache_{cfg['study_type']}_{cat}.json"340 if cache_path.exists():341 try:342 cache_path.unlink()343 except Exception:344 pass345 346 print(f"[ASSIGN] Recorded completions for {cat}: indices {indices} "347 f"(user {user_id[:8]})")348 349 350# ── Prolific status polling ───────────────────────────────────────────────────351 352def _prolific_returned_pids(cfg: dict) -> set:353 """354 Query Prolific for participants who have RETURNED or TIMED-OUT from the355 active study. Returns a set of their PIDs. Cached for PROLIFIC_POLL_CACHE_TTL.356 """357 token = cfg.get("prolific_api_token", "")358 study_id = cfg.get("prolific_study_id", "")359 if not token or not study_id:360 return set()361 362 cache_path = _data_dir(cfg) / "prolific_returned_cache.json"363 now = time.time()364 365 if cache_path.exists():366 try:367 with open(cache_path) as f:368 c = json.load(f)369 if now - c.get("timestamp", 0) < PROLIFIC_POLL_CACHE_TTL:370 return set(c.get("returned_pids", []))371 except Exception:372 pass373 374 returned = set()375 try:376 import requests377 url = f"https://api.prolific.com/api/v1/studies/{study_id}/submissions/"378 headers = {"Authorization": f"Token {token}"}379 resp = requests.get(url, headers=headers, timeout=10)380 resp.raise_for_status()381 for sub in resp.json().get("results", []):382 status = sub.get("status", "")383 if status in ("RETURNED", "TIMED-OUT", "TIMED_OUT"):384 pid = sub.get("participant_id") or sub.get("participant", "")385 if pid:386 returned.add(pid)387 print(f"[PROLIFIC] Found {len(returned)} returned/timed-out participants")388 except Exception as e:389 print(f"[PROLIFIC] Could not query API: {e}")390 391 try:392 with open(cache_path, "w") as f:393 json.dump({"timestamp": now, "returned_pids": list(returned)}, f)394 except Exception:395 pass396 397 return returned398 399 400def _release_returned_reservations(reservations: dict, cfg: dict) -> None:401 """402 Remove reservations held by Prolific participants who have RETURNED or403 TIMED-OUT. Mutates the reservations dict in place.404 """405 returned_pids = _prolific_returned_pids(cfg)406 if not returned_pids:407 return408 409 released = []410 for idx, r in list(reservations.items()):411 pid = r.get("prolific_pid", "")412 if pid and pid in returned_pids:413 released.append(idx)414 del reservations[idx]415 if released:416 print(f"[ASSIGN] Released {len(released)} reservations from returned/timed-out participants: {released}")417 418 419def all_items_covered(cfg: dict) -> bool:420 """421 Returns True if every item in every category has been accepted at least once.422 Used for auto-pausing the Prolific study.423 """424 for cat_cfg in cfg["categories"]:425 cat = cat_cfg["name"]426 pool = _load_pool(str(_pool_path(cat, cfg)))427 counts = _get_accepted_counts(cat, cfg)428 for i in range(len(pool)):429 if counts.get(str(i), 0) < 1:430 return False431 return True432 433 434def pause_prolific_study(cfg: dict) -> bool:435 """436 Call Prolific's API to pause the study. Returns True on success.437 Requires prolific_api_token (env PROLIFIC_API_TOKEN) and prolific_study_id.438 Idempotent — safe to call multiple times (Prolific treats repeated pauses as no-ops).439 """440 token = cfg.get("prolific_api_token", "")441 study_id = cfg.get("prolific_study_id", "")442 if not token or not study_id:443 print("[PROLIFIC] Cannot auto-pause: no API token or study_id configured")444 return False445 446 # Idempotency marker so we don't spam the API on every completion after447 # the first time all items are covered.448 paused_marker = _data_dir(cfg) / ".prolific_paused"449 if paused_marker.exists():450 return True451 452 try:453 import requests454 url = f"https://api.prolific.com/api/v1/studies/{study_id}/transition/"455 headers = {"Authorization": f"Token {token}", "Content-Type": "application/json"}456 resp = requests.post(url, headers=headers, json={"action": "PAUSE"}, timeout=10)457 resp.raise_for_status()458 paused_marker.touch()459 print(f"[PROLIFIC] ✅ Study {study_id} paused automatically — all items covered.")460 return True461 except Exception as e:462 print(f"[PROLIFIC] Could not auto-pause study: {e}")463 return False464 465 466# ── Core assignment ───────────────────────────────────────────────────────────467 468def _assign_from_category(category: str, n: int, user_id: str, cfg: dict) -> list:469 """470 Assign n items using least-coverage-first strategy.471 472 Priority order (via sort key):473 1. Uncovered + unreserved (count=0, not reserved)474 2. Uncovered + reserved by other (count=0, reserved)475 3. Covered + unreserved (count>0, not reserved)476 4. Covered + reserved by other (count>0, reserved)477 478 Reservations are ONLY created for participants who come via Prolific479 (i.e. have a non-empty prolific_pid in the URL). Non-Prolific visitors480 (testers, previewers, direct-URL visitors) still get items assigned so481 they can run through the study, but they don't hold reservations.482 483 Reservations from participants who have RETURNED/TIMED-OUT on Prolific484 are released BEFORE the sort, so their items are treated as unreserved.485 """486 pool = _load_pool(str(_pool_path(category, cfg)))487 accepted_counts = _get_accepted_counts(category, cfg)488 lock = FileLock(str(_reservation_lock_path(cfg)), timeout=10)489 490 # Capture prolific_pid early so we can decide whether to reserve.491 # Read from query_params directly — session_state.study_state doesn't492 # exist yet during init_state, which is what calls this function.493 prolific_pid = ""494 try:495 params = st.query_params496 prolific_pid = params.get("PROLIFIC_PID", "") or ""497 except Exception:498 pass499 is_prolific = bool(prolific_pid)500 501 with lock:502 reservations = _load_reservations(cfg)503 _expire_reservations(reservations)504 _release_returned_reservations(reservations, cfg)505 506 # If this Prolific PID already has reservations (e.g. they refreshed507 # the tab, got a new user_id, and came back), release the old ones508 # before creating new ones. Prevents the same participant from509 # accumulating multiple reservations.510 if is_prolific:511 stale = [512 idx for idx, r in list(reservations.items())513 if r.get("prolific_pid") == prolific_pid514 ]515 for idx in stale:516 del reservations[idx]517 if stale:518 print(f"[ASSIGN] Released {len(stale)} prior reservations "519 f"for returning PID {prolific_pid}")520 521 def is_reserved_by_other(i):522 r = reservations.get(str(i))523 return r is not None and r["user_id"] != user_id524 525 def sort_key(i):526 count = accepted_counts.get(str(i), 0)527 reserved = int(is_reserved_by_other(i))528 return (count, reserved)529 530 all_indices = sorted(range(len(pool)), key=sort_key)531 selected_indices = all_indices[:n]532 533 # Only reserve if this is a Prolific participant — keeps the534 # admin "in progress" count accurate and stops testers/bouncers535 # from blocking items for real users.536 if is_prolific:537 expiry = time.time() + RESERVATION_TTL538 for i in selected_indices:539 reservations[str(i)] = {540 "user_id": user_id,541 "prolific_pid": prolific_pid,542 "expiry": expiry,543 }544 _save_reservations(reservations, cfg)545 print(f"[ASSIGN] Reserved for Prolific PID {prolific_pid}")546 else:547 print(f"[ASSIGN] Non-Prolific visitor — no reservation created")548 549 selected = []550 for i in selected_indices:551 item = dict(pool[i])552 item["_pool_index"] = i553 item["_pool_category"] = category554 selected.append(item)555 556 print(f"[ASSIGN] {category}: assigned indices {selected_indices} "557 f"(counts: {[accepted_counts.get(str(i), 0) for i in selected_indices]})")558 return selected559 560 561# ── Variant assignment ────────────────────────────────────────────────────────562 563def _assign_variants(cfg: dict, n: int) -> list:564 variants = cfg.get("model_variants")565 if not variants:566 return [{"name": "default",567 "model_name": cfg["model_name"],568 "prompt_variant": cfg["prompt_variant"]}] * n569 570 if len(variants) == 1:571 return [variants[0]] * n572 573 lock = FileLock(str(_data_dir(cfg) / "variant_counter.lock"), timeout=10)574 with lock:575 counter_path = _data_dir(cfg) / "variant_counter.txt"576 ctr = int(counter_path.read_text().strip()) if counter_path.exists() else 0577 counter_path.write_text(str(ctr + 1))578 579 v0, v1 = variants[0], variants[1]580 if ctr % 2 == 1:581 v0, v1 = v1, v0582 583 from itertools import zip_longest584 interleaved = []585 for a, b in zip_longest([v0] * v0["count"], [v1] * v1["count"]):586 if a: interleaved.append(a)587 if b: interleaved.append(b)588 589 print(f"[VARIANTS] user {ctr}: {[v['name'] for v in interleaved]}")590 return interleaved591 592 593# ── Category count computation ────────────────────────────────────────────────594 595def _compute_counts(cfg: dict) -> dict:596 cats = cfg["categories"]597 n = cfg["pairs_per_user"]598 599 if len(cats) == 1:600 return {cats[0]["name"]: n}601 602 lock = FileLock(str(_data_dir(cfg) / "alternation_counter.lock"), timeout=10)603 with lock:604 path = _data_dir(cfg) / "alternation_counter.txt"605 ctr = int(path.read_text().strip()) if path.exists() else 0606 path.write_text(str(ctr + 1))607 608 base = {c["name"]: c["count"] for c in cats}609 if sum(base.values()) != n:610 base = {}611 for i, c in enumerate(cats):612 base[c["name"]] = n // len(cats) + (1 if i < n % len(cats) else 0)613 return base614 615 if ctr % 2 == 1:616 names = [c["name"] for c in cats]617 base[names[0]], base[names[1]] = base[names[1]], base[names[0]]618 619 return base620 621 622def assign_items(cfg: dict, user_id: str) -> list:623 counts = _compute_counts(cfg)624 items = []625 for cat_name, n in counts.items():626 items.extend(_assign_from_category(cat_name, n, user_id, cfg))627 random.shuffle(items)628 return items629 630 631# ── Item slot construction ────────────────────────────────────────────────────632 633def _make_item_slot(item: dict, study_type: str) -> dict:634 base = {635 "_pool_index": item.get("_pool_index"),636 "_pool_category": item.get("_pool_category", item.get("category", "")),637 "conversation": {638 "system_prompt": "",639 "closing_message": "",640 "turns": [],641 "num_turns": 0,642 },643 "reflection": {},644 "pre_rating": None,645 "post_rating": None,646 "rating_delta": None,647 }648 if study_type == "preference":649 base.update({650 "pair_id": item.get("pair_id", str(uuid.uuid4())),651 "category": item.get("category", ""),652 "product_a": item.get("product_a", {}),653 "product_b": item.get("product_b", {}),654 "familiarity_a": None,655 "familiarity_b": None,656 })657 else:658 base.update({659 "item_id": item.get("item_id", str(uuid.uuid4())),660 "category": item.get("category", ""),661 "product": item,662 "familiarity": None,663 })664 return base665 666 667# ── Session-state construction ────────────────────────────────────────────────668 669def init_state(cfg: dict) -> dict:670 """Build the initial session-state dict for a new participant."""671 n = cfg["pairs_per_user"]672 user_id = str(uuid.uuid4())673 variants = _assign_variants(cfg, n)674 items = assign_items(cfg, user_id)[:n]675 676 slots = [_make_item_slot(it, cfg["study_type"]) for it in items]677 for slot, variant in zip(slots, variants):678 slot["model_name"] = variant["model_name"]679 slot["prompt_variant"] = variant["prompt_variant"]680 slot["sampler_path"] = variant.get("sampler_path", "")681 682 for i, slot in enumerate(slots):683 print(f"[ITEM {i}] category={slot.get('category')} "684 f"pool_index={slot.get('_pool_index')} "685 f"model={slot.get('model_name')} "686 f"personalization={slot.get('prompt_variant', {}).get('personalization')}")687 688 try:689 params = st.query_params690 except Exception:691 params = {}692 693 return {694 "submission_id": str(uuid.uuid4()),695 "user_id": user_id,696 "prolific_pid": params.get("PROLIFIC_PID", ""),697 "study_id": params.get("STUDY_ID", ""),698 "session_id": params.get("SESSION_ID", ""),699 "start_time": time.time(),700 "study_type": cfg["study_type"],701 "demographics": {},702 "background": {},703 "items": slots,704 "current_index": 0,705 "screen": "welcome",706 "meta": {},707 }