Spaces:

bearda
/

croissantminer

Sleeping

Update to paper systems
394445e verified about 1 month ago
41.1 kB
#!/usr/bin/env python3
"""CroissantMiner Agentic V2 pipeline.
Phases:
  0. PDF parse (PyPDF2 + canonical clean_text + UTF-8 sanitize).
  1. Load Gemini Flash triage (from data/agentic/phase1/<ds_id>.json).
  2. Full-context extraction that also asks for evidence quotes per field.
  2.5 Self-correction pass over fields that are null but triage says present.
  3. Cross-document enrichment (HuggingFace + Semantic Scholar) for gaps.
  4. Schema validation + confidence scoring.
  5. Write canonical output to data/extractions/agentic_v2_<backbone_slug>/.
Output schema (canonical):
  {"dataset_id", "model", "extraction" (flat 30 fields), "usage", "valid",
   "_meta": {provenance, pipeline="agentic_v2", extraction_details: {field:
             {confidence, status, source, evidence}}, paper_stats, ...}}
Usage:
  python scripts/agentic_v2.py --backbone sonnet-4-5
  python scripts/agentic_v2.py --backbone gpt-5.4 --dev-only
  python scripts/agentic_v2.py --backbone sonnet-4-5 --batch   # Anthropic only
"""

from __future__ import annotations

import argparse
import json
import re
import time
import urllib.parse
import urllib.request
from pathlib import Path

from dotenv import load_dotenv

from _agentic_helpers import (
    ROOT,
    SYSTEM_PROMPT,
    CANONICAL_FIELDS,
    call_llm,
    estimate_cost,
    fill_canonical,
    get_paper_text,
    meta_block,
    normalize_field_names,
    output_dir_for,
    parse_json_response,
    resolve_backbone,
    strip_references,
)

load_dotenv(ROOT / ".env")

MAX_PAPER_CHARS = 200_000
MAX_TOKENS_EXTRACT = 16384  # was 8192; bumped to clear Gemini 3.x truncation failures (2026-04-30)
MAX_TOKENS_CORRECT = 8192   # was 4096; bumped proportionally
SCRIPT_NAME = "scripts/agentic_v2.py"

PHASE1_DIR = ROOT / "data" / "agentic" / "phase1"

with open(ROOT / "data" / "agentic" / "dev_test_split.json") as f:
    SPLIT = json.load(f)


# ── Field groups (used by triage + self-correction) ───────────────
EASY_FIELDS = [
    "name", "description", "url", "license", "creator",
    "publisher", "datePublished", "inLanguage", "citeAs", "isLiveDataset",
]
HARD_FIELDS = [
    "rai:dataCollection", "rai:dataCollectionType", "rai:dataCollectionMissingData",
    "rai:dataCollectionRawData", "rai:dataCollectionTimeframe",
    "rai:dataImputationProtocol", "rai:dataManipulationProtocol",
    "rai:dataPreprocessingProtocol", "rai:dataAnnotationProtocol",
    "rai:dataAnnotationPlatform", "rai:dataAnnotationAnalysis",
    "rai:annotationsPerItem", "rai:annotatorDemographics",
    "rai:machineAnnotationTools", "rai:dataReleaseMaintenancePlan",
    "rai:personalSensitiveInformation", "rai:dataSocialImpact",
    "rai:dataBiases", "rai:dataLimitations", "rai:dataUseCases",
]
ALL_FIELDS = EASY_FIELDS + HARD_FIELDS

TRIAGE_GROUPS = {
    "G3_collection": [
        "rai:dataCollection", "rai:dataCollectionType",
        "rai:dataCollectionMissingData", "rai:dataCollectionRawData",
        "rai:dataCollectionTimeframe",
    ],
    "G4_annotation": [
        "rai:dataAnnotationProtocol", "rai:dataAnnotationPlatform",
        "rai:dataAnnotationAnalysis", "rai:annotationsPerItem",
        "rai:annotatorDemographics", "rai:machineAnnotationTools",
    ],
    "G5_rai": [
        "rai:dataBiases", "rai:dataLimitations", "rai:dataSocialImpact",
        "rai:personalSensitiveInformation", "rai:dataUseCases",
        "rai:dataReleaseMaintenancePlan",
    ],
    "G6_processing": [
        "rai:dataImputationProtocol", "rai:dataManipulationProtocol",
        "rai:dataPreprocessingProtocol",
    ],
}


