feat: per-annotator document assignments with configurable overlap

- annotator_config.yaml: defines annotators and overlap % (default 10%)
- generate_assignments.py: auto-distributes docs with overlap, seed-based
- Documents API reads config, filters docs by logged-in user
- 10 shared docs for inter-annotator agreement
- Fallback: show all docs if user not in config

app/api/documents/route.js

@@ -1,7 +1,43 @@
 import { HF_DATASET_BASE_URL, MAX_DOCS_TO_SCAN } from '../../../utils/config.js';
+import yaml from 'js-yaml';
 
-export async function GET() {
+/**
+ * Fetch annotator_config.yaml and return the doc list for a given user.
+ * Returns null if no config or user not found (show all docs).
+ */
+async function getUserAssignedDocs(username) {
+    if (!username) return null;
+
+    try {
+        const configUrl = `${HF_DATASET_BASE_URL}/raw/main/annotation_data/annotator_config.yaml`;
+        const res = await fetch(configUrl, {
+            headers: { 'Authorization': `Bearer ${process.env.HF_TOKEN}` },
+            next: { revalidate: 300 } // cache 5 min
+        });
+        if (!res.ok) return null;
+
+        const text = await res.text();
+        const config = yaml.load(text);
+
+        const annotator = (config.annotators || []).find(a => a.username === username);
+        if (!annotator || !annotator.docs || annotator.docs.length === 0) return null;
+
+        return new Set(annotator.docs);
+    } catch (e) {
+        console.warn('Could not load annotator_config.yaml:', e.message);
+        return null;
+    }
+}
+
+export async function GET(request) {
     try {
+        // Get username from query param
+        const { searchParams } = new URL(request.url);
+        const username = searchParams.get('user');
+
+        // Fetch user's assigned docs (if configured)
+        const assignedDocs = await getUserAssignedDocs(username);
+
         // Fetch the index file from HF Datasets
         const linksUrl = `${HF_DATASET_BASE_URL}/raw/main/annotation_data/wbg_data/wbg_pdf_links.json`;
         const linksRes = await fetch(linksUrl, {
@@ -21,12 +57,17 @@ export async function GET() {
 
         const links = await linksRes.json();
 
-        // Filter to docs with revalidation data and take the first N
-        const successLinks = links
-            .filter(l => l.status === 'success' && l.has_revalidation === true)
-            .slice(0, MAX_DOCS_TO_SCAN);
+        // Filter to docs with revalidation data, then by user assignment if available
+        let successLinks = links
+            .filter(l => l.status === 'success' && l.has_revalidation === true);
+
+        if (assignedDocs) {
+            successLinks = successLinks.filter(l => assignedDocs.has(l.index));
+        }
+
+        successLinks = successLinks.slice(0, MAX_DOCS_TO_SCAN);
 
-        // Parallel fetch — much faster than sequential scanning
+        // Parallel fetch
         const results = await Promise.allSettled(
             successLinks.map(async (link) => {
                 const docUrl = `${HF_DATASET_BASE_URL}/raw/main/annotation_data/wbg_extractions/doc_${link.index}/raw/doc_${link.index}_direct_judged.jsonl`;
@@ -44,7 +85,7 @@ export async function GET() {
                 if (annotatablePages.length === 0) return null;
 
                 const pdfUrl = link.direct_pdf_url;
-                if (!pdfUrl) return null; // no PDF URL at all
+                if (!pdfUrl) return null;
 
                 return {
                     index: link.index,
@@ -63,7 +104,7 @@ export async function GET() {
             status: 200,
             headers: {
                 'Content-Type': 'application/json',
-                'Cache-Control': 'public, s-maxage=3600, stale-while-revalidate=59'
+                'Cache-Control': 'public, s-maxage=300, stale-while-revalidate=59'
             }
         });
     } catch (error) {

app/page.js

@@ -48,9 +48,14 @@ export default function Home() {
     const annotatablePages = currentDoc?.annotatable_pages ?? [];
     const currentPageNumber = annotatablePages[pageIdx] ?? null;
 
-    // Load documents on mount
+    // Load documents (re-fetches when annotatorName changes to get user-specific assignment)
     useEffect(() => {
-        fetch('/api/documents')
+        setLoading(true);
+        const url = annotatorName
+            ? `/api/documents?user=${encodeURIComponent(annotatorName)}`
+            : '/api/documents';
+
+        fetch(url)
             .then(res => res.json())
             .then(data => {
                 setDocuments(data);
@@ -75,7 +80,7 @@ export default function Home() {
                 console.error("Failed to load documents", err);
                 setLoading(false);
             });
-    }, []);
+    }, [annotatorName]);
 
     // Read HF OAuth cookie for annotator identity
     useEffect(() => {

generate_assignments.py

@@ -0,0 +1,156 @@
+#!/usr/bin/env python3
+"""
+generate_assignments.py
+
+Reads annotator_config.yaml, distributes available docs across annotators
+with configurable overlap, and writes back the updated config.
+
+Usage:
+    python3 generate_assignments.py                 # Generate and save
+    python3 generate_assignments.py --dry-run        # Preview only
+    python3 generate_assignments.py --upload          # Upload config to HF
+
+Requires: pyyaml, huggingface_hub (for --upload)
+"""
+
+import argparse
+import json
+import random
+import sys
+from pathlib import Path
+
+try:
+    import yaml
+except ImportError:
+    print("❌ pyyaml required: uv pip install pyyaml")
+    sys.exit(1)
+
+CONFIG_PATH = Path(__file__).parent / "annotation_data" / "annotator_config.yaml"
+LINKS_PATH = Path(__file__).parent / "annotation_data" / "wbg_data" / "wbg_pdf_links.json"
+
+
+def load_config():
+    return yaml.safe_load(CONFIG_PATH.read_text())
+
+
+def save_config(config):
+    CONFIG_PATH.write_text(yaml.dump(config, default_flow_style=False, sort_keys=False))
+
+
+def get_available_docs():
+    """Get list of active English doc indices."""
+    links = json.loads(LINKS_PATH.read_text())
+    return sorted([
+        l["index"] for l in links
+        if l.get("has_revalidation") and l.get("status") == "success"
+    ])
+
+
+def generate_assignments(config, seed=42):
+    """Distribute docs across annotators with overlap."""
+    settings = config.get("settings", {})
+    overlap_pct = settings.get("overlap_percent", 10)
+    annotators = config.get("annotators", [])
+
+    if not annotators:
+        print("❌ No annotators defined in config.")
+        return config
+
+    all_docs = get_available_docs()
+    n_docs = len(all_docs)
+    n_annotators = len(annotators)
+
+    # Calculate overlap
+    n_overlap = max(1, round(n_docs * overlap_pct / 100))
+
+    # Shuffle docs deterministically
+    rng = random.Random(seed)
+    shuffled = all_docs.copy()
+    rng.shuffle(shuffled)
+
+    # Pick overlap docs (shared by ALL annotators)
+    overlap_docs = sorted(shuffled[:n_overlap])
+    remaining = shuffled[n_overlap:]
+
+    # Split remaining docs evenly across annotators
+    per_annotator = len(remaining) // n_annotators
+    extra = len(remaining) % n_annotators
+
+    print(f"\n📊 Assignment Summary:")
+    print(f"  Total docs:       {n_docs}")
+    print(f"  Annotators:       {n_annotators}")
+    print(f"  Overlap ({overlap_pct}%):   {n_overlap} docs shared by all")
+    print(f"  Per annotator:    ~{per_annotator + n_overlap} docs each")
+    print(f"  Overlap docs:     {overlap_docs}")
+    print()
+
+    start = 0
+    for i, ann in enumerate(annotators):
+        # Distribute remaining: first `extra` annotators get 1 more
+        count = per_annotator + (1 if i < extra else 0)
+        exclusive = sorted(remaining[start:start + count])
+        start += count
+
+        ann["docs"] = sorted(overlap_docs + exclusive)
+        print(f"  {ann['username']}: {len(ann['docs'])} docs "
+              f"({n_overlap} overlap + {len(exclusive)} exclusive)")
+
+    return config
+
+
+def upload_config():
+    """Upload annotator_config.yaml to HF."""
+    try:
+        from huggingface_hub import HfApi
+        import os
+
+        token = os.environ.get("HF_TOKEN")
+        if not token:
+            env_path = Path(__file__).parent / ".env"
+            if env_path.exists():
+                for line in env_path.read_text().splitlines():
+                    if line.startswith("HF_TOKEN="):
+                        token = line.split("=", 1)[1].strip()
+
+        if not token:
+            print("❌ No HF_TOKEN found.")
+            return
+
+        api = HfApi(token=token)
+        api.upload_file(
+            path_or_fileobj=str(CONFIG_PATH),
+            path_in_repo="annotation_data/annotator_config.yaml",
+            repo_id="ai4data/annotation_data",
+            repo_type="dataset",
+            commit_message="Update annotator assignments",
+        )
+        print("✅ Uploaded annotator_config.yaml to HF")
+    except ImportError:
+        print("❌ huggingface_hub required: uv pip install huggingface_hub")
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Generate document assignments")
+    parser.add_argument("--dry-run", action="store_true", help="Preview only")
+    parser.add_argument("--upload", action="store_true", help="Upload config to HF")
+    parser.add_argument("--seed", type=int, default=42, help="Random seed")
+    args = parser.parse_args()
+
+    config = load_config()
+    config = generate_assignments(config, seed=args.seed)
+
+    if args.dry_run:
+        print("\n[DRY RUN] Would save:")
+        print(yaml.dump(config, default_flow_style=False, sort_keys=False))
+    else:
+        save_config(config)
+        print(f"\n💾 Saved to {CONFIG_PATH}")
+
+        if args.upload:
+            upload_config()
+
+    print("\n✅ Done!")
+
+
+if __name__ == "__main__":
+    main()

package-lock.json

@@ -9,6 +9,7 @@
       "version": "1.0.0",
       "dependencies": {
         "@huggingface/hub": "^2.10.3",
+        "js-yaml": "^4.1.1",
         "next": "14.2.14",
         "react": "^18",
         "react-dom": "^18",
@@ -286,6 +287,12 @@
         "node": ">=8"
       }
     },
+    "node_modules/argparse": {
+      "version": "2.0.1",
+      "resolved": "https://registry.npmjs.org/argparse/-/argparse-2.0.1.tgz",
+      "integrity": "sha512-8+9WqebbFzpX9OR+Wa6O29asIogeRMzcGtAINdpMHHyAg10f05aSFVBbcEqGf/PXw1EjAZ+q2/bEBg3DvurK3Q==",
+      "license": "Python-2.0"
+    },
     "node_modules/bail": {
       "version": "2.0.2",
       "resolved": "https://registry.npmjs.org/bail/-/bail-2.0.2.tgz",
@@ -634,6 +641,18 @@
       "integrity": "sha512-RdJUflcE3cUzKiMqQgsCu06FPu9UdIJO0beYbPhHN4k6apgJtifcoCtT9bcxOpYBtpD2kCM6Sbzg4CausW/PKQ==",
       "license": "MIT"
     },
+    "node_modules/js-yaml": {
+      "version": "4.1.1",
+      "resolved": "https://registry.npmjs.org/js-yaml/-/js-yaml-4.1.1.tgz",
+      "integrity": "sha512-qQKT4zQxXl8lLwBtHMWwaTcGfFOZviOJet3Oy/xmGk2gZH677CJM9EvtfdSkgWcATZhj/55JZ0rmy3myCT5lsA==",
+      "license": "MIT",
+      "dependencies": {
+        "argparse": "^2.0.1"
+      },
+      "bin": {
+        "js-yaml": "bin/js-yaml.js"
+      }
+    },
     "node_modules/longest-streak": {
       "version": "3.1.0",
       "resolved": "https://registry.npmjs.org/longest-streak/-/longest-streak-3.1.0.tgz",

package.json

@@ -9,6 +9,7 @@
   },
   "dependencies": {
     "@huggingface/hub": "^2.10.3",
+    "js-yaml": "^4.1.1",
     "next": "14.2.14",
     "react": "^18",
     "react-dom": "^18",