Spaces:
Sleeping
Sleeping
| """Fixed-group assignment and legacy balanced assignment: pure logic, | |
| no FastAPI/HF imports. | |
| Everything here takes plain data in and returns plain data out, so it can be | |
| unit-tested without touching the network. The I/O layer (reading/writing | |
| annotation and assignment records from the HF dataset repo, locking) lives | |
| in main.py and calls into these functions. | |
| Legacy assignment settings are env-overridable, matching the existing HF_* | |
| variable convention in main.py. Fixed-group size and slots per group are | |
| function parameters, defaulting to at most 50 videos and 3 participants. | |
| Saved groups are reused rather than rebuilt when these defaults change. | |
| """ | |
| import os | |
| import re | |
| import secrets | |
| import string | |
| from datetime import datetime | |
| # Legacy round configuration; new groups default to at most 50 videos and 3 slots. | |
| VIDEOS_PER_ANNOTATOR = int(os.getenv("VIDEOS_PER_ANNOTATOR", "1000")) | |
| TARGET_ANNOTATIONS_PER_VIDEO = int( | |
| os.getenv("TARGET_ANNOTATIONS_PER_VIDEO", "3") | |
| ) | |
| # Abandonment TTL for both legacy reservations and fixed-group assignments | |
| # that have not saved any answers. | |
| RESERVATION_TTL_SECONDS = int(os.getenv("ASSIGNMENT_RESERVATION_TTL_SECONDS", str(24 * 3600))) | |
| COMPLETION_CODE_ALPHABET = string.ascii_uppercase + string.digits | |
| COMPLETION_CODE_LENGTH = 12 | |
| def build_allocation_state(catalog_video_ids, group_size=50, slots_per_group=3): | |
| """Build initial allocation_state.json data: groups and their member slots. | |
| Uses actual resolved video IDs rather than a configured catalog size. | |
| Call only when no saved allocation state exists; reuse it thereafter. | |
| Members map annotator IDs to in_progress/completed status. | |
| Existing assignments continue using the legacy functions below. | |
| """ | |
| for name, value in (("group_size", group_size), ("slots_per_group", slots_per_group)): | |
| if type(value) is not int or value <= 0: | |
| raise ValueError(f"{name} must be a positive integer") | |
| ids = list(catalog_video_ids) | |
| if not ids or any(not isinstance(v, str) or not v.strip() for v in ids): | |
| raise ValueError("catalog must contain non-empty video IDs") | |
| if len(ids) != len(set(ids)): | |
| raise ValueError("catalog contains duplicate video IDs") | |
| group_count = (len(ids) + group_size - 1) // group_size | |
| smaller_size, larger_group_count = divmod(len(ids), group_count) | |
| group_sizes = [smaller_size + 1] * larger_group_count | |
| group_sizes += [smaller_size] * (group_count - larger_group_count) | |
| groups = [] | |
| start = 0 | |
| for index, size in enumerate(group_sizes, start=1): | |
| groups.append({ | |
| "group_id": f"group_{index:04d}", | |
| "video_ids": ids[start:start + size], | |
| "members": {}, | |
| }) | |
| start += size | |
| return { | |
| "schema_version": 1, | |
| "group_size": group_size, | |
| "slots_per_group": slots_per_group, | |
| "total_videos": len(ids), | |
| "groups": groups, | |
| } | |
| def build_workload_allocation_state(group_specs, slots_per_group=2): | |
| """Build persisted variable-size groups for the workload-duration pilot.""" | |
| if not isinstance(group_specs, list) or not group_specs: | |
| raise ValueError("group_specs must be a non-empty list") | |
| if type(slots_per_group) is not int or slots_per_group <= 0: | |
| raise ValueError("slots_per_group must be a positive integer") | |
| groups = [] | |
| all_ids = [] | |
| group_ids = set() | |
| for spec in group_specs: | |
| if not isinstance(spec, dict): | |
| raise ValueError("each group spec must be an object") | |
| group_id = spec.get("group_id") | |
| video_ids = spec.get("video_ids") | |
| if (not isinstance(group_id, str) or not group_id | |
| or not isinstance(video_ids, list) or not video_ids | |
| or any(not isinstance(v, str) or not v.strip() for v in video_ids)): | |
| raise ValueError("each group must have an ID and non-empty video IDs") | |
| if group_id in group_ids: | |
| raise ValueError("workload group IDs must be unique") | |
| group_ids.add(group_id) | |
| if len(video_ids) != len(set(video_ids)): | |
| raise ValueError(f"group {group_id} contains duplicate video IDs") | |
| all_ids.extend(video_ids) | |
| groups.append({ | |
| "group_id": group_id, | |
| "video_ids": list(video_ids), | |
| "members": {}, | |
| }) | |
| if len(all_ids) != len(set(all_ids)): | |
| raise ValueError("a video cannot appear in more than one workload group") | |
| return { | |
| "schema_version": 1, | |
| "allocation_mode": "workload_test", | |
| "group_size": None, | |
| "group_sizes": [len(group["video_ids"]) for group in groups], | |
| "slots_per_group": slots_per_group, | |
| "total_videos": len(all_ids), | |
| "groups": groups, | |
| } | |
| def pick_available_group(state, rng): | |
| """Choose a group ID from saved state, or None if every group is full. | |
| Both in-progress and completed members occupy slots. Does not mutate | |
| state or claim a slot; the I/O caller must persist claims atomically | |
| and re-read/recompute after a commit conflict. | |
| """ | |
| available = [ | |
| group for group in state["groups"] | |
| if len(group["members"]) + group.get("legacy_reserved_slots", 0) < state["slots_per_group"] | |
| ] | |
| if not available: | |
| return None | |
| lowest = min(len(group["members"]) + group.get("legacy_reserved_slots", 0) for group in available) | |
| return rng.choice([ | |
| group["group_id"] for group in available | |
| if len(group["members"]) + group.get("legacy_reserved_slots", 0) == lowest | |
| ]) | |
| def generate_completion_code(): | |
| """Cryptographically random completion code, format T2AV-XXXXXXXXXXXX. | |
| Uses `secrets`, not `random` - this is a credential the annotator submits | |
| elsewhere (e.g. an MTurk HIT) as proof of completion, not a UI cosmetic.""" | |
| suffix = "".join(secrets.choice(COMPLETION_CODE_ALPHABET) for _ in range(COMPLETION_CODE_LENGTH)) | |
| return f"T2AV-{suffix}" | |
| def build_user_state(round_record, latest_annotations=None, completion=None): | |
| """Build users/{annotator_id}.json from an assignment and latest records. | |
| latest_annotations maps video IDs to metadata (id, created_at, | |
| annotation_path), never full answer bodies. This also supports migration | |
| of existing rounds; group_id is None for legacy assignments. | |
| """ | |
| import copy | |
| latest = copy.deepcopy(latest_annotations or {}) | |
| video_ids = set(round_record["video_ids"]) | |
| if not set(latest) <= video_ids: | |
| raise ValueError("annotation video is outside the current assignment") | |
| for record in latest.values(): | |
| if not all(record.get(field) for field in ("id", "created_at", "annotation_path")): | |
| raise ValueError("latest annotation metadata is incomplete") | |
| if completion is not None: | |
| if (completion["annotator_id"] != round_record["annotator_id"] | |
| or completion["assignment_id"] != round_record["assignment_id"]): | |
| raise ValueError("completion belongs to a different assignment") | |
| if not video_ids or set(latest) != video_ids: | |
| raise ValueError("cannot complete a round with unfinished videos") | |
| return { | |
| "schema_version": 1, | |
| "annotator_id": round_record["annotator_id"], | |
| "group_id": round_record.get("group_id"), | |
| "assignment": copy.deepcopy(round_record), | |
| "completed_video_ids": sorted(latest), | |
| "latest_annotations": latest, | |
| "status": "completed" if completion is not None else "in_progress", | |
| "completion_code": completion["completion_code"] if completion is not None else None, | |
| "completed_at": completion["completed_at"] if completion is not None else None, | |
| } | |
| def normalize_participant_id(value): | |
| """Trim input without silently replacing invalid characters or merging IDs. | |
| Keep legacy canonical names usable; this is syntax validation, not | |
| verification of a Prolific account. | |
| """ | |
| value = value.strip() if isinstance(value, str) else "" | |
| if not re.fullmatch(r"[A-Za-z0-9_][A-Za-z0-9_.-]{0,127}", value): | |
| raise ValueError("Enter your Prolific ID without spaces or special characters") | |
| return value.lower() | |
| def slugify(value): | |
| """Canonical annotator_id from a raw display name. Lowercased (unlike | |
| main.py's filename slugify) so "Adi" and "adi" resolve to one identity - | |
| this is the identity key used for completion/assignment lookups, not a | |
| filename, so case-insensitivity is the more correct default here.""" | |
| slug = re.sub(r"[^A-Za-z0-9_.-]+", "-", value or "").strip("-").lower() | |
| return slug or "unknown" | |
| def parse_iso(value): | |
| return datetime.fromisoformat(value.replace("Z", "+00:00")) | |
| def fixed_assignment_is_abandoned(state, now, ttl_seconds=RESERVATION_TTL_SECONDS): | |
| """Return whether an untouched fixed-group assignment may release its slot. | |
| Any saved answer preserves the assignment indefinitely. Completed and | |
| legacy assignments are never expired by this rule. | |
| """ | |
| if (not state or state.get("group_id") is None | |
| or state.get("status") != "in_progress" | |
| or state.get("completed_video_ids") | |
| or state.get("latest_annotations")): | |
| return False | |
| try: | |
| created_at = parse_iso(state["assignment"]["created_at"]) | |
| except (KeyError, ValueError, AttributeError, TypeError): | |
| return False | |
| return (now - created_at).total_seconds() > ttl_seconds | |
| def dedupe_latest(annotation_records): | |
| """A "completed annotation" is one row per unique (annotator_id, video_id) | |
| pair - repeated saves/autosaves/updates from the same annotator for the | |
| same video collapse into a single record, keeping the most recent | |
| created_at. Records missing a required field are skipped (never crash on | |
| older/malformed data) and reported separately. | |
| Returns (deduped: dict[(annotator_id, video_id)] -> record, skipped: int). | |
| """ | |
| latest = {} | |
| skipped = 0 | |
| for record in annotation_records: | |
| try: | |
| video_id = record["video_id"] | |
| user_raw = record["user"] | |
| created_at = record["created_at"] | |
| if not video_id or not user_raw or not created_at: | |
| raise ValueError("empty required field") | |
| parse_iso(created_at) # validate shape before trusting string comparisons below | |
| except (KeyError, TypeError, ValueError, AttributeError): | |
| skipped += 1 | |
| continue | |
| key = (slugify(user_raw), video_id) | |
| existing = latest.get(key) | |
| if existing is None or created_at > existing["created_at"]: | |
| latest[key] = record | |
| return latest, skipped | |
| def completed_by_video(deduped): | |
| result = {} | |
| for annotator_id, video_id in deduped: | |
| result.setdefault(video_id, set()).add(annotator_id) | |
| return result | |
| def completed_by_annotator(deduped): | |
| result = {} | |
| for annotator_id, video_id in deduped: | |
| result.setdefault(annotator_id, set()).add(video_id) | |
| return result | |
| def reserved_by_video(assignment_records, completed_by_annotator_map, now, ttl_seconds, exclude_annotator=None): | |
| """Open (not-yet-completed, not-yet-expired) reservations per video, | |
| across every OTHER annotator's assignment. A reservation stops counting | |
| once its assignment is older than ttl_seconds ("abandoned") - this never | |
| deletes or mutates the stored assignment record, it only stops | |
| contributing to the live coverage snapshot used for balancing.""" | |
| result = {} | |
| for record in assignment_records: | |
| annotator_id = record.get("annotator_id") | |
| if not annotator_id or annotator_id == exclude_annotator: | |
| continue | |
| try: | |
| created_at = parse_iso(record["created_at"]) | |
| except (KeyError, ValueError, AttributeError, TypeError): | |
| continue | |
| if (now - created_at).total_seconds() > ttl_seconds: | |
| continue | |
| completed = completed_by_annotator_map.get(annotator_id, set()) | |
| for video_id in record.get("video_ids", []): | |
| if video_id in completed: | |
| continue | |
| result[video_id] = result.get(video_id, 0) + 1 | |
| return result | |
| def pick_balanced(candidates, coverage_fn, count, rng, target=TARGET_ANNOTATIONS_PER_VIDEO): | |
| """Picks up to `count` videos, lowest coverage (completed + reserved) | |
| first, random tie-break among equal coverage. A video already at | |
| `target` is never picked. Returns fewer than `count` when fewer | |
| under-target videos remain.""" | |
| if count <= 0: | |
| return [] | |
| under_target = [v for v in candidates if coverage_fn(v) < target] | |
| ordered = sorted(under_target, key=lambda video_id: (coverage_fn(video_id), rng.random())) | |
| return ordered[:count] | |
| def build_round( | |
| annotator_id, | |
| annotator_raw, | |
| catalog_video_ids, | |
| already_completed, | |
| completed_by_video_map, | |
| reserved_by_video_map, | |
| rng, | |
| round_number, | |
| videos_per_annotator=VIDEOS_PER_ANNOTATOR, | |
| target=TARGET_ANNOTATIONS_PER_VIDEO, | |
| ): | |
| """Builds a new round: up to `videos_per_annotator` videos this | |
| annotator has never completed before. Videos already completed are | |
| excluded. Videos already at the coverage target are excluded too, so a | |
| short round is returned instead of giving a video a fourth rater.""" | |
| candidates = [v for v in catalog_video_ids if v not in already_completed] | |
| def coverage(video_id): | |
| return len(completed_by_video_map.get(video_id, set())) + reserved_by_video_map.get(video_id, 0) | |
| video_ids = pick_balanced(candidates, coverage, videos_per_annotator, rng, target) | |
| return { | |
| "annotator_id": annotator_id, | |
| "annotator_name_raw": annotator_raw, | |
| "round_number": round_number, | |
| "video_ids": video_ids, | |
| "requested_size": videos_per_annotator, | |
| "actual_size": len(video_ids), | |
| } | |
| def summarize( | |
| catalog_video_ids, | |
| completed_by_video_map, | |
| reserved_by_video_map, | |
| total_completed_preserved, | |
| skipped_records, | |
| duplicate_records_deduped, | |
| assignment_count, | |
| target=TARGET_ANNOTATIONS_PER_VIDEO, | |
| ): | |
| below = sum(1 for v in catalog_video_ids if len(completed_by_video_map.get(v, set())) < target) | |
| return { | |
| "total_videos": len(catalog_video_ids), | |
| "target_per_video": target, | |
| "videos_per_annotator": VIDEOS_PER_ANNOTATOR, | |
| "completed_counts": {v: len(completed_by_video_map.get(v, set())) for v in catalog_video_ids}, | |
| "reserved_counts": {v: reserved_by_video_map.get(v, 0) for v in catalog_video_ids}, | |
| "videos_below_target": below, | |
| "videos_at_or_above_target": len(catalog_video_ids) - below, | |
| "total_existing_annotations_preserved": total_completed_preserved, | |
| "skipped_malformed_records": skipped_records, | |
| "duplicate_records_deduplicated": duplicate_records_deduped, | |
| "total_assignments": assignment_count, | |
| } | |