# ── V2-specific prompt (asks for {value, evidence} per field) ─────
USER_PROMPT_WITH_EVIDENCE = '''Extract metadata from the following academic paper by matching text to the field descriptions above.
For each non-null field, also provide a brief "evidence" quote: the key sentence(s) from the paper that support your extraction. This helps verify accuracy.
Your output must conform exactly to the following schema. For each field, provide "value" and "evidence":
SCHEMA:
{
  "name": {"value": "string", "evidence": "quote from paper"},
  "description": {"value": "string", "evidence": "quote from paper"},
  "url": {"value": "string", "evidence": "quote from paper"},
  "license": {"value": "string or null", "evidence": "quote or null"},
  "creator": {"value": "string", "evidence": "quote from paper"},
  "publisher": {"value": "string or null", "evidence": "quote or null"},
  "datePublished": {"value": "string (YYYY or YYYY-MM-DD)", "evidence": "quote from paper"},
  "inLanguage": {"value": "string (ISO codes)", "evidence": "quote or null"},
  "citeAs": {"value": "string or null", "evidence": "quote or null"},
  "isLiveDataset": {"value": "Yes/No or null", "evidence": "quote or null"},
  "rai:dataCollection": {"value": "string or null", "evidence": "quote or null"},
  "rai:dataCollectionType": {"value": "string (from recommended values) or null", "evidence": "quote or null"},
  "rai:dataCollectionMissingData": {"value": "string or null", "evidence": "quote or null"},
  "rai:dataCollectionRawData": {"value": "string or null", "evidence": "quote or null"},
  "rai:dataCollectionTimeframe": {"value": "string or null", "evidence": "quote or null"},
  "rai:dataImputationProtocol": {"value": "string or null", "evidence": "quote or null"},
  "rai:dataManipulationProtocol": {"value": "string or null", "evidence": "quote or null"},
  "rai:dataPreprocessingProtocol": {"value": "string or null", "evidence": "quote or null"},
  "rai:dataAnnotationProtocol": {"value": "string or null", "evidence": "quote or null"},
  "rai:dataAnnotationPlatform": {"value": "string or null", "evidence": "quote or null"},
  "rai:dataAnnotationAnalysis": {"value": "string or null", "evidence": "quote or null"},
  "rai:annotationsPerItem": {"value": "string or null", "evidence": "quote or null"},
  "rai:annotatorDemographics": {"value": "string or null", "evidence": "quote or null"},
  "rai:machineAnnotationTools": {"value": "string or null", "evidence": "quote or null"},
  "rai:dataReleaseMaintenancePlan": {"value": "string or null", "evidence": "quote or null"},
  "rai:personalSensitiveInformation": {"value": "string or null", "evidence": "quote or null"},
  "rai:dataSocialImpact": {"value": "string or null", "evidence": "quote or null"},
  "rai:dataBiases": {"value": "string or null", "evidence": "quote or null"},
  "rai:dataLimitations": {"value": "string or null", "evidence": "quote or null"},
  "rai:dataUseCases": {"value": "string or null", "evidence": "quote or null"}
}
For null fields, use: {"value": null, "evidence": null}
PAPER TEXT:
%s
Return ONLY valid JSON matching the schema above. No markdown, no explanations.'''


CORRECTION_PROMPT_TEMPLATE = """You previously extracted metadata from this paper but returned null for the following fields:
{null_fields}
The paper likely discusses these topics based on keyword analysis. Please re-read the paper carefully and try to extract these specific fields. If the information is genuinely not in the paper, return null.
Return ONLY valid JSON with just these keys: {field_keys}
PAPER TEXT:
{paper_text}"""


# ── Prompt variant v2: anti-null-bias with few-shot null examples ─
# Phase 3 Fix A from the 2026-05-02 V2/LEV iteration plan.
# Targets failure mode F1 (null-gold hallucination, ~70% of round-0 score=3 cells).
USER_PROMPT_WITH_EVIDENCE_V2_PREFIX = '''CRITICAL — DEFAULT TO NULL WHEN UNCERTAIN.
Many Croissant RAI fields are rarely documented in dataset papers. The
following fields are MOST OFTEN null and you should default to null
unless the paper has an explicit discussion you can quote:
  - rai:dataImputationProtocol      (most papers have no missing data)
  - rai:personalSensitiveInformation (only if paper explicitly lists collected attributes)
  - rai:dataAnnotationPlatform      (only if a specific platform is named: MTurk, Label Studio, etc.)
  - rai:dataReleaseMaintenancePlan  (only if paper discusses versioning or update cadence)
  - rai:annotationsPerItem          (only if a specific count is stated)
  - rai:annotatorDemographics       (only if demographics are explicitly described)
  - rai:dataSocialImpact            (only if a discussion of social implications exists)
  - rai:machineAnnotationTools      (only if specific software tools are named)
  - rai:dataManipulationProtocol    (only if explicit post-preprocessing transforms are described)
RULES:
  1. DO NOT generate plausible-sounding content based on what a similar
     paper might say. If THIS paper does not discuss it, return null.
  2. DO NOT infer from related content. A paper discussing biases does
     NOT automatically document rai:personalSensitiveInformation.
  3. EVERY non-null extraction MUST have an evidence quote that is a
     near-verbatim sentence from the paper. If you cannot quote a
     sentence supporting your extraction, the value should be null.
EXAMPLES:
EXAMPLE A — paper does NOT discuss the field, so the answer is null:
  Paper text: "We crowdsourced 10K questions and verified accuracy..."
  Field: rai:dataAnnotationPlatform
  CORRECT output: {"value": null, "evidence": null}
  WRONG output: {"value": "Amazon Mechanical Turk", "evidence": "..."}
    (wrong because paper says "crowdsourced" but never names MTurk)
EXAMPLE B — paper DOES discuss the field, so extract with evidence:
  Paper text: "Annotations were collected via Amazon Mechanical Turk with 3
  workers per item. Each worker had >97%% HIT acceptance rate."
  Field: rai:dataAnnotationPlatform
  CORRECT output: {"value": "Amazon Mechanical Turk", "evidence":
    "Annotations were collected via Amazon Mechanical Turk..."}
EXAMPLE C — paper has no missing data, so dataImputationProtocol is null:
  Paper text: "Our dataset contains 50K complete examples, no
  preprocessing was needed beyond normalization."
  Field: rai:dataImputationProtocol
  CORRECT output: {"value": null, "evidence": null}
  WRONG output: extracting "data was filled using mean imputation"
    (would be hallucinated; paper says nothing about imputation)
Now extract metadata from the paper below, applying these rules strictly.
'''
USER_PROMPT_WITH_EVIDENCE_V2 = USER_PROMPT_WITH_EVIDENCE_V2_PREFIX + USER_PROMPT_WITH_EVIDENCE


# Prompt variant v5: v4 (Fix A + Fix C) + Fix D (multi-aspect emphasis).
# Phase 3 iteration #4 of the V2/LEV plan, on top of accepted v4.
USER_PROMPT_WITH_EVIDENCE_V5_PREFIX = '''ADDITIONAL GUIDANCE FOR MULTI-ASPECT FIELDS
Several Croissant fields capture MULTIPLE aspects of a single concept.
When the paper discusses multiple aspects, be EXHAUSTIVE and list them ALL,
not just the primary or most prominent one.
  - rai:dataLimitations: capture ALL discussed limitations
    (temporal, geographic, quality, scope, generalization, exclusion of specific
    cases, modality, etc.). If the paper mentions multiple types of
    limitations, include ALL of them.
  - rai:dataBiases: capture ALL discussed biases (selection, demographic,
    contamination, source-medium, stylistic, evaluation, etc.).
  - rai:dataUseCases: list ALL stated use cases (training, testing, validation,
    fine-tuning, benchmarking, evaluation).
  - rai:dataCollectionType: list ALL applicable types from the controlled
    vocabulary. Many datasets combine 2 or 3 types (e.g., "Web Scraping" +
    "Manual Human Curator"). Pick ALL that apply.
Format: use " | " separator OR continuous text covering all aspects.
Do not omit aspects to keep the answer brief.
'''
USER_PROMPT_WITH_EVIDENCE_V5 = (
    USER_PROMPT_WITH_EVIDENCE_V2_PREFIX
    + USER_PROMPT_WITH_EVIDENCE_V5_PREFIX
    + USER_PROMPT_WITH_EVIDENCE
)


LICENSE_MAP = {
    "cc-by-4.0": "CC-BY-4.0", "cc by 4.0": "CC-BY-4.0",
    "cc-by-sa-4.0": "CC-BY-SA-4.0", "cc by-sa 4.0": "CC-BY-SA-4.0",
    "cc-by-nc-4.0": "CC-BY-NC-4.0", "cc by-nc 4.0": "CC-BY-NC-4.0",
    "cc-by-nc-sa-4.0": "CC-BY-NC-SA-4.0",
    "cc0-1.0": "CC0-1.0", "cc0": "CC0-1.0", "public domain": "CC0-1.0",
    "mit": "MIT", "mit license": "MIT",
    "apache-2.0": "Apache-2.0", "apache 2.0": "Apache-2.0",
}


# ── Phase 0: PDF text ─────────────────────────────────────────────
def extract_paper_text(ds_id: str) -> str | None:
    text = get_paper_text(ds_id)
    if text is None:
        return None
    text = strip_references(text).strip()
    if len(text) > MAX_PAPER_CHARS:
        text = text[:MAX_PAPER_CHARS]
    return text


# ── Phase 1: triage ───────────────────────────────────────────────
def load_triage(ds_id: str) -> dict:
    path = PHASE1_DIR / f"{ds_id}.json"
    if not path.exists():
        return {}
    with open(path) as f:
        return json.load(f).get("triage", {})


# ── Phase 2: extract + evidence ───────────────────────────────────
def extract_phase2(cfg, paper_text, with_evidence=True, prompt_variant="v1"):
    if with_evidence:
        # v2 / v4 / v6 use the anti-null-bias prompt prefix.
        # v5 adds multi-aspect enumeration on top of v2's prefix.
        # v6 uses v4's prompt but stricter post-hoc verification (in process_paper).
        if prompt_variant == "v5":
            user_content = USER_PROMPT_WITH_EVIDENCE_V5 % paper_text
        elif prompt_variant in ("v2", "v4", "v6"):
            user_content = USER_PROMPT_WITH_EVIDENCE_V2 % paper_text
        else:
            user_content = USER_PROMPT_WITH_EVIDENCE % paper_text
        raw, usage = call_llm(cfg, SYSTEM_PROMPT, user_content, MAX_TOKENS_EXTRACT)
        parsed = parse_json_response(raw)
        extraction: dict = {}
        evidence: dict = {}
        for field in ALL_FIELDS:
            entry = parsed.get(field, {})
            if isinstance(entry, dict):
                extraction[field] = entry.get("value")
                evidence[field] = entry.get("evidence")
            else:
                extraction[field] = entry
                evidence[field] = None
        return extraction, evidence, usage

    # Ablation baseline: reuse single-pass user prompt (flat schema)
    from config import USER_PROMPT_TEMPLATE
    user_content = USER_PROMPT_TEMPLATE % paper_text
    raw, usage = call_llm(cfg, SYSTEM_PROMPT, user_content, MAX_TOKENS_EXTRACT)
    parsed = normalize_field_names(parse_json_response(raw))
    return parsed, {f: None for f in ALL_FIELDS}, usage


# ── Phase 2.5: self-correction ────────────────────────────────────
def _is_null(val) -> bool:
    if val is None:
        return True
    if isinstance(val, str) and val.strip().lower() in ("", "null", "none"):
        return True
    return False


def find_null_but_expected(extraction: dict, triage: dict) -> list[str]:
    missing = []
    for field in ALL_FIELDS:
        if not _is_null(extraction.get(field)):
            continue
        if field in EASY_FIELDS:
            missing.append(field)
            continue
        for group_name, group_fields in TRIAGE_GROUPS.items():
            if field not in group_fields:
                continue
            info = triage.get(group_name, {})
            if info.get("presence") in ("likely", "certain"):
                missing.append(field)
                break
    return missing


def self_correct(cfg, paper_text, extraction, triage):
    null_fields = find_null_but_expected(extraction, triage)
    if not null_fields:
        return extraction, 0, {"input_tokens": 0, "output_tokens": 0}

    prompt = CORRECTION_PROMPT_TEMPLATE.format(
        null_fields="\n".join(f"- {f}" for f in null_fields),
        field_keys=", ".join(null_fields),
        paper_text=paper_text,
    )

    try:
        raw, usage = call_llm(cfg, SYSTEM_PROMPT, prompt, MAX_TOKENS_CORRECT)
        corrected = parse_json_response(raw)
    except Exception as e:
        print(f"    Self-correction failed: {str(e)[:80]}")
        return extraction, 0, {"input_tokens": 0, "output_tokens": 0}

    applied = 0
    for field in null_fields:
        val = corrected.get(field)
        if not _is_null(val):
            extraction[field] = val
            applied += 1
    return extraction, applied, usage


# ── Phase 3: cross-document enrichment ────────────────────────────
def _enrich_huggingface(ds_id: str, extraction: dict) -> dict:
    enrichments: dict = {}
    hf_id = ds_id.replace("_", "/", 1)
    try:
        url = f"https://huggingface.co/api/datasets/{hf_id}"
        req = urllib.request.Request(url, headers={"User-Agent": "CroissantMiner/2.0"})
        resp = urllib.request.urlopen(req, timeout=10)
        data = json.loads(resp.read())
        card = data.get("cardData", {}) or {}

        if _is_null(extraction.get("license")):
            lic = card.get("license")
            if lic:
                enrichments["license"] = lic
        if _is_null(extraction.get("url")):
            enrichments["url"] = f"https://huggingface.co/datasets/{hf_id}"
        if _is_null(extraction.get("inLanguage")):
            lang = card.get("language")
            if lang:
                enrichments["inLanguage"] = ", ".join(lang) if isinstance(lang, list) else str(lang)
    except Exception:
        pass
    return enrichments


def _enrich_semantic_scholar(ds_id: str, extraction: dict) -> dict:
    if not _is_null(extraction.get("citeAs")):
        return {}
    name = extraction.get("name") or ds_id
    try:
        q = urllib.parse.quote(str(name))
        url = (f"https://api.semanticscholar.org/graph/v1/paper/search"
               f"?query={q}&limit=1&fields=title,citationStyles")
        req = urllib.request.Request(url, headers={"User-Agent": "CroissantMiner/2.0"})
        resp = urllib.request.urlopen(req, timeout=10)
        data = json.loads(resp.read())
        papers = data.get("data", [])
        if papers:
            bibtex = papers[0].get("citationStyles", {}).get("bibtex", "")
            if bibtex:
                return {"citeAs": bibtex}
    except Exception:
        pass
    return {}


def run_enrichment(ds_id: str, extraction: dict) -> tuple[dict, dict]:
    enrichments: dict = {}
    enrichments.update(_enrich_huggingface(ds_id, extraction))
    enrichments.update(_enrich_semantic_scholar(ds_id, extraction))
    for field, val in enrichments.items():
        extraction[field] = val
    return extraction, enrichments


# ── Phase 4: validation + confidence ──────────────────────────────
def validate_and_score(extraction, evidence, enrichments, triage) -> dict:
    details: dict = {}
    for field in ALL_FIELDS:
        val = extraction.get(field)

        if isinstance(val, str) and val.strip().lower() in (
            "", "null", "none", "n/a", "not applicable",
            "[null]", "[null - not found in paper]",
        ):
            val = None
        if isinstance(val, list):
            val = ", ".join(str(x) for x in val)
        if field == "license" and val:
            norm = LICENSE_MAP.get(str(val).strip().lower())
            if norm:
                val = norm
        if field == "url" and val and not re.match(r"https?://", str(val)):
            val = f"https://{val}"

        extraction[field] = val

        if val is None:
            details[field] = {"confidence": 0.0, "status": "NULL", "source": "absent", "evidence": None}
            continue

        confidence = 0.8
        source = "extraction"
        field_ev = evidence.get(field)
        has_ev = field_ev is not None and str(field_ev).strip() not in ("", "null", "None")
        if has_ev:
            confidence = min(confidence + 0.15, 1.0)
        if field in enrichments:
            source = "enrichment"
            confidence = 0.7

        if field in HARD_FIELDS:
            for group_name, group_fields in TRIAGE_GROUPS.items():
                if field in group_fields:
                    if triage.get(group_name, {}).get("presence") == "unlikely":
                        confidence = min(confidence, 0.5)
                    break

        status = ("VERIFIED" if confidence >= 0.7
                  else "UNCERTAIN" if confidence >= 0.4
                  else "LOW_CONFIDENCE")
        details[field] = {
            "confidence": round(confidence, 2),
            "status": status,
            "source": source,
            "evidence": field_ev if has_ev else None,
        }
    return details


# ── Phase 5: save canonical output ────────────────────────────────
def _paper_stats(details: dict) -> dict:
    extracted = sum(1 for d in details.values() if d["status"] != "NULL")
    verified = sum(1 for d in details.values() if d["status"] == "VERIFIED")
    uncertain = sum(1 for d in details.values() if d["status"] == "UNCERTAIN")
    ev = sum(1 for d in details.values() if d.get("evidence"))
    confs = [d["confidence"] for d in details.values() if d["status"] != "NULL"]
    avg = sum(confs) / len(confs) if confs else 0.0
    return {
        "fields_extracted": extracted,
        "fields_null": 30 - extracted,
        "fields_verified": verified,
        "fields_uncertain": uncertain,
        "evidence_count": ev,
        "evidence_coverage": round(ev / max(extracted, 1) * 100, 1),
        "avg_confidence": round(avg, 3),
    }


def save_paper(out_dir: Path, ds_id: str, cfg, extraction, details, usage,
               mode: str, extra_meta: dict) -> dict:
    stats = _paper_stats(details)
    extra = {
        "pipeline": "agentic_v2",
        "extraction_details": details,
        "paper_stats": stats,
        "cost_usd": estimate_cost(cfg, usage, batch_discount=(mode == "batch")),
        "split": "dev" if ds_id in SPLIT["dev"] else "test",
    }
    extra.update(extra_meta)
    meta = meta_block(cfg, SCRIPT_NAME, mode, extra=extra)
    from _agentic_helpers import write_canonical_output
    write_canonical_output(out_dir, ds_id, extraction, usage, cfg, meta)
    return stats


# ── Sequential processing ─────────────────────────────────────────
def verify_evidence_grounding(extraction: dict, evidence: dict, paper_text: str,
                              min_match_len: int = 30) -> tuple[dict, dict, int]:
    """Phase 3 Fix C: null RAI prose extractions whose evidence quote does
    not appear in the paper text. Restricted to LONG_TEXT_RAI_FIELDS — the
    fields where the failure-mode taxonomy showed hallucination dominates.
    Match rule: lowercase + collapse whitespace; require at least one
    min_match_len-char contiguous substring of the evidence quote to appear
    in the normalized paper. This catches verbatim and near-verbatim
    quotes; hallucinated quotes (which the model confabulates rather than
    pulls from the paper) typically fail.
    Tier 1 fields are not verified here — their gold values are short and
    enforced by the Tier 1 canonicalizer at scoring time.
    """
    from evaluation.field_metrics import LONG_TEXT_RAI_FIELDS

    if not paper_text:
        return extraction, evidence, 0
    norm_paper = re.sub(r"\s+", " ", paper_text.lower())
    nullified = 0
    for field in LONG_TEXT_RAI_FIELDS:
        v = extraction.get(field)
        if v is None or (isinstance(v, str) and not v.strip()):
            continue  # already null
        ev = evidence.get(field) if evidence else None
        if not ev or not isinstance(ev, str):
            # No evidence quote: cannot verify. Default behaviour is to
            # null the extraction since this fix's purpose is "if you
            # cannot ground it in the paper, do not keep it."
            extraction[field] = None
            if evidence:
                evidence[field] = None
            nullified += 1
            continue
        norm_ev = re.sub(r"\s+", " ", ev.lower()).strip()
        if len(norm_ev) < min_match_len:
            # Quote too short to require a substring match — keep as-is.
            continue
        matched = any(
            norm_ev[i:i + min_match_len] in norm_paper
            for i in range(len(norm_ev) - min_match_len + 1)
        )
        if not matched:
            extraction[field] = None
            if evidence:
                evidence[field] = None
            nullified += 1
    return extraction, evidence, nullified


def process_paper(cfg, ds_id: str, args, out_dir: Path) -> dict | None:
    t0 = time.time()
    paper_text = extract_paper_text(ds_id)
    if not paper_text:
        print(f"  SKIP {ds_id}: no PDF")
        return None

    triage = load_triage(ds_id)

    usage = {"input_tokens": 0, "output_tokens": 0}

    with_evidence = not args.no_verification
    prompt_variant = getattr(args, "prompt_variant", "v1")
    extraction, evidence, ext_usage = extract_phase2(cfg, paper_text, with_evidence,
                                                     prompt_variant=prompt_variant)
    usage["input_tokens"] += ext_usage["input_tokens"]
    usage["output_tokens"] += ext_usage["output_tokens"]

    corrections = 0
    if not args.no_correction:
        extraction, corrections, cusage = self_correct(cfg, paper_text, extraction, triage)
        usage["input_tokens"] += cusage["input_tokens"]
        usage["output_tokens"] += cusage["output_tokens"]

    enrichments: dict = {}
    if not args.no_enrichment:
        extraction, enrichments = run_enrichment(ds_id, extraction)

    # Phase 3 Fix C: evidence-grounding verification (variants v3, v4, v5, v6).
    # v6 uses a stricter 50-char substring match (vs 30 in v3/v4/v5) to null
    # more borderline cases more aggressively.
    nullified_by_verify = 0
    pv = getattr(args, "prompt_variant", "v1")
    if pv in ("v3", "v4", "v5"):
        extraction, evidence, nullified_by_verify = verify_evidence_grounding(
            extraction, evidence, paper_text, min_match_len=30)
    elif pv == "v6":
        # Same prompt as v4 (A+C compound) but stricter post-hoc verification.
        extraction, evidence, nullified_by_verify = verify_evidence_grounding(
            extraction, evidence, paper_text, min_match_len=50)

    if args.no_verification:
        details = {
            f: {"confidence": 0.8 if extraction.get(f) is not None else 0.0,
                "status": "EXTRACTED" if extraction.get(f) is not None else "NULL",
                "source": "extraction", "evidence": None}
            for f in ALL_FIELDS
        }
    else:
        details = validate_and_score(extraction, evidence, enrichments, triage)

    extra_meta = {
        "corrections_applied": corrections,
        "enrichments": list(enrichments.keys()),
        "elapsed_seconds": round(time.time() - t0, 1),
        "flags": {
            "correction": not args.no_correction,
            "enrichment": not args.no_enrichment,
            "verification": not args.no_verification,
        },
        "prompt_variant": getattr(args, "prompt_variant", "v1"),
        "evidence_verify_nullified": nullified_by_verify,
    }
    return save_paper(out_dir, ds_id, cfg, extraction, details, usage, "realtime", extra_meta)


def run_sequential(cfg, papers, args, out_dir):
    ok = fail = total_extracted = total_verified = 0
    total_usage = {"input_tokens": 0, "output_tokens": 0}

    for i, ds_id in enumerate(papers):
        safe = ds_id.replace("/", "__")
        out_path = out_dir / f"{safe}.json"
        if out_path.exists() and not args.paper:
            ok += 1
            if (i + 1) % 25 == 0:
                print(f"  [{i+1}/{len(papers)}] {ok} done (skipping existing)")
            continue

        try:
            stats = process_paper(cfg, ds_id, args, out_dir)
            if stats is None:
                fail += 1
                continue
            ok += 1
            total_extracted += stats["fields_extracted"]
            total_verified += stats["fields_verified"]
            with open(out_path) as f:
                saved = json.load(f)
            total_usage["input_tokens"] += saved["usage"]["input_tokens"]
            total_usage["output_tokens"] += saved["usage"]["output_tokens"]
            print(f"  [{i+1}/{len(papers)}] {ds_id}: "
                  f"{stats['fields_extracted']}/30 extracted, "
                  f"{stats['fields_verified']} verified, "
                  f"${saved['_meta']['cost_usd']:.3f}, "
                  f"{saved['_meta']['elapsed_seconds']:.0f}s")
        except Exception as e:
            print(f"  [{i+1}/{len(papers)}] {ds_id}: FAILED ({str(e)[:80]})")
            fail += 1
        time.sleep(0.4)

    total_cost = estimate_cost(cfg, total_usage, batch_discount=False)
    total_fields = ok * 30
    print(f"\n{'=' * 70}")
    print(f"SUMMARY — V2 sequential ({cfg['name']})")
    print(f"{'=' * 70}")
    print(f"Papers: {ok} OK, {fail} failed")
    print(f"Fields: {total_extracted}/{total_fields} "
          f"({total_extracted / max(total_fields, 1) * 100:.1f}%)")
    print(f"Verified: {total_verified}/{total_fields}")
    print(f"Tokens: {total_usage['input_tokens']:,} in, "
          f"{total_usage['output_tokens']:,} out")
    print(f"Cost: ${total_cost:.2f}")
    print(f"Output: {out_dir}/")

    summary = {
        "backbone": cfg["model_id"],
        "papers_ok": ok, "papers_failed": fail,
        "total_extracted": total_extracted,
        "total_verified": total_verified,
        "fill_rate": round(total_extracted / max(total_fields, 1) * 100, 1),
        "usage": total_usage, "cost_usd": total_cost,
    }
    with open(out_dir / "_summary.json", "w") as f:
        json.dump(summary, f, indent=2)


# ── Anthropic-only batch mode ─────────────────────────────────────
def _require_anthropic(cfg):
    if cfg["provider"] != "anthropic":
        raise SystemExit(
            f"Batch mode currently supports provider=anthropic only "
            f"(backbone '{cfg['model_id']}' uses provider={cfg['provider']}). "
            f"Run without --batch for sequential execution."
        )


def _sanitize_custom_id(ds_id: str) -> str:
    return re.sub(r"[^a-zA-Z0-9_-]", "_", ds_id.replace("/", "__"))[:64]


def run_batch(cfg, papers, args, out_dir):
    _require_anthropic(cfg)
    import anthropic
    from anthropic.types.message_create_params import MessageCreateParamsNonStreaming
    from anthropic.types.messages.batch_create_params import Request

    client = anthropic.Anthropic()
    wave_dir = out_dir / "_batch"
    wave_dir.mkdir(parents=True, exist_ok=True)

    # ── Wave 1: base extraction ──
    wave1_path = wave_dir / "wave1.json"
    if wave1_path.exists():
        wave1 = json.loads(wave1_path.read_text())
    else:
        requests = []
        id_map: dict = {}
        for ds_id in papers:
            safe = ds_id.replace("/", "__")
            if (out_dir / f"{safe}.json").exists():
                continue
            text = extract_paper_text(ds_id)
            if not text:
                print(f"  skip (no PDF): {ds_id}")
                continue
            cid = _sanitize_custom_id(ds_id)
            id_map[cid] = ds_id
            requests.append(Request(
                custom_id=cid,
                params=MessageCreateParamsNonStreaming(
                    model=cfg["model_id"],
                    max_tokens=MAX_TOKENS_EXTRACT,
                    temperature=0.0,
                    system=SYSTEM_PROMPT,
                    messages=[{"role": "user", "content": USER_PROMPT_WITH_EVIDENCE % text}],
                ),
            ))
        if not requests:
            print("  No papers to submit for Wave 1.")
            return
        batch = client.messages.batches.create(requests=requests)
        wave1 = {"batch_id": batch.id, "id_map": id_map,
                 "submitted_at": time.strftime("%Y-%m-%d %H:%M:%S")}
        wave1_path.write_text(json.dumps(wave1, indent=2))
        print(f"  Wave 1 submitted: {batch.id} ({len(requests)} requests)")

    # Poll Wave 1
    while True:
        info = client.messages.batches.retrieve(wave1["batch_id"])
        counts = info.request_counts
        print(f"    Wave 1 {info.processing_status}: "
              f"proc={counts.processing} ok={counts.succeeded} err={counts.errored}")
        if info.processing_status == "ended":
            break
        time.sleep(30)

    # Collect Wave 1 parsed results
    parsed: dict = {}
    for result in client.messages.batches.results(wave1["batch_id"]):
        ds_id = wave1["id_map"].get(result.custom_id, result.custom_id)
        if result.result.type != "succeeded":
            print(f"    Wave 1 FAILED {ds_id}: {result.result.type}")
            continue
        msg = result.result.message
        raw = "".join(getattr(b, "text", "") for b in msg.content)
        usage = {"input_tokens": msg.usage.input_tokens,
                 "output_tokens": msg.usage.output_tokens}
        try:
            data = parse_json_response(raw)
        except Exception as exc:
            print(f"    Wave 1 parse failed {ds_id}: {str(exc)[:60]}")
            continue
        extraction, evidence = {}, {}
        for field in ALL_FIELDS:
            entry = data.get(field, {})
            if isinstance(entry, dict):
                extraction[field] = entry.get("value")
                evidence[field] = entry.get("evidence")
            else:
                extraction[field] = entry
                evidence[field] = None
        parsed[ds_id] = {"extraction": extraction, "evidence": evidence, "usage": usage}

    print(f"  Wave 1 parsed: {len(parsed)} papers")

    # ── Wave 2: corrections batch (optional) ──
    wave2_path = wave_dir / "wave2.json"
    corrections: dict = {}
    if not args.no_correction:
        need: dict = {}
        for ds_id, rec in parsed.items():
            triage = load_triage(ds_id)
            nulls = find_null_but_expected(rec["extraction"], triage)
            if nulls:
                need[ds_id] = nulls
        print(f"  Papers needing corrections: {len(need)}/{len(parsed)}")

        if need and not wave2_path.exists():
            reqs = []
            id_map2: dict = {}
            for ds_id, nulls in need.items():
                text = extract_paper_text(ds_id)
                if not text:
                    continue
                prompt = CORRECTION_PROMPT_TEMPLATE.format(
                    null_fields="\n".join(f"- {f}" for f in nulls),
                    field_keys=", ".join(nulls),
                    paper_text=text,
                )
                cid = _sanitize_custom_id(ds_id)
                id_map2[cid] = ds_id
                reqs.append(Request(
                    custom_id=cid,
                    params=MessageCreateParamsNonStreaming(
                        model=cfg["model_id"],
                        max_tokens=MAX_TOKENS_CORRECT,
                        temperature=0.0,
                        system=SYSTEM_PROMPT,
                        messages=[{"role": "user", "content": prompt}],
                    ),
                ))
            if reqs:
                batch = client.messages.batches.create(requests=reqs)
                wave2 = {"batch_id": batch.id, "id_map": id_map2, "need": need,
                         "submitted_at": time.strftime("%Y-%m-%d %H:%M:%S")}
                wave2_path.write_text(json.dumps(wave2, indent=2))
                print(f"  Wave 2 submitted: {batch.id} ({len(reqs)} requests)")

        if wave2_path.exists():
            wave2 = json.loads(wave2_path.read_text())
            while True:
                info = client.messages.batches.retrieve(wave2["batch_id"])
                counts = info.request_counts
                print(f"    Wave 2 {info.processing_status}: "
                      f"proc={counts.processing} ok={counts.succeeded} err={counts.errored}")
                if info.processing_status == "ended":
                    break
                time.sleep(30)
            for result in client.messages.batches.results(wave2["batch_id"]):
                ds_id = wave2["id_map"].get(result.custom_id, result.custom_id)
                if result.result.type != "succeeded":
                    print(f"    Wave 2 FAILED {ds_id}")
                    continue
                msg = result.result.message
                raw = "".join(getattr(b, "text", "") for b in msg.content)
                usage = {"input_tokens": msg.usage.input_tokens,
                         "output_tokens": msg.usage.output_tokens}
                try:
                    corrected = parse_json_response(raw)
                except Exception:
                    continue
                corrections[ds_id] = {"corrected": corrected,
                                      "nulls": wave2["need"].get(ds_id, []),
                                      "usage": usage}

    # ── Apply corrections + enrichment + validation + save ──
    ok = 0
    for ds_id, rec in parsed.items():
        extraction = rec["extraction"]
        evidence = rec["evidence"]
        usage = dict(rec["usage"])
        applied = 0
        if ds_id in corrections:
            c = corrections[ds_id]
            usage["input_tokens"] += c["usage"]["input_tokens"]
            usage["output_tokens"] += c["usage"]["output_tokens"]
            for field in c["nulls"]:
                val = c["corrected"].get(field)
                if not _is_null(val):
                    extraction[field] = val
                    applied += 1

        enrichments: dict = {}
        if not args.no_enrichment:
            extraction, enrichments = run_enrichment(ds_id, extraction)

        triage = load_triage(ds_id)
        details = validate_and_score(extraction, evidence, enrichments, triage)

        extra_meta = {
            "corrections_applied": applied,
            "enrichments": list(enrichments.keys()),
            "flags": {
                "correction": not args.no_correction,
                "enrichment": not args.no_enrichment,
                "verification": True,
            },
        }
        save_paper(out_dir, ds_id, cfg, extraction, details, usage, "batch", extra_meta)
        ok += 1

    print(f"  Batch pipeline done: {ok} papers written to {out_dir}")


# ── CLI ───────────────────────────────────────────────────────────
def main():
    p = argparse.ArgumentParser(description="CroissantMiner Agentic V2 pipeline")
    p.add_argument("--backbone", default="sonnet-4-5",
                   help="Backbone key from scripts/_agentic_helpers.py::MODELS")
    p.add_argument("--dev-only", action="store_true", help="Process only the 15 dev papers")
    p.add_argument("--paper", type=str, help="Process a single paper by dataset_id")
    p.add_argument("--batch", action="store_true",
                   help="Use Anthropic batch API (sonnet-4-5 only)")
    p.add_argument("--no-enrichment", action="store_true", help="Ablation: skip Phase 3")
    p.add_argument("--no-verification", action="store_true", help="Ablation: skip Phase 4")
    p.add_argument("--no-correction", action="store_true", help="Ablation: skip Phase 2.5")
    p.add_argument("--prompt-variant", default="v1", choices=["v1", "v2", "v3", "v4", "v5", "v6"],
                   help="v1=round-0; v2=Fix A; v3=Fix C; v4=A+C (accepted +0.052); v5=A+C+D; v6=A+C with stricter 50-char evidence verification.")
    args = p.parse_args()

    if not any(PHASE1_DIR.glob("*.json")):
        raise SystemExit("Triage + Critique reads the stored Gemini 2.5 Flash triage in data/agentic/phase1/, "
                         "which is missing (it ships with the repository; rebuild it with scripts/agentic_phase1.py).")
    cfg = resolve_backbone(args.backbone)
    out_dir = output_dir_for("v2", args.backbone)
    if args.prompt_variant != "v1":
        out_dir = out_dir.parent / f"{out_dir.name}_{args.prompt_variant}"
    out_dir.mkdir(parents=True, exist_ok=True)

    if args.paper:
        papers = [args.paper]
    elif args.dev_only:
        papers = sorted(SPLIT["dev"])
    else:
        papers = sorted(SPLIT["dev"] + SPLIT["test"])

    print(f"{'=' * 70}")
    print(f"CROISSANTMINER AGENTIC V2 — backbone={args.backbone} ({cfg['name']})")
    print(f"Papers: {len(papers)}  Output: {out_dir}")
    print(f"Flags: correction={'ON' if not args.no_correction else 'OFF'}, "
          f"enrichment={'ON' if not args.no_enrichment else 'OFF'}, "
          f"verification={'ON' if not args.no_verification else 'OFF'}")
    print(f"{'=' * 70}")

    if args.batch:
        run_batch(cfg, papers, args, out_dir)
    else:
        run_sequential(cfg, papers, args, out_dir)


if __name__ == "__main__":
    main()