diff --git a/README.md b/README.md index 6e9003b..970c8ee 100644 --- a/README.md +++ b/README.md @@ -37,7 +37,7 @@ Read our blog post [here](https://openai.com/index/scaling-social-science-resear Most of the evidence social scientists and analysts care about lives in unstructured formats: interviews, speeches, transcripts, product photos, archival scans. Modern GPT models can judge attributes, extract facts, and reason about this material with high fidelity, but building robust pipelines is still tedious. GABRIEL provides: - 🧠 **Human-level comprehension on demand** – express the attribute the way you would brief a human coder; GABRIEL packages the prompt, context, and retries for you. -- 📊 **Quantitative outputs** – ratings (0–100), grounded comparisons, classifications, and structured extractions return as tidy DataFrames with reproducible settings. +- 📊 **Quantitative outputs** – ratings (0–100), pairwise comparisons, classifications, and structured extractions return as tidy DataFrames with reproducible settings. - ⚙️ **Operational tooling** – automatic parallelism (hundreds of concurrent calls), resumable runs, raw response logs, and helper UIs make it safe to scale to very large corpora. - 🧰 **Extensibility** – swap instructions with `additional_instructions`, bring your own templates, or drop down to `gabriel.whatever` + custom `response_fn` for bespoke prompts while still reusing the infrastructure. @@ -50,7 +50,7 @@ The tutorial notebook walks through these ideas step-by-step—from setting up a | Function | Purpose & Output Scale | Example Use | | --- | --- | --- | | `gabriel.rate` | Asks GPT to score each text / image / audio / item on natural language attributes. Output = 0-100 rating. | Measure “populist rhetoric” in a speech; “toxicity” of tweets; “luxury” in ad images. | -| `gabriel.rank` | Pairwise comparisons between texts yields ELO-like attribute ratings. Output = grounded, relative z scores for each text. | Rank technologies by “bulkiness” or artworks by “fine brushwork”. | +| `gabriel.rank` | Pairwise comparisons yield Bradley–Terry attribute ratings. Output = model-derived z-scores standardized within connected comparison components. | Rank technologies by “bulkiness” or artworks by “fine brushwork”. | | `gabriel.classify` | Classifies texts / images / audio / items on whether provided labels apply. Output = one or more classes per item. | Tag news articles, product photos, or interview clips into topical categories. | | `gabriel.extract` | Structured fact or entity extraction. One source can return one row or expand into several extracted-entity rows. | Extract every product and its name, price, description, and color from each catalog page. | | `gabriel.discover` | Discovers natural language features which discriminate two classes of data. | Identify what distinguishes 5 star vs. 1 star reviews or successful vs. failed startups. | @@ -73,7 +73,7 @@ The tutorial notebook walks through these ideas step-by-step—from setting up a | `gabriel.bucket` | Proposes a compact taxonomy from many terms. Output = bucket names and definitions, not row-level assignments. | Develop candidate categories for technologies, artworks, or HR complaints, then review and apply them downstream. | | `gabriel.seed` | Enforces a representative distribution / diversity of seeds. | Initialize unique personas that match US population distribution. | | `gabriel.poll` | Seeds personas, expands them into full biographies, and surveys them. | Simulate a synthetic opinion poll on policy, values, and open-ended attitudes. | -| `gabriel.ideate` | Generates many novel scientific theories and filters the cream of the crop. | Procure novel theories on inflation for potential research. | +| `gabriel.ideate` | Generates and filters candidate scientific theories for further evaluation. | Procure novel theories on inflation for potential research. | | `gabriel.debias` | Post-process measurements to remove inference bias. | Ensure GPT isn't guessing climate opinions in speeches based on general political lean. | | `gabriel.load` | Prepares a folder of text / image / audio files into a spreadsheet for use in GABRIEL. | Image directory converted into spreadsheet of file paths. | | `gabriel.view` | UI to view sample texts with ratings / passage coding. | Spot-check classify / rating outputs; view coded passages. | @@ -183,7 +183,7 @@ The tutorial notebook covers full projects end-to-end. The list below matches it ### Measurement primitives - **`gabriel.rate`** – assign 0–100 scores per attribute across text, entities, images, audio, or web-sourced context. -- **`gabriel.rank`** – pairwise tournaments that surface relative winners with grounded z-scores. +- **`gabriel.rank`** – pairwise tournaments fitted with a Bradley–Terry model, with symmetric pseudo-comparison regularization on observed edges when `learning_rate > 0`; a zero value requires Ford's directed strong-connectivity condition within each nontrivial observed component for a finite fit. Draws are encoded as half-wins in a binary Bradley–Terry working objective (one comparison total), not as a three-outcome tie likelihood. `insufficient_signal_policy="tie"` makes the modeling assumption that missing direct evidence is equality; use `"abstain"` to condition fitting on retained judgments, noting that informative abstention can still bias estimates. In non-recursive runs, z-scores are standardized separately within connected comparison components (with a neutral zero fallback when a component has no score dispersion), and `_component` columns identify which values are mutually comparable; isolated items remain unranked with `NaN` scores and standard errors. `_se` is an asymptotic, model-based sandwich standard error for the component-centered raw log-skill `_raw`, computed while treating the realized comparison graph as fixed. Observed comparison weights determine the meat under the binary-BT working variance, while fixed pseudo-comparisons add curvature to the bread. The standard error excludes shrinkage bias, shared-judge dependence, misspecification, adaptive-selection uncertainty, and uncertainty for the reported z-score. A non-recursive `_diagnostics.csv` sidecar records decoded outcome counts, effective comparison weight, graph components, isolates, and the count of finite standard errors. Resume metadata fingerprints the effective judge settings and input payloads; `judge_version` is required for custom judges and must change whenever their hidden logic changes. Planned rounds and Batch jobs are recovered transactionally rather than silently mixed or resubmitted. - **`gabriel.classify`** – single- or multi-label tagging with label definitions and consensus columns. - **`gabriel.extract`** – extract typed fields or multiple entities from each passage, page, image, PDF, or audio item; one source may intentionally produce several rows. - **`gabriel.discover`** – contrast two labeled corpora to learn discriminating features. diff --git a/gabriel_tutorial_notebook.ipynb b/gabriel_tutorial_notebook.ipynb index 45242d6..e595a30 100644 --- a/gabriel_tutorial_notebook.ipynb +++ b/gabriel_tutorial_notebook.ipynb @@ -137,7 +137,7 @@ " \n", " \n", " gabriel.rateAsks GPT to score each text / image / audio / item on natural language attributes. Output = 0-100 rating.Measure \"populist rhetoric\" in a speech; \"toxicity\" of tweets; \"luxury\" in ad images.\n", - " gabriel.rankPairwise comparisons between texts yields ELO-like attribute ratings. Output = grounded, relative z scores for each text.Rank technologies by \"bulkiness\" or artworks by \"fine brushwork\".\n", + " gabriel.rankPairwise comparisons yield Bradley–Terry ratings. Output = model-derived z-scores standardized within connected comparison components.Rank technologies by \"bulkiness\" or artworks by \"fine brushwork\".\n", " gabriel.classifyClassifies texts / images / audio / items on whether provided labels apply. Output = one or more classes per item.Tag news articles, product photos, or interview clips into topical categories.\n", " gabriel.extractStructured fact extraction on each item. Output = string / numeric values.For each product, provide the \"company\", \"CEO\", and \"year of invention\".\n", " gabriel.discoverDiscovers natural language features which discriminate two classes of data.Identify what distinguishes 5 star vs. 1 star reviews or successful vs. failed startups.\n", @@ -184,7 +184,7 @@ " gabriel.bucketBuilds taxonomies from many terms. Output = bucket/cluster labels.Group technologies, artworks, or HR complaints into emergent categories.\n", " gabriel.seedEnforces a representative distribution / diversity of seeds.Initialize unique personas that match US population distribution.\n", " gabriel.pollSeeds personas, expands them into full biographies, and surveys them.Simulate a synthetic opinion poll on policy, values, and open-ended attitudes.\n", - " gabriel.ideateGenerates many novel scientific theories and filters the cream of the crop.Procure novel theories on inflation for potential research.\n", + " gabriel.ideateGenerates and filters candidate scientific theories for further evaluation.Procure novel theories on inflation for potential research.\n", " gabriel.debiasPost-process measurements to remove inference bias.Ensure GPT isn't guessing climate opinions in speeches based on general political lean.\n", " gabriel.loadPrepares a folder of text / image / audio files into a spreadsheet for use in GABRIEL.Image directory converted into spreadsheet of file paths.\n", " gabriel.viewUI to view sample texts with ratings / passage coding.Spot-check classify / rating outputs; view coded passages.\n", @@ -3160,7 +3160,7 @@ "\n", "For example, say you have a dataset of speeches. Your attributes might be `\"populism\"` or `\"critical of the supreme court\"` or `\"adversarial\"` or anything. You'll get a rating for each speech, on each attribute, 0 to 100. Same if you had interviews with people about their childhood upbringing: you could rate `\"emphasizes family over friends\"` or any other features of interest.\n", "\n", - "Run the toy example below. **We recommend you try out `gabriel.rank` too** - it also generates ratings, but by comparing the text to each other rather that rating them abstractly, making it more grounded.\n", + "Run the toy example below. **We recommend you try out `gabriel.rank` too** - it generates relative, model-derived ratings by comparing texts to one another rather than rating each one independently.\n", "\n", "Ratings are as simple as `results = await gabriel.rate(your_data, column_name, attributes_to_rate, save_folder)`!" ], @@ -7112,7 +7112,7 @@ "source": [ "# Ranking texts using pairwise comparisons.\n", "\n", - "Get grounded z scores as attribute ratings using a bunch of pairwise comparisons between texts." + "Estimate component-relative Bradley–Terry z-scores using pairwise comparisons between texts." ], "metadata": { "id": "C1sWtm66CnVQ" @@ -7121,15 +7121,15 @@ { "cell_type": "markdown", "source": [ - "Here we cover how to use `gabriel.rank` to get z scores for each text on how much it manifests each attribute relative to the other texts. We do this using many pairwise comparisons between samples of two texts (or entities, images, audio; see multimodal section). Like a chess tournament, this creates ordered, ELO-like scores for each text. Unlike `gabriel.rate`, **these z scores are clearly grounded because they stem from comparisons**, not ratings in a vacuum.\n", + "Here we cover how to use `gabriel.rank` to estimate how strongly each text manifests an attribute relative to texts in the same connected comparison component. We use repeated pairwise comparisons between samples of two texts (or entities, images, audio; see multimodal section) and fit Bradley–Terry log-skills, then standardize them within each connected component. These are model-derived comparative measurements, not ground truth.\n", "\n", "The use case is similar to `gabriel.rate`: we want continuous numerical measurements on the texts. The difference here is our z scores will be **relative**, not absolute.\n", "\n", "In this case, we do not pre-define a ratings scale like 0-100; we do many pairwise comparisons of texts over multiple rounds to get an ordering of where each text is in the relative distribution, separately for each attribute.\n", "\n", - "The `gabriel.rank` method call is analogous to `gabriel.rate`. `gabriel.rank` is more expensive and takes more time, since each entity is involved in many rounds of comparisons instead of one-shot ratings. Instead of `n_runs`, use the `n_rounds` and `matches_per_round` parameters to control how many iterations of the round-robin are run. More is better!\n", + "The `gabriel.rank` method call is analogous to `gabriel.rate`. `gabriel.rank` is more expensive and takes more time, since each entity is involved in many rounds of comparisons instead of one-shot ratings. Instead of `n_runs`, use the `n_rounds` and `matches_per_round` parameters to control the comparison budget. More comparisons generally improve model-based precision within a connected component, at higher cost; they do not remove prompt or judge bias.\n", "\n", - "**We believe `gabriel.rate` is best for most use cases.** But if subtle relative differences and grounded ratings are needed, `gabriel.rank` is for you. Set `recursive` to `True` if you want to identify the absolute top performers on an attribute in a large dataset!" + "**We believe `gabriel.rate` is best for most use cases.** But if subtle relative differences are important, `gabriel.rank` may be useful. Set `recursive` to `True` to progressively retain candidates with higher estimated scores on one selected attribute. At Rank-produced recursive stages, pruning requires a connected comparison graph for that attribute; a preliminary Rate stage has no comparison graph." ], "metadata": { "id": "D6LlwRM3CpQ8" @@ -7140,7 +7140,7 @@ "source": [ "## Toy rankings example.\n", "\n", - "Here's a simple exercise using rankings to measure a distribution of z scores across texts. We compare a series of short tweet-like texts to each other on a range of features." + "Here's a simple exercise using rankings to measure a distribution of Bradley–Terry z-scores across texts. We compare short tweet-like texts on several features. Draws are encoded as half-wins in a binary Bradley–Terry working objective (one comparison total), not as a three-outcome tie likelihood. By default, `insufficient signal` makes the modeling assumption that missing direct evidence is equality; set `insufficient_signal_policy=\"abstain\"` to exclude those judgments, noting that informative abstention can bias the retained sample. Z-scores and raw scores are component-relative, `_component` columns show which values are comparable, zero-dispersion components use a neutral z-score fallback of zero, and isolated items are unranked. `_se` is an asymptotic, model-based sandwich standard error for the component-centered raw log-skill, computed while treating the realized graph as fixed. Observed comparison weights determine binary-BT working variance; positive pseudo-comparison weight adds curvature. It excludes shrinkage bias, shared-judge dependence, misspecification, adaptive-selection uncertainty, and uncertainty for the z-score. `power_matching=True` uses an uncertainty-guided scheduling heuristic rather than exact expected information gain. `matches_per_round` is a target degree capped at the number of opponents, and realized degrees can be higher when an item is also selected as another item's opponent. Setting `learning_rate=0` requires Ford's directed strong-connectivity condition within every nontrivial observed component for a finite fit. The non-recursive diagnostics sidecar records decoded outcome counts, effective comparison weight, topology, isolates, and the count of finite standard errors." ], "metadata": { "id": "Fx-lOgEsDL8u" @@ -7187,14 +7187,14 @@ " toy_data,\n", " column_name = \"text\", # name of column with posts to rank\n", " attributes = attributes, # attributes to rank by\n", - " n_rounds = 3, # number of ELO rounds (more = more pairwise judgments)\n", - " matches_per_round = 5, # 5 * 3 means 15 total matchups for every entry\n", + " n_rounds = 3, # number of Bradley–Terry rounds (more = more pairwise judgments)\n", + " matches_per_round = 5, # target about 5 scheduled matchups per item in each round\n", " save_dir = \"toy_rank\", # directory to save results\n", " model = \"gpt-5.6-luna\", # model used for pairwise comparisons\n", " modality = \"text\", # input modality can also be \"entity\" (for terms like 'apple pie', not full texts), \"pdf\", \"image\", or \"audio\" (see multimodal section)\n", " n_parallels = 650, # max parallel threads (reduce if many errors, increase for higher speed)\n", " reset_files = False, # rerunning the cell loads from save / picks up from checkpoint\n", - " recursive = False, # set to True to recursively cut low performers on a single attribute and rerank until revealing the cream of the crop (isolates and orders the true top performers)\n", + " recursive = False, # set to True to progressively retain candidates with higher estimated scores on one attribute\n", ")" ], "metadata": { @@ -21500,7 +21500,7 @@ "source": [ "# Using GPT as an ideator.\n", "\n", - "We can use GPT to come up with thousands of novel research theories, then filter itself to find the cream of the crop." + "We can use GPT to propose many research theories, then filter them into a smaller candidate set for further evaluation." ], "metadata": { "id": "M59xsTZKL-9E" @@ -21511,7 +21511,7 @@ "source": [ "The `gabriel.ideate` method takes a research topic and brainstorms hundreds of novel research theories in parallel.\n", "\n", - "Then we use `gabriel.rank` with `recursive = True` to find the absolute best theories proposed by GPT." + "Then we use `gabriel.rank` with `recursive = True` to retain theories with higher estimated model scores through successive stages." ], "metadata": { "id": "6OB91xcqMc0z" @@ -21532,7 +21532,7 @@ " n_ideas = 500, # the number of initial ideas to generate\n", " scientific_theory = True, # set to False if you want general ideation to problem solve on any topic\n", " save_dir = \"automation_ideate\",\n", - " evaluation_mode = \"recursive_rank\", # uses gabriel.rank recursively to find cream of the crop (set to 'rate' or 'rank' for faster runs)\n", + " evaluation_mode = \"recursive_rank\", # progressively retains higher-scoring candidates (set to 'rate' or 'rank' for faster runs)\n", " model = \"gpt-5.6-sol\",\n", " ranking_model = \"gpt-5.6-sol\",\n", " reasoning_effort = \"high\",\n", diff --git a/src/gabriel/api.py b/src/gabriel/api.py index e41f422..4de7f8c 100644 --- a/src/gabriel/api.py +++ b/src/gabriel/api.py @@ -3,6 +3,7 @@ import os import re import inspect +import json import warnings import pandas as pd from dataclasses import fields @@ -917,7 +918,7 @@ async def ideate( embedding_fn: Optional[Callable[..., Awaitable[Any]]] = None, get_all_embeddings_fn: Optional[Callable[..., Awaitable[Dict[str, List[float]]]]] = None, ) -> pd.DataFrame: - """Generates many novel scientific theories and filters the cream of the crop. + """Generate and filter candidate scientific theories for further evaluation. Example Use ----------- @@ -1196,6 +1197,7 @@ async def rank( power_matching: bool = True, return_raw_scores: bool = False, learning_rate: float = 0.1, + insufficient_signal_policy: str = "tie", n_parallels: int = 650, n_attributes_per_run: Optional[int] = None, file_name: str = "rankings", @@ -1219,12 +1221,13 @@ async def rank( primer_scores: Optional[Dict[str, Dict[str, float]]] = None, primer_scale: float = 1.0, primer_center: bool = True, + judge_version: Optional[str] = None, id_column: Optional[str] = None, response_fn: Optional[Callable[..., Awaitable[Any]]] = None, get_all_responses_fn: Optional[Callable[..., Awaitable[pd.DataFrame]]] = None, **cfg_kwargs, ) -> pd.DataFrame: - """Pairwise comparisons between texts yields ELO-like attribute ratings. Output = grounded, relative z scores for each text. + """Fit model-derived Bradley–Terry ratings from pairwise comparisons. Example Use ----------- @@ -1249,7 +1252,30 @@ async def rank( when this package version shipped; model IDs change, so verify the exact slug in the official OpenAI model catalog before overriding it. n_rounds, matches_per_round, power_matching, learning_rate: - Parameters controlling the Elo-style tournament mechanics. + Parameters controlling the Bradley–Terry tournament mechanics. + ``matches_per_round`` is a per-item target capped at the number of + possible opponents. Realized degrees may exceed it because an item can + also be selected as another item's opponent; this is especially common + in random mode. ``power_matching=True`` uses an uncertainty-guided + heuristic, not exact expected information gain; ``False`` uses random + pairing. ``learning_rate`` is the total symmetric pseudo-comparison + weight applied to each observed pair. A zero value requires Ford's + directed strong-connectivity condition within every nontrivial observed + component for a finite maximum-likelihood fit. + ``_se`` is an asymptotic, model-based sandwich standard error + for the component-centered raw log-skill, computed while treating the + realized comparison graph as fixed. Observed comparison weights + determine the meat under binary-BT working variance; positive fixed + pseudo-comparison weight adds curvature to the bread. It excludes + shrinkage bias, shared-judge dependence, misspecification, + adaptive-selection uncertainty, and uncertainty for the reported + z-score. + insufficient_signal_policy: + ``"tie"`` makes the modeling assumption that lack of direct signal is + equality and gives each item a half-win. ``"abstain"`` excludes that + judgment from fitting; informative abstention can still bias the + retained sample. True draws use half-wins in a binary Bradley–Terry + working objective, not a three-outcome tie likelihood. n_parallels: Concurrency ceiling, not a fixed worker count. Keep the default unchanged: GABRIEL ramps up and adapts to rate limits and repeated errors. Temporary @@ -1273,16 +1299,22 @@ async def rank( Settings for recursive pruning (fraction kept, minimum remaining, etc.). initial_rating_pass: Whether to run a preliminary rating stage before comparisons. Enabled by - default to give the tournament grounded starting scores; set to + default to give the tournament model-derived starting scores; set to ``False`` to skip the rating seed. rate_kwargs: Additional configuration forwarded to the preliminary rating stage. A legacy ``max_output_tokens`` entry is accepted, warned about, and - ignored. + ignored. ``attributes``, ``save_dir``, and ``file_name`` are owned by + Rank and cannot be overridden here. primer_scores, primer_scale, primer_center: Optional seed ratings to prime the Bradley–Terry loop. Scores are centred per attribute when ``primer_center`` is ``True`` and scaled by ``primer_scale``. + judge_version: + Stable label required when ``response_fn`` or + ``get_all_responses_fn`` supplies a custom judge. Change it whenever + hidden judgment logic or configuration changes so resumed tournaments + reject mixed measurement regimes. id_column: Optional existing identifier column; otherwise hashes of ``column_name`` are generated. @@ -1300,12 +1332,86 @@ async def rank( Returns ------- pandas.DataFrame - Ranked outputs. The CSV written to ``save_dir`` always contains raw - scores and standard errors, but the returned DataFrame hides those - columns unless ``return_raw_scores`` is ``True``. + Ranked outputs. For non-recursive runs, the CSV written to ``save_dir`` + contains raw scores, model-based sandwich standard errors, and + comparison-graph component labels. Z-scores and raw scores are only + comparable within a component; components with no score dispersion use + a neutral z-score fallback of zero, while isolated items have ``NaN`` + scores and standard errors. Raw scores and standard errors are hidden unless + ``return_raw_scores`` is ``True``. Recursive runs instead return + stage-relative z-scores, ``exit_stage``, and ``overall_rank``. At + Rank-produced recursive stages, pruning stops with an error if the cut + attribute's comparison graph is disconnected; a preliminary Rate stage + has no comparison graph. Non-recursive runs also save a diagnostics sidecar with + outcome counts, effective comparison weight, graph topology, and the + count of finite standard errors. """ save_dir = os.path.expandvars(os.path.expanduser(save_dir)) os.makedirs(save_dir, exist_ok=True) + + def _persisted_rank_attribute_keys() -> Optional[List[str]]: + base_name = os.path.splitext(file_name)[0] + for path in ( + Path(save_dir) / f"{base_name}_attrs.json", + Path(save_dir) / "attributes.json", + ): + try: + with path.open(encoding="utf-8") as attribute_file: + payload = json.load(attribute_file) + except Exception: + continue + if isinstance(payload, dict) and all( + isinstance(key, str) for key in payload + ): + return list(payload) + if isinstance(payload, list) and all( + isinstance(key, str) for key in payload + ): + return payload + return None + + def _project_public_rank_columns(result: pd.DataFrame) -> pd.DataFrame: + if return_raw_scores or recursive: + return result + # Infer persisted Rank attributes from the output schema rather than + # from this invocation's arguments. Cached results may have been + # produced with a different in-memory attribute object, but the public + # projection must never leak hidden raw estimates or standard errors. + schema_attr_keys = [ + column[: -len("_component")] + for column in result.columns + if column.endswith("_component") + and column[: -len("_component")] in result.columns + and f"{column[: -len('_component')]}_raw" in result.columns + and f"{column[: -len('_component')]}_se" in result.columns + ] + caller_attr_keys = ( + list(attributes.keys()) + if isinstance(attributes, dict) + else list(attributes) + ) + persisted_attr_keys = _persisted_rank_attribute_keys() + # Persisted attributes are authoritative and avoid hiding unrelated + # input columns that happen to share Rank's suffix convention. Caller + # and schema inference retain safe behavior for legacy cache layouts. + persisted_matches = [ + attr + for attr in (persisted_attr_keys or []) + if f"{attr}_raw" in result.columns and f"{attr}_se" in result.columns + ] + attr_keys = ( + persisted_matches + if persisted_matches + else list(dict.fromkeys([*schema_attr_keys, *caller_attr_keys])) + ) + drop_cols = [ + column + for attr in attr_keys + for column in (f"{attr}_raw", f"{attr}_se") + if column in result.columns + ] + return result.drop(columns=drop_cols) if drop_cols else result + if df is None: if recursive: base_folder = os.path.join(save_dir, f"{file_name}_recursive") @@ -1313,7 +1419,9 @@ async def rank( else: base_name = os.path.splitext(file_name)[0] final_path = os.path.join(save_dir, f"{base_name}_final.csv") - return _load_cached_dataframe(final_path, task_name="Rank") + return _project_public_rank_columns( + _load_cached_dataframe(final_path, task_name="Rank") + ) cfg_kwargs, response_kwargs = _split_cfg_and_response_kwargs( RankConfig, dict(cfg_kwargs), @@ -1327,6 +1435,7 @@ async def rank( matches_per_round=matches_per_round, power_matching=power_matching, learning_rate=learning_rate, + insufficient_signal_policy=insufficient_signal_policy, model=model, n_parallels=n_parallels, n_attributes_per_run=n_attributes_per_run, @@ -1351,6 +1460,7 @@ async def rank( primer_scores=primer_scores, primer_scale=primer_scale, primer_center=primer_center, + judge_version=judge_version, **cfg_kwargs, ) result_df = await Rank(cfg, template_path=template_path).run( @@ -1363,26 +1473,9 @@ async def rank( **response_kwargs, ) - # By default only expose the z-score columns (attribute names without suffixes) - # to API callers while keeping the raw/SE columns persisted in the CSV output. - if return_raw_scores: - return result_df - - if isinstance(attributes, dict): - attr_keys: List[str] = list(attributes.keys()) - else: - attr_keys = list(attributes) - drop_cols: List[str] = [] - for attr in attr_keys: - raw_col = f"{attr}_raw" - se_col = f"{attr}_se" - if raw_col in result_df.columns: - drop_cols.append(raw_col) - if se_col in result_df.columns: - drop_cols.append(se_col) - if drop_cols: - result_df = result_df.drop(columns=drop_cols) - return result_df + # Keep raw estimates and uncertainty persisted while presenting the same + # public projection for freshly computed and cached results. + return _project_public_rank_columns(result_df) async def codify( diff --git a/src/gabriel/tasks/_attribute_utils.py b/src/gabriel/tasks/_attribute_utils.py index 97d0f28..eccc065 100644 --- a/src/gabriel/tasks/_attribute_utils.py +++ b/src/gabriel/tasks/_attribute_utils.py @@ -11,14 +11,15 @@ def load_persisted_attributes( task_name: str, item_name: str = "attributes", legacy_filename: Optional[str] = None, + persist_missing: bool = True, ) -> Dict[str, Any]: """Load attributes/labels from disk for reproducibility. Preference order: 1) ``attributes.json`` in ``save_dir`` 2) ``legacy_filename`` (e.g., ``ratings_attrs.json``) in ``save_dir`` - When neither exists, ``incoming`` is written to both paths (when - applicable) for future runs. + When neither exists and ``persist_missing`` is true, ``incoming`` is + written to both paths (when applicable) for future runs. """ primary_path = os.path.join(save_dir, "attributes.json") @@ -57,13 +58,14 @@ def load_persisted_attributes( print(message) return loaded - for path in candidate_paths: - if not path: - continue - try: - with open(path, "w") as f: - json.dump(incoming, f, indent=2) - except Exception: - pass + if persist_missing: + for path in candidate_paths: + if not path: + continue + try: + with open(path, "w") as f: + json.dump(incoming, f, indent=2) + except Exception: + pass return incoming diff --git a/src/gabriel/tasks/_run_utils.py b/src/gabriel/tasks/_run_utils.py index 579ac52..6680574 100644 --- a/src/gabriel/tasks/_run_utils.py +++ b/src/gabriel/tasks/_run_utils.py @@ -4,6 +4,7 @@ import json import os import re +import tempfile from typing import Any, Dict, List, Optional, Sequence, Tuple, TypeVar import pandas as pd @@ -44,7 +45,21 @@ def load_run_metadata(save_dir: str, base_name: str, *, reset_files: bool) -> Di return payload if isinstance(payload, dict) else {} -def update_run_metadata(save_dir: str, base_name: str, **updates: Any) -> None: +def update_run_metadata( + save_dir: str, + base_name: str, + *, + strict: bool = False, + **updates: Any, +) -> None: + """Merge and atomically persist task metadata. + + Most tasks historically treated metadata as best-effort bookkeeping, so + the default retains that behavior. Tasks whose checkpoint correctness + depends on a commit marker can set ``strict=True`` to surface read or write + failures rather than silently advancing with ambiguous state. + """ + path = run_metadata_path(save_dir, base_name) payload: Dict[str, Any] = {} try: @@ -53,16 +68,37 @@ def update_run_metadata(save_dir: str, base_name: str, **updates: Any) -> None: loaded = json.load(f) if isinstance(loaded, dict): payload.update(loaded) + elif strict: + raise ValueError(f"Run metadata {path!r} is not a JSON object") except Exception: + if strict: + raise payload = {} payload.setdefault("metadata_version", 1) payload.update(updates) + temporary_path: Optional[str] = None try: os.makedirs(save_dir, exist_ok=True) - with open(path, "w", encoding="utf-8") as f: + file_descriptor, temporary_path = tempfile.mkstemp( + prefix=f".{base_name}_run_metadata.", + suffix=".tmp", + dir=save_dir, + text=True, + ) + with os.fdopen(file_descriptor, "w", encoding="utf-8") as f: json.dump(payload, f, indent=2) + f.flush() + os.fsync(f.fileno()) + os.replace(temporary_path, path) + temporary_path = None except Exception: - pass + if temporary_path is not None: + try: + os.unlink(temporary_path) + except OSError: + pass + if strict: + raise def hash_identifier(value: Any, *, bits: int = DEFAULT_IDENTIFIER_HASH_BITS) -> str: @@ -303,11 +339,14 @@ def write_task_run_metadata( identifier_hash_bits: int, n_attributes_per_run: Optional[int], attribute_batches: Sequence[Sequence[Tuple[str, Any]]], + strict: bool = False, + **task_metadata: Any, ) -> None: attr_batches = metadata_attribute_batches(attribute_batches) update_run_metadata( save_dir, base_name, + strict=strict, task=task_name, output_base_name=base_name, model=model, @@ -315,4 +354,5 @@ def write_task_run_metadata( n_attributes_per_run=n_attributes_per_run, attribute_count=sum(len(batch) for batch in attr_batches), attribute_batches=attr_batches, + **task_metadata, ) diff --git a/src/gabriel/tasks/rank.py b/src/gabriel/tasks/rank.py index 3491105..a7d3862 100644 --- a/src/gabriel/tasks/rank.py +++ b/src/gabriel/tasks/rank.py @@ -20,35 +20,42 @@ * Support for the new rankings prompt (``rankings_prompt.jinja2``) which allows the large language model to return one of four outcomes for each attribute: ``"circle"``, ``"square"``, ``"draw`` - or ``"insufficient signal"``. ``draw`` and ``insufficient signal`` - are both treated as a tie and contribute equally to both items when - fitting the BT model. + or ``"insufficient signal"``. A draw contributes one half-win to each + item. ``insufficient signal`` can use the same equality interpretation + or be treated as an abstention via ``insufficient_signal_policy``. * A cleaned up asynchronous ``run`` method that accepts a pandas ``DataFrame`` and the name of the column containing the text to be ranked. Each row receives a stable identifier derived from a hash of its contents; no external ``id_col`` argument is required. The method - produces a DataFrame with one row per input passage, a numeric - rating for each attribute, along with z‑scores and standard errors, - and writes the results to disk under ``save_dir``. - -The core ranking logic remains largely unchanged from ``elo.py`` -because the underlying mathematics of the BT model and the pairing -strategies continue to work well. However, comments have been added -throughout the code to clarify intent and to highlight areas where -further experimentation (e.g. alternative information gain metrics) can -be incorporated. + produces a DataFrame with one row per input passage and writes results + under ``save_dir``. Non-recursive runs report component-relative raw + scores, z-scores, standard errors, and graph labels; recursive runs report + stage-relative scores, exit stages, and an overall rank. + +The Bradley–Terry implementation fits each connected observed comparison +component, regularizes only observed edges when ``learning_rate > 0``, and +reports component-aware scores and model-based uncertainty diagnostics. +Persisted runs include estimator metadata so outcomes are never silently +replayed under different missing-signal or regularization semantics. """ from __future__ import annotations import os +import json from pathlib import Path import random +import re import math import copy +import shutil +import tempfile +import warnings +from numbers import Integral, Real from dataclasses import dataclass, field, fields -from typing import Any, Callable, Dict, List, Optional, Sequence, Set, Tuple, Union +from typing import Any, Callable, Dict, List, Mapping, Optional, Sequence, Set, Tuple, Union +from urllib.parse import urlsplit, urlunsplit import numpy as np import pandas as pd @@ -72,12 +79,278 @@ from ._run_utils import ( hash_identifier, load_run_metadata, + run_metadata_path, resolve_attribute_batches, resolve_identifier_hash_bits, + update_run_metadata, write_task_run_metadata, ) +WeightedOutcome = Tuple[str, str, float] +_RANK_ESTIMATOR_VERSION = 2 +_RANK_OWNED_RATE_KEYS = frozenset({"attributes", "save_dir", "file_name"}) +_RANK_OWNED_RESPONSE_KEYS = frozenset( + { + "prompts", + "identifiers", + "prompt_images", + "prompt_audio", + "prompt_pdfs", + "prompt_web_search_filters", + "model", + "n_parallels", + "json_mode", + "save_path", + "reset_files", + "use_dummy", + "max_retries", + "reasoning_effort", + } +) +_RANK_OPERATIONAL_RESPONSE_KEYS = frozenset( + { + "api_key", + "max_output_tokens", + "estimated_output_tokens_per_prompt", + "print_example_prompt", + "skip_tail_fails", + "ramp_up_seconds", + "ramp_up_start_fraction", + "timeout", + "timeout_factor", + "max_timeout", + "dynamic_timeout", + "timeout_burst_window", + "timeout_burst_cooldown", + "timeout_burst_max_restarts", + "background_mode", + "background_poll_interval", + "cancel_existing_batch", + "use_batch", + "batch_completion_window", + "batch_poll_interval", + "batch_wait_for_completion", + "max_batch_requests", + "max_batch_file_bytes", + "save_every_x_responses", + "verbose", + "quiet", + "global_cooldown", + "rate_limit_window", + "connection_error_window", + "token_sample_size", + "status_report_interval", + "planning_rate_limit_buffer", + "logging_level", + "service_tier", + "include", + "return_raw", + "request_phase_callback", + } +) + + +def _canonical_rank_spec_value(value: Any) -> Any: + """Return a deterministic, secret-free representation of judge settings.""" + + if value is None or isinstance(value, (str, bool)): + return value + if isinstance(value, Integral) and not isinstance(value, bool): + return int(value) + if isinstance(value, Real) and not isinstance(value, bool): + number = float(value) + if not math.isfinite(number): + raise ValueError("Rank judgment settings must contain finite numbers") + return number + if isinstance(value, Path): + return str(value) + if isinstance(value, Mapping): + if any(not isinstance(key, str) for key in value): + raise ValueError("Rank judgment-setting mappings must use string keys") + return { + key: _canonical_rank_spec_value(value[key]) + for key in sorted(value) + } + if isinstance(value, (list, tuple)): + return [_canonical_rank_spec_value(item) for item in value] + if isinstance(value, (set, frozenset)): + items = [_canonical_rank_spec_value(item) for item in value] + return sorted( + items, + key=lambda item: json.dumps( + item, sort_keys=True, ensure_ascii=False, separators=(",", ":") + ), + ) + if callable(value): + target = value if hasattr(value, "__qualname__") else type(value) + return { + "callable_module": getattr(target, "__module__", ""), + "callable_qualname": getattr(target, "__qualname__", ""), + } + raise ValueError( + "Rank cannot fingerprint a non-serializable judgment setting of type " + f"{type(value).__name__!r}; pass a stable, serializable value instead" + ) + + +def _sanitized_rank_endpoint(value: Any) -> str: + """Canonicalize an API endpoint without persisting credentials or queries.""" + + parsed = urlsplit(str(value)) + host = parsed.hostname or "" + if parsed.port is not None: + host = f"{host}:{parsed.port}" + return urlunsplit((parsed.scheme.lower(), host.lower(), parsed.path, "", "")) + + +def _rank_batch_state_has_external_work(path: Path) -> bool: + """Fail closed on malformed state and identify any durable Batch job.""" + + try: + with path.open(encoding="utf-8") as state_file: + state = json.load(state_file) + except Exception as exc: + raise ValueError( + f"Could not safely inspect Rank Batch API state {str(path)!r}" + ) from exc + batches = state.get("batches", []) if isinstance(state, dict) else None + if not isinstance(batches, list): + raise ValueError(f"Rank Batch API state {str(path)!r} is malformed") + return bool(batches) or bool(state.get("batch_id")) + + +def _validate_rank_attribute_names( + attributes: Union[Dict[str, Any], List[str]], *, recursive: bool +) -> List[str]: + """Validate every column namespace generated by effective Rank attributes.""" + + if not isinstance(attributes, (dict, list)) or not attributes: + raise ValueError("attributes must contain at least one attribute") + attribute_names = ( + list(attributes.keys()) if isinstance(attributes, dict) else list(attributes) + ) + if any( + not isinstance(attribute, str) or not attribute.strip() + for attribute in attribute_names + ): + raise ValueError("attribute names must be nonblank strings") + canonical_attributes = [ + attribute.strip().lower() for attribute in attribute_names + ] + if len(canonical_attributes) != len(set(canonical_attributes)): + raise ValueError( + "attribute names must be unique after case and whitespace normalization" + ) + generated_column_owners: Dict[str, str] = {} + for attribute in attribute_names: + for generated_column in ( + attribute, + f"{attribute}_raw", + f"{attribute}_se", + f"{attribute}_component", + ): + owner = generated_column_owners.get(generated_column) + if owner is not None and owner != attribute: + raise ValueError( + "Rank attribute output namespaces overlap: " + f"{owner!r} and {attribute!r} both generate " + f"{generated_column!r}" + ) + generated_column_owners[generated_column] = attribute + if recursive: + reserved_recursive_names = {"identifier", "overall_rank", "exit_stage"} + invalid_recursive_names = [ + attribute + for attribute in attribute_names + if attribute.strip().lower() in reserved_recursive_names + or re.match(r"^stage\d+_", attribute.strip(), flags=re.IGNORECASE) + ] + if invalid_recursive_names: + raise ValueError( + "Recursive Rank attribute names cannot use internal output " + "names or the stage_ prefix: " + + ", ".join(repr(name) for name in invalid_recursive_names) + ) + return attribute_names + + +def _validate_rank_resume_metadata( + metadata: Dict[str, Any], + *, + artifacts_exist: bool, + reset_files: bool, + insufficient_signal_policy: str, + learning_rate: float, + modality: str, + measurement_spec_fingerprint: str, +) -> None: + """Fail closed when persisted judgments use incompatible semantics.""" + + if reset_files or not artifacts_exist: + return + + mismatches: List[str] = [] + saved_version = metadata.get("rank_estimator_version") + if type(saved_version) is not int or saved_version != _RANK_ESTIMATOR_VERSION: + mismatches.append( + "rank_estimator_version " + f"(saved={saved_version!r}, required={_RANK_ESTIMATOR_VERSION})" + ) + + saved_policy = metadata.get("insufficient_signal_policy") + if saved_policy != insufficient_signal_policy: + mismatches.append( + "insufficient_signal_policy " + f"(saved={saved_policy!r}, requested={insufficient_signal_policy!r})" + ) + + saved_learning_rate = metadata.get("learning_rate") + valid_saved_rate = isinstance(saved_learning_rate, (int, float)) and not isinstance( + saved_learning_rate, bool + ) + try: + saved_rate_float = float(saved_learning_rate) + valid_saved_rate = valid_saved_rate and math.isfinite(saved_rate_float) + except (TypeError, ValueError, OverflowError): + saved_rate_float = math.nan + valid_saved_rate = False + if not valid_saved_rate or saved_rate_float != float(learning_rate): + mismatches.append( + "learning_rate " + f"(saved={saved_learning_rate!r}, requested={learning_rate!r})" + ) + + last_completed_round = metadata.get("last_completed_round") + if ( + type(last_completed_round) is not int + or last_completed_round < -1 + ): + mismatches.append( + "last_completed_round " + f"(saved={last_completed_round!r}, required=an integer >= -1)" + ) + + if not isinstance(metadata.get("input_fingerprints"), dict): + mismatches.append("input_fingerprints (missing or malformed)") + saved_modality = metadata.get("modality") + if saved_modality != modality: + mismatches.append( + f"modality (saved={saved_modality!r}, requested={modality!r})" + ) + saved_measurement_spec = metadata.get("measurement_spec_fingerprint") + if saved_measurement_spec != measurement_spec_fingerprint: + mismatches.append("measurement_spec_fingerprint (changed or missing)") + + if mismatches: + raise ValueError( + "Existing Rank artifacts are incompatible with the current " + "Bradley-Terry estimator configuration: " + + "; ".join(mismatches) + + ". Use reset_files=True or a new save_dir to recompute them." + ) + + def _is_missing_scalar(value: Any) -> bool: if value is None: return True @@ -88,6 +361,12 @@ def _is_missing_scalar(value: Any) -> bool: return bool(result) if isinstance(result, (bool, np.bool_)) else False +def _is_valid_identifier(value: Any) -> bool: + """Return whether a user identifier can be persisted unambiguously.""" + + return not _is_missing_scalar(value) and bool(str(value).strip()) + + def _hash_text_identifier( value: Any, *, @@ -110,6 +389,96 @@ def _hash_text_identifier( return hash_identifier(text, bits=bits) +def _rank_input_fingerprints( + df: pd.DataFrame, + *, + id_column: str, + payload_column: str, + modality: str, +) -> Dict[str, str]: + """Fingerprint the effective comparison payload for each stable ID. + + Local media are encoded before hashing so replacing a file in place cannot + silently reuse judgments about its previous contents. Remote locators are + hashed as locators because fetching them here would duplicate model I/O. + """ + + fingerprints: Dict[str, str] = {} + for item_id, payload in zip(df[id_column], df[payload_column]): + if modality == "image": + effective_payload: Any = load_image_inputs(payload) + elif modality == "audio": + effective_payload = load_audio_inputs(payload) + elif modality == "pdf": + effective_payload = load_pdf_inputs(payload) + else: + effective_payload = payload + serialized = json.dumps( + {"modality": modality, "payload": effective_payload}, + sort_keys=True, + ensure_ascii=False, + separators=(",", ":"), + default=repr, + ) + fingerprints[str(item_id)] = hash_identifier(serialized, bits=160) + return fingerprints + + +def _current_rank_input_fingerprints( + df: pd.DataFrame, + *, + payload_column: str, + id_column: Optional[str], + modality: str, + identifier_hash_bits: int, +) -> Dict[str, str]: + """Prepare valid IDs and fingerprint current effective payloads.""" + + if payload_column not in df.columns: + raise ValueError(f"column_name '{payload_column}' not found in DataFrame") + prepared = df.reset_index(drop=True).copy() + if id_column is not None: + if id_column not in prepared.columns: + raise ValueError(f"id_column '{id_column}' not found in DataFrame") + valid_mask = pd.Series( + [_is_valid_identifier(value) for value in prepared[id_column]], + index=prepared.index, + dtype=bool, + ) + prepared = prepared.loc[valid_mask].copy().reset_index(drop=True) + prepared["_id"] = prepared[id_column].astype(str) + else: + strict = modality in {"text", "entity", "web"} + hashed_ids = prepared[payload_column].map( + lambda value: _hash_text_identifier( + value, + strict=strict, + bits=identifier_hash_bits, + ) + ) + valid_mask = hashed_ids.notna() + prepared = prepared.loc[valid_mask].copy().reset_index(drop=True) + prepared["_id"] = ( + hashed_ids.loc[valid_mask].astype(str).reset_index(drop=True) + ) + duplicate_ids = prepared.loc[ + prepared["_id"].duplicated(keep=False), "_id" + ].unique() + if len(duplicate_ids) > 0: + preview = ", ".join(repr(value) for value in duplicate_ids[:3]) + raise ValueError( + "Rank requires a unique identifier for every row; duplicate " + f"identifier(s): {preview}. Provide a unique id_column when " + "ranking duplicate content." + ) + return _rank_input_fingerprints( + prepared, + id_column="_id", + payload_column=payload_column, + modality=modality, + ) + + @dataclass class RankConfig: """User‑visible configuration for :class:`Rank`. @@ -129,15 +498,35 @@ class RankConfig: n_rounds: Number of rounds of pairwise comparisons to perform. matches_per_round: - Number of matches per item per round. + Target matches per item per round, capped at the number of possible + opponents. Realized degrees may exceed the target because deduplicated + undirected edges can also select an item as another item's opponent. power_matching: - Whether to use an information‑theoretic pairing heuristic. The - resulting rankings always include per‑attribute z‑scores alongside the - raw Bradley–Terry estimates (``"_raw"``) and their - standard errors (``"_se"``). + Whether to use an uncertainty-guided pairing heuristic. If ``False``, + pairs are sampled randomly; random mode can exceed the target when an + item is chosen by several other items. Non-recursive rankings include + per‑attribute z‑scores alongside the raw Bradley–Terry estimates + (``"_raw"``) and their standard errors + (``"_se"``). learning_rate: - Pseudo‑count used by the BT model to regularise the win/loss - matrix. A larger value makes updates more conservative. + Total symmetric pseudo‑comparison weight added to each observed + pair by the BT model. A larger value shrinks pairwise estimates + more strongly toward a draw. Reported standard errors are asymptotic, + model-based sandwich standard errors for component-centered raw + log-skills, computed while treating the realized comparison graph as + fixed. They exclude shrinkage bias, adaptive-selection uncertainty, and + are not standard errors for the reported z-scores. Pseudo-comparison + curvature is present only when ``learning_rate`` is positive. A zero + value requires Ford's directed strong-connectivity condition within + every nontrivial observed component for a finite maximum-likelihood fit. + insufficient_signal_policy: + How to use an ``"insufficient signal"`` judgment. ``"tie"`` + records one half‑win per item, making the modeling assumption that + neither item manifests the attribute more strongly. ``"abstain"`` + records no comparison when absence of direct evidence should be + treated as missingness; informative abstention can still bias the + retained sample. A true draw uses half-wins in a binary Bradley–Terry + working objective rather than a separate three-outcome likelihood. model: Name of the language model to call via ``get_all_responses``. n_parallels: @@ -164,7 +553,7 @@ class RankConfig: recursive_rate_first_round: If ``True`` perform a :class:`Rate` sweep before the first recursive stage and seed subsequent rounds with those scores. - This is enabled by default so the initial culling uses grounded + This is enabled by default so the initial culling uses model-derived single-pass ratings; set to ``False`` to skip. recursive_rewrite_func, recursive_rewrite_text_col: Optional hook to rewrite surviving passages between stages and @@ -181,10 +570,17 @@ class RankConfig: rate_kwargs: Optional dictionary of overrides forwarded to the rating task whenever it is invoked (either as a seed or during recursion). + Rank owns ``attributes``, ``save_dir``, and ``file_name`` so those + keys cannot be overridden here. primer_scores, primer_scale, primer_center: Optional manual primers to seed the Bradley–Terry rating state. Scores are centred per attribute when ``primer_center`` is ``True`` and scaled by ``primer_scale``. + judge_version: + Stable label required when a custom response function or external judge + is supplied. Change it whenever that judge's hidden prompt, logic, or + configuration changes so persisted rounds cannot silently mix + measurement regimes. """ attributes: Union[Dict[str, str], List[str]] @@ -192,6 +588,7 @@ class RankConfig: matches_per_round: int = 5 power_matching: bool = True learning_rate: float = 0.1 + insufficient_signal_policy: str = "tie" model: str = "gpt-5.6-luna" n_parallels: int = 650 use_dummy: bool = False @@ -224,11 +621,79 @@ class RankConfig: primer_scores: Optional[Dict[str, Dict[str, float]]] = None primer_scale: float = 1.0 primer_center: bool = True + judge_version: Optional[str] = None def __post_init__(self) -> None: + _validate_rank_attribute_names(self.attributes, recursive=self.recursive) + if not isinstance(self.rate_kwargs, dict): + raise ValueError("rate_kwargs must be a dictionary") + conflicting_rate_keys = sorted( + _RANK_OWNED_RATE_KEYS & self.rate_kwargs.keys() + ) + if conflicting_rate_keys: + raise ValueError( + "rate_kwargs cannot override Rank-owned Rate setting(s): " + + ", ".join(conflicting_rate_keys) + ) if self.additional_instructions is not None: cleaned = str(self.additional_instructions).strip() self.additional_instructions = cleaned or None + if self.insufficient_signal_policy not in {"tie", "abstain"}: + raise ValueError( + "insufficient_signal_policy must be either 'tie' or 'abstain'" + ) + for name in ("n_rounds", "matches_per_round"): + value = getattr(self, name) + if ( + not isinstance(value, Integral) + or isinstance(value, bool) + or int(value) < 1 + ): + raise ValueError(f"{name} must be an integer of at least 1") + setattr(self, name, int(value)) + if not isinstance(self.learning_rate, Real) or isinstance( + self.learning_rate, bool + ): + raise ValueError("learning_rate must be a finite non-negative number") + self.learning_rate = float(self.learning_rate) + if not math.isfinite(self.learning_rate) or self.learning_rate < 0: + raise ValueError("learning_rate must be a finite non-negative number") + if not isinstance(self.primer_scale, Real) or isinstance( + self.primer_scale, bool + ): + raise ValueError("primer_scale must be a finite number") + self.primer_scale = float(self.primer_scale) + if not math.isfinite(self.primer_scale): + raise ValueError("primer_scale must be a finite number") + if type(self.primer_center) is not bool: + raise ValueError("primer_center must be a boolean") + if self.judge_version is not None: + if not isinstance(self.judge_version, str) or not self.judge_version.strip(): + raise ValueError("judge_version must be a nonblank string or None") + self.judge_version = self.judge_version.strip() + if not isinstance(self.recursive_fraction, Real) or isinstance( + self.recursive_fraction, bool + ): + raise ValueError("recursive_fraction must be between 0 and 1") + self.recursive_fraction = float(self.recursive_fraction) + if ( + not math.isfinite(self.recursive_fraction) + or self.recursive_fraction <= 0 + or self.recursive_fraction >= 1 + ): + raise ValueError("recursive_fraction must be between 0 and 1") + for name in ( + "recursive_min_remaining", + "recursive_final_round_multiplier", + ): + value = getattr(self, name) + if ( + not isinstance(value, Integral) + or isinstance(value, bool) + or int(value) < 1 + ): + raise ValueError(f"{name} must be an integer of at least 1") + setattr(self, name, int(value)) class Rank: @@ -238,8 +703,10 @@ class Rank: of sampling pairs, calling a language model to adjudicate which passage better exhibits each attribute, and then fitting a Bradley–Terry model to those outcomes. Standard errors and - z‑scores are computed for every attribute. Results are persisted to disk - after the final round. + component-aware z‑scores are computed for every attribute in + non-recursive mode. A component with no score dispersion receives the + neutral z-score fallback of zero; a singleton remains unranked. Results + are persisted to disk after the final round. """ def __init__( @@ -280,6 +747,7 @@ def __init__( # place holders for multiway rankings and aggregated standard errors self.history_multi: Dict[str, List[List[str]]] = {} self._last_se_agg: Optional[Dict[str, float]] = None + self._warned_disconnected_graph = False # internal constants for the pairing and BT algorithms. These # values are deliberately not exposed through the public API as @@ -291,14 +759,9 @@ def __init__( self._HIGH_SE_FRAC = 0.25 # fraction of high‑uncertainty items self._MAX_ITER = 1000 # maximum iterations for BT optimisation self._TOL = 1e-6 # convergence tolerance for BT - # A small ridge term stabilises the inversion of the Fisher information - # matrix when computing standard errors. The previous value (1e‑9) - # occasionally led to extremely large uncertainties for items with - # limited or contradictory comparisons. Increasing this value - # regularises the covariance estimate and prevents unreasonably - # large standard errors. If you observe inflated SE values, - # consider increasing this further (e.g. to 1e‑4). - self._SE_RIDGE = 1e-5 + # Relative eigenvalue tolerance for the Fisher-information pseudoinverse. + # This is a numerical cutoff only; it must not add synthetic information. + self._SE_EIGEN_TOL = 1e-12 # The maximum number of candidate pairs to consider per pairing round. # When the number of items becomes very large (e.g. tens of thousands), # evaluating all possible pairs is intractable. We therefore cap the @@ -309,6 +772,81 @@ def __init__( # 200 000, each item will only consider approximately 20 neighbours. self._MAX_CANDIDATE_PAIRS_PER_ROUND = 200_000 + def _measurement_spec_fingerprint( + self, runtime_kwargs: Optional[Mapping[str, Any]] = None + ) -> str: + """Hash judgment settings that must not change within a tournament.""" + + effective_runtime: Dict[str, Any] = { + "web_search": self.cfg.modality == "web", + "image_detail": None, + "n": 1, + "temperature": 0.9, + "expected_schema": None, + "tools": None, + "tool_choice": None, + "reasoning_summary": None, + "dummy_responses": None, + "response_fn": None, + "get_all_responses_fn": None, + } + for key, value in dict(runtime_kwargs or {}).items(): + if key in _RANK_OPERATIONAL_RESPONSE_KEYS: + continue + effective_runtime[key] = value + effective_runtime["web_search"] = bool( + effective_runtime.get("web_search") + ) + if effective_runtime.get("image_detail") in {None, "none"}: + effective_runtime["image_detail"] = None + if effective_runtime.get("web_search"): + effective_runtime.setdefault("web_search_filters", None) + context_size = effective_runtime.get("search_context_size", "medium") + effective_runtime["search_context_size"] = { + "small": "low", + "large": "high", + }.get(context_size, context_size) + else: + effective_runtime.pop("web_search_filters", None) + effective_runtime.pop("search_context_size", None) + endpoint = effective_runtime.get("base_url") or os.getenv("OPENAI_BASE_URL") + if endpoint: + effective_runtime["base_url"] = _sanitized_rank_endpoint(endpoint) + else: + effective_runtime.pop("base_url", None) + + semantic_rate_kwargs = { + key: value + for key, value in self.cfg.rate_kwargs.items() + if key not in _RANK_OPERATIONAL_RESPONSE_KEYS + and key != "n_parallels" + } + payload = _canonical_rank_spec_value( + { + "model": self.cfg.model, + "template": self.template.text, + "attributes": self.cfg.attributes, + "recursive": self.cfg.recursive, + "recursive_rate_first_round": self.cfg.recursive_rate_first_round, + "recursive_rewrite_func": self.cfg.recursive_rewrite_func, + "initial_rating_pass": self.cfg.initial_rating_pass, + "rate_kwargs": semantic_rate_kwargs, + "additional_instructions": self.cfg.additional_instructions, + "circle_first": self.cfg.circle_first, + "reasoning_effort": self.cfg.reasoning_effort, + "use_dummy": self.cfg.use_dummy, + "judge_version": self.cfg.judge_version, + "runtime_judgment": effective_runtime, + } + ) + serialized = json.dumps( + payload, + sort_keys=True, + ensure_ascii=False, + separators=(",", ":"), + ) + return hash_identifier(serialized, bits=160) + # ------------------------------------------------------------------ def _apply_primer( self, @@ -333,7 +871,9 @@ def _apply_primer( for attr in attr_keys: if attr in amap and amap[attr] is not None: try: - attr_to_vals[attr].append(float(amap[attr])) + value = float(amap[attr]) + if math.isfinite(value): + attr_to_vals[attr].append(value) except Exception: continue @@ -343,7 +883,7 @@ def _apply_primer( if vals: attr_offset[attr] = float(np.mean(vals)) - scale = self.cfg.primer_scale or 1.0 + scale = self.cfg.primer_scale for ident, amap in primer.items(): if ident not in ratings: continue @@ -351,8 +891,11 @@ def _apply_primer( if attr not in amap or amap[attr] is None: continue try: - val = (float(amap[attr]) - attr_offset[attr]) * scale - ratings[ident][attr] = val + value = float(amap[attr]) + if math.isfinite(value): + ratings[ident][attr] = ( + value - attr_offset[attr] + ) * scale except Exception: continue @@ -375,12 +918,20 @@ def _attributes_as_dict(self) -> Dict[str, str]: return dict(self.cfg.attributes) return {attr: "" for attr in self.cfg.attributes} - def _split_rate_kwargs(self, overrides: Optional[Dict[str, Any]] = None) -> Tuple[Dict[str, Any], Dict[str, Any]]: + def _split_rate_kwargs( + self, overrides: Optional[Dict[str, Any]] = None + ) -> Tuple[Dict[str, Any], Dict[str, Any]]: merged: Dict[str, Any] = {} if self.cfg.rate_kwargs: merged.update(self.cfg.rate_kwargs) if overrides: merged.update(overrides) + conflicting_rate_keys = sorted(_RANK_OWNED_RATE_KEYS & merged.keys()) + if conflicting_rate_keys: + raise ValueError( + "rate_kwargs cannot override Rank-owned Rate setting(s): " + + ", ".join(conflicting_rate_keys) + ) config_fields = {f.name for f in fields(RateConfig)} cfg_kwargs: Dict[str, Any] = {} run_kwargs: Dict[str, Any] = {} @@ -448,10 +999,11 @@ def _seed_ratings_from_rate( if id_column and id_column in rate_df.columns: key_series = rate_df[id_column].astype(str) elif text_column in rate_df.columns: + strict_payload = self.cfg.modality in {"text", "entity", "web"} key_series = rate_df[text_column].map( lambda x: _hash_text_identifier( x, - strict=True, + strict=strict_payload, bits=identifier_hash_bits, ) ) @@ -474,13 +1026,1087 @@ def _seed_ratings_from_rate( # Only retain seeds for items that will appear in the ranking loop return {item_id: seeds[item_id] for item_id in item_ids if item_id in seeds} - # ------------------------------------------------------------------ - # BT / PL fitting utilities - # ------------------------------------------------------------------ + # ------------------------------------------------------------------ + # BT / PL fitting utilities + # ------------------------------------------------------------------ + @staticmethod + def _matches_outcome_label(value: str, aliases: Sequence[str]) -> bool: + """Match only explicit, complete labels from the response contract.""" + + return value in aliases + + def _decode_pairwise_outcome( + self, + id_a: str, + id_b: str, + raw_value: Any, + ) -> Tuple[str, List[WeightedOutcome]]: + """Decode one model judgment into evidence for the BT model. + + The prompt distinguishes a true draw from a lack of direct signal. + Keeping those labels separate here lets callers choose whether the + latter is an equality judgment or an abstention. Fractional wins make + every non-abstaining response contribute exactly one comparison. + """ + + value = raw_value + if isinstance(value, dict): + winner_key = next( + (key for key in value if str(key).strip().lower() == "winner"), + None, + ) + value = value.get(winner_key) if winner_key is not None else None + if not isinstance(value, str): + return "invalid", [] + + label = " ".join(value.strip().lower().split()) + if self._matches_outcome_label(label, ("insufficient signal", "insufficient")): + if self.cfg.insufficient_signal_policy == "abstain": + return "insufficient_signal", [] + return "insufficient_signal", [ + (id_a, id_b, 0.5), + (id_b, id_a, 0.5), + ] + if self._matches_outcome_label(label, ("draw", "tie", "equal")): + return "draw", [(id_a, id_b, 0.5), (id_b, id_a, 0.5)] + if self._matches_outcome_label( + label, ("circle", "circle wins", "text a", "text a wins") + ): + return "circle", [(id_a, id_b, 1.0)] + if self._matches_outcome_label( + label, ("square", "square wins", "text b", "text b wins") + ): + return "square", [(id_b, id_a, 1.0)] + return "invalid", [] + + def _record_pairwise_outcome( + self, + history: List[WeightedOutcome], + id_a: str, + id_b: str, + raw_value: Any, + ) -> str: + """Append a decoded response to ``history`` and return its category.""" + + category, outcomes = self._decode_pairwise_outcome(id_a, id_b, raw_value) + history.extend(outcomes) + return category + + @staticmethod + def _successful_response_mask(frame: pd.DataFrame) -> pd.Series: + """Identify rows with both a successful request and a usable response.""" + + if "Response" not in frame.columns: + return pd.Series(False, index=frame.index, dtype=bool) + responses = frame["Response"] + response_mask = responses.notna() & responses.astype(str).str.strip().ne("") + if "Successful" not in frame.columns: + return response_mask.astype(bool) + + success_raw = frame["Successful"] + success_mask = pd.Series(False, index=frame.index, dtype=bool) + try: + success_mask |= success_raw.astype("boolean").fillna(False) + except (TypeError, ValueError): + pass + success_mask |= ( + success_raw.astype(str) + .str.strip() + .str.lower() + .isin({"true", "1", "yes", "y", "completed", "succeeded", "success"}) + ) + return (response_mask & success_mask).astype(bool) + + @classmethod + def _validate_requested_responses( + cls, + frame: pd.DataFrame, + requested_ids: Sequence[str], + *, + context: str, + allow_extra: bool, + ) -> pd.Series: + """Require one successful, nonblank response for every requested ID.""" + + required_columns = {"Identifier", "Response"} + missing_columns = required_columns - set(frame.columns) + if missing_columns: + raise ValueError( + f"{context} response table is missing columns: " + + ", ".join(sorted(missing_columns)) + ) + identifiers = frame["Identifier"].astype(str) + requested_set = set(requested_ids) + requested_mask = identifiers.isin(requested_set) + returned_requested = identifiers.loc[requested_mask] + if ( + returned_requested.duplicated().any() + or set(returned_requested) != requested_set + or (not allow_extra and int(requested_mask.sum()) != len(frame)) + ): + raise ValueError( + f"{context} did not return exactly one row for every requested " + "judgment" + ) + successful = cls._successful_response_mask(frame) + failed_requested = requested_mask & ~successful + if failed_requested.any(): + failed_preview = ", ".join( + repr(value) + for value in identifiers.loc[failed_requested].head(3) + ) + raise ValueError( + f"{context} returned failed or blank judgment response(s): " + f"{failed_preview}. The round was not committed and can be retried." + ) + return requested_mask + + @staticmethod + async def _coerce_response_dict(raw: Any) -> Dict[str, Any]: + """Decode the response container forms accepted by the task API.""" + + obj = await safest_json(raw) + if isinstance(obj, dict): + return obj + if isinstance(obj, str): + nested = await safest_json(obj) + if isinstance(nested, dict): + return nested + if isinstance(obj, list) and obj: + nested = await safest_json(obj[0]) + if isinstance(nested, dict): + return nested + return {} + + async def _validate_pairwise_response_payloads( + self, + frame: pd.DataFrame, + meta_map: Dict[str, Tuple[int, int, str, str]], + attr_batches: Sequence[Sequence[str]], + *, + context: str, + retry_path: Optional[str] = None, + ) -> None: + """Require complete, recognized outcomes before committing a round.""" + + invalid: Dict[str, str] = {} + for identifier, raw_response in zip( + frame["Identifier"].astype(str), frame["Response"] + ): + meta = meta_map.get(identifier) + if meta is None: + invalid[identifier] = "missing prompt metadata" + continue + batch_idx, _, id_a, id_b = meta + if batch_idx < 0 or batch_idx >= len(attr_batches): + invalid[identifier] = "invalid attribute batch" + continue + response_obj = await self._coerce_response_dict(raw_response) + expected = { + str(attribute).strip().lower(): str(attribute) + for attribute in attr_batches[batch_idx] + } + normalized: Dict[str, Any] = {} + duplicate_key = False + for raw_key, value in response_obj.items(): + key = str(raw_key).strip().lower() + if key in normalized: + duplicate_key = True + normalized[key] = value + if duplicate_key or set(normalized) != set(expected): + invalid[identifier] = ( + "response must contain exactly the requested attributes" + ) + continue + for key, value in normalized.items(): + category, _ = self._decode_pairwise_outcome(id_a, id_b, value) + if category == "invalid": + invalid[identifier] = ( + f"unrecognized outcome for {expected[key]!r}" + ) + break + + if not invalid: + return + if retry_path is not None: + retry_frame = frame.copy() + if "Successful" not in retry_frame.columns: + retry_frame["Successful"] = True + retry_mask = retry_frame["Identifier"].astype(str).isin(invalid) + retry_frame.loc[retry_mask, "Successful"] = False + if "Error Log" not in retry_frame.columns: + retry_frame["Error Log"] = "" + retry_frame.loc[retry_mask, "Error Log"] = retry_frame.loc[ + retry_mask, "Identifier" + ].astype(str).map(invalid) + self._write_dataframe_atomically(retry_frame, retry_path) + preview = "; ".join( + f"{identifier!r}: {reason}" + for identifier, reason in list(invalid.items())[:3] + ) + raise ValueError( + f"{context} returned semantically invalid judgment payload(s): " + f"{preview}. The round was not committed and can be retried." + ) + + @staticmethod + def _read_rank_checkpoint(path: str, n_attribute_batches: int) -> pd.DataFrame: + """Read and validate one committed v2 round without coercing IDs.""" + + try: + frame = pd.read_csv(path, dtype=str, keep_default_na=False) + except Exception as exc: + raise ValueError(f"Could not read Rank checkpoint {path!r}") from exc + + required_columns = { + "Identifier", + "Response", + "Batch", + "Pair", + "IdA", + "IdB", + } + missing_columns = required_columns - set(frame.columns) + if missing_columns: + raise ValueError( + f"Rank checkpoint {path!r} is missing structured columns: " + + ", ".join(sorted(missing_columns)) + ) + if frame.empty: + raise ValueError(f"Rank checkpoint {path!r} contains no judgments") + if not Rank._successful_response_mask(frame).all(): + raise ValueError( + f"Rank checkpoint {path!r} contains failed or blank responses" + ) + # Checkpoints written before the response collector exposed an explicit + # success flag are valid when every response is nonblank. Normalize + # them before a staged catch-up is concatenated; otherwise pandas adds + # ``Successful`` only for the new rows and the legacy rows become NaN, + # causing the atomic checkpoint validator to reject the merge. + if "Successful" not in frame.columns: + frame["Successful"] = True + identifiers = frame["Identifier"].str.strip() + if identifiers.eq("").any() or identifiers.duplicated().any(): + raise ValueError( + f"Rank checkpoint {path!r} has blank or duplicate identifiers" + ) + if frame["IdA"].str.strip().eq("").any() or frame[ + "IdB" + ].str.strip().eq("").any(): + raise ValueError(f"Rank checkpoint {path!r} has blank item IDs") + + batch_numbers = pd.to_numeric(frame["Batch"], errors="coerce") + pair_numbers = pd.to_numeric(frame["Pair"], errors="coerce") + valid_batches = ( + batch_numbers.notna() + & np.isfinite(batch_numbers) + & (batch_numbers >= 0) + & (batch_numbers < n_attribute_batches) + & (batch_numbers == np.floor(batch_numbers)) + ) + valid_pairs = ( + pair_numbers.notna() + & np.isfinite(pair_numbers) + & (pair_numbers >= 0) + & (pair_numbers == np.floor(pair_numbers)) + ) + if not valid_batches.all() or not valid_pairs.all(): + raise ValueError( + f"Rank checkpoint {path!r} has invalid Batch or Pair values" + ) + frame["Batch"] = batch_numbers.astype(int) + frame["Pair"] = pair_numbers.astype(int) + return frame + + @classmethod + def _write_rank_checkpoint( + cls, + frame: pd.DataFrame, + path: str, + n_attribute_batches: int, + ) -> None: + """Validate and atomically replace one structured round checkpoint.""" + + save_dir = os.path.dirname(path) or "." + os.makedirs(save_dir, exist_ok=True) + file_descriptor, temporary_path = tempfile.mkstemp( + prefix=f".{Path(path).name}.", + suffix=".tmp", + dir=save_dir, + text=True, + ) + os.close(file_descriptor) + try: + frame.to_csv(temporary_path, index=False) + cls._read_rank_checkpoint(temporary_path, n_attribute_batches) + os.replace(temporary_path, path) + except Exception: + try: + os.unlink(temporary_path) + except OSError: + pass + raise + + @staticmethod + def _write_dataframe_atomically(frame: pd.DataFrame, path: str) -> None: + """Durably replace a CSV staging checkpoint.""" + + save_dir = os.path.dirname(path) or "." + os.makedirs(save_dir, exist_ok=True) + file_descriptor, temporary_path = tempfile.mkstemp( + prefix=f".{Path(path).name}.", + suffix=".tmp", + dir=save_dir, + text=True, + ) + try: + with os.fdopen( + file_descriptor, "w", encoding="utf-8", newline="" + ) as temporary_file: + frame.to_csv(temporary_file, index=False) + temporary_file.flush() + os.fsync(temporary_file.fileno()) + os.replace(temporary_path, path) + except Exception: + try: + os.unlink(temporary_path) + except OSError: + pass + raise + + @staticmethod + def _write_json_atomically(payload: Dict[str, Any], path: str) -> None: + """Durably replace a JSON sidecar without exposing a partial file.""" + + save_dir = os.path.dirname(path) or "." + os.makedirs(save_dir, exist_ok=True) + file_descriptor, temporary_path = tempfile.mkstemp( + prefix=f".{Path(path).name}.", + suffix=".tmp", + dir=save_dir, + text=True, + ) + try: + with os.fdopen(file_descriptor, "w", encoding="utf-8") as handle: + json.dump(payload, handle, indent=2) + handle.flush() + os.fsync(handle.fileno()) + os.replace(temporary_path, path) + except Exception: + try: + os.unlink(temporary_path) + except OSError: + pass + raise + + @staticmethod + def _comparison_components(n_ij: np.ndarray) -> List[np.ndarray]: + """Return deterministic components of the observed comparison graph.""" + + n = n_ij.shape[0] + if n_ij.shape != (n, n): + raise ValueError("n_ij must be a square matrix") + visited = np.zeros(n, dtype=bool) + components: List[np.ndarray] = [] + for start in range(n): + if visited[start]: + continue + stack = [start] + visited[start] = True + component: List[int] = [] + while stack: + node = stack.pop() + component.append(node) + for neighbour in np.flatnonzero(n_ij[node] > 0): + neighbour_int = int(neighbour) + if not visited[neighbour_int]: + visited[neighbour_int] = True + stack.append(neighbour_int) + components.append(np.array(sorted(component), dtype=int)) + return components + + @classmethod + def _comparison_component_labels(cls, n_ij: np.ndarray) -> np.ndarray: + """Map every item to its observed comparison-graph component.""" + + labels = np.empty(n_ij.shape[0], dtype=int) + for label, component in enumerate(cls._comparison_components(n_ij)): + labels[component] = label + return labels + + @staticmethod + def _component_zscores( + values: np.ndarray, component_labels: np.ndarray + ) -> np.ndarray: + """Standardize scores within identifiable comparison components.""" + + if values.shape != component_labels.shape: + raise ValueError("values and component_labels must have the same shape") + zscores = np.full(len(values), np.nan, dtype=float) + for component in np.unique(component_labels): + members = component_labels == component + local_values = values[members] + if int(np.sum(members)) == 1: + continue + local_std = float(local_values.std(ddof=0)) + if local_std > 0: + zscores[members] = (local_values - local_values.mean()) / local_std + else: + zscores[members] = 0.0 + return zscores + + @staticmethod + def _is_strongly_connected( + observed_wins: np.ndarray, component: np.ndarray + ) -> bool: + """Check Ford's directed-win condition within one component.""" + + if len(component) < 2: + return False + local = observed_wins[np.ix_(component, component)] > 0 + + def reaches_all(adjacency: np.ndarray) -> bool: + visited = np.zeros(len(component), dtype=bool) + stack = [0] + visited[0] = True + while stack: + node = stack.pop() + for neighbour in np.flatnonzero(adjacency[node]): + neighbour_int = int(neighbour) + if not visited[neighbour_int]: + visited[neighbour_int] = True + stack.append(neighbour_int) + return bool(np.all(visited)) + + return reaches_all(local) and reaches_all(local.T) + + @staticmethod + def _build_bt_win_matrices( + item_ids: Sequence[str], + outcomes: Sequence[Union[Tuple[str, str], WeightedOutcome]], + pseudo: float, + ) -> Tuple[np.ndarray, np.ndarray]: + """Build separate observed and regularized directed-win matrices. + + ``pseudo`` is the total symmetric pseudo-comparison mass for each + observed undirected edge. Half is assigned as a win in each direction. + Unobserved edges and the diagonal receive no artificial comparisons. + """ + + if len(set(item_ids)) != len(item_ids): + raise ValueError("item_ids must contain unique values") + if not np.isfinite(pseudo) or pseudo < 0: + raise ValueError("pseudo must be a finite non-negative number") + + n = len(item_ids) + idx = {item: i for i, item in enumerate(item_ids)} + observed_wins = np.zeros((n, n), dtype=float) + for outcome in outcomes: + if len(outcome) == 2: + winner, loser = outcome + weight = 1.0 + elif len(outcome) == 3: + winner, loser, raw_weight = outcome + try: + weight = float(raw_weight) + except (TypeError, ValueError) as exc: + raise ValueError("outcome weights must be numeric") from exc + else: + raise ValueError( + "outcomes must be (winner, loser) or (winner, loser, weight)" + ) + if not np.isfinite(weight) or weight < 0: + raise ValueError("outcome weights must be finite and non-negative") + if weight == 0: + continue + if winner not in idx or loser not in idx: + raise ValueError( + "outcome identifiers must be present in item_ids; " + f"received ({winner!r}, {loser!r})" + ) + if winner == loser: + continue + observed_wins[idx[winner], idx[loser]] += weight + + np.fill_diagonal(observed_wins, 0.0) + observed_matches = observed_wins + observed_wins.T + edge_mask = observed_matches > 0 + fit_wins = observed_wins + 0.5 * pseudo * edge_mask + np.fill_diagonal(fit_wins, 0.0) + return observed_wins, fit_wins + + @staticmethod + def _bt_logit_spanning_tree_start( + local_fit_wins: np.ndarray, + ) -> Optional[np.ndarray]: + """Construct a scale-stable empirical-logit start for a BT component. + + Each bidirectionally observed edge supplies the empirical equation + ``s_i - s_j = log(w_ij) - log(w_ji)``. Integrating those equations + along a maximum-curvature spanning tree avoids squaring the condition + number in normal equations. On a tree this is the exact regularized + optimum, even when pseudo mass is 1e-300 or a long score chain lies far + outside the range in which abilities can be exponentiated. + """ + + size = int(local_fit_wins.shape[0]) + edge_data: List[Tuple[int, int, float, float]] = [] + for left in range(size): + for right in range(left + 1, size): + left_wins = float(local_fit_wins[left, right]) + right_wins = float(local_fit_wins[right, left]) + if left_wins <= 0 or right_wins <= 0: + continue + high = max(left_wins, right_wins) + low = min(left_wins, right_wins) + weight = low / (1.0 + low / high) + log_odds = math.log(left_wins) - math.log(right_wins) + if not np.isfinite(weight) or not np.isfinite(log_odds): + return None + edge_data.append((left, right, weight, log_odds)) + if not edge_data: + return None + + parent = list(range(size)) + + def find(node: int) -> int: + while parent[node] != node: + parent[node] = parent[parent[node]] + node = parent[node] + return node + + adjacency: List[List[Tuple[int, float]]] = [ + [] for _ in range(size) + ] + tree_edges = 0 + for left, right, _weight, log_odds in sorted( + edge_data, key=lambda edge: edge[2], reverse=True + ): + left_root = find(left) + right_root = find(right) + if left_root == right_root: + continue + parent[right_root] = left_root + adjacency[left].append((right, -log_odds)) + adjacency[right].append((left, log_odds)) + tree_edges += 1 + if tree_edges == size - 1: + break + if tree_edges != size - 1: + return None + + scores = np.full(size, np.nan, dtype=float) + scores[0] = 0.0 + stack = [0] + while stack: + node = stack.pop() + for neighbor, offset in adjacency[node]: + if np.isfinite(scores[neighbor]): + continue + scores[neighbor] = scores[node] + offset + stack.append(neighbor) + scores -= scores.mean() + if np.any(~np.isfinite(scores)): + return None + return scores + + @staticmethod + def _bt_strong_components(local_observed_wins: np.ndarray) -> List[np.ndarray]: + """Return deterministic SCCs of an observed directed-win graph. + + Tiny edge regularization can put SCC-wide location contrasts hundreds + of orders of magnitude below the within-SCC curvature. Keeping the + observed SCCs explicit lets the numerical fallback optimize those weak + contrasts without subtracting the O(1) information internal to an SCC. + """ + + size = int(local_observed_wins.shape[0]) + adjacency = local_observed_wins > 0 + + def reachable(start: int, graph: np.ndarray) -> Set[int]: + visited = {start} + stack = [start] + while stack: + node = stack.pop() + for neighbor in np.flatnonzero(graph[node]): + neighbor_int = int(neighbor) + if neighbor_int not in visited: + visited.add(neighbor_int) + stack.append(neighbor_int) + return visited + + remaining = set(range(size)) + components: List[np.ndarray] = [] + while remaining: + start = min(remaining) + component = reachable(start, adjacency) & reachable(start, adjacency.T) + ordered = np.array(sorted(component), dtype=int) + components.append(ordered) + remaining.difference_update(component) + return components + + @classmethod + def _bt_scc_coordinate_directions( + cls, + local_observed_wins: np.ndarray, + include_within_scc: bool = True, + ) -> List[np.ndarray]: + """Build a basis of SCC-wide and within-SCC score contrasts.""" + + size = int(local_observed_wins.shape[0]) + components = cls._bt_strong_components(local_observed_wins) + directions: List[np.ndarray] = [] + if len(components) > 1: + reference = components[-1] + for component in components[:-1]: + direction = np.zeros(size, dtype=float) + direction[component] = 1.0 / len(component) + direction[reference] = -1.0 / len(reference) + directions.append(direction) + if include_within_scc: + for component in components: + reference = int(component[-1]) + for member in component[:-1]: + direction = np.zeros(size, dtype=float) + direction[int(member)] = 1.0 + direction[reference] = -1.0 + directions.append(direction) + expected_count = size - 1 if include_within_scc else len(components) - 1 + if len(directions) != expected_count: + raise RuntimeError("Could not construct Bradley-Terry SCC coordinates") + return directions + + @staticmethod + def _bt_stable_edge_gradient( + scores: np.ndarray, + local_fit_wins: np.ndarray, + ) -> np.ndarray: + """Evaluate pairwise BT score residuals without tail cancellation.""" + + differences = scores[:, None] - scores[None, :] + tail_probability = np.exp(-np.abs(differences)) + winner_probability = np.where( + differences >= 0, + 1.0 / (1.0 + tail_probability), + tail_probability / (1.0 + tail_probability), + ) + loser_probability = np.where( + differences >= 0, + tail_probability / (1.0 + tail_probability), + 1.0 / (1.0 + tail_probability), + ) + edge_gradient = ( + -local_fit_wins * loser_probability + + local_fit_wins.T * winner_probability + ) + # Enforce pairwise cancellation exactly enough that summing over an SCC + # removes its O(1) internal score terms. Otherwise roundoff from a + # balanced internal edge can overwhelm a 1e-300 cross-SCC residual. + return 0.5 * (edge_gradient - edge_gradient.T) + + @classmethod + def _bt_stable_gradient( + cls, + scores: np.ndarray, + local_fit_wins: np.ndarray, + ) -> np.ndarray: + """Evaluate BT score equations without saturated-probability loss.""" + + return cls._bt_stable_edge_gradient(scores, local_fit_wins).sum(axis=1) + + @classmethod + def _bt_stable_directional_gradient( + cls, + scores: np.ndarray, + local_fit_wins: np.ndarray, + direction: np.ndarray, + ) -> float: + """Contract pairwise residuals before summing over score directions.""" + + edge_gradient = cls._bt_stable_edge_gradient(scores, local_fit_wins) + direction_differences = direction[:, None] - direction[None, :] + return float(0.5 * np.sum(edge_gradient * direction_differences)) + + @classmethod + def _bt_scc_solution_is_certified( + cls, + scores: np.ndarray, + local_observed_wins: np.ndarray, + local_fit_wins: np.ndarray, + tol: float, + include_within_scc: bool = True, + ) -> bool: + """Check that every SCC-coordinate optimum is within score tolerance.""" + + convergence_tol = max(float(tol), 1e-10) + for direction in cls._bt_scc_coordinate_directions( + local_observed_wins, + include_within_scc=include_within_scc, + ): + radius = convergence_tol / float(np.max(np.abs(direction))) + lower_derivative = cls._bt_stable_directional_gradient( + scores - radius * direction, + local_fit_wins, + direction, + ) + upper_derivative = cls._bt_stable_directional_gradient( + scores + radius * direction, + local_fit_wins, + direction, + ) + if ( + not np.isfinite(lower_derivative) + or not np.isfinite(upper_derivative) + or lower_derivative > 0 + or upper_derivative < 0 + ): + return False + return True + + @classmethod + def _refine_bt_scc_offsets_with_newton( + cls, + initial_scores: np.ndarray, + local_observed_wins: np.ndarray, + local_fit_wins: np.ndarray, + tol: float, + max_iter: int = 2000, + ) -> np.ndarray: + """Refine only SCC-wide offsets on the condensation graph. + + Internal SCC comparisons are constant under these offsets and are + removed before forming the objective and Hessian. The remaining + Newton system therefore retains regularized cross-SCC curvature even + when it is hundreds of orders of magnitude below internal curvature. + """ + + scores = np.asarray(initial_scores, dtype=float).copy() + scores -= scores.mean() + components = cls._bt_strong_components(local_observed_wins) + if len(components) <= 1: + return scores + + labels = np.empty(len(scores), dtype=int) + for component_index, component in enumerate(components): + labels[component] = component_index + cross_mask = labels[:, None] != labels[None, :] + cross_fit_wins = np.where(cross_mask, local_fit_wins, 0.0) + cross_n_ij = cross_fit_wins + cross_fit_wins.T + convergence_tol = max(float(tol), 1e-10) + + def objective(values: np.ndarray) -> float: + directed_log_losses = np.logaddexp( + 0.0, values[None, :] - values[:, None] + ) + return float(np.sum(cross_fit_wins * directed_log_losses)) + + for _ in range(max_iter): + if cls._bt_scc_solution_is_certified( + scores, + local_observed_wins, + local_fit_wins, + tol, + include_within_scc=False, + ): + return scores + + edge_gradient = cls._bt_stable_edge_gradient( + scores, cross_fit_wins + ) + item_gradient = edge_gradient.sum(axis=1) + group_gradient = np.array( + [item_gradient[component].sum() for component in components] + ) + + differences = scores[:, None] - scores[None, :] + tail_probability = np.exp(-np.abs(differences)) + variance = tail_probability / np.square(1.0 + tail_probability) + item_curvature = cross_n_ij * variance + group_curvature = np.zeros( + (len(components), len(components)), dtype=float + ) + for left, left_component in enumerate(components): + for right in range(left + 1, len(components)): + right_component = components[right] + curvature = float( + item_curvature[ + np.ix_(left_component, right_component) + ].sum() + ) + group_curvature[left, right] = curvature + group_curvature[right, left] = curvature + hessian = ( + np.diag(group_curvature.sum(axis=1)) - group_curvature + ) + scale = float(np.max(np.abs(hessian))) + if not np.isfinite(scale) or scale <= 0: + break + ground = int(np.argmax(np.diag(hessian))) + keep = [ + index for index in range(len(components)) if index != ground + ] + try: + reduced_direction = np.linalg.solve( + hessian[np.ix_(keep, keep)] / scale, + group_gradient[keep] / scale, + ) + except np.linalg.LinAlgError: + break + group_direction = np.zeros(len(components), dtype=float) + group_direction[keep] = reduced_direction + item_direction = group_direction[labels] + item_direction -= item_direction.mean() + if np.any(~np.isfinite(item_direction)): + break + direction_norm = float(np.max(np.abs(item_direction))) + decrement = float(np.dot(group_gradient, group_direction)) + if ( + direction_norm <= convergence_tol + or not np.isfinite(decrement) + or decrement <= 0 + ): + break + + current_objective = objective(scores) + step_size = 1.0 + accepted = False + for _ in range(60): + candidate = scores - step_size * item_direction + candidate -= candidate.mean() + candidate_objective = objective(candidate) + if ( + np.isfinite(candidate_objective) + and candidate_objective + <= current_objective - 1e-4 * step_size * decrement + ): + scores = candidate + accepted = True + break + step_size *= 0.5 + if not accepted: + break + + raise RuntimeError( + "Bradley-Terry SCC-offset optimization failed to converge" + ) + + @classmethod + def _fit_bt_component_with_scc_coordinates( + cls, + initial_scores: np.ndarray, + local_observed_wins: np.ndarray, + local_fit_wins: np.ndarray, + tol: float, + max_sweeps: int = 4000, + include_within_scc: bool = True, + ) -> np.ndarray: + """Finish an ill-conditioned BT fit by exact coordinate minimization. + + The coordinates consist of SCC-wide mean contrasts plus within-SCC + contrasts, and therefore span the component's sum-to-zero space. Each + one-dimensional convex subproblem is solved by bracketing its stable + score equation. This avoids a global Hessian inversion when a weak + regularized cut coexists with high-information comparisons inside an + SCC, while still optimizing the original regularized likelihood. + """ + + scores = np.asarray(initial_scores, dtype=float).copy() + scores -= scores.mean() + convergence_tol = max(float(tol), 1e-10) + directions = cls._bt_scc_coordinate_directions( + local_observed_wins, + include_within_scc=include_within_scc, + ) + + def directional_gradient(direction: np.ndarray, shift: float) -> float: + candidate = scores + shift * direction + if np.any(~np.isfinite(candidate)): + return math.copysign(math.inf, shift) + return cls._bt_stable_directional_gradient( + candidate, + local_fit_wins, + direction, + ) + + for _ in range(max_sweeps): + max_update = 0.0 + for direction in directions: + coordinate_tol = max( + convergence_tol * 1e-3, + 8.0 + * np.finfo(float).eps + * max(1.0, float(np.max(np.abs(scores)))), + ) + certification_radius = convergence_tol / float( + np.max(np.abs(direction)) + ) + lower_derivative = directional_gradient( + direction, -certification_radius + ) + upper_derivative = directional_gradient( + direction, certification_radius + ) + if ( + np.isfinite(lower_derivative) + and np.isfinite(upper_derivative) + and lower_derivative <= 0 + and upper_derivative >= 0 + ): + continue + derivative = directional_gradient(direction, 0.0) + if not np.isfinite(derivative): + raise RuntimeError( + "Bradley-Terry SCC coordinate gradient was non-finite" + ) + if derivative == 0.0: + continue + + if derivative > 0: + lower, upper = -1.0, 0.0 + for _ in range(2048): + if directional_gradient(direction, lower) <= 0: + break + lower *= 2.0 + else: + raise RuntimeError( + "Could not bracket a Bradley-Terry SCC coordinate" + ) + else: + lower, upper = 0.0, 1.0 + for _ in range(2048): + if directional_gradient(direction, upper) >= 0: + break + upper *= 2.0 + else: + raise RuntimeError( + "Could not bracket a Bradley-Terry SCC coordinate" + ) + + for _ in range(256): + midpoint = 0.5 * (lower + upper) + midpoint_derivative = directional_gradient(direction, midpoint) + if midpoint_derivative > 0: + upper = midpoint + elif midpoint_derivative < 0: + lower = midpoint + else: + lower = upper = midpoint + break + if upper - lower <= coordinate_tol: + break + shift = 0.5 * (lower + upper) + scores += shift * direction + scores -= scores.mean() + max_update = max( + max_update, + abs(shift) * float(np.max(np.abs(direction))), + ) + if max_update <= convergence_tol: + return scores + + raise RuntimeError( + "Bradley-Terry SCC coordinate optimization failed to converge" + ) + + @staticmethod + def _refine_bt_component_with_newton( + initial_scores: np.ndarray, + local_fit_wins: np.ndarray, + local_n_ij: np.ndarray, + tol: float, + max_iter: int = 2000, + ) -> np.ndarray: + """Polish a slow MM fit with constrained, damped Newton steps. + + The Bradley–Terry negative log-likelihood is convex in log-skills and + has one location null direction. Each Newton system is therefore + solved in the positive-eigenvalue subspace of its Laplacian Hessian, + preserving the component sum-to-zero constraint. Backtracking keeps + every accepted step likelihood-improving. + """ + + scores = np.asarray(initial_scores, dtype=float).copy() + scores -= scores.mean() + convergence_tol = max(float(tol), 1e-10) + + def objective(values: np.ndarray) -> float: + directed_log_losses = np.logaddexp( + 0.0, values[None, :] - values[:, None] + ) + return float(np.sum(local_fit_wins * directed_log_losses)) + + for _ in range(max_iter): + differences = scores[:, None] - scores[None, :] + tail_probability = np.exp(-np.abs(differences)) + # Form the score equations from directed edge residuals. Writing + # them as ``n_ij * p_ij - wins_i`` subtracts nearly equal O(1) + # terms after separation and loses pseudo-counts as small as + # 1e-300 to cancellation. + gradient = Rank._bt_stable_gradient(scores, local_fit_wins) + variance = tail_probability / np.square(1.0 + tail_probability) + curvature = local_n_ij * variance + hessian = np.diag(curvature.sum(axis=1)) - curvature + hessian = 0.5 * (hessian + hessian.T) + scale = float(np.max(np.abs(hessian))) + if not np.isfinite(scale) or scale <= 0: + break + # Ground the highest-curvature vertex and solve the reduced + # Laplacian. Unlike a full eigendecomposition, this preserves a + # weak cut next to a high-information subgraph instead of losing + # its tiny positive eigenvalue to cancellation with the location + # null direction. + ground = int(np.argmax(np.diag(hessian))) + keep = [index for index in range(len(scores)) if index != ground] + reduced_hessian = hessian[np.ix_(keep, keep)] / scale + reduced_gradient = gradient[keep] / scale + try: + reduced_direction = np.linalg.solve( + reduced_hessian, reduced_gradient + ) + except np.linalg.LinAlgError: + break + direction = np.zeros(len(scores), dtype=float) + direction[keep] = reduced_direction + if np.any(~np.isfinite(direction)): + break + direction -= direction.mean() + direction_norm = float(np.max(np.abs(direction))) + if direction_norm <= convergence_tol: + return scores + decrement = float(np.dot(gradient, direction)) + if not np.isfinite(decrement) or decrement <= 0: + break + + current_objective = objective(scores) + step_size = 1.0 + accepted = False + for _ in range(60): + candidate = scores - step_size * direction + candidate -= candidate.mean() + candidate_objective = objective(candidate) + if ( + np.isfinite(candidate_objective) + and candidate_objective + <= current_objective - 1e-4 * step_size * decrement + ): + scores = candidate + accepted = True + break + step_size *= 0.5 + if not accepted: + break + if step_size * direction_norm <= convergence_tol: + return scores + + raise RuntimeError( + "Bradley-Terry optimization failed to converge; no scores or " + "standard errors were produced" + ) + def _fit_bt( self, item_ids: List[str], - outcomes: List[Tuple[str, str]], + outcomes: Sequence[Union[Tuple[str, str], WeightedOutcome]], pseudo: float, max_iter: int, tol: float, @@ -493,13 +2119,12 @@ def _fit_bt( item_ids: List of unique item identifiers. outcomes: - List of tuples ``(winner, loser)`` representing outcomes of - pairwise matches. Ties can be represented by including - both ``(a, b)`` and ``(b, a)`` in the list; each entry - contributes a single increment to the win matrix. + Pairwise ``(winner, loser)`` tuples, or weighted + ``(winner, loser, weight)`` tuples. A tie is represented by two + directions with weight ``0.5`` each. pseudo: - Pseudo count added to both win and total match counts. Acts - as a smoothing prior. + Total symmetric pseudo-comparison weight added to each observed + edge. Unobserved pairs receive no pseudo-comparisons. max_iter, tol: Control convergence of the iterative fixed‑point updates. return_info: @@ -511,76 +2136,242 @@ def _fit_bt( scores : dict Mapping from item identifier to estimated log‑skill. (scores, n_ij, p_ij) : tuple - When ``return_info`` is ``True``, also return the total - match counts and predicted win probabilities for each pair. + When ``return_info`` is ``True``, also return total observed match + counts and predicted win probabilities for observed components. + Pseudo-comparisons are intentionally excluded from the counts; + cross-component probabilities are ``NaN`` because they are not + identified by the comparison data. """ + if max_iter < 1: + raise ValueError("max_iter must be at least 1") + if not np.isfinite(tol) or tol <= 0: + raise ValueError("tol must be a finite positive number") + + observed_wins, fit_wins = self._build_bt_win_matrices( + item_ids, outcomes, pseudo + ) n = len(item_ids) - idx = {item: i for i, item in enumerate(item_ids)} - # win matrix; wins[i,j] counts how many times i beat j - wins = np.zeros((n, n), dtype=float) - for w, l in outcomes: - if w in idx and l in idx: - wins[idx[w], idx[l]] += 1.0 - # total matches between each pair - n_ij = wins + wins.T - # total wins for each item - w_i = wins.sum(axis=1) - # add pseudo counts - n_ij += pseudo - w_i += pseudo - # initialise skill parameters uniformly - p = np.ones(n, dtype=float) - for _ in range(max_iter): - # denominator for each player in the fixed point update - denom = (n_ij / (p[:, None] + p[None, :])).sum(axis=1) - p_new = w_i / denom - if np.max(np.abs(p_new - p)) < tol: - p = p_new - break - p = p_new - # convert to log space and centre at zero mean - s = np.log(p) - s -= s.mean() + observed_n_ij = observed_wins + observed_wins.T + fit_n_ij = fit_wins + fit_wins.T + scores = np.zeros(n, dtype=float) + + for component in self._comparison_components(observed_n_ij): + if len(component) == 1: + continue + if pseudo == 0 and not self._is_strongly_connected( + observed_wins, component + ): + raise ValueError( + "learning_rate must be positive when the directed win graph " + "is not strongly connected" + ) + local_observed_wins = observed_wins[np.ix_(component, component)] + local_fit_wins = fit_wins[np.ix_(component, component)] + local_wins = local_fit_wins.sum(axis=1) + local_n_ij = fit_n_ij[np.ix_(component, component)] + requires_scc_coordinates = ( + len(self._bt_strong_components(local_observed_wins)) > 1 + ) + if np.any(local_wins <= 0): + raise ValueError( + "learning_rate must be positive when the observed win graph " + "does not have a finite Bradley-Terry maximum-likelihood estimate" + ) + + def scc_solution_is_certified(candidate_scores: np.ndarray) -> bool: + return not requires_scc_coordinates or ( + self._bt_scc_solution_is_certified( + candidate_scores, + local_observed_wins, + local_fit_wins, + tol, + ) + ) + + def finish_with_scc_coordinates( + initial_scores: np.ndarray, + ) -> np.ndarray: + if requires_scc_coordinates: + try: + offset_scores = self._refine_bt_scc_offsets_with_newton( + initial_scores, + local_observed_wins, + local_fit_wins, + tol, + ) + except RuntimeError: + try: + offset_scores = ( + self._fit_bt_component_with_scc_coordinates( + initial_scores, + local_observed_wins, + local_fit_wins, + tol, + include_within_scc=False, + ) + ) + except RuntimeError: + offset_scores = initial_scores + if scc_solution_is_certified(offset_scores): + return offset_scores + try: + polished_scores = ( + self._refine_bt_component_with_newton( + offset_scores, + local_fit_wins, + local_n_ij, + tol, + ) + ) + except RuntimeError: + polished_scores = offset_scores + if scc_solution_is_certified(polished_scores): + return polished_scores + else: + polished_scores = initial_scores + return self._fit_bt_component_with_scc_coordinates( + polished_scores, + local_observed_wins, + local_fit_wins, + tol, + ) + + def refine_component( + initial_scores: np.ndarray, + fallback_on_newton_failure: bool, + ) -> np.ndarray: + try: + refined_scores = self._refine_bt_component_with_newton( + initial_scores, + local_fit_wins, + local_n_ij, + tol, + ) + except RuntimeError: + if not fallback_on_newton_failure: + raise + return finish_with_scc_coordinates(initial_scores) + if not scc_solution_is_certified(refined_scores): + return finish_with_scc_coordinates(refined_scores) + return refined_scores + + log_abilities = np.zeros(len(component), dtype=float) + converged = False + logit_start = self._bt_logit_spanning_tree_start(local_fit_wins) + if logit_start is not None: + try: + log_abilities = refine_component( + logit_start, + fallback_on_newton_failure=False, + ) + converged = True + except RuntimeError: + # Retain Hunter's globally convergent MM update as the + # conservative path when direct refinement is unresolved + # from the log-odds projection. + pass + + abilities = np.ones(len(component), dtype=float) + for _ in range(max_iter): + if converged: + break + denom = ( + local_n_ij + / (abilities[:, None] + abilities[None, :]) + ).sum(axis=1) + if np.any(~np.isfinite(denom)) or np.any(denom <= 0): + raise FloatingPointError( + "Bradley-Terry update encountered an invalid denominator" + ) + next_abilities = local_wins / denom + if np.any(~np.isfinite(next_abilities)) or np.any( + next_abilities <= 0 + ): + raise FloatingPointError( + "Bradley-Terry update encountered an invalid ability" + ) + + next_log_abilities = np.log(next_abilities) + next_log_abilities -= next_log_abilities.mean() + next_abilities = np.exp(next_log_abilities) + delta = float( + np.max(np.abs(next_log_abilities - log_abilities)) + ) + abilities = next_abilities + log_abilities = next_log_abilities + if delta < tol: + converged = True + break + if not converged: + log_abilities = refine_component( + log_abilities, + fallback_on_newton_failure=True, + ) + elif not scc_solution_is_certified(log_abilities): + log_abilities = finish_with_scc_coordinates(log_abilities) + scores[component] = log_abilities + if not return_info: - return {item: float(val) for item, val in zip(item_ids, s)} - # predicted win probabilities between each pair - exp_s = np.exp(s) - p_ij = exp_s[:, None] / (exp_s[:, None] + exp_s[None, :]) - return {item: float(val) for item, val in zip(item_ids, s)}, n_ij, p_ij + return {item: float(val) for item, val in zip(item_ids, scores)} + p_ij = np.full((n, n), np.nan, dtype=float) + for component in self._comparison_components(observed_n_ij): + local_scores = scores[component] + score_diff = np.clip( + local_scores[:, None] - local_scores[None, :], -700, 700 + ) + p_ij[np.ix_(component, component)] = 1.0 / ( + 1.0 + np.exp(-score_diff) + ) + return ( + {item: float(val) for item, val in zip(item_ids, scores)}, + observed_n_ij, + p_ij, + ) def _bt_standard_errors( self, s: np.ndarray, n_ij: np.ndarray, p_ij: np.ndarray, - ridge: float, + rcond: float, + regularization_strength: float = 0.0, ) -> np.ndarray: - """Estimate standard errors for BT skill parameters. + """Estimate sandwich standard errors for regularized BT scores. - The observed Fisher information for the Bradley–Terry model is given by + The observed information for the Bradley–Terry working model is given by ``I = diag(q 1) - q`` where ``q = n_ij * p_ij * (1 - p_ij)`` encodes the - uncertainty contributed by each pairwise comparison (Ford, 1957). The - estimates satisfy a sum-to-zero constraint, so the Fisher information is - rank deficient with the all-ones vector in its null space. Instead of - selecting an arbitrary reference item (which previously produced - inflated standard errors for that reference when it received few - comparisons), we project the matrix onto the constrained subspace and - take its Moore–Penrose pseudoinverse. A small ridge term stabilises the - inversion for sparse comparison graphs. The standard error for item ``i`` - is the square root of the ``i``-th diagonal entry of the resulting - covariance matrix. + model-based binary-BT variability determined by observed comparison + weights. Because + the point estimator includes fixed pseudo-comparisons, its frequentist + covariance uses the penalized-information matrix as the sandwich bread + and observed information as the meat. Thus regularization affects the + estimator's sampling variance without being treated as random data. + The calculation treats the realized comparison graph as fixed. This + model-based variance excludes regularization bias, shared-judge + dependence, likelihood misspecification, and adaptive-selection + uncertainty. It is therefore not a total-error or mean-squared-error + estimate. + + The calculation is performed separately under a sum-to-zero constraint + within each observed component. Isolated items remain ``NaN``. Draws + use the package's fractional-binomial half-win approximation rather than + a separate three-outcome tie likelihood. Parameters ---------- s : np.ndarray Array of estimated log-skills for each item. n_ij : np.ndarray - Matrix of total match counts between items (wins + losses). + Matrix of observed total match weights between items. p_ij : np.ndarray Matrix of predicted win probabilities between items. - ridge : float - Small constant added to the diagonal of the projected Fisher - information matrix for numerical stability. + rcond : float + Relative eigenvalue tolerance used when taking the constrained + inverse. It does not add information to the Fisher matrix. + regularization_strength : float + Total fixed pseudo-comparison weight added to each observed edge + by the point estimator. This enters the sandwich bread, not the + observed-data meat. Returns ------- @@ -591,25 +2382,90 @@ def _bt_standard_errors( n = len(s) if n == 0: return np.array([], dtype=float) - if n == 1: - return np.zeros(1, dtype=float) - - q_ij = n_ij * p_ij * (1 - p_ij) - diag = q_ij.sum(axis=1) - I = np.diag(diag) - q_ij - I = np.nan_to_num(I) - ones = np.ones((n, 1)) - proj = np.eye(n) - ones @ ones.T / n - I_proj = proj @ I @ proj - I_proj[np.diag_indices(n)] += ridge - try: - cov = np.linalg.pinv(I_proj, rcond=1e-12) - except np.linalg.LinAlgError: - cov = np.linalg.pinv(np.nan_to_num(I_proj), rcond=1e-12) - cov = proj @ cov @ proj - cov = 0.5 * (cov + cov.T) - se = np.sqrt(np.clip(np.diag(cov), 0.0, None)) - return np.nan_to_num(se) + if n_ij.shape != (n, n) or p_ij.shape != (n, n): + raise ValueError("n_ij and p_ij must be square matrices matching s") + if not np.isfinite(rcond) or rcond < 0: + raise ValueError("rcond must be a finite non-negative number") + if ( + not np.isfinite(regularization_strength) + or regularization_strength < 0 + ): + raise ValueError( + "regularization_strength must be a finite non-negative number" + ) + components = self._comparison_components(n_ij) + if ( + len(components) > 1 + and np.any(n_ij > 0) + and not self._warned_disconnected_graph + ): + warnings.warn( + "The observed comparison graph is disconnected. Scores and " + "standard errors are component-relative; inspect the saved " + "'_component' columns before comparing items.", + RuntimeWarning, + stacklevel=2, + ) + self._warned_disconnected_graph = True + + se = np.full(n, np.nan, dtype=float) + for component in components: + if len(component) == 1: + continue + local_n_ij = n_ij[np.ix_(component, component)] + local_s = s[component] + local_p_ij = p_ij[np.ix_(component, component)] + if np.any(~np.isfinite(local_p_ij)): + continue + # exp(-|delta|) / (1 + exp(-|delta|))^2 is algebraically + # p * (1 - p), but remains symmetric when a direct sigmoid rounds + # to exactly zero or one at an extreme score gap. + tail_probability = np.exp( + -np.abs(local_s[:, None] - local_s[None, :]) + ) + variance = tail_probability / np.square(1.0 + tail_probability) + observed_q = local_n_ij * variance + regularized_n_ij = local_n_ij + regularization_strength * ( + local_n_ij > 0 + ) + bread_q = regularized_n_ij * variance + meat = np.diag(observed_q.sum(axis=1)) - observed_q + bread = np.diag(bread_q.sum(axis=1)) - bread_q + meat = 0.5 * (meat + meat.T) + bread = 0.5 * (bread + bread.T) + scale = float(np.max(np.abs(bread))) + if not np.isfinite(scale) or scale <= 0: + continue + scaled_bread = bread / scale + scaled_meat = meat / scale + try: + eigenvalues, eigenvectors = np.linalg.eigh(scaled_bread) + except np.linalg.LinAlgError: + continue + max_eigenvalue = float(np.max(eigenvalues)) + threshold = max( + float(rcond), np.finfo(float).eps * len(component) + ) * max(0.0, max_eigenvalue) + identified = eigenvalues > threshold + if int(np.sum(identified)) != len(component) - 1: + continue + basis = eigenvectors[:, identified] + local_meat = basis.T @ scaled_meat @ basis + local_meat = 0.5 * (local_meat + local_meat.T) + try: + meat_eigenvalues, meat_eigenvectors = np.linalg.eigh(local_meat) + except np.linalg.LinAlgError: + continue + meat_root = meat_eigenvectors * np.sqrt( + np.clip(meat_eigenvalues, 0.0, None) + )[None, :] + covariance_factor = basis @ ( + meat_root / eigenvalues[identified, None] + ) + se[component] = ( + np.hypot.reduce(covariance_factor, axis=1) / np.sqrt(scale) + ) + return se def _fit_pl( self, @@ -719,22 +2575,31 @@ def _pairs_info_gain( se_agg: Dict[str, float], mpr: int, ) -> List[Tuple[Tuple[str, str], Tuple[str, str]]]: - """Select pairs by maximising expected information gain while ensuring - that every item participates in the prescribed number of matches. - - This implementation differs from the original heuristics by - considering a bounded set of candidate pairs that scales with the - number of items. Each pair is assigned a score based on the - expected reduction in uncertainty (estimated from the current - ratings and aggregated standard errors). Pairs with larger - scores are chosen first, subject to the constraint that each - item is matched exactly ``mpr`` times. If some items remain - unmatched after exhausting the scored pairs, additional pairs - are filled in randomly to satisfy the per‑item quota. + """Select pairs with an uncertainty-guided scheduling heuristic. + + A bounded candidate set scales with the number of items. The heuristic + favors uncertain, similarly rated pairs using current model-based + standard errors and binary-BT outcome variance; it is not an exact + expected-information-gain calculation. ``mpr`` is a capped target + degree, not a guarantee: an item can exceed it when selected as another + item's opponent. """ n = len(item_ids) if n < 2: return [] + finite_se = [ + float(value) + for value in se_agg.values() + if np.isfinite(value) and float(value) >= 0 + ] + unidentified_se = 2.0 * max(finite_se, default=1.0) + + def item_se(item_id: str) -> float: + value = se_agg.get(item_id, unidentified_se) + if not np.isfinite(value) or float(value) < 0: + return unidentified_se + return float(value) + max_pairs = max(1, self._MAX_CANDIDATE_PAIRS_PER_ROUND) desired_neighbors = max_pairs // max(1, n) candidate_neighbors = max( @@ -751,7 +2616,7 @@ def logistic_clip(x: float) -> float: ids_sorted = sorted(item_ids, key=lambda i: current_ratings[i]) idx_of = {i_id: k for k, i_id in enumerate(ids_sorted)} num_high_se = max(1, int(self._HIGH_SE_FRAC * n)) - high_se_ids = sorted(item_ids, key=lambda i: se_agg.get(i, 1.0), reverse=True)[ + high_se_ids = sorted(item_ids, key=item_se, reverse=True)[ :num_high_se ] candidate_pairs_set: Set[Tuple[str, str]] = set() @@ -801,13 +2666,11 @@ def logistic_clip(x: float) -> float: diff = current_ratings[a] - current_ratings[b] p = logistic_clip(diff) outcome_var = p * (1 - p) - var_a = se_agg.get(a, 1.0) ** 2 - var_b = se_agg.get(b, 1.0) ** 2 + var_a = item_se(a) ** 2 + var_b = item_se(b) ** 2 param_unc = var_a + var_b - # Encourage comparisons between similarly‑rated items (high information - # gain) while still prioritising uncertain pairs. The closeness term - # dampens pairings with large rating gaps to tease out subtle ordering - # differences. + # Favor uncertain, similarly rated pairs. The closeness term dampens + # pairings with large rating gaps to probe subtle ordering differences. closeness = 1.0 / (1.0 + abs(diff)) score = outcome_var * param_unc * closeness scored_pairs.append((score, a, b)) @@ -830,16 +2693,23 @@ def logistic_clip(x: float) -> float: break # Choose an item that still needs matches a = self.rng.choice(ids_needing) - # Try to pair it with any other item (not just those needing matches) to avoid self‑pairs - potential = [x for x in item_ids if x != a] + # Prefer an unscheduled partner that also needs a comparison. If + # none exists, allow one endpoint to exceed its target by one. + potential = [ + x + for x in item_ids + if x != a and tuple(sorted((a, x))) not in pairs_seen + ] if not potential: - # Degenerate case: only one item exists; cannot form a valid pair break - b = self.rng.choice(potential) + partners_needing = [x for x in potential if needed[x] > 0] + b = self.rng.choice(partners_needing or potential) tup = (a, b) if a < b else (b, a) pairs_selected.append(tup) + pairs_seen.add(tup) needed[a] -= 1 - needed[b] -= 1 + if needed[b] > 0: + needed[b] -= 1 return [((a, texts_by_id[a]), (b, texts_by_id[b])) for a, b in pairs_selected] def _generate_pairs( @@ -850,8 +2720,11 @@ def _generate_pairs( se_agg: Optional[Dict[str, float]], ) -> List[Tuple[Tuple[str, str], Tuple[str, str]]]: """Dispatch to the appropriate pairing strategy.""" - mpr = max(1, self.cfg.matches_per_round) - # Always use information gain pairing to guarantee exact match counts + if len(item_ids) < 2: + return [] + mpr = min(self.cfg.matches_per_round, len(item_ids) - 1) + if not self.cfg.power_matching: + return self._pairs_random(item_ids, texts_by_id, mpr) if current_ratings is None: current_ratings = {i: 0.0 for i in item_ids} if se_agg is None or len(se_agg) != len(item_ids): @@ -864,7 +2737,7 @@ def _generate_pairs( async def _catch_up_existing_rounds( self, - new_ids: List[str], + candidate_ids: List[str], round_indices: List[int], item_ids: List[str], texts_by_id: Dict[str, str], @@ -873,9 +2746,11 @@ async def _catch_up_existing_rounds( pdfs_by_id: Dict[str, List[Dict[str, str]]], attr_batches: List[List[str]], attr_keys: List[str], - history_pairs: Dict[str, List[Tuple[str, str]]], + history_pairs: Dict[str, List[WeightedOutcome]], + outcome_counts: Dict[str, Dict[str, int]], ratings: Dict[str, Dict[str, float]], se_store: Dict[str, Dict[str, float]], + component_store: Dict[str, Dict[str, int]], base_name: str, df_proc: pd.DataFrame, _write_checkpoint: Callable[[], None], @@ -885,127 +2760,417 @@ async def _catch_up_existing_rounds( identifier_hash_bits: int, **kwargs: Any, ) -> None: - if not new_ids: + if not candidate_ids: return for rnd in round_indices: round_path = os.path.join(self.cfg.save_dir, f"{base_name}_round{rnd}.csv") if not os.path.exists(round_path): continue - try: - df_round = pd.read_csv(round_path) - except Exception: - continue + df_round = self._read_rank_checkpoint( + round_path, len(attr_batches) + ) counts: Dict[str, int] = {} + current_id_set = set(item_ids) + committed_pairs: Set[Tuple[str, str]] = set() if {"IdA", "IdB"}.issubset(df_round.columns): - for a, b in zip(df_round["IdA"], df_round["IdB"]): - counts[str(a)] = counts.get(str(a), 0) + 1 - counts[str(b)] = counts.get(str(b), 0) + 1 + counts_by_batch: Dict[int, Dict[str, int]] = { + batch_idx: {} for batch_idx in range(len(attr_batches)) + } + for batch_raw, a, b in zip( + df_round.get("Batch", pd.Series(dtype=float)), + df_round["IdA"], + df_round["IdB"], + ): + try: + batch_idx = int(batch_raw) + except (TypeError, ValueError): + continue + if batch_idx not in counts_by_batch: + continue + id_a = str(a) + id_b = str(b) + if id_a not in current_id_set or id_b not in current_id_set: + continue + committed_pairs.add(tuple(sorted((id_a, id_b)))) + batch_counts = counts_by_batch[batch_idx] + batch_counts[id_a] = batch_counts.get(id_a, 0) + 1 + batch_counts[id_b] = batch_counts.get(id_b, 0) + 1 + for item_id in item_ids: + counts[item_id] = min( + ( + batch_counts.get(item_id, 0) + for batch_counts in counts_by_batch.values() + ), + default=0, + ) else: for ident in df_round.get("Identifier", []): parts = str(ident).split("|") if len(parts) != 5: continue _, _, _, id_a, id_b = parts + if id_a not in current_id_set or id_b not in current_id_set: + continue + committed_pairs.add(tuple(sorted((id_a, id_b)))) counts[id_a] = counts.get(id_a, 0) + 1 counts[id_b] = counts.get(id_b, 0) + 1 - pairs_needed: List[Tuple[str, str]] = [] - for nid in new_ids: - needed = self.cfg.matches_per_round - counts.get(nid, 0) - if needed <= 0: - continue - opponents = [i for i in item_ids if i != nid] - self.rng.shuffle(opponents) - for opp in opponents[:needed]: - pairs_needed.append((nid, opp)) - if not pairs_needed: - continue - announce_prompt_rendering( - "Rank:catchup", len(attr_batches) * len(pairs_needed) + catchup_path = os.path.join( + self.cfg.save_dir, f".{base_name}_catchup_round{rnd}.csv" + ) + plan_path = os.path.join( + self.cfg.save_dir, f".{base_name}_catchup_round{rnd}_plan.json" ) + batch_state_path = f"{catchup_path}.batch_state.json" + committed_response_ids = set(df_round["Identifier"].astype(str)) + expected_target_matches = min( + self.cfg.matches_per_round, max(0, len(item_ids) - 1) + ) + canonical_item_ids = sorted(item_ids) + plan_records: Optional[List[Dict[str, Any]]] = None + if os.path.exists(plan_path): + try: + with open(plan_path, encoding="utf-8") as plan_file: + plan_payload = json.load(plan_file) + except Exception as exc: + raise ValueError( + f"Could not read Rank catch-up plan {plan_path!r}" + ) from exc + planned_item_ids = ( + plan_payload.get("item_ids") + if isinstance(plan_payload, dict) + else None + ) + if ( + not isinstance(plan_payload, dict) + or plan_payload.get("version") != 1 + or plan_payload.get("round") != rnd + or plan_payload.get("attribute_batches") != attr_batches + or not isinstance(planned_item_ids, list) + or any(not isinstance(item_id, str) for item_id in planned_item_ids) + or len(planned_item_ids) != len(set(planned_item_ids)) + or sorted(planned_item_ids) != canonical_item_ids + or plan_payload.get("target_matches") + != expected_target_matches + or not isinstance(plan_payload.get("records"), list) + or not plan_payload["records"] + ): + raise ValueError( + f"Rank catch-up plan {plan_path!r} is incompatible or malformed" + ) + candidate_records = plan_payload["records"] + required_record_keys = { + "identifier", + "batch", + "pair", + "id_a", + "id_b", + "circle_first", + } + if any( + not isinstance(record, dict) + or not required_record_keys.issubset(record) + for record in candidate_records + ): + raise ValueError( + f"Rank catch-up plan {plan_path!r} has malformed records" + ) + for record in candidate_records: + if ( + type(record["batch"]) is not int + or type(record["pair"]) is not int + or type(record["circle_first"]) is not bool + or not str(record["id_a"]).strip() + or not str(record["id_b"]).strip() + ): + raise ValueError( + f"Rank catch-up plan {plan_path!r} has invalid fields" + ) + expected_identifier = hash_identifier( + "catchup|" + f"{rnd}|{record['batch']}|{record['pair']}|" + f"{record['id_a']}|{record['id_b']}", + bits=identifier_hash_bits, + ) + if str(record["identifier"]) != expected_identifier: + raise ValueError( + f"Rank catch-up plan {plan_path!r} has an invalid identifier" + ) + planned_ids = [str(record["identifier"]) for record in candidate_records] + if len(planned_ids) != len(set(planned_ids)): + raise ValueError( + f"Rank catch-up plan {plan_path!r} has duplicate identifiers" + ) + committed_plan_ids = set(planned_ids) & committed_response_ids + if committed_plan_ids and committed_plan_ids != set(planned_ids): + raise ValueError( + "Rank found a partially committed catch-up plan. Use " + "reset_files=True or a new save_dir to recompute safely." + ) + if committed_plan_ids == set(planned_ids): + for completed_artifact in ( + batch_state_path, + catchup_path, + plan_path, + ): + try: + Path(completed_artifact).unlink(missing_ok=True) + except OSError: + pass + else: + planned_endpoints = { + str(record[key]) + for record in candidate_records + for key in ("id_a", "id_b") + } + if not planned_endpoints.issubset(current_id_set): + raise ValueError( + "Rank has an unfinished catch-up plan whose items are " + "not all present. Restore the same input set before " + "resuming, or use reset_files=True after confirming no " + "external batch remains." + ) + plan_records = candidate_records + + if os.path.exists(batch_state_path) and not os.path.exists( + catchup_path + ): + try: + with open(batch_state_path, encoding="utf-8") as state_file: + batch_state = json.load(state_file) + except Exception as exc: + raise ValueError( + "Could not read Rank catch-up Batch API state " + f"{batch_state_path!r}" + ) from exc + batches = ( + batch_state.get("batches", []) + if isinstance(batch_state, dict) + else None + ) + if not isinstance(batches, list): + raise ValueError( + f"Rank catch-up Batch API state {batch_state_path!r} " + "is malformed" + ) + unresolved_submission = any( + not isinstance(batch, dict) + or batch.get("status") == "submitting" + or not batch.get("batch_id") + for batch in batches + ) + has_active_batch = bool(batch_state.get("batch_id")) or bool( + batches + ) + if unresolved_submission: + raise ValueError( + "Rank found an unresolved catch-up Batch API submission. " + "The server may already have accepted a paid job; " + "reconcile or cancel it before resetting local state." + ) + if not has_active_batch or plan_records is None: + raise ValueError( + "Rank found catch-up Batch API state without a durable " + "response checkpoint and recoverable plan. Paid results " + "may already exist; use reset_files=True only after " + "confirming that it is safe to discard the external state." + ) + # A submitted batch ID plus the durable plan is sufficient to + # reconstruct the request set. The collector will poll that + # batch and atomically create the staging CSV when rows arrive. + + if plan_records is None: + deficits = { + item_id: max( + 0, expected_target_matches - counts.get(item_id, 0) + ) + for item_id in candidate_ids + } + pairs_needed: List[Tuple[str, str]] = [] + scheduled_pairs: Set[Tuple[str, str]] = set(committed_pairs) + while any(deficit > 0 for deficit in deficits.values()): + max_deficit = max(deficits.values()) + focal_candidates = [ + item_id + for item_id, deficit in deficits.items() + if deficit == max_deficit + ] + id_a = self.rng.choice(focal_candidates) + available = [ + item_id + for item_id in item_ids + if item_id != id_a + and tuple(sorted((id_a, item_id))) not in scheduled_pairs + ] + if not available: + break + partners_needing = [ + item_id + for item_id in available + if deficits.get(item_id, 0) > 0 + ] + id_b = self.rng.choice(partners_needing or available) + pair = tuple(sorted((id_a, id_b))) + scheduled_pairs.add(pair) + pairs_needed.append((id_a, id_b)) + deficits[id_a] -= 1 + if deficits.get(id_b, 0) > 0: + deficits[id_b] -= 1 + if not pairs_needed: + continue + plan_records = [] + for batch_idx, _batch in enumerate(attr_batches): + for pair_idx, (id_a, id_b) in enumerate(pairs_needed): + raw_ident = ( + f"catchup|{rnd}|{batch_idx}|{pair_idx}|{id_a}|{id_b}" + ) + hashed_ident = hash_identifier( + raw_ident, bits=identifier_hash_bits + ) + circle_first_flag = ( + self.cfg.circle_first + if self.cfg.circle_first is not None + else self.rng.random() < 0.5 + ) + plan_records.append( + { + "identifier": hashed_ident, + "batch": batch_idx, + "pair": pair_idx, + "id_a": id_a, + "id_b": id_b, + "circle_first": circle_first_flag, + } + ) + self._write_json_atomically( + { + "version": 1, + "round": rnd, + "attribute_batches": attr_batches, + # Pair prompts and plan endpoints are ID-keyed. Store a + # canonical set representation so harmless DataFrame row + # reordering does not strand already-paid staged rows. + "item_ids": canonical_item_ids, + "target_matches": expected_target_matches, + "records": plan_records, + }, + plan_path, + ) + + announce_prompt_rendering("Rank:catchup", len(plan_records)) prompts: List[str] = [] ids: List[str] = [] pair_images: Dict[str, List[str]] = {} pair_audio: Dict[str, List[Dict[str, str]]] = {} pair_pdfs: Dict[str, List[Dict[str, str]]] = {} meta_map: Dict[str, Tuple[int, int, str, str]] = {} - id_to_circle_first: Dict[str, bool] = {} - for batch_idx, batch in enumerate(attr_batches): + for record in plan_records: + batch_idx = int(record["batch"]) + pair_idx = int(record["pair"]) + id_a = str(record["id_a"]) + id_b = str(record["id_b"]) + hashed_ident = str(record["identifier"]) + circle_first_flag = bool(record["circle_first"]) + if batch_idx < 0 or batch_idx >= len(attr_batches): + raise ValueError( + f"Rank catch-up plan {plan_path!r} has an invalid batch" + ) + batch = attr_batches[batch_idx] attr_def_map = ( {a: self.cfg.attributes[a] for a in batch} if isinstance(self.cfg.attributes, dict) else {a: "" for a in batch} ) - for pair_idx, (id_a, id_b) in enumerate(pairs_needed): - raw_ident = f"catchup|{rnd}|{batch_idx}|{pair_idx}|{id_a}|{id_b}" - hashed_ident = hash_identifier(raw_ident, bits=identifier_hash_bits) - circle_first_flag = ( - self.cfg.circle_first - if self.cfg.circle_first is not None - else self.rng.random() < 0.5 - ) - id_to_circle_first[hashed_ident] = circle_first_flag - prompts.append( - self.template.render( - entry_circle=texts_by_id[id_a], - entry_square=texts_by_id[id_b], - attributes=attr_def_map, - additional_instructions=self.cfg.additional_instructions or "", - modality=self.cfg.modality, - circle_first=circle_first_flag, - ) + prompts.append( + self.template.render( + entry_circle=texts_by_id[id_a], + entry_square=texts_by_id[id_b], + attributes=attr_def_map, + additional_instructions=self.cfg.additional_instructions or "", + modality=self.cfg.modality, + circle_first=circle_first_flag, ) - ids.append(hashed_ident) - meta_map[hashed_ident] = (batch_idx, pair_idx, id_a, id_b) - if images_by_id: - imgs = [] - ia = images_by_id.get(id_a, []) - ib = images_by_id.get(id_b, []) - if circle_first_flag: - if ia: - imgs.extend(ia) - if ib: - imgs.extend(ib) - else: - if ib: - imgs.extend(ib) - if ia: - imgs.extend(ia) - if imgs: - pair_images[hashed_ident] = imgs - if audio_by_id: - auds = [] - aa = audio_by_id.get(id_a, []) - ab = audio_by_id.get(id_b, []) - if circle_first_flag: - if aa: - auds.extend(aa) - if ab: - auds.extend(ab) - else: - if ab: - auds.extend(ab) - if aa: - auds.extend(aa) - if auds: - pair_audio[hashed_ident] = auds - if pdfs_by_id: - pdfs: List[Dict[str, str]] = [] - pa = pdfs_by_id.get(id_a, []) - pb = pdfs_by_id.get(id_b, []) - if circle_first_flag: - if pa: - pdfs.extend(pa) - if pb: - pdfs.extend(pb) - else: - if pb: - pdfs.extend(pb) - if pa: - pdfs.extend(pa) - if pdfs: - pair_pdfs[hashed_ident] = pdfs + ) + ids.append(hashed_ident) + meta_map[hashed_ident] = (batch_idx, pair_idx, id_a, id_b) + if images_by_id: + imgs = [] + ia = images_by_id.get(id_a, []) + ib = images_by_id.get(id_b, []) + if circle_first_flag: + if ia: + imgs.extend(ia) + if ib: + imgs.extend(ib) + else: + if ib: + imgs.extend(ib) + if ia: + imgs.extend(ia) + if imgs: + pair_images[hashed_ident] = imgs + if audio_by_id: + auds = [] + aa = audio_by_id.get(id_a, []) + ab = audio_by_id.get(id_b, []) + if circle_first_flag: + if aa: + auds.extend(aa) + if ab: + auds.extend(ab) + else: + if ab: + auds.extend(ab) + if aa: + auds.extend(aa) + if auds: + pair_audio[hashed_ident] = auds + if pdfs_by_id: + pdfs: List[Dict[str, str]] = [] + pa = pdfs_by_id.get(id_a, []) + pb = pdfs_by_id.get(id_b, []) + if circle_first_flag: + if pa: + pdfs.extend(pa) + if pb: + pdfs.extend(pb) + else: + if pb: + pdfs.extend(pb) + if pa: + pdfs.extend(pa) + if pdfs: + pair_pdfs[hashed_ident] = pdfs if not prompts: continue + if len(ids) != len(set(ids)): + raise ValueError( + "Rank prompt identifier collision; use reset_files=True " + "or a save_dir configured for 64-bit identifiers" + ) + response_kwargs = dict(kwargs) + # Rank commits a round only when every planned judgment is + # present. The generic collector's large-run tail shortcut is + # therefore incompatible with transactional tournament replay. + response_kwargs["skip_tail_fails"] = False + if self.cfg.use_dummy: + response_kwargs.setdefault( + "dummy_responses", + { + identifier: { + "responses": [ + json.dumps( + { + attribute: "draw" + for attribute in attr_batches[ + meta_map[identifier][0] + ] + } + ) + ] + } + for identifier in ids + }, + ) resp_df = await get_all_responses( prompts=prompts, identifiers=ids, @@ -1015,13 +3180,33 @@ async def _catch_up_existing_rounds( n_parallels=self.cfg.n_parallels, model=self.cfg.model, json_mode=self.cfg.modality != "audio", - save_path=round_path, + # Never let response collection append raw rows to an already + # committed structured round. The staging checkpoint can resume + # interrupted or failed requests, then the validated rows are + # merged into the committed round atomically below. + save_path=catchup_path, reset_files=reset_files, use_dummy=self.cfg.use_dummy, max_retries=1, reasoning_effort=self.cfg.reasoning_effort, - **kwargs, + **response_kwargs, + ) + new_mask = self._validate_requested_responses( + resp_df, + ids, + context="Rank catch-up", + allow_extra=True, ) + resp_df = resp_df.loc[new_mask].copy() + await self._validate_pairwise_response_payloads( + resp_df, + meta_map, + attr_batches, + context="Rank catch-up", + retry_path=catchup_path, + ) + if "Successful" not in resp_df.columns: + resp_df["Successful"] = True resp_df["Batch"] = resp_df.Identifier.map( lambda x: meta_map.get(str(x), (np.nan, np.nan, "", ""))[0] ) @@ -1034,7 +3219,22 @@ async def _catch_up_existing_rounds( resp_df["IdB"] = resp_df.Identifier.map( lambda x: meta_map.get(str(x), (np.nan, np.nan, "", ""))[3] ) - resp_df.to_csv(round_path, index=False) + combined_round = pd.concat([df_round, resp_df], ignore_index=True) + combined_round = combined_round.drop_duplicates( + subset=["Identifier"], keep="last" + ) + self._write_rank_checkpoint( + combined_round, round_path, len(attr_batches) + ) + for completed_artifact in ( + batch_state_path, + catchup_path, + plan_path, + ): + try: + Path(completed_artifact).unlink(missing_ok=True) + except OSError: + pass async def _coerce_dict(raw: Any) -> Dict[str, Any]: obj = await safest_json(raw) @@ -1065,22 +3265,10 @@ async def _coerce_dict(raw: Any) -> Dict[str, Any]: if attr_key_l not in batch_attr_map: continue real_attr = batch_attr_map[attr_key_l] - val = winner_raw - if isinstance(val, dict) and "winner" in val: - val = val.get("winner") - if isinstance(val, str): - v = val.strip().lower() - else: - v = "" - if v.startswith(("cir", "c", "left", "text a")): - history_pairs[real_attr].append((id_a, id_b)) - elif v.startswith(("squ", "b", "right", "text b")): - history_pairs[real_attr].append((id_b, id_a)) - elif v.startswith("draw") or v.startswith("insufficient"): - history_pairs[real_attr].append((id_a, id_b)) - history_pairs[real_attr].append((id_b, id_a)) - else: - continue + category = self._record_pairwise_outcome( + history_pairs[real_attr], id_a, id_b, winner_raw + ) + outcome_counts[real_attr][category] += 1 se_agg_next: Dict[str, float] = {i: 0.0 for i in item_ids} se_agg_counts: Dict[str, int] = {i: 0 for i in item_ids} for attr in attr_keys: @@ -1102,12 +3290,17 @@ async def _coerce_dict(raw: Any) -> Dict[str, Any]: s=s_vec, n_ij=n_ij, p_ij=p_ij, - ridge=self._SE_RIDGE, + rcond=self._SE_EIGEN_TOL, + regularization_strength=self.cfg.learning_rate, ) + component_labels = self._comparison_component_labels(n_ij) for i, se_val in zip(item_ids, se_vec): se_store[attr][i] = float(se_val) - se_agg_next[i] += float(se_val) - se_agg_counts[i] += 1 + if np.isfinite(se_val): + se_agg_next[i] += float(se_val) + se_agg_counts[i] += 1 + for i, component in zip(item_ids, component_labels): + component_store[attr][i] = int(component) for i in item_ids: if se_agg_counts[i] > 0: se_agg_next[i] /= se_agg_counts[i] @@ -1143,14 +3336,33 @@ async def _run_recursive( cut_side = (self.cfg.recursive_cut_side or "top").lower() if cut_side not in {"top", "bottom"}: raise ValueError("recursive_cut_side must be 'top' or 'bottom'") + reserved_input_columns = [ + str(column) + for column in df.columns + if str(column).strip().lower() + in {"identifier", "overall_rank", "exit_stage"} + or re.match( + r"^stage\d+_", str(column).strip(), flags=re.IGNORECASE + ) + ] + if reserved_input_columns: + raise ValueError( + "Recursive Rank input columns cannot use internal output names " + "or the stage_ prefix: " + + ", ".join(repr(name) for name in reserved_input_columns) + ) work_df = df.reset_index(drop=True).copy() strict_text_mode = self.cfg.modality in {"text", "entity", "web"} if id_column is not None: if id_column not in work_df.columns: raise ValueError(f"id_column '{id_column}' not found in DataFrame") - valid_mask = ~work_df[id_column].map(_is_missing_scalar) - dropped = int((~valid_mask).sum()) + valid_mask = pd.Series( + [_is_valid_identifier(value) for value in work_df[id_column]], + index=work_df.index, + dtype=bool, + ) + dropped = len(work_df) - int(valid_mask.sum()) if dropped > 0: total = len(work_df) pct = (dropped / total * 100.0) if total else 0.0 @@ -1178,33 +3390,61 @@ async def _run_recursive( work_df = work_df.loc[valid_mask].copy().reset_index(drop=True) work_df["identifier"] = hashed.loc[valid_mask].astype(str).reset_index(drop=True) if text_column != "text": - work_df = work_df.rename(columns={text_column: "text"}) + # Keep the caller's original column for the returned DataFrame and + # use an internal text view for prompt rendering. + work_df["text"] = work_df[text_column] rewrite_col = self.cfg.recursive_rewrite_text_col or "text" if rewrite_col not in work_df.columns: work_df[rewrite_col] = work_df["text"] work_df["identifier"] = work_df["identifier"].astype(str) + duplicate_ids = work_df.loc[ + work_df["identifier"].duplicated(keep=False), "identifier" + ].unique() + if len(duplicate_ids) > 0: + preview = ", ".join(repr(value) for value in duplicate_ids[:3]) + raise ValueError( + "Rank requires a unique identifier for every row; duplicate " + f"identifier(s): {preview}. Provide a unique id_column when " + "ranking duplicate content." + ) + + base_folder = os.path.join( + self.cfg.save_dir, f"{self.cfg.file_name}_recursive" + ) + base_folder_path = Path(base_folder) + existing_recursive_artifacts = base_folder_path.exists() and any( + path.is_file() for path in base_folder_path.rglob("*") + ) if work_df.empty: empty_out = work_df[[c for c in df.columns if c in work_df.columns]].copy() for attr in attr_list: empty_out[attr] = pd.Series(dtype="float64") - empty_out[f"{attr}_raw"] = pd.Series(dtype="float64") - empty_out[f"{attr}_se"] = pd.Series(dtype="float64") empty_out["overall_rank"] = pd.Series(dtype="float64") empty_out["exit_stage"] = pd.Series(dtype="float64") - final_file = os.path.join(self.cfg.save_dir, f"{self.cfg.file_name}_recursive_final.csv") + if existing_recursive_artifacts and not reset_files: + return empty_out + os.makedirs(base_folder, exist_ok=True) + final_file = os.path.join(base_folder, "recursive_final.csv") empty_out.to_csv(final_file, index=False) return empty_out - original_cols = [c for c in df.columns if c in work_df.columns] + # A requested attribute replaces any same-named input column. Exclude + # those columns here so the final column selection cannot duplicate the + # generated stage-relative score. Recursive Rank does not emit the + # non-recursive raw/SE/component suffixes, so unrelated user suffix + # columns remain untouched. + generated_recursive_cols = set(attr_list) + original_cols = [ + c + for c in df.columns + if c in work_df.columns and c not in generated_recursive_cols + ] original_df = work_df[original_cols + ["identifier"]].copy() latest_text: Dict[str, str] = { ident: txt for ident, txt in zip(work_df["identifier"], work_df["text"]) } - base_folder = os.path.join( - self.cfg.save_dir, f"{self.cfg.file_name}_recursive" - ) os.makedirs(base_folder, exist_ok=True) def _compute_stage_zscores( @@ -1214,33 +3454,86 @@ def _compute_stage_zscores( scales: Dict[str, float] = {attr: 1.0 for attr in attr_list} for attr in attr_list: raw_col = f"{attr}_raw" - source_col = raw_col if raw_col in stage_df.columns else attr + component_col = f"{attr}_component" + # Rank already returns component-wise z-scores. Preserve them + # instead of globally re-standardizing arbitrary component + # locations. Rate stages do not have component labels and are + # normalized here as before. + if attr in stage_df.columns and component_col in stage_df.columns: + source_col = attr + else: + source_col = raw_col if raw_col in stage_df.columns else attr if source_col not in stage_df.columns: continue series = pd.to_numeric(stage_df[source_col], errors="coerce") - mean = series.mean() - std = series.std(ddof=0) - if std == 0 or np.isnan(std): - normed = pd.Series([0.0] * len(series), index=stage_df.index) - scales[attr] = 1.0 + if source_col == attr and component_col in stage_df.columns: + normed = series + if raw_col in stage_df.columns: + raw_std = pd.to_numeric( + stage_df[raw_col], errors="coerce" + ).std(ddof=0) + scales[attr] = ( + float(raw_std) + if np.isfinite(raw_std) and raw_std > 0 + else 1.0 + ) else: - normed = (series - mean) / std - scales[attr] = float(std) if raw_col in stage_df.columns else 1.0 + mean = series.mean() + std = series.std(ddof=0) + if std == 0 or np.isnan(std): + normed = pd.Series( + [0.0] * len(series), index=stage_df.index + ) + scales[attr] = 1.0 + else: + normed = (series - mean) / std + scales[attr] = ( + float(std) if raw_col in stage_df.columns else 1.0 + ) for ident, val in zip(stage_df["identifier"], normed): zscores[attr][str(ident)] = float(val) return zscores, scales - def _select_next_ids(active_ids: Sequence[str], stage_zs: Dict[str, Dict[str, float]]) -> List[str]: + def _require_identifiable_cut_scores(stage_df: pd.DataFrame) -> None: + component_col = f"{cut_attr}_component" + if component_col not in stage_df.columns: + return + components = pd.to_numeric( + stage_df[component_col], errors="coerce" + ).dropna() + if components.nunique() > 1: + raise ValueError( + "Recursive Rank cannot prune or globally order a disconnected " + f"comparison graph for '{cut_attr}'. Component-relative " + "Bradley-Terry scores have no identified ordering across " + "components. Increase n_rounds or matches_per_round, or use " + "a comparison design that connects every surviving item." + ) + + def _select_next_ids( + active_ids: Sequence[str], + stage_zs: Dict[str, Dict[str, float]], + ) -> List[str]: n = len(active_ids) if n <= self.cfg.recursive_min_remaining: return list(active_ids) - keep_n = max( - int(math.ceil(n * self.cfg.recursive_fraction)), - self.cfg.recursive_min_remaining, + keep_n = min( + n - 1, + max( + int(math.ceil(n * self.cfg.recursive_fraction)), + self.cfg.recursive_min_remaining, + ), ) - scores = {i: stage_zs.get(cut_attr, {}).get(i, 0.0) for i in active_ids} + scores = { + item_id: stage_zs.get(cut_attr, {}).get(item_id, 0.0) + for item_id in active_ids + } ascending = cut_side == "bottom" - ranked = sorted(active_ids, key=lambda x: scores.get(x, 0.0), reverse=not ascending) + ranked = sorted( + active_ids, + key=lambda item_id: scores.get(item_id, 0.0), + reverse=not ascending, + ) return ranked[:keep_n] def _maybe_rewrite_texts( @@ -1271,7 +3564,6 @@ def _maybe_rewrite_texts( return df_local stage_idx = 0 - final_stage_idx: Optional[int] = None final_stage_df: Optional[pd.DataFrame] = None stage_z_history: Dict[int, Dict[str, Dict[str, float]]] = {} exit_stage: Dict[str, Optional[int]] = {ident: None for ident in work_df["identifier"]} @@ -1285,9 +3577,12 @@ def _maybe_rewrite_texts( if n_current <= self.cfg.recursive_min_remaining: is_final_stage = True else: - next_keep = max( - int(math.ceil(n_current * self.cfg.recursive_fraction)), - self.cfg.recursive_min_remaining, + next_keep = min( + n_current - 1, + max( + int(math.ceil(n_current * self.cfg.recursive_fraction)), + self.cfg.recursive_min_remaining, + ), ) if next_keep <= self.cfg.recursive_min_remaining: is_final_stage = True @@ -1336,6 +3631,7 @@ def _maybe_rewrite_texts( **kwargs, ) + _require_identifiable_cut_scores(stage_df_out) stage_zs, stage_scales = _compute_stage_zscores(stage_df_out) stage_z_history[stage_idx] = stage_zs @@ -1343,7 +3639,6 @@ def _maybe_rewrite_texts( for ident in current_ids: exit_stage[ident] = stage_idx final_stage_df = stage_df_out - final_stage_idx = stage_idx break next_ids = _select_next_ids(current_ids, stage_zs) @@ -1362,8 +3657,6 @@ def _maybe_rewrite_texts( if final_stage_df is None: final_stage_df = work_df[work_df["identifier"].isin(current_ids)].copy() - if stage_idx: - final_stage_idx = stage_idx # Build final output stage_cols: Dict[str, List[Optional[float]]] = {} @@ -1382,7 +3675,7 @@ def _maybe_rewrite_texts( for attr in attr_list: col_name = f"stage{stage}_{attr}" stage_cols.setdefault(col_name, []).append(zs.get(attr, {}).get(ident)) - if final_stage_idx is not None and stage == final_stage_idx: + if ident_stage is not None and stage == ident_stage: final_attr_vals[attr] = zs.get(attr, {}).get(ident) for attr in attr_list: final_attr_cols[attr].append(final_attr_vals[attr]) @@ -1479,27 +3772,213 @@ async def run( Returns ------- pandas.DataFrame - A DataFrame with one row per input passage. For each + In non-recursive mode, a DataFrame with one row per input passage. For each attribute the DataFrame contains a ``""`` column holding the z‑score, a ``"_raw"`` column with the centred Bradley–Terry estimate, and a ``"_se"`` - column with the standard error. The DataFrame is also written - to ``save_dir``. + column with its model-based sandwich standard error. The + ``"_component"`` column identifies disconnected + comparison-graph components; z-scores and raw scores are only + comparable within the same component. Recursive mode instead + returns stage-relative scores, exit stages, and an overall rank. + The DataFrame is also written to ``save_dir``. """ + has_custom_judge = any( + kwargs.get(key) is not None + for key in ("response_fn", "get_all_responses_fn") + ) or any( + self.cfg.rate_kwargs.get(key) is not None + for key in ("response_fn", "get_all_responses_fn") + ) + if has_custom_judge and self.cfg.judge_version is None: + raise ValueError( + "judge_version is required when response_fn or " + "get_all_responses_fn supplies a custom Rank judge" + ) + conflicting_response_keys = sorted( + _RANK_OWNED_RESPONSE_KEYS & kwargs.keys() + ) + if conflicting_response_keys: + raise TypeError( + "Rank owns response setting(s) that cannot be overridden at " + "runtime: " + ", ".join(conflicting_response_keys) + ) + kwargs.setdefault("web_search", self.cfg.modality == "web") base_name = os.path.splitext(self.cfg.file_name)[0] + final_path = os.path.join(self.cfg.save_dir, f"{base_name}_final.csv") + recursive_base_folder = Path(self.cfg.save_dir) / ( + f"{self.cfg.file_name}_recursive" + ) + initial_rate_folder = Path(self.cfg.save_dir) / f"{base_name}_initial_rate" + if reset_files: + if self.cfg.recursive: + reset_batch_states = ( + list(recursive_base_folder.rglob("*.batch_state.json")) + if recursive_base_folder.exists() + else [] + ) + else: + root = Path(self.cfg.save_dir) + reset_batch_states = [] + for pattern in ( + f"{base_name}_round*.csv.batch_state.json", + f".{base_name}_round*.csv.batch_state.json", + f".{base_name}_catchup_round*.csv.batch_state.json", + ): + reset_batch_states.extend(root.glob(pattern)) + if initial_rate_folder.exists(): + reset_batch_states.extend( + initial_rate_folder.rglob("*.batch_state.json") + ) + for state_path in dict.fromkeys(reset_batch_states): + try: + has_external_work = _rank_batch_state_has_external_work( + state_path + ) + except Exception as exc: + raise ValueError( + "Rank cannot safely reset unreadable Batch API state " + f"{str(state_path)!r}; reconcile the external job first." + ) from exc + if has_external_work: + raise ValueError( + "Rank cannot reset while durable Batch API state may " + f"reference external work ({str(state_path)!r}). Resume, " + "reconcile, or cancel that batch before resetting files." + ) + if initial_rate_folder.exists(): + try: + shutil.rmtree(initial_rate_folder) + except OSError as exc: + raise OSError( + "Could not clear stale Rank initial-rating artifacts in " + f"{str(initial_rate_folder)!r}" + ) from exc + if reset_files and self.cfg.recursive and recursive_base_folder.exists(): + try: + shutil.rmtree(recursive_base_folder) + except OSError as exc: + raise OSError( + "Could not clear stale recursive Rank artifacts in " + f"{str(recursive_base_folder)!r}" + ) from exc + if reset_files and not self.cfg.recursive: + stale_paths = list( + Path(self.cfg.save_dir).glob(f"{base_name}_round*.csv") + ) + stale_paths.extend( + Path(self.cfg.save_dir).glob( + f"{base_name}_round*.csv.batch_state.json" + ) + ) + stale_paths.extend( + Path(self.cfg.save_dir).glob(f".{base_name}_round*.csv") + ) + stale_paths.extend( + Path(self.cfg.save_dir).glob( + f".{base_name}_round*.csv.batch_state.json" + ) + ) + stale_paths.extend( + Path(self.cfg.save_dir).glob( + f".{base_name}_round*_plan.json" + ) + ) + stale_paths.extend( + Path(self.cfg.save_dir).glob( + f".{base_name}_catchup_round*.csv" + ) + ) + stale_paths.extend( + Path(self.cfg.save_dir).glob( + f".{base_name}_catchup_round*.csv.batch_state.json" + ) + ) + stale_paths.extend( + Path(self.cfg.save_dir).glob( + f".{base_name}_catchup_round*_plan.json" + ) + ) + stale_paths.extend( + [ + Path(final_path), + Path(self.cfg.save_dir) / f"{base_name}_diagnostics.csv", + ] + ) + for stale_path in stale_paths: + try: + stale_path.unlink(missing_ok=True) + except OSError as exc: + raise OSError( + f"Could not clear stale Rank artifact {str(stale_path)!r}" + ) from exc checkpoint_paths = [ str(path) for path in Path(self.cfg.save_dir).glob(f"{base_name}_round*.csv") ] + normal_staging_paths = [ + str(path) + for path in Path(self.cfg.save_dir).glob(f".{base_name}_round*.csv") + ] + normal_plan_paths = [ + str(path) + for path in Path(self.cfg.save_dir).glob( + f".{base_name}_round*_plan.json" + ) + ] + round_batch_state_paths = [ + str(path) + for path in Path(self.cfg.save_dir).glob( + f"{base_name}_round*.csv.batch_state.json" + ) + ] + round_batch_state_paths.extend( + str(path) + for path in Path(self.cfg.save_dir).glob( + f".{base_name}_round*.csv.batch_state.json" + ) + ) run_metadata = load_run_metadata( self.cfg.save_dir, base_name, reset_files=reset_files ) + recursive_artifacts_exist = self.cfg.recursive and ( + recursive_base_folder.exists() + and any(path.is_file() for path in recursive_base_folder.rglob("*")) + ) + initial_rate_artifacts_exist = ( + initial_rate_folder.exists() + and any(path.is_file() for path in initial_rate_folder.rglob("*")) + ) + rank_resume_artifacts_exist = ( + bool(checkpoint_paths) + or bool(normal_staging_paths) + or bool(normal_plan_paths) + or bool(round_batch_state_paths) + or os.path.exists(final_path) + or os.path.exists(run_metadata_path(self.cfg.save_dir, base_name)) + or recursive_artifacts_exist + or initial_rate_artifacts_exist + ) identifier_hash_bits = resolve_identifier_hash_bits( task_name="Rank", metadata=run_metadata, reset_files=reset_files, checkpoint_paths=checkpoint_paths, ) + attribute_sidecar_paths = [ + Path(self.cfg.save_dir) / "attributes.json", + Path(self.cfg.save_dir) / f"{base_name}_attrs.json", + ] + has_readable_attribute_sidecar = False + if rank_resume_artifacts_exist and not reset_files: + for attribute_path in attribute_sidecar_paths: + try: + with attribute_path.open(encoding="utf-8") as attribute_file: + json.load(attribute_file) + has_readable_attribute_sidecar = True + break + except Exception: + continue self.cfg.attributes = load_persisted_attributes( save_dir=self.cfg.save_dir, incoming=self.cfg.attributes, @@ -1507,10 +3986,242 @@ async def run( task_name="Rank", item_name="attributes", legacy_filename=f"{base_name}_attrs.json", + persist_missing=reset_files or not rank_resume_artifacts_exist, + ) + _validate_rank_attribute_names( + self.cfg.attributes, recursive=self.cfg.recursive + ) + measurement_spec_fingerprint = self._measurement_spec_fingerprint(kwargs) + _validate_rank_resume_metadata( + run_metadata, + artifacts_exist=rank_resume_artifacts_exist, + reset_files=reset_files, + insufficient_signal_policy=self.cfg.insufficient_signal_policy, + learning_rate=self.cfg.learning_rate, + modality=self.cfg.modality, + measurement_spec_fingerprint=measurement_spec_fingerprint, ) + if ( + rank_resume_artifacts_exist + and not reset_files + and not has_readable_attribute_sidecar + ): + # Restore missing/corrupt attribute sidecars only after metadata has + # proved that the effective measurement definition is compatible. + self.cfg.attributes = load_persisted_attributes( + save_dir=self.cfg.save_dir, + incoming=self.cfg.attributes, + reset_files=False, + task_name="Rank", + item_name="attributes", + legacy_filename=f"{base_name}_attrs.json", + persist_missing=True, + ) + _validate_rank_attribute_names( + self.cfg.attributes, recursive=self.cfg.recursive + ) + if not reset_files: + for batch_state_path in round_batch_state_paths: + try: + with open(batch_state_path, encoding="utf-8") as state_file: + batch_state = json.load(state_file) + except Exception as exc: + raise ValueError( + f"Could not read Rank Batch API state {batch_state_path!r}" + ) from exc + if not isinstance(batch_state, dict): + raise ValueError( + f"Rank Batch API state {batch_state_path!r} is malformed" + ) + round_checkpoint_path = batch_state_path.removesuffix( + ".batch_state.json" + ) + batches = batch_state.get("batches", []) + if not isinstance(batches, list): + raise ValueError( + f"Rank Batch API state {batch_state_path!r} is malformed" + ) + has_active_batch = bool(batch_state.get("batch_id")) or bool( + batches + ) + unresolved_submission = any( + not isinstance(batch, dict) + or batch.get("status") == "submitting" + or not batch.get("batch_id") + for batch in batches + ) + if unresolved_submission: + raise ValueError( + "Rank found an unresolved Batch API submission. The " + "server may already have accepted a paid job; reconcile " + "or cancel it before resetting local state." + ) + if has_active_batch: + normal_plan_path = ( + f"{round_checkpoint_path.removesuffix('.csv')}_plan.json" + ) + is_journaled_staging = Path(round_checkpoint_path).name.startswith( + f".{base_name}_round" + ) and os.path.exists(normal_plan_path) + if not is_journaled_staging: + raise ValueError( + "Rank found active Batch API state without a " + "recoverable round plan. Reconcile or cancel the " + "external batch before resetting local state." + ) + # The durable plan reconstructs the exact prompt IDs, so + # get_all_responses can safely poll this submitted batch. + continue + if not os.path.exists(round_checkpoint_path): + raise ValueError( + "Rank found completed/empty Batch API state without a " + "durable response checkpoint. Paid results may already " + "exist; use reset_files=True only after confirming that " + "it is safe to discard that state." + ) kwargs.setdefault("web_search", self.cfg.modality == "web") if self.cfg.recursive: + current_input_fingerprints = _current_rank_input_fingerprints( + df, + payload_column=column_name, + id_column=id_column, + modality=self.cfg.modality, + identifier_hash_bits=identifier_hash_bits, + ) + if ( + rank_resume_artifacts_exist + and not reset_files + and run_metadata.get("recursive_input_fingerprint_version") != 1 + ): + raise ValueError( + "Existing recursive Rank artifacts lack content-aware input " + "fingerprints. Use reset_files=True or a new save_dir to " + "recompute them safely." + ) + saved_input_fingerprints = ( + run_metadata.get("input_fingerprints", {}) + if rank_resume_artifacts_exist and not reset_files + else {} + ) + if not isinstance(saved_input_fingerprints, dict): + saved_input_fingerprints = {} + recursive_transaction_plans: List[Path] = [] + if recursive_base_folder.exists(): + normal_stage_plan = re.compile( + rf"^\.{re.escape(base_name)}_round(\d+)_plan\.json$" + ) + catchup_stage_plan = re.compile( + rf"^\.{re.escape(base_name)}_catchup_round(\d+)_plan\.json$" + ) + for plan_path in recursive_base_folder.rglob("*_plan.json"): + stage_metadata = load_run_metadata( + str(plan_path.parent), base_name, reset_files=False + ) + marker = stage_metadata.get("last_completed_round", -1) + marker = marker if type(marker) is int else -1 + normal_match = normal_stage_plan.match(plan_path.name) + if normal_match is not None: + if int(normal_match.group(1)) > marker: + recursive_transaction_plans.append(plan_path) + continue + catchup_match = catchup_stage_plan.match(plan_path.name) + if catchup_match is None: + recursive_transaction_plans.append(plan_path) + continue + round_index = int(catchup_match.group(1)) + round_path = plan_path.parent / ( + f"{base_name}_round{round_index}.csv" + ) + try: + plan_payload = json.loads(plan_path.read_text()) + records = plan_payload["records"] + if not isinstance(records, list) or not records: + raise ValueError("malformed catch-up plan records") + planned_response_ids = { + str(record["identifier"]) for record in records + } + if len(planned_response_ids) != len(records): + raise ValueError("duplicate catch-up plan identifiers") + committed_response_ids = set( + pd.read_csv( + round_path, + usecols=["Identifier"], + dtype=str, + keep_default_na=False, + )["Identifier"].astype(str) + ) + except Exception: + recursive_transaction_plans.append(plan_path) + continue + committed_plan_ids = ( + planned_response_ids & committed_response_ids + ) + if committed_plan_ids == planned_response_ids: + continue + if committed_plan_ids: + raise ValueError( + "Recursive Rank found a partially committed stage " + "catch-up plan" + ) + recursive_transaction_plans.append(plan_path) + recursive_batch_states = ( + list(recursive_base_folder.rglob("*.batch_state.json")) + if recursive_base_folder.exists() + else [] + ) + if initial_rate_folder.exists(): + recursive_batch_states.extend( + initial_rate_folder.rglob("*.batch_state.json") + ) + has_active_recursive_batch = any( + _rank_batch_state_has_external_work(path) + for path in recursive_batch_states + ) + if ( + (recursive_transaction_plans or has_active_recursive_batch) + and set(current_input_fingerprints) + != set(saved_input_fingerprints) + ): + raise ValueError( + "Recursive Rank has an unfinished stage transaction. " + "Restore the exact top-level input set before resuming." + ) + changed_ids = [ + item_id + for item_id, fingerprint in current_input_fingerprints.items() + if item_id in saved_input_fingerprints + and saved_input_fingerprints[item_id] != fingerprint + ] + if changed_ids: + preview = ", ".join(repr(value) for value in changed_ids[:3]) + raise ValueError( + "Recursive Rank comparison payload changed for persisted " + f"identifier(s): {preview}. Existing stage judgments are " + "no longer valid; use reset_files=True or a new save_dir." + ) + merged_input_fingerprints = { + **saved_input_fingerprints, + **current_input_fingerprints, + } + update_run_metadata( + self.cfg.save_dir, + base_name, + strict=True, + task="Rank", + output_base_name=base_name, + model=self.cfg.model, + identifier_hash_bits=identifier_hash_bits, + rank_estimator_version=_RANK_ESTIMATOR_VERSION, + insufficient_signal_policy=self.cfg.insufficient_signal_policy, + learning_rate=self.cfg.learning_rate, + last_completed_round=-1, + modality=self.cfg.modality, + input_fingerprints=merged_input_fingerprints, + recursive_input_fingerprint_version=1, + measurement_spec_fingerprint=measurement_spec_fingerprint, + recursive=True, + ) return await self._run_recursive( df, column_name, @@ -1521,7 +4232,6 @@ async def run( ) # prepare file paths - final_path = os.path.join(self.cfg.save_dir, f"{base_name}_final.csv") if n_runs is not None: print( "Parameter 'n_runs' is ignored. Use 'n_rounds' to control the number of iterations. " @@ -1534,8 +4244,12 @@ async def run( if id_column is not None: if id_column not in df_proc.columns: raise ValueError(f"id_column '{id_column}' not found in DataFrame") - valid_mask = ~df_proc[id_column].map(_is_missing_scalar) - dropped = int((~valid_mask).sum()) + valid_mask = pd.Series( + [_is_valid_identifier(value) for value in df_proc[id_column]], + index=df_proc.index, + dtype=bool, + ) + dropped = len(df_proc) - int(valid_mask.sum()) if dropped > 0: total = len(df_proc) pct = (dropped / total * 100.0) if total else 0.0 @@ -1573,8 +4287,347 @@ async def run( df_out[attr] = pd.Series(dtype="float64") df_out[f"{attr}_raw"] = pd.Series(dtype="float64") df_out[f"{attr}_se"] = pd.Series(dtype="float64") + df_out[f"{attr}_component"] = pd.Series(dtype="int64") + if rank_resume_artifacts_exist and not reset_files: + # An empty view of an existing tournament must not erase its + # checkpoints, fingerprints, or committed-round marker. + return df_out + empty_attr_items = [(attr, attr) for attr in attr_keys_empty] + empty_batches, effective_empty_batch_size = resolve_attribute_batches( + task_name="Rank", + items=empty_attr_items, + requested_n=self.cfg.n_attributes_per_run, + metadata=run_metadata, + reset_files=reset_files, + checkpoint_paths=checkpoint_paths, + ) + write_task_run_metadata( + save_dir=self.cfg.save_dir, + base_name=base_name, + task_name="Rank", + model=self.cfg.model, + identifier_hash_bits=identifier_hash_bits, + n_attributes_per_run=effective_empty_batch_size, + attribute_batches=empty_batches, + strict=True, + rank_estimator_version=_RANK_ESTIMATOR_VERSION, + insufficient_signal_policy=self.cfg.insufficient_signal_policy, + learning_rate=self.cfg.learning_rate, + last_completed_round=-1, + modality=self.cfg.modality, + input_fingerprints={}, + measurement_spec_fingerprint=measurement_spec_fingerprint, + ) df_out.to_csv(final_path, index=False) + pd.DataFrame( + { + "attribute": attr_keys_empty, + "circle_count": 0, + "square_count": 0, + "draw_count": 0, + "insufficient_signal_count": 0, + "invalid_count": 0, + "effective_comparison_weight": 0.0, + "comparison_components": 0, + "isolated_items": 0, + "finite_standard_errors": 0, + } + ).to_csv( + os.path.join(self.cfg.save_dir, f"{base_name}_diagnostics.csv"), + index=False, + ) return df_out + + duplicate_ids = df_proc.loc[ + df_proc["_id"].duplicated(keep=False), "_id" + ].unique() + if len(duplicate_ids) > 0: + preview = ", ".join(repr(value) for value in duplicate_ids[:3]) + raise ValueError( + "Rank requires a unique identifier for every row; duplicate " + f"identifier(s): {preview}. Provide a unique id_column when " + "ranking duplicate content." + ) + current_input_fingerprints = _rank_input_fingerprints( + df_proc, + id_column="_id", + payload_column=column_name, + modality=self.cfg.modality, + ) + saved_input_fingerprints = ( + run_metadata.get("input_fingerprints", {}) + if rank_resume_artifacts_exist + else {} + ) + if not isinstance(saved_input_fingerprints, dict): + saved_input_fingerprints = {} + initial_rate_batch_states = ( + list(initial_rate_folder.rglob("*.batch_state.json")) + if initial_rate_folder.exists() + else [] + ) + if ( + any( + _rank_batch_state_has_external_work(path) + for path in initial_rate_batch_states + ) + and set(current_input_fingerprints) != set(saved_input_fingerprints) + ): + raise ValueError( + "Rank has an unfinished initial-rating Batch transaction. " + "Restore the exact input set before resuming." + ) + completed_marker = run_metadata.get("last_completed_round", -1) + completed_marker = ( + completed_marker if type(completed_marker) is int else -1 + ) + persisted_endpoint_ids: Set[str] = set() + current_id_set = set(current_input_fingerprints) + saved_attribute_batches = run_metadata.get("attribute_batches") + saved_batch_count = ( + len(saved_attribute_batches) + if isinstance(saved_attribute_batches, list) + and saved_attribute_batches + else 0 + ) + + # A durable plan is a transaction intent. Validate its input universe + # before merging new fingerprints into metadata; otherwise a rejected + # accidental superset can permanently bind never-judged payloads. + normal_plan_pattern = re.compile( + rf"^\.{re.escape(base_name)}_round(\d+)_plan\.json$" + ) + for raw_plan_path in normal_plan_paths: + plan_path = Path(raw_plan_path) + match = normal_plan_pattern.match(plan_path.name) + if match is None: + raise ValueError(f"Rank round plan {str(plan_path)!r} is noncanonical") + round_index = int(match.group(1)) + try: + with plan_path.open(encoding="utf-8") as plan_file: + plan = json.load(plan_file) + except Exception as exc: + raise ValueError( + f"Could not read Rank round plan {str(plan_path)!r}" + ) from exc + planned_ids = plan.get("item_ids") if isinstance(plan, dict) else None + records = plan.get("records") if isinstance(plan, dict) else None + if ( + not isinstance(plan, dict) + or plan.get("version") != 1 + or plan.get("round") != round_index + or not isinstance(planned_ids, list) + or any(not isinstance(item_id, str) for item_id in planned_ids) + or len(planned_ids) != len(set(planned_ids)) + or not isinstance(records, list) + or not records + ): + raise ValueError( + f"Rank round plan {str(plan_path)!r} is incompatible or malformed" + ) + round_path = Path(self.cfg.save_dir) / ( + f"{base_name}_round{round_index}.csv" + ) + promotion_safe = False + if round_path.exists() and round_index <= completed_marker: + promotion_safe = True + elif round_path.exists() and round_index == completed_marker + 1: + if saved_batch_count == 0: + raise ValueError( + "Rank cannot validate an uncommitted round without its " + "persisted attribute-batch schema" + ) + try: + self._read_rank_checkpoint(str(round_path), saved_batch_count) + except Exception as exc: + raise ValueError( + "Rank found an uncommitted round checkpoint that is not " + "safe to promote" + ) from exc + promotion_safe = True + if promotion_safe: + continue + if ( + round_index != completed_marker + 1 + or set(planned_ids) != current_id_set + ): + raise ValueError( + "Rank has an unfinished round plan for a different input " + "set. Restore the exact planned items before resuming." + ) + for record in records: + if not isinstance(record, dict): + raise ValueError( + f"Rank round plan {str(plan_path)!r} has malformed records" + ) + for key in ("id_a", "id_b"): + endpoint = record.get(key) + if not isinstance(endpoint, str) or endpoint not in current_id_set: + raise ValueError( + f"Rank round plan {str(plan_path)!r} has invalid endpoints" + ) + persisted_endpoint_ids.add(endpoint) + + catchup_plan_pattern = re.compile( + rf"^\.{re.escape(base_name)}_catchup_round(\d+)_plan\.json$" + ) + catchup_plan_paths = list( + Path(self.cfg.save_dir).glob( + f".{base_name}_catchup_round*_plan.json" + ) + ) + for plan_path in catchup_plan_paths: + match = catchup_plan_pattern.match(plan_path.name) + if match is None: + raise ValueError( + f"Rank catch-up plan {str(plan_path)!r} is noncanonical" + ) + round_index = int(match.group(1)) + try: + with plan_path.open(encoding="utf-8") as plan_file: + plan = json.load(plan_file) + except Exception as exc: + raise ValueError( + f"Could not read Rank catch-up plan {str(plan_path)!r}" + ) from exc + planned_ids = plan.get("item_ids") if isinstance(plan, dict) else None + records = plan.get("records") if isinstance(plan, dict) else None + if ( + not isinstance(plan, dict) + or plan.get("version") != 1 + or plan.get("round") != round_index + or not isinstance(planned_ids, list) + or any(not isinstance(item_id, str) for item_id in planned_ids) + or len(planned_ids) != len(set(planned_ids)) + or not isinstance(records, list) + or not records + ): + raise ValueError( + f"Rank catch-up plan {str(plan_path)!r} is incompatible or malformed" + ) + round_path = Path(self.cfg.save_dir) / ( + f"{base_name}_round{round_index}.csv" + ) + try: + committed_ids = set( + pd.read_csv( + round_path, + usecols=["Identifier"], + dtype=str, + keep_default_na=False, + )["Identifier"].astype(str) + ) + except Exception as exc: + raise ValueError( + "Rank catch-up plan does not have a readable committed " + f"round checkpoint: {str(plan_path)!r}" + ) from exc + plan_response_ids = { + str(record.get("identifier")) + for record in records + if isinstance(record, dict) + } + if len(plan_response_ids) != len(records) or "None" in plan_response_ids: + raise ValueError( + f"Rank catch-up plan {str(plan_path)!r} has malformed records" + ) + already_committed = plan_response_ids & committed_ids + if already_committed == plan_response_ids: + continue + if already_committed: + raise ValueError("Rank found a partially committed catch-up plan") + if set(planned_ids) != current_id_set: + raise ValueError( + "Rank has an unfinished catch-up plan for a different input " + "set. Restore the exact planned items before resuming." + ) + for record in records: + for key in ("id_a", "id_b"): + endpoint = record.get(key) + if not isinstance(endpoint, str) or endpoint not in current_id_set: + raise ValueError( + f"Rank catch-up plan {str(plan_path)!r} has invalid endpoints" + ) + persisted_endpoint_ids.add(endpoint) + + round_name_pattern = re.compile( + rf"^{re.escape(base_name)}_round(\d+)\.csv$" + ) + for checkpoint_path in checkpoint_paths: + match = round_name_pattern.match(Path(checkpoint_path).name) + if match is None or int(match.group(1)) > completed_marker: + continue + try: + endpoint_frame = pd.read_csv( + checkpoint_path, + usecols=["IdA", "IdB"], + dtype=str, + keep_default_na=False, + ) + except Exception as exc: + raise ValueError( + f"Could not verify persisted Rank endpoints in " + f"{checkpoint_path!r}" + ) from exc + persisted_endpoint_ids.update(endpoint_frame["IdA"].astype(str)) + persisted_endpoint_ids.update(endpoint_frame["IdB"].astype(str)) + missing_or_malformed_fingerprints = sorted( + item_id + for item_id in persisted_endpoint_ids + if not isinstance(saved_input_fingerprints.get(item_id), str) + or re.fullmatch( + r"[0-9a-f]{40}", str(saved_input_fingerprints.get(item_id, "")) + ) + is None + ) + if missing_or_malformed_fingerprints: + preview = ", ".join( + repr(value) for value in missing_or_malformed_fingerprints[:3] + ) + raise ValueError( + "Rank metadata lacks valid content fingerprints for persisted " + f"checkpoint identifier(s): {preview}. Use reset_files=True or " + "a new save_dir to recompute the tournament safely." + ) + changed_ids = [ + item_id + for item_id, fingerprint in current_input_fingerprints.items() + if item_id in saved_input_fingerprints + and saved_input_fingerprints[item_id] != fingerprint + ] + if changed_ids: + preview = ", ".join(repr(value) for value in changed_ids[:3]) + raise ValueError( + "Rank comparison payload changed for persisted identifier(s): " + f"{preview}. Existing pairwise judgments are no longer valid; " + "use reset_files=True or a new save_dir." + ) + merged_input_fingerprints = { + **saved_input_fingerprints, + **current_input_fingerprints, + } + update_run_metadata( + self.cfg.save_dir, + base_name, + strict=True, + task="Rank", + output_base_name=base_name, + model=self.cfg.model, + identifier_hash_bits=identifier_hash_bits, + rank_estimator_version=_RANK_ESTIMATOR_VERSION, + insufficient_signal_policy=self.cfg.insufficient_signal_policy, + learning_rate=self.cfg.learning_rate, + last_completed_round=( + run_metadata["last_completed_round"] + if type(run_metadata.get("last_completed_round")) is int + and run_metadata["last_completed_round"] >= -1 + else -1 + ), + modality=self.cfg.modality, + input_fingerprints=merged_input_fingerprints, + measurement_spec_fingerprint=measurement_spec_fingerprint, + recursive=False, + ) # Determine how many rounds have already been processed when # `reset_files` is False. We look for files named # ``_round.csv`` to infer progress. If a final @@ -1590,51 +4643,107 @@ async def run( if fname.startswith(f"{base_name}_round") and fname.endswith( ".csv" ): + idx_str = fname[ + len(base_name) + 6 : -4 + ] # len("_round") == 6 try: - idx_str = fname[ - len(base_name) + 6 : -4 - ] # len("_round") == 6 rnd_idx = int(idx_str) - existing_rounds.append(rnd_idx) - except Exception: + except (TypeError, ValueError): continue + if fname != f"{base_name}_round{rnd_idx}.csv": + raise ValueError( + "Existing Rank checkpoint has a noncanonical " + f"filename: {fname!r}. Rename or remove it, " + "or use reset_files=True." + ) + existing_rounds.append(rnd_idx) + except ValueError: + raise except Exception: existing_rounds = [] + completed_round = ( + int(run_metadata.get("last_completed_round", -1)) + if existing_rounds or os.path.exists(final_path) + else -1 + ) if existing_rounds: - last_completed = max(existing_rounds) - if os.path.exists(final_path): + existing_rounds = sorted(set(existing_rounds)) + if existing_rounds != list(range(existing_rounds[-1] + 1)): + raise ValueError( + "Existing Rank round checkpoints are not contiguous from " + "round 0. Use reset_files=True or a new save_dir to " + "recompute them." + ) + if completed_round > existing_rounds[-1]: + raise ValueError( + "Rank metadata marks a completed round whose checkpoint " + "file is missing. Use reset_files=True or a new save_dir " + "to recompute the run." + ) + if existing_rounds[-1] > completed_round + 1: + raise ValueError( + "Rank has more than one uncommitted round checkpoint. " + "Use reset_files=True or a new save_dir to recompute the " + "run." + ) + if existing_rounds[-1] == completed_round + 1: + uncommitted_path = Path(self.cfg.save_dir) / ( + f"{base_name}_round{completed_round + 1}.csv" + ) + saved_attribute_batches = run_metadata.get("attribute_batches") + if ( + not isinstance(saved_attribute_batches, list) + or not saved_attribute_batches + ): + raise ValueError( + "Rank found an uncommitted round but cannot validate its " + "attribute-batch schema. Use reset_files=True or a new " + "save_dir to recompute the run." + ) try: - final_df = pd.read_csv(final_path) - identifier_col = ( - id_column - if id_column and id_column in final_df.columns - else column_name + self._read_rank_checkpoint( + str(uncommitted_path), len(saved_attribute_batches) ) - if identifier_col not in final_df.columns: - raise ValueError( - "Existing ranking output is missing identifier column " - f"'{identifier_col}'." - ) - if id_column: - final_ids = set(final_df[identifier_col].astype(str)) - else: - final_ids = set( - final_df[identifier_col] - .map( - lambda x: _hash_text_identifier( - x, - strict=True, - bits=identifier_hash_bits, - ) - ) - .dropna() - .astype(str) - ) - if last_completed >= self.cfg.n_rounds - 1 and set(df_proc["_id"]) <= final_ids: - return final_df - except Exception: - pass - start_round = last_completed + 1 + except Exception as exc: + raise ValueError( + "Rank found an uncommitted round checkpoint that is not " + "safe to promote. Use reset_files=True or a new save_dir " + "to recompute the run." + ) from exc + completed_round += 1 + update_run_metadata( + self.cfg.save_dir, + base_name, + strict=True, + last_completed_round=completed_round, + ) + print( + "[Rank] Promoted a complete uncommitted round checkpoint " + f"without repeating its judgments (round {completed_round})." + ) + for completed_artifact in ( + Path(self.cfg.save_dir) + / f".{base_name}_round{completed_round}.csv.batch_state.json", + Path(self.cfg.save_dir) + / f".{base_name}_round{completed_round}.csv", + Path(self.cfg.save_dir) + / f".{base_name}_round{completed_round}_plan.json", + ): + completed_artifact.unlink(missing_ok=True) + if completed_round + 1 > self.cfg.n_rounds: + raise ValueError( + "Rank cannot resume with fewer n_rounds than are already " + f"committed (requested={self.cfg.n_rounds}, " + f"committed={completed_round + 1}). Use reset_files=True " + "or a new save_dir to recompute the shorter tournament." + ) + start_round = completed_round + 1 + elif completed_round >= 0: + raise ValueError( + "Rank metadata marks completed rounds, but no round checkpoint " + "files exist. Use reset_files=True or a new save_dir to " + "recompute the run." + ) # extract contents and build lookup if self.cfg.modality in {"image", "audio", "pdf"}: texts = list(zip(df_proc["_id"], ["" for _ in df_proc[column_name]])) @@ -1673,7 +4782,7 @@ async def run( rate_seed: Dict[str, Dict[str, float]] = {} if self.cfg.primer_scores: self._apply_primer(ratings, self.cfg.primer_scores, attr_keys) - if self.cfg.initial_rating_pass and attr_keys: + if self.cfg.initial_rating_pass and attr_keys and len(item_ids) > 1: print( "[Rank] Running initial rating pass to seed pairwise comparisons " "(disable with initial_rating_pass=False)." @@ -1706,11 +4815,28 @@ async def run( ratings[item_id][attr] = val has_seed_ratings = bool(rate_seed) # maintain a history of pairwise outcomes for each attribute - history_pairs: Dict[str, List[Tuple[str, str]]] = {a: [] for a in attr_keys} + history_pairs: Dict[str, List[WeightedOutcome]] = { + a: [] for a in attr_keys + } + outcome_categories = ( + "circle", + "square", + "draw", + "insufficient_signal", + "invalid", + ) + outcome_counts: Dict[str, Dict[str, int]] = { + attr: {category: 0 for category in outcome_categories} + for attr in attr_keys + } # store per‑attribute standard errors across items se_store: Dict[str, Dict[str, float]] = { a: {i: np.nan for i in item_ids} for a in attr_keys } + component_store: Dict[str, Dict[str, int]] = { + a: {item_id: index for index, item_id in enumerate(item_ids)} + for a in attr_keys + } # Define attribute batches once to reuse across replay and new rounds attr_items = [(attr, attr) for attr in attr_keys] attr_batch_items, effective_n_attributes_per_run = resolve_attribute_batches( @@ -1744,35 +4870,64 @@ async def run( identifier_hash_bits=identifier_hash_bits, n_attributes_per_run=effective_n_attributes_per_run, attribute_batches=attr_batch_items, + strict=True, + rank_estimator_version=_RANK_ESTIMATOR_VERSION, + insufficient_signal_policy=self.cfg.insufficient_signal_policy, + learning_rate=self.cfg.learning_rate, + last_completed_round=completed_round, + modality=self.cfg.modality, + input_fingerprints=merged_input_fingerprints, + measurement_spec_fingerprint=measurement_spec_fingerprint, ) # Helper function to write the current results to the final CSV. This # builds the output DataFrame from the current ``df_proc`` and # ``ratings``/``se_store``/``zscores`` and writes it to # ``final_path``. + checkpoint_result: Optional[pd.DataFrame] = None + def _write_checkpoint() -> None: - # Compute z‑scores for each attribute so that we can expose centred - # BT scores alongside their normalised variants. + nonlocal checkpoint_result + # Compute z-scores within observed comparison components. A global + # normalization would make even within-component values change when + # an unrelated disconnected group is added. zscores_local: Dict[str, Dict[str, float]] = {} for attr in attr_keys: vals = np.array([ratings[i][attr] for i in item_ids]) - mean = vals.mean() - std = vals.std(ddof=0) - if std == 0: - zscores_local[attr] = {i: 0.0 for i in item_ids} - else: - zscores_local[attr] = { - i: float((ratings[i][attr] - mean) / std) for i in item_ids - } + components = np.array( + [component_store[attr].get(i, -1) for i in item_ids], + dtype=int, + ) + zscores = self._component_zscores(vals, components) + zscores_local[attr] = { + i: float(value) for i, value in zip(item_ids, zscores) + } # Merge computed results back into the original DataFrame copy. for attr in attr_keys: raw_col = f"{attr}_raw" # ratings - val_map = {i: ratings[i][attr] for i in item_ids} + component_sizes: Dict[int, int] = {} + for component in component_store[attr].values(): + component_sizes[component] = component_sizes.get(component, 0) + 1 + val_map = { + i: ( + ratings[i][attr] + if component_sizes.get(component_store[attr].get(i, -1), 0) + > 1 + else np.nan + ) + for i in item_ids + } df_proc[raw_col] = df_proc["_id"].map(val_map) # standard errors se_map = {i: se_store[attr].get(i, np.nan) for i in item_ids} df_proc[f"{attr}_se"] = df_proc["_id"].map(se_map) + # Connected-component labels make it explicit when scores and + # standard errors are only comparable within a graph component. + component_map = { + i: component_store[attr].get(i, -1) for i in item_ids + } + df_proc[f"{attr}_component"] = df_proc["_id"].map(component_map) # z‑scores z_map = zscores_local.get(attr, {i: np.nan for i in item_ids}) df_proc[attr] = df_proc["_id"].map(z_map) @@ -1780,21 +4935,67 @@ def _write_checkpoint() -> None: # Reorder columns: original user columns first (excluding the internal ``_id``), # then for each attribute the z‑score column followed by raw scores and # standard errors. + generated_rank_cols = { + generated + for attr in attr_keys + for generated in ( + attr, + f"{attr}_raw", + f"{attr}_se", + f"{attr}_component", + ) + } original_cols = [ - c for c in df.columns - ] # preserve the order provided by the user + c for c in df.columns if c not in generated_rank_cols + ] # preserve the order of unaffected user columns new_cols: List[str] = [] for attr in attr_keys: new_cols.append(attr) new_cols.append(f"{attr}_raw") new_cols.append(f"{attr}_se") + new_cols.append(f"{attr}_component") final_cols = original_cols + new_cols final_cols = [c for c in final_cols if c in df_proc.columns] df_out_local = df_proc[final_cols].copy() + checkpoint_result = df_out_local # Write the final results to disk in CSV format. Using CSV avoids # Excel row limits and unnecessary overhead. df_out_local.to_csv(final_path, index=False) + diagnostic_rows: List[Dict[str, Any]] = [] + for attr in attr_keys: + components = list(component_store[attr].values()) + component_sizes: Dict[int, int] = {} + for component in components: + component_sizes[component] = component_sizes.get(component, 0) + 1 + counts = outcome_counts[attr] + diagnostic_rows.append( + { + "attribute": attr, + **{f"{category}_count": counts[category] for category in outcome_categories}, + "effective_comparison_weight": float( + sum(weight for _, _, weight in history_pairs[attr]) + ), + "comparison_components": len(component_sizes), + "isolated_items": sum( + size == 1 for size in component_sizes.values() + ), + "finite_standard_errors": sum( + np.isfinite(se_store[attr].get(item_id, np.nan)) + for item_id in item_ids + ), + } + ) + pd.DataFrame(diagnostic_rows).to_csv( + os.path.join(self.cfg.save_dir, f"{base_name}_diagnostics.csv"), + index=False, + ) + + if len(item_ids) < 2: + _write_checkpoint() + assert checkpoint_result is not None + return checkpoint_result.copy() + # If there are completed rounds and we're resuming, replay them to # reconstruct the ratings and uncertainties. After each replayed # round we write a checkpoint to ``final_path``. @@ -1806,13 +5007,38 @@ def _write_checkpoint() -> None: if not os.path.exists(round_path): break try: - # Load existing responses for this round - df_round = pd.read_csv(round_path) + df_round = self._read_rank_checkpoint( + round_path, len(attr_batches) + ) df_round["Response"] = df_round["Response"].apply( lambda x: None if pd.isna(x) else x ) - except Exception: - continue + except Exception as exc: + raise ValueError( + f"Could not replay Rank checkpoint {round_path!r}. " + "Use reset_files=True or a new save_dir to recompute it." + ) from exc + replay_meta = { + str(identifier): ( + int(batch_idx), + int(pair_idx), + str(id_a), + str(id_b), + ) + for identifier, batch_idx, pair_idx, id_a, id_b in zip( + df_round["Identifier"], + df_round["Batch"], + df_round["Pair"], + df_round["IdA"], + df_round["IdB"], + ) + } + await self._validate_pairwise_response_payloads( + df_round, + replay_meta, + attr_batches, + context=f"Rank checkpoint {round_path!r}", + ) # Parse each response to build history_pairs async def _coerce_dict_replay(raw: Any) -> Dict[str, Any]: @@ -1842,6 +5068,10 @@ async def _coerce_dict_replay(raw: Any) -> Dict[str, Any]: batch_idx = int(batch_idx_raw) except (TypeError, ValueError): continue + id_a = str(id_a) + id_b = str(id_b) + if id_a not in ratings or id_b not in ratings: + continue if batch_idx < 0 or batch_idx >= len(attr_batches): continue batch = attr_batches[batch_idx] @@ -1854,22 +5084,10 @@ async def _coerce_dict_replay(raw: Any) -> Dict[str, Any]: if attr_key_l not in batch_attr_map: continue real_attr = batch_attr_map[attr_key_l] - val = winner_raw - if isinstance(val, dict) and "winner" in val: - val = val.get("winner") - if isinstance(val, str): - v = val.strip().lower() - else: - v = "" - if v.startswith(("cir", "c", "left", "text a")): - history_pairs[real_attr].append((id_a, id_b)) - elif v.startswith(("squ", "b", "right", "text b")): - history_pairs[real_attr].append((id_b, id_a)) - elif v.startswith("draw") or v.startswith("insufficient"): - history_pairs[real_attr].append((id_a, id_b)) - history_pairs[real_attr].append((id_b, id_a)) - else: - continue + category = self._record_pairwise_outcome( + history_pairs[real_attr], id_a, id_b, winner_raw + ) + outcome_counts[real_attr][category] += 1 else: for ident, resp_raw in zip( df_round["Identifier"], df_round["Response"] @@ -1878,10 +5096,14 @@ async def _coerce_dict_replay(raw: Any) -> Dict[str, Any]: if len(parts) != 5: continue _, batch_idx_str, _, id_a, id_b = parts + id_a = str(id_a) + id_b = str(id_b) try: batch_idx = int(batch_idx_str) except (TypeError, ValueError): continue + if id_a not in ratings or id_b not in ratings: + continue if batch_idx < 0 or batch_idx >= len(attr_batches): continue batch = attr_batches[batch_idx] @@ -1894,22 +5116,10 @@ async def _coerce_dict_replay(raw: Any) -> Dict[str, Any]: if attr_key_l not in batch_attr_map: continue real_attr = batch_attr_map[attr_key_l] - val = winner_raw - if isinstance(val, dict) and "winner" in val: - val = val.get("winner") - if isinstance(val, str): - v = val.strip().lower() - else: - v = "" - if v.startswith(("cir", "c", "left", "text a")): - history_pairs[real_attr].append((id_a, id_b)) - elif v.startswith(("squ", "b", "right", "text b")): - history_pairs[real_attr].append((id_b, id_a)) - elif v.startswith("draw") or v.startswith("insufficient"): - history_pairs[real_attr].append((id_a, id_b)) - history_pairs[real_attr].append((id_b, id_a)) - else: - continue + category = self._record_pairwise_outcome( + history_pairs[real_attr], id_a, id_b, winner_raw + ) + outcome_counts[real_attr][category] += 1 # After parsing all pairs for this round, update ratings se_agg_next: Dict[str, float] = {i: 0.0 for i in item_ids} se_agg_counts: Dict[str, int] = {i: 0 for i in item_ids} @@ -1932,12 +5142,17 @@ async def _coerce_dict_replay(raw: Any) -> Dict[str, Any]: s=s_vec, n_ij=n_ij, p_ij=p_ij, - ridge=self._SE_RIDGE, + rcond=self._SE_EIGEN_TOL, + regularization_strength=self.cfg.learning_rate, ) + component_labels = self._comparison_component_labels(n_ij) for i, se_val in zip(item_ids, se_vec): se_store[attr][i] = float(se_val) - se_agg_next[i] += float(se_val) - se_agg_counts[i] += 1 + if np.isfinite(se_val): + se_agg_next[i] += float(se_val) + se_agg_counts[i] += 1 + for i, component in zip(item_ids, component_labels): + component_store[attr][i] = int(component) for i in item_ids: if se_agg_counts[i] > 0: se_agg_next[i] /= se_agg_counts[i] @@ -1954,14 +5169,8 @@ async def _coerce_dict_replay(raw: Any) -> Dict[str, Any]: _write_checkpoint() # Determine if any new items were added and need to catch up on existing rounds - seen_ids: Set[str] = set() - for pair_list in history_pairs.values(): - for a, b in pair_list: - seen_ids.add(a) - seen_ids.add(b) - new_ids = [i for i in item_ids if i not in seen_ids] await self._catch_up_existing_rounds( - new_ids=new_ids, + candidate_ids=item_ids, round_indices=list(range(start_round)), item_ids=item_ids, texts_by_id=texts_by_id, @@ -1971,8 +5180,10 @@ async def _coerce_dict_replay(raw: Any) -> Dict[str, Any]: attr_batches=attr_batches, attr_keys=attr_keys, history_pairs=history_pairs, + outcome_counts=outcome_counts, ratings=ratings, se_store=se_store, + component_store=component_store, base_name=base_name, df_proc=df_proc, _write_checkpoint=_write_checkpoint, @@ -1992,102 +5203,272 @@ async def _coerce_dict_replay(raw: Any) -> Dict[str, Any]: se_agg_local = self._last_se_agg use_current = rnd > 0 or start_round > 0 or has_seed_ratings se_source = se_agg_local if (rnd > 0 or start_round > 0 or se_agg_local is not None) else None - pairs = self._generate_pairs( - item_ids=item_ids, - texts_by_id=texts_by_id, - current_ratings=current_agg if use_current else None, - se_agg=se_source, + round_path = os.path.join( + self.cfg.save_dir, f"{base_name}_round{rnd}.csv" ) - if not pairs: - break - announce_prompt_rendering( - "Rank", len(attr_batches) * len(pairs) + staging_path = os.path.join( + self.cfg.save_dir, f".{base_name}_round{rnd}.csv" + ) + plan_path = os.path.join( + self.cfg.save_dir, f".{base_name}_round{rnd}_plan.json" ) + batch_state_path = f"{staging_path}.batch_state.json" + plan_records: Optional[List[Dict[str, Any]]] = None + canonical_item_ids = sorted(item_ids) + if os.path.exists(plan_path): + try: + with open(plan_path, encoding="utf-8") as plan_file: + plan_payload = json.load(plan_file) + except Exception as exc: + raise ValueError( + f"Could not read Rank round plan {plan_path!r}" + ) from exc + planned_item_ids = ( + plan_payload.get("item_ids") + if isinstance(plan_payload, dict) + else None + ) + if ( + not isinstance(plan_payload, dict) + or plan_payload.get("version") != 1 + or plan_payload.get("round") != rnd + or plan_payload.get("attribute_batches") != attr_batches + or not isinstance(planned_item_ids, list) + or any( + not isinstance(item_id, str) + for item_id in planned_item_ids + ) + or len(planned_item_ids) != len(set(planned_item_ids)) + or sorted(planned_item_ids) != canonical_item_ids + or not isinstance(plan_payload.get("records"), list) + or not plan_payload["records"] + ): + raise ValueError( + f"Rank round plan {plan_path!r} is incompatible or malformed" + ) + candidate_records = plan_payload["records"] + required_record_keys = { + "identifier", + "batch", + "pair", + "id_a", + "id_b", + "circle_first", + } + pair_endpoints: Dict[int, Tuple[str, str]] = {} + pair_indices_by_batch: Dict[int, Set[int]] = { + batch_idx: set() for batch_idx in range(len(attr_batches)) + } + planned_identifiers: List[str] = [] + for record in candidate_records: + if ( + not isinstance(record, dict) + or not required_record_keys.issubset(record) + or type(record["batch"]) is not int + or type(record["pair"]) is not int + or type(record["circle_first"]) is not bool + ): + raise ValueError( + f"Rank round plan {plan_path!r} has malformed records" + ) + batch_idx = record["batch"] + pair_idx = record["pair"] + id_a = str(record["id_a"]) + id_b = str(record["id_b"]) + if ( + batch_idx not in pair_indices_by_batch + or pair_idx < 0 + or id_a == id_b + or id_a not in texts_by_id + or id_b not in texts_by_id + ): + raise ValueError( + f"Rank round plan {plan_path!r} has invalid endpoints" + ) + endpoints = (id_a, id_b) + if ( + pair_idx in pair_endpoints + and pair_endpoints[pair_idx] != endpoints + ): + raise ValueError( + f"Rank round plan {plan_path!r} changes pair endpoints" + ) + pair_endpoints[pair_idx] = endpoints + if pair_idx in pair_indices_by_batch[batch_idx]: + raise ValueError( + f"Rank round plan {plan_path!r} duplicates a pair" + ) + pair_indices_by_batch[batch_idx].add(pair_idx) + expected_identifier = hash_identifier( + f"{rnd}|{batch_idx}|{pair_idx}|{id_a}|{id_b}", + bits=identifier_hash_bits, + ) + if str(record["identifier"]) != expected_identifier: + raise ValueError( + f"Rank round plan {plan_path!r} has an invalid identifier" + ) + planned_identifiers.append(expected_identifier) + if ( + len(planned_identifiers) != len(set(planned_identifiers)) + or len(set(map(frozenset, pair_indices_by_batch.values()))) + != 1 + ): + raise ValueError( + f"Rank round plan {plan_path!r} is incomplete or colliding" + ) + plan_records = candidate_records + else: + pairs = self._generate_pairs( + item_ids=item_ids, + texts_by_id=texts_by_id, + current_ratings=current_agg if use_current else None, + se_agg=se_source, + ) + if not pairs: + break + plan_records = [] + for batch_idx, _batch in enumerate(attr_batches): + for pair_idx, ((id_a, _), (id_b, _)) in enumerate(pairs): + raw_identifier = ( + f"{rnd}|{batch_idx}|{pair_idx}|{id_a}|{id_b}" + ) + plan_records.append( + { + "identifier": hash_identifier( + raw_identifier, bits=identifier_hash_bits + ), + "batch": batch_idx, + "pair": pair_idx, + "id_a": id_a, + "id_b": id_b, + "circle_first": ( + self.cfg.circle_first + if self.cfg.circle_first is not None + else self.rng.random() < 0.5 + ), + } + ) + self._write_json_atomically( + { + "version": 1, + "round": rnd, + "attribute_batches": attr_batches, + "item_ids": canonical_item_ids, + "records": plan_records, + }, + plan_path, + ) + + announce_prompt_rendering("Rank", len(plan_records)) prompts: List[str] = [] ids: List[str] = [] pair_images: Dict[str, List[str]] = {} pair_audio: Dict[str, List[Dict[str, str]]] = {} pair_pdfs: Dict[str, List[Dict[str, str]]] = {} meta_map: Dict[str, Tuple[int, int, str, str]] = {} - id_to_circle_first: Dict[str, bool] = {} - for batch_idx, batch in enumerate(attr_batches): + for record in plan_records: + batch_idx = int(record["batch"]) + pair_idx = int(record["pair"]) + id_a = str(record["id_a"]) + id_b = str(record["id_b"]) + hashed_ident = str(record["identifier"]) + circle_first_flag = bool(record["circle_first"]) + batch = attr_batches[batch_idx] attr_def_map = ( {a: self.cfg.attributes[a] for a in batch} if isinstance(self.cfg.attributes, dict) else {a: "" for a in batch} ) - for pair_idx, ((id_a, t_a), (id_b, t_b)) in enumerate(pairs): - raw_ident = f"{rnd}|{batch_idx}|{pair_idx}|{id_a}|{id_b}" - hashed_ident = hash_identifier(raw_ident, bits=identifier_hash_bits) - circle_first_flag = ( - self.cfg.circle_first - if self.cfg.circle_first is not None - else self.rng.random() < 0.5 + prompts.append( + self.template.render( + entry_circle=texts_by_id[id_a], + entry_square=texts_by_id[id_b], + attributes=attr_def_map, + additional_instructions=self.cfg.additional_instructions + or "", + modality=self.cfg.modality, + circle_first=circle_first_flag, ) - id_to_circle_first[hashed_ident] = circle_first_flag - prompts.append( - self.template.render( - entry_circle=t_a, - entry_square=t_b, - attributes=attr_def_map, - additional_instructions=self.cfg.additional_instructions - or "", - modality=self.cfg.modality, - circle_first=circle_first_flag, - ) - ) - ids.append(hashed_ident) - meta_map[hashed_ident] = (batch_idx, pair_idx, id_a, id_b) - if images_by_id: - imgs = [] - ia = images_by_id.get(id_a, []) - ib = images_by_id.get(id_b, []) - if circle_first_flag: - if ia: - imgs.extend(ia) - if ib: - imgs.extend(ib) - else: - if ib: - imgs.extend(ib) - if ia: - imgs.extend(ia) - if imgs: - pair_images[hashed_ident] = imgs - if audio_by_id: - auds = [] - aa = audio_by_id.get(id_a, []) - ab = audio_by_id.get(id_b, []) - if circle_first_flag: - if aa: - auds.extend(aa) - if ab: - auds.extend(ab) - else: - if ab: - auds.extend(ab) - if aa: - auds.extend(aa) - if auds: - pair_audio[hashed_ident] = auds - if pdfs_by_id: - pdfs: List[Dict[str, str]] = [] - pa = pdfs_by_id.get(id_a, []) - pb = pdfs_by_id.get(id_b, []) - if circle_first_flag: - if pa: - pdfs.extend(pa) - if pb: - pdfs.extend(pb) - else: - if pb: - pdfs.extend(pb) - if pa: - pdfs.extend(pa) - if pdfs: - pair_pdfs[hashed_ident] = pdfs + ) + ids.append(hashed_ident) + meta_map[hashed_ident] = (batch_idx, pair_idx, id_a, id_b) + if images_by_id: + imgs = [] + ia = images_by_id.get(id_a, []) + ib = images_by_id.get(id_b, []) + if circle_first_flag: + if ia: + imgs.extend(ia) + if ib: + imgs.extend(ib) + else: + if ib: + imgs.extend(ib) + if ia: + imgs.extend(ia) + if imgs: + pair_images[hashed_ident] = imgs + if audio_by_id: + auds = [] + aa = audio_by_id.get(id_a, []) + ab = audio_by_id.get(id_b, []) + if circle_first_flag: + if aa: + auds.extend(aa) + if ab: + auds.extend(ab) + else: + if ab: + auds.extend(ab) + if aa: + auds.extend(aa) + if auds: + pair_audio[hashed_ident] = auds + if pdfs_by_id: + pdfs: List[Dict[str, str]] = [] + pa = pdfs_by_id.get(id_a, []) + pb = pdfs_by_id.get(id_b, []) + if circle_first_flag: + if pa: + pdfs.extend(pa) + if pb: + pdfs.extend(pb) + else: + if pb: + pdfs.extend(pb) + if pa: + pdfs.extend(pa) + if pdfs: + pair_pdfs[hashed_ident] = pdfs # obtain responses from the language model for this round - round_path = os.path.join(self.cfg.save_dir, f"{base_name}_round{rnd}.csv") + if len(ids) != len(set(ids)): + raise ValueError( + "Rank prompt identifier collision; use reset_files=True " + "or a save_dir configured for 64-bit identifiers" + ) + response_kwargs = dict(kwargs) + # A partial tail cannot be committed or safely treated as a + # completed tournament round, even for very large collections. + response_kwargs["skip_tail_fails"] = False + if self.cfg.use_dummy: + response_kwargs.setdefault( + "dummy_responses", + { + identifier: { + "responses": [ + json.dumps( + { + attribute: "draw" + for attribute in attr_batches[ + meta_map[identifier][0] + ] + } + ) + ] + } + for identifier in ids + }, + ) resp_df = await get_all_responses( prompts=prompts, identifiers=ids, @@ -2097,12 +5478,25 @@ async def _coerce_dict_replay(raw: Any) -> Dict[str, Any]: n_parallels=self.cfg.n_parallels, model=self.cfg.model, json_mode=self.cfg.modality != "audio", - save_path=round_path, + save_path=staging_path, reset_files=reset_files, use_dummy=self.cfg.use_dummy, max_retries=1, reasoning_effort=self.cfg.reasoning_effort, - **kwargs, + **response_kwargs, + ) + self._validate_requested_responses( + resp_df, + ids, + context="Rank", + allow_extra=False, + ) + await self._validate_pairwise_response_payloads( + resp_df, + meta_map, + attr_batches, + context="Rank", + retry_path=staging_path, ) # attach metadata columns and overwrite the round CSV resp_df["Batch"] = resp_df.Identifier.map( @@ -2117,7 +5511,7 @@ async def _coerce_dict_replay(raw: Any) -> Dict[str, Any]: resp_df["IdB"] = resp_df.Identifier.map( lambda x: meta_map.get(str(x), (np.nan, np.nan, "", ""))[3] ) - resp_df.to_csv(round_path, index=False) + self._write_rank_checkpoint(resp_df, round_path, len(attr_batches)) # parse each response # reuse the _coerce_dict function defined in the original implementation @@ -2150,22 +5544,10 @@ async def _coerce_dict(raw: Any) -> Dict[str, Any]: if attr_key_l not in batch_attr_map: continue real_attr = batch_attr_map[attr_key_l] - val = winner_raw - if isinstance(val, dict) and "winner" in val: - val = val.get("winner") - if isinstance(val, str): - v = val.strip().lower() - else: - v = "" - if v.startswith(("cir", "c", "left", "text a")): - history_pairs[real_attr].append((id_a, id_b)) - elif v.startswith(("squ", "b", "right", "text b")): - history_pairs[real_attr].append((id_b, id_a)) - elif v.startswith("draw") or v.startswith("insufficient"): - history_pairs[real_attr].append((id_a, id_b)) - history_pairs[real_attr].append((id_b, id_a)) - else: - continue + category = self._record_pairwise_outcome( + history_pairs[real_attr], id_a, id_b, winner_raw + ) + outcome_counts[real_attr][category] += 1 # update ratings using the BT model for this round se_agg_next: Dict[str, float] = {i: 0.0 for i in item_ids} se_agg_counts: Dict[str, int] = {i: 0 for i in item_ids} @@ -2188,12 +5570,17 @@ async def _coerce_dict(raw: Any) -> Dict[str, Any]: s=s_vec, n_ij=n_ij, p_ij=p_ij, - ridge=self._SE_RIDGE, + rcond=self._SE_EIGEN_TOL, + regularization_strength=self.cfg.learning_rate, ) + component_labels = self._comparison_component_labels(n_ij) for i, se_val in zip(item_ids, se_vec): se_store[attr][i] = float(se_val) - se_agg_next[i] += float(se_val) - se_agg_counts[i] += 1 + if np.isfinite(se_val): + se_agg_next[i] += float(se_val) + se_agg_counts[i] += 1 + for i, component in zip(item_ids, component_labels): + component_store[attr][i] = int(component) for i in item_ids: if se_agg_counts[i] > 0: se_agg_next[i] /= se_agg_counts[i] @@ -2208,6 +5595,24 @@ async def _coerce_dict(raw: Any) -> Dict[str, Any]: ratings[i][attr] -= mean_val # Write checkpoint after this new round _write_checkpoint() + update_run_metadata( + self.cfg.save_dir, + base_name, + strict=True, + last_completed_round=rnd, + ) + for completed_artifact in ( + batch_state_path, + staging_path, + plan_path, + ): + try: + Path(completed_artifact).unlink(missing_ok=True) + except OSError: + pass # After processing all rounds, return the final DataFrame # The checkpoint has already been written in the final iteration - return pd.read_csv(final_path) + if checkpoint_result is None: + _write_checkpoint() + assert checkpoint_result is not None + return checkpoint_result.copy() diff --git a/src/gabriel/utils/openai_utils.py b/src/gabriel/utils/openai_utils.py index 9e3351f..e4b1070 100644 --- a/src/gabriel/utils/openai_utils.py +++ b/src/gabriel/utils/openai_utils.py @@ -4295,8 +4295,88 @@ async def get_all_responses( cols.insert(7, "Reasoning Summary") df = pd.DataFrame(columns=cols) done = set() - written_identifiers: Set[str] = set(df["Identifier"].astype(str)) if not df.empty else set() + # Successful checkpoint rows are immutable within this invocation. Failed + # rows must remain writable so a later successful retry can replace them. + written_identifiers: Set[str] = set(done) requested_identifiers = [str(i) for i in identifiers] + + def _persist_response_rows( + batch_df: pd.DataFrame, *, atomic: bool = False + ) -> None: + """Persist new rows, replacing any previously failed identifiers.""" + + nonlocal df, csv_header_written, written_identifiers + if batch_df.empty: + return + if "Web Search Sources" not in batch_df.columns: + batch_df["Web Search Sources"] = pd.NA + batch_df = batch_df.drop_duplicates(subset=["Identifier"], keep="last") + batch_df = batch_df[ + ~batch_df["Identifier"].astype(str).isin(written_identifiers) + ] + if batch_df.empty: + return + + batch_identifiers = set(batch_df["Identifier"].astype(str)) + existing_identifiers = ( + set(df["Identifier"].astype(str)) if not df.empty else set() + ) + replaces_failed_rows = bool(batch_identifiers & existing_identifiers) + if atomic or replaces_failed_rows: + merged = pd.concat([df, batch_df], ignore_index=True) + merged = merged.drop_duplicates( + subset=["Identifier"], keep="last" + ).reset_index(drop=True) + to_save = merged.copy() + for col in ("Response", "Error Log", "Web Search Sources"): + if col in to_save: + to_save[col] = to_save[col].apply(_ser) + file_descriptor, temporary_path = tempfile.mkstemp( + prefix=f".{Path(save_path).name}.", + suffix=".tmp", + dir=str(save_dir), + text=True, + ) + try: + with os.fdopen( + file_descriptor, + "w", + encoding="utf-8", + newline="", + ) as temporary_file: + to_save.to_csv( + temporary_file, + index=False, + quoting=csv.QUOTE_MINIMAL, + ) + temporary_file.flush() + os.fsync(temporary_file.fileno()) + os.replace(temporary_path, save_path) + except Exception: + with contextlib.suppress(OSError): + os.unlink(temporary_path) + raise + df = merged + csv_header_written = True + else: + to_save = batch_df.copy() + for col in ("Response", "Error Log", "Web Search Sources"): + if col in to_save: + to_save[col] = to_save[col].apply(_ser) + to_save.to_csv( + save_path, + mode="a" if csv_header_written else "w", + header=not csv_header_written, + index=False, + quoting=csv.QUOTE_MINIMAL, + ) + csv_header_written = True + if df.empty: + df = batch_df.reset_index(drop=True) + else: + df = pd.concat([df, batch_df], ignore_index=True) + written_identifiers.update(batch_identifiers) + # Helper to calculate and report final run cost def _report_cost() -> None: nonlocal df @@ -4607,34 +4687,35 @@ def _select_ceiling(limit_val: Optional[float], remaining_val: Optional[float]) if use_batch: state_path = save_path + ".batch_state.json" + def _write_batch_state(payload: Dict[str, Any]) -> None: + """Atomically persist every externally visible Batch transition.""" + + file_descriptor, temporary_path = tempfile.mkstemp( + prefix=f".{Path(state_path).name}.", + suffix=".tmp", + dir=str(save_dir), + text=True, + ) + try: + with os.fdopen( + file_descriptor, "w", encoding="utf-8" + ) as temporary_file: + json.dump(payload, temporary_file) + temporary_file.flush() + os.fsync(temporary_file.fileno()) + os.replace(temporary_path, state_path) + except Exception: + with contextlib.suppress(OSError): + os.unlink(temporary_path) + raise + # Helper to append batch rows def _append_results(rows: List[Dict[str, Any]]) -> None: - nonlocal df, csv_header_written, written_identifiers if not rows: return - batch_df = pd.DataFrame(rows) - if "Web Search Sources" not in batch_df.columns: - batch_df["Web Search Sources"] = pd.NA - batch_df = batch_df[~batch_df["Identifier"].astype(str).isin(written_identifiers)] - if batch_df.empty: - return - to_save = batch_df.copy() - for col in ("Response", "Error Log", "Web Search Sources"): - if col in to_save: - to_save[col] = to_save[col].apply(_ser) - to_save.to_csv( - save_path, - mode="a" if csv_header_written else "w", - header=not csv_header_written, - index=False, - quoting=csv.QUOTE_MINIMAL, - ) - csv_header_written = True - if df.empty: - df = batch_df.reset_index(drop=True) - else: - df = pd.concat([df, batch_df], ignore_index=True) - written_identifiers.update(batch_df["Identifier"].astype(str)) + # A batch ID is removed from state only after this atomic checkpoint + # succeeds, so a crash cannot lose already-paid response rows. + _persist_response_rows(pd.DataFrame(rows), atomic=True) client = _get_client(base_url) # Load existing state @@ -4655,6 +4736,24 @@ def _append_results(rows: List[Dict[str, Any]]) -> None: } ] } + if not isinstance(state.get("batches", []), list): + raise ValueError(f"Malformed Batch API state in {state_path!r}") + unresolved_submissions = [ + batch_state + for batch_state in state.get("batches", []) + if not isinstance(batch_state, dict) + or ( + batch_state.get("status") == "submitting" + or not batch_state.get("batch_id") + ) + ] + if unresolved_submissions: + raise RuntimeError( + "Found an unresolved Batch API submission intent. The server " + "may already have accepted a paid job, so automatic resubmission " + "is disabled. Reconcile or cancel the external batch, then " + "remove/reset the local state only when safe." + ) # Cancel unfinished batches if requested if cancel_existing_batch and state.get("batches"): logger.info("Cancelling unfinished batch jobs...") @@ -4798,29 +4897,33 @@ def _append_results(rows: List[Dict[str, Any]]) -> None: uploaded = await client.files.create( file=open(input_filename, "rb"), purpose="batch" ) + pending_batch: Dict[str, Any] = { + "batch_id": None, + "status": "submitting", + "input_file_id": uploaded.id, + "total": len(batch_tasks), + "submitted_at": int(time.time()), + } + state["batches"].append(pending_batch) + # Record intent before the irreversible external call. A + # connection loss after server acceptance is ambiguous and + # must fail closed rather than submit the same work again. + _write_batch_state(state) batch = await client.batches.create( input_file_id=uploaded.id, endpoint="/v1/responses", completion_window=batch_completion_window, ) - state["batches"].append( - { - "batch_id": batch.id, - "input_file_id": uploaded.id, - "total": len(batch_tasks), - "submitted_at": int(time.time()), - } - ) + pending_batch["batch_id"] = batch.id + pending_batch["status"] = "submitted" + _write_batch_state(state) logger.info( f"Submitted batch {batch.id} with {len(batch_tasks)} requests." ) - with open(state_path, "w") as f: - json.dump(state, f) # Return immediately if not waiting for completion if not batch_wait_for_completion: return df unfinished_batches: List[Dict[str, Any]] = list(state.get("batches", [])) - completed_rows: List[Dict[str, Any]] = [] while unfinished_batches: for b in list(unfinished_batches): bid = b.get("batch_id") @@ -4831,6 +4934,7 @@ def _append_results(rows: List[Dict[str, Any]]) -> None: continue status = job.status if status == "completed": + batch_completed_rows: List[Dict[str, Any]] = [] output_file_id = job.output_file_id error_file_id = job.error_file_id logger.info(f"Batch {bid} completed. Downloading results...") @@ -4928,7 +5032,7 @@ def _append_results(rows: List[Dict[str, Any]]) -> None: } if reasoning_summary is not None: row["Reasoning Summary"] = None - completed_rows.append(row) + batch_completed_rows.append(row) continue resp_obj = rec["response"] resp_text: Optional[str] = None @@ -5041,15 +5145,15 @@ def _append_results(rows: List[Dict[str, Any]]) -> None: } if reasoning_summary is not None: row["Reasoning Summary"] = summary_text - completed_rows.append(row) + batch_completed_rows.append(row) + _append_results(batch_completed_rows) unfinished_batches.remove(b) state["batches"] = [ bb for bb in state.get("batches", []) if bb.get("batch_id") != bid ] - with open(state_path, "w") as f: - json.dump(state, f) + _write_batch_state(state) elif status in {"failed", "cancelled", "expired"}: logger.warning(f"Batch {bid} finished with status {status}.") unfinished_batches.remove(b) @@ -5058,8 +5162,7 @@ def _append_results(rows: List[Dict[str, Any]]) -> None: for bb in state.get("batches", []) if bb.get("batch_id") != bid ] - with open(state_path, "w") as f: - json.dump(state, f) + _write_batch_state(state) else: rc = job.request_counts logger.info( @@ -5067,8 +5170,8 @@ def _append_results(rows: List[Dict[str, Any]]) -> None: ) if unfinished_batches: await asyncio.sleep(batch_poll_interval) - # Append and return - _append_results(completed_rows) + # All completed rows were durably checkpointed before their batch IDs + # were removed from state. _report_cost() return df # Non‑batch path @@ -5716,30 +5819,9 @@ def _maybe_trigger_threshold_refresh() -> None: _maybe_refresh_estimates("parallel-cap reached", force=True) async def flush() -> None: - nonlocal results, df, processed, csv_header_written, written_identifiers + nonlocal results, processed if results: - batch_df = pd.DataFrame(results) - if "Web Search Sources" not in batch_df.columns: - batch_df["Web Search Sources"] = pd.NA - batch_df = batch_df[~batch_df["Identifier"].astype(str).isin(written_identifiers)] - if not batch_df.empty: - to_save = batch_df.copy() - for col in ("Response", "Error Log", "Web Search Sources"): - if col in to_save: - to_save[col] = to_save[col].apply(_ser) - to_save.to_csv( - save_path, - mode="a" if csv_header_written else "w", - header=not csv_header_written, - index=False, - quoting=csv.QUOTE_MINIMAL, - ) - csv_header_written = True - if df.empty: - df = batch_df.reset_index(drop=True) - else: - df = pd.concat([df, batch_df], ignore_index=True) - written_identifiers.update(batch_df["Identifier"].astype(str)) + _persist_response_rows(pd.DataFrame(results)) results = [] if logger.isEnabledFor(logging.INFO) and processed: logger.info( diff --git a/tests/test_basic.py b/tests/test_basic.py index 3667d69..cda1efe 100644 --- a/tests/test_basic.py +++ b/tests/test_basic.py @@ -1420,6 +1420,7 @@ def test_rank_outputs_zscores_and_raw_columns(tmp_path): assert attr in df.columns assert f"{attr}_raw" in df.columns assert f"{attr}_se" in df.columns + assert f"{attr}_component" in df.columns assert np.isfinite(df[f"{attr}_se"].fillna(0.0)).all() @@ -1459,7 +1460,20 @@ def test_rank_drops_malformed_rows_in_text_mode(tmp_path, capsys): assert set(df["text"]) == {"good", "great"} -def test_recursive_rank_drops_malformed_rows_in_text_mode(tmp_path, capsys): +def test_recursive_rank_drops_malformed_rows_in_text_mode( + monkeypatch, tmp_path, capsys +): + async def fake_get_all_responses(*, identifiers, **kwargs): + return pd.DataFrame( + { + "Identifier": identifiers, + "Response": ['{"clarity": "draw"}'] * len(identifiers), + } + ) + + monkeypatch.setattr( + "gabriel.tasks.rank.get_all_responses", fake_get_all_responses + ) cfg = RankConfig( attributes={"clarity": ""}, save_dir=str(tmp_path), @@ -1491,7 +1505,18 @@ def test_rank_primer_centering(): assert ratings["b"]["clarity"] == 10.0 -def test_recursive_rank_outputs(tmp_path): +def test_recursive_rank_outputs(monkeypatch, tmp_path): + async def fake_get_all_responses(*, identifiers, **kwargs): + return pd.DataFrame( + { + "Identifier": identifiers, + "Response": ['{"clarity": "draw"}'] * len(identifiers), + } + ) + + monkeypatch.setattr( + "gabriel.tasks.rank.get_all_responses", fake_get_all_responses + ) cfg = RankConfig( attributes={"clarity": ""}, save_dir=str(tmp_path), @@ -1536,11 +1561,13 @@ def test_api_rank_hides_raw_columns(tmp_path): assert attr in df.columns assert f"{attr}_raw" not in df.columns assert f"{attr}_se" not in df.columns + assert f"{attr}_component" in df.columns final_path = tmp_path / "rankings_final.csv" saved = pd.read_csv(final_path) for attr in ("clarity", "originality"): assert f"{attr}_raw" in saved.columns assert f"{attr}_se" in saved.columns + assert f"{attr}_component" in saved.columns def test_api_rank_can_return_raw_scores_when_requested(tmp_path): @@ -1565,6 +1592,153 @@ def test_api_rank_can_return_raw_scores_when_requested(tmp_path): assert f"{attr}_se" in df.columns +@pytest.mark.parametrize("recursive", [False, True]) +def test_api_rank_replaces_same_named_input_column_without_duplicates( + tmp_path, recursive +): + data = pd.DataFrame( + {"text": ["first", "second"], "quality": ["old-a", "old-b"]} + ) + result = asyncio.run( + gabriel.rank( + data, + "text", + attributes={"quality": ""}, + save_dir=str(tmp_path), + file_name="rankings.csv", + use_dummy=True, + n_rounds=1, + matches_per_round=1, + n_parallels=4, + initial_rating_pass=False, + recursive=recursive, + recursive_rate_first_round=False, + ) + ) + + assert list(result.columns).count("quality") == 1 + assert list(result.columns).count("text") == 1 + + +@pytest.mark.parametrize("persisted_attributes", [{"quality": ""}, ["quality"]]) +def test_api_rank_cached_load_respects_raw_score_projection( + tmp_path, persisted_attributes +): + pd.DataFrame( + { + "quality": [0.0], + "quality_raw": [1.2], + "quality_se": [0.3], + "quality_component": [0], + "temperature": [20.0], + "temperature_raw": [19.5], + "temperature_se": [0.1], + "temperature_component": [7], + } + ).to_csv(tmp_path / "rankings_final.csv", index=False) + (tmp_path / "attributes.json").write_text(json.dumps({"unrelated": ""})) + (tmp_path / "rankings_attrs.json").write_text(json.dumps(persisted_attributes)) + + public = asyncio.run( + gabriel.rank( + None, + "text", + attributes={"quality": ""}, + save_dir=str(tmp_path), + return_raw_scores=False, + ) + ) + detailed = asyncio.run( + gabriel.rank( + None, + "text", + attributes={"quality": ""}, + save_dir=str(tmp_path), + return_raw_scores=True, + ) + ) + mismatched = asyncio.run( + gabriel.rank( + None, + "text", + attributes={"different": ""}, + save_dir=str(tmp_path), + return_raw_scores=False, + ) + ) + + assert "quality_raw" not in public + assert "quality_se" not in public + assert "quality_component" in public + assert "quality_raw" not in mismatched + assert "quality_se" not in mismatched + assert {"temperature_raw", "temperature_se"}.issubset(public.columns) + assert {"quality_raw", "quality_se"}.issubset(detailed.columns) + + +def test_api_recursive_cache_preserves_user_suffix_columns(tmp_path): + recursive_dir = tmp_path / "rankings_recursive" + recursive_dir.mkdir() + pd.DataFrame( + { + "text": ["sample"], + "quality": [0.0], + "quality_raw": ["user value"], + "quality_se": ["user value"], + "overall_rank": [1], + "exit_stage": [1], + } + ).to_csv(recursive_dir / "recursive_final.csv", index=False) + + result = asyncio.run( + gabriel.rank( + None, + "text", + attributes={"quality": ""}, + save_dir=str(tmp_path), + recursive=True, + return_raw_scores=False, + ) + ) + assert {"quality_raw", "quality_se"}.issubset(result.columns) + + +def test_empty_recursive_rank_matches_nonempty_public_schema(tmp_path): + result = asyncio.run( + gabriel.rank( + pd.DataFrame({"id": [], "text": []}), + "text", + id_column="id", + attributes={"quality": ""}, + save_dir=str(tmp_path), + recursive=True, + recursive_rate_first_round=False, + use_dummy=True, + reset_files=True, + ) + ) + assert "quality" in result.columns + assert "quality_raw" not in result.columns + assert "quality_se" not in result.columns + + +@pytest.mark.parametrize("override", [{"max_retries": 2}, {"json_mode": False}]) +def test_api_rank_rejects_response_settings_owned_by_transaction(override, tmp_path): + with pytest.raises(TypeError, match=next(iter(override))): + asyncio.run( + gabriel.rank( + pd.DataFrame({"text": ["first", "second"]}), + "text", + attributes={"quality": ""}, + save_dir=str(tmp_path), + use_dummy=True, + n_rounds=1, + matches_per_round=1, + **override, + ) + ) + + def test_deidentifier_dummy(tmp_path): cfg = DeidentifyConfig(save_dir=str(tmp_path), file_name="deid.csv", use_dummy=True) task = Deidentifier(cfg) @@ -2393,7 +2567,7 @@ async def fake_get_all_responses(*, prompts, identifiers, **kwargs): assert any(call.get("prompt_pdfs") for call in calls) -def test_rank_resume_ignores_nan_batch_rows(tmp_path): +def test_rank_resume_rejects_nan_batch_rows(tmp_path): cfg = RankConfig( attributes={"clarity": ""}, save_dir=str(tmp_path), @@ -2409,17 +2583,31 @@ def test_rank_resume_ignores_nan_batch_rows(tmp_path): "Identifier": ["x"], "Response": ["{}"], "Batch": [np.nan], + "Pair": [0], "IdA": ["a"], "IdB": ["b"], } ).to_csv(round0, index=False) - task = Rank(cfg) - data = pd.DataFrame({"text": ["first", "second"]}) - df = asyncio.run(task.run(data, column_name="text", reset_files=False)) + (tmp_path / "rankings_run_metadata.json").write_text( + json.dumps( + { + "rank_estimator_version": 2, + "insufficient_signal_policy": "tie", + "learning_rate": 0.1, + "last_completed_round": 0, + "modality": "text", + "input_fingerprints": {"a": "0" * 40, "b": "1" * 40}, + "measurement_spec_fingerprint": task._measurement_spec_fingerprint(), + } + ) + ) - assert "clarity" in df.columns - assert len(df) == 2 + data = pd.DataFrame({"text": ["first", "second"]}) + before = round0.read_bytes() + with pytest.raises(ValueError, match="Could not replay Rank checkpoint"): + asyncio.run(task.run(data, column_name="text", reset_files=False)) + assert round0.read_bytes() == before def test_seed_api_passes_embedding_overrides(monkeypatch, tmp_path): diff --git a/tests/test_rank_statistics.py b/tests/test_rank_statistics.py new file mode 100644 index 0000000..1d030a0 --- /dev/null +++ b/tests/test_rank_statistics.py @@ -0,0 +1,3121 @@ +import asyncio +import importlib +import json +from pathlib import Path + +import numpy as np +import pandas as pd +import pytest +from scipy.optimize import minimize + +from gabriel.tasks.rank import Rank, RankConfig + + +def _rank(tmp_path, *, insufficient_signal_policy="tie"): + return Rank( + RankConfig( + attributes={"quality": ""}, + save_dir=str(tmp_path), + initial_rating_pass=False, + insufficient_signal_policy=insufficient_signal_policy, + ) + ) + + +def _fit(rank, item_ids, outcomes, *, pseudo=0.1): + return rank._fit_bt( + item_ids=list(item_ids), + outcomes=outcomes, + pseudo=pseudo, + max_iter=20_000, + tol=1e-12, + return_info=True, + ) + + +def _standard_errors( + rank, item_ids, scores, n_ij, p_ij, *, regularization_strength=0.1 +): + return rank._bt_standard_errors( + s=np.array([scores[item] for item in item_ids]), + n_ij=n_ij, + p_ij=p_ij, + rcond=rank._SE_EIGEN_TOL, + regularization_strength=regularization_strength, + ) + + +def test_outcome_decoder_uses_unit_weight_draws_and_bounded_labels(tmp_path): + rank = _rank(tmp_path) + + assert rank._decode_pairwise_outcome("a", "b", "draw") == ( + "draw", + [("a", "b", 0.5), ("b", "a", 0.5)], + ) + assert rank._decode_pairwise_outcome("a", "b", {"Winner": "circle"}) == ( + "circle", + [("a", "b", 1.0)], + ) + assert rank._decode_pairwise_outcome("a", "b", "square wins") == ( + "square", + [("b", "a", 1.0)], + ) + assert rank._decode_pairwise_outcome("a", "b", "cannot determine") == ( + "invalid", + [], + ) + assert rank._decode_pairwise_outcome("a", "b", "both lack evidence") == ( + "invalid", + [], + ) + for ambiguous in ( + "a tie", + "a draw", + "A: insufficient signal", + "b: cannot determine", + "c: cannot determine", + "circle is worse", + "square loses", + "draw but square is stronger", + "left", + "right", + ): + assert rank._decode_pairwise_outcome("a", "b", ambiguous) == ( + "invalid", + [], + ) + + +def test_insufficient_signal_policy_is_explicit(tmp_path): + tie_rank = _rank(tmp_path / "tie", insufficient_signal_policy="tie") + abstain_rank = _rank(tmp_path / "abstain", insufficient_signal_policy="abstain") + + assert tie_rank._decode_pairwise_outcome( + "a", "b", "insufficient signal" + ) == ( + "insufficient_signal", + [("a", "b", 0.5), ("b", "a", 0.5)], + ) + assert abstain_rank._decode_pairwise_outcome( + "a", "b", "insufficient signal" + ) == ("insufficient_signal", []) + + with pytest.raises(ValueError, match="insufficient_signal_policy"): + RankConfig( + attributes={"quality": ""}, + insufficient_signal_policy="guess", + ) + for invalid_rate in (True, "0.1", np.inf, -0.1): + with pytest.raises(ValueError, match="learning_rate"): + RankConfig( + attributes={"quality": ""}, learning_rate=invalid_rate + ) + normalized = RankConfig( + attributes={"quality": ""}, learning_rate=np.float64(0.25) + ) + assert normalized.learning_rate == 0.25 + for invalid_scale in (True, "1", np.nan, np.inf, -np.inf): + with pytest.raises(ValueError, match="primer_scale"): + RankConfig(attributes={"quality": ""}, primer_scale=invalid_scale) + with pytest.raises(ValueError, match="primer_center"): + RankConfig(attributes={"quality": ""}, primer_center=1) + with pytest.raises(ValueError, match="judge_version"): + RankConfig(attributes={"quality": ""}, judge_version=" ") + + zero_primer = RankConfig( + attributes={"quality": ""}, primer_scale=0, primer_center=True + ) + ratings = {"a": {"quality": 0.0}, "b": {"quality": 0.0}} + Rank(zero_primer)._apply_primer( + ratings, + {"a": {"quality": 10.0}, "b": {"quality": 0.0}}, + ["quality"], + ) + assert ratings == {"a": {"quality": 0.0}, "b": {"quality": 0.0}} + for suffix in ("raw", "se", "component"): + with pytest.raises(ValueError, match="namespaces overlap"): + RankConfig(attributes={"foo": "", f"foo_{suffix}": ""}) + for reserved in ("identifier", "overall_rank", "exit_stage", "stage1_quality"): + with pytest.raises(ValueError, match="Recursive Rank attribute names"): + RankConfig(attributes={reserved: ""}, recursive=True) + for reserved_key in ("attributes", "save_dir", "file_name"): + with pytest.raises(ValueError, match="Rank-owned"): + RankConfig( + attributes={"quality": ""}, + rate_kwargs={reserved_key: "not-allowed"}, + ) + mutable_config = RankConfig(attributes={"quality": ""}) + mutable_config.rate_kwargs["save_dir"] = str(tmp_path / "outside") + with pytest.raises(ValueError, match="Rank-owned"): + Rank(mutable_config)._split_rate_kwargs() + assert type(normalized.learning_rate) is float + for field_name in ("n_rounds", "matches_per_round"): + with pytest.raises(ValueError, match=field_name): + RankConfig(attributes={"quality": ""}, **{field_name: 0}) + with pytest.raises(ValueError, match=field_name): + RankConfig(attributes={"quality": ""}, **{field_name: True}) + for empty_attributes in ({}, []): + with pytest.raises(ValueError, match="at least one attribute"): + RankConfig(attributes=empty_attributes) + for invalid_attributes in ([""], [1], ["Quality", " quality "]): + with pytest.raises(ValueError, match="attribute names"): + RankConfig(attributes=invalid_attributes) + for invalid_fraction in (True, 0, 1, -0.1, np.inf): + with pytest.raises(ValueError, match="recursive_fraction"): + RankConfig( + attributes={"quality": ""}, + recursive_fraction=invalid_fraction, + ) + for field_name in ( + "recursive_min_remaining", + "recursive_final_round_multiplier", + ): + for invalid_value in (True, 0, 1.5): + with pytest.raises(ValueError, match=field_name): + RankConfig( + attributes={"quality": ""}, + **{field_name: invalid_value}, + ) + + +def test_persisted_rank_attributes_cannot_bypass_namespace_validation(tmp_path): + tmp_path.mkdir(exist_ok=True) + (tmp_path / "attributes.json").write_text( + json.dumps({"foo": "", "foo_raw": ""}) + ) + rank = Rank( + RankConfig( + attributes={"quality": ""}, + save_dir=str(tmp_path), + initial_rating_pass=False, + ) + ) + + with pytest.raises(ValueError, match="namespaces overlap"): + asyncio.run(rank.run(pd.DataFrame({"text": ["A", "B"]}), "text")) + + +@pytest.mark.parametrize("reserved_column", ["overall_rank", "exit_stage", "stage2_x"]) +def test_recursive_rank_rejects_reserved_input_columns_before_judging( + tmp_path, reserved_column +): + rank = Rank( + RankConfig( + attributes={"quality": ""}, + save_dir=str(tmp_path), + recursive=True, + initial_rating_pass=False, + ) + ) + data = pd.DataFrame({"text": ["A", "B"], reserved_column: [0, 1]}) + + with pytest.raises(ValueError, match="input columns cannot use internal"): + asyncio.run(rank.run(data, "text")) + + +def test_power_matching_flag_selects_random_scheduler_and_caps_degree( + monkeypatch, tmp_path +): + rank = Rank( + RankConfig( + attributes={"quality": ""}, + save_dir=str(tmp_path), + initial_rating_pass=False, + power_matching=False, + matches_per_round=99, + ) + ) + observed = {} + + def fake_random(item_ids, texts_by_id, mpr): + observed["args"] = (item_ids, texts_by_id, mpr) + return [(('a', 'A'), ('b', 'B'))] + + def fail_info_gain(*args, **kwargs): + raise AssertionError("power_matching=False must not use the heuristic") + + monkeypatch.setattr(rank, "_pairs_random", fake_random) + monkeypatch.setattr(rank, "_pairs_info_gain", fail_info_gain) + result = rank._generate_pairs( + ["a", "b", "c"], + {"a": "A", "b": "B", "c": "C"}, + current_ratings={"a": 0.0, "b": 0.0, "c": 0.0}, + se_agg={"a": 1.0, "b": 1.0, "c": 1.0}, + ) + + assert result == [(('a', 'A'), ('b', 'B'))] + assert observed["args"][2] == 2 + + +def test_initial_rate_seed_uses_media_identifier_semantics(tmp_path): + rank_module = importlib.import_module("gabriel.tasks.rank") + payloads = [["a.png", "b.png"], ["c.png"]] + item_ids = [ + rank_module._hash_text_identifier(payload, strict=False, bits=64) + for payload in payloads + ] + rank = Rank( + RankConfig( + attributes={"quality": ""}, + modality="image", + save_dir=str(tmp_path), + ) + ) + + seeds = rank._seed_ratings_from_rate( + pd.DataFrame({"image": payloads, "quality": [20.0, 80.0]}), + id_column=None, + text_column="image", + item_ids=item_ids, + attr_keys=["quality"], + identifier_hash_bits=64, + ) + + assert set(seeds) == set(item_ids) + assert seeds[item_ids[0]]["quality"] == pytest.approx(-30.0) + assert seeds[item_ids[1]]["quality"] == pytest.approx(30.0) + + +def test_regularization_is_derived_from_a_coherent_win_matrix(tmp_path): + rank = _rank(tmp_path) + observed, fitted = rank._build_bt_win_matrices( + ["a", "b", "c"], + [("a", "b", 1.0), ("b", "a", 0.5)], + pseudo=0.1, + ) + + np.testing.assert_allclose( + observed, + [[0.0, 1.0, 0.0], [0.5, 0.0, 0.0], [0.0, 0.0, 0.0]], + ) + np.testing.assert_allclose( + fitted, + [[0.0, 1.05, 0.0], [0.55, 0.0, 0.0], [0.0, 0.0, 0.0]], + ) + np.testing.assert_allclose(np.diag(fitted), 0.0) + assert fitted[0, 2] == fitted[2, 0] == 0.0 + np.testing.assert_allclose(fitted + fitted.T, (fitted + fitted.T).T) + np.testing.assert_allclose(fitted.sum(axis=1), [1.05, 0.55, 0.0]) + assert fitted.sum() == pytest.approx( + np.triu(fitted + fitted.T, k=1).sum() + ) + + +def test_two_item_fit_matches_regularized_closed_form(tmp_path): + rank = _rank(tmp_path) + outcomes = [("a", "b")] * 9 + [("b", "a")] + + scores, n_ij, p_ij = _fit(rank, ["a", "b"], outcomes) + + expected_gap = np.log((9.0 + 0.05) / (1.0 + 0.05)) + assert scores["a"] - scores["b"] == pytest.approx(expected_gap, abs=1e-10) + np.testing.assert_allclose(n_ij, [[0.0, 10.0], [10.0, 0.0]]) + assert p_ij[0, 1] / p_ij[1, 0] == pytest.approx(np.exp(expected_gap)) + + +@pytest.mark.parametrize("n_items", [2, 5, 10, 50, 100]) +def test_balanced_extensions_do_not_distort_existing_contrasts(tmp_path, n_items): + rank = _rank(tmp_path) + item_ids = ["a", "b"] + [f"u{i}" for i in range(n_items - 2)] + outcomes = [("a", "b")] * 9 + [("b", "a")] + for item in item_ids[2:]: + outcomes.extend([(item, "b")] * 5) + outcomes.extend([("b", item)] * 5) + + scores, _, _ = _fit(rank, item_ids, outcomes) + + expected_gap = np.log((9.0 + 0.05) / (1.0 + 0.05)) + assert scores["a"] - scores["b"] == pytest.approx(expected_gap, abs=1e-8) + for item in item_ids[2:]: + assert scores[item] - scores["b"] == pytest.approx(0.0, abs=1e-8) + + +def test_fit_is_permutation_equivariant(tmp_path): + rank = _rank(tmp_path) + outcomes = [ + ("a", "b", 3.0), + ("b", "a", 1.0), + ("b", "c", 2.0), + ("c", "b", 1.0), + ("c", "a", 1.5), + ("a", "c", 0.5), + ] + + reference, _, _ = _fit(rank, ["a", "b", "c"], outcomes) + permuted, _, _ = _fit(rank, ["c", "a", "b"], outcomes) + + for item in reference: + assert permuted[item] == pytest.approx(reference[item], abs=1e-10) + + +def test_reversing_every_outcome_negates_scores(tmp_path): + rank = _rank(tmp_path) + outcomes = [ + ("a", "b", 4.0), + ("b", "a", 1.0), + ("b", "c", 3.0), + ("c", "b", 2.0), + ("c", "a", 1.0), + ] + reversed_outcomes = [(loser, winner, weight) for winner, loser, weight in outcomes] + + scores, _, _ = _fit(rank, ["a", "b", "c"], outcomes) + reversed_scores, _, _ = _fit( + rank, ["a", "b", "c"], reversed_outcomes + ) + + for item in scores: + assert reversed_scores[item] == pytest.approx(-scores[item], abs=1e-9) + + +def test_bt_fit_matches_constrained_likelihood_optimization(tmp_path): + rank = _rank(tmp_path) + item_ids = ["a", "b", "c", "d"] + outcomes = [ + ("a", "b", 4.0), + ("b", "a", 1.0), + ("a", "c", 1.0), + ("c", "a", 3.0), + ("b", "c", 2.0), + ("c", "b", 2.0), + ("b", "d", 3.0), + ("d", "b", 1.0), + ("c", "d", 1.0), + ("d", "c", 2.0), + ] + fitted_scores, _, _ = _fit(rank, item_ids, outcomes) + _, fitted_wins = rank._build_bt_win_matrices(item_ids, outcomes, pseudo=0.1) + + def objective(free_scores): + scores = np.append(free_scores, -np.sum(free_scores)) + value = 0.0 + for i in range(len(item_ids)): + for j in range(len(item_ids)): + if i != j and fitted_wins[i, j] > 0: + value += fitted_wins[i, j] * np.logaddexp( + 0.0, -(scores[i] - scores[j]) + ) + return value + + result = minimize( + objective, + np.zeros(len(item_ids) - 1), + method="BFGS", + options={"gtol": 1e-10, "maxiter": 10_000}, + ) + assert result.success or np.linalg.norm(result.jac) < 1e-6 + optimizer_scores = np.append(result.x, -np.sum(result.x)) + + np.testing.assert_allclose( + [fitted_scores[item] for item in item_ids], optimizer_scores, atol=2e-6 + ) + + +def test_draw_has_one_comparison_in_fractional_binomial_working_model(tmp_path): + rank = _rank(tmp_path) + _, outcomes = rank._decode_pairwise_outcome("a", "b", "draw") + scores, n_ij, p_ij = _fit(rank, ["a", "b"], outcomes, pseudo=0.0) + se = _standard_errors( + rank, + ["a", "b"], + scores, + n_ij, + p_ij, + regularization_strength=0.0, + ) + + assert n_ij[0, 1] == n_ij[1, 0] == pytest.approx(1.0) + np.testing.assert_allclose(se, [1.0, 1.0], atol=1e-10) + + doubled_scores, doubled_n, doubled_p = _fit( + rank, ["a", "b"], outcomes + outcomes, pseudo=0.0 + ) + doubled_se = _standard_errors( + rank, + ["a", "b"], + doubled_scores, + doubled_n, + doubled_p, + regularization_strength=0.0, + ) + np.testing.assert_allclose(doubled_se, se / np.sqrt(2.0), atol=1e-10) + + +def test_standard_errors_use_penalized_estimator_sandwich(tmp_path): + rank = _rank(tmp_path) + outcomes = [("a", "b", 5.0), ("b", "a", 5.0)] + + low_scores, low_n, low_p = _fit(rank, ["a", "b"], outcomes, pseudo=0.01) + high_scores, high_n, high_p = _fit(rank, ["a", "b"], outcomes, pseudo=10.0) + low_se = _standard_errors( + rank, + ["a", "b"], + low_scores, + low_n, + low_p, + regularization_strength=0.01, + ) + high_se = _standard_errors( + rank, + ["a", "b"], + high_scores, + high_n, + high_p, + regularization_strength=10.0, + ) + + np.testing.assert_allclose(low_n, high_n) + np.testing.assert_allclose(low_se, np.sqrt(10.0) / 10.01, atol=1e-12) + np.testing.assert_allclose(high_se, np.sqrt(10.0) / 20.0, atol=1e-12) + assert np.all(high_se < low_se) + + +def test_standard_errors_remain_identified_after_sigmoid_saturation(tmp_path): + rank = _rank(tmp_path) + pseudo = 1e-18 + scores, n_ij, p_ij = _fit( + rank, ["a", "b"], [("a", "b")], pseudo=pseudo + ) + se = _standard_errors( + rank, + ["a", "b"], + scores, + n_ij, + p_ij, + regularization_strength=pseudo, + ) + + gap = abs(scores["a"] - scores["b"]) + tail = np.exp(-gap) + variance = tail / (1.0 + tail) ** 2 + expected = np.sqrt(n_ij[0, 1]) / ( + 2.0 * (n_ij[0, 1] + pseudo) * np.sqrt(variance) + ) + + assert p_ij[0, 1] == 1.0 + assert np.isfinite(se).all() + np.testing.assert_allclose(se, [expected, expected], rtol=1e-12) + + +@pytest.mark.parametrize("pseudo", [1e-18, 1e-100, 1e-300]) +@pytest.mark.parametrize("n_items", [3, 5, 20]) +def test_tiny_pseudo_chain_has_stable_closed_form_fit_and_se( + tmp_path, pseudo, n_items +): + rank = _rank(tmp_path) + item_ids = [str(index) for index in range(n_items)] + outcomes = list(zip(item_ids[:-1], item_ids[1:])) + + scores, n_ij, p_ij = rank._fit_bt( + item_ids, + outcomes, + pseudo=pseudo, + max_iter=1, + tol=1e-6, + return_info=True, + ) + expected_gap = np.log1p(pseudo / 2.0) - np.log(pseudo / 2.0) + fitted_gaps = np.array( + [scores[item_ids[i]] - scores[item_ids[i + 1]] for i in range(n_items - 1)] + ) + np.testing.assert_allclose(fitted_gaps, expected_gap, atol=2e-8, rtol=0) + + se = _standard_errors( + rank, + item_ids, + scores, + n_ij, + p_ij, + regularization_strength=pseudo, + ) + assert np.isfinite(se).all() + + +@pytest.mark.parametrize("pseudo", [1e-18, 1e-100, 1e-300]) +def test_tiny_pseudo_mixed_curvature_component_converges(tmp_path, pseudo): + rank = _rank(tmp_path) + item_ids = ["0", "1", "2"] + outcomes = [ + ("1", "0", 1.0), + ("0", "2", 3.0), + ("2", "0", 1.0), + ("1", "2", 1.0), + ] + scores = rank._fit_bt( + item_ids, + outcomes, + pseudo=pseudo, + max_iter=1000, + tol=1e-6, + ) + + assert all(np.isfinite(list(scores.values()))) + assert scores["1"] > scores["0"] > scores["2"] + assert scores["0"] - scores["2"] == pytest.approx(np.log(3.0), abs=1e-10) + + +@pytest.mark.parametrize("pseudo", [1e-16, 1e-100, 1e-300]) +def test_tiny_pseudo_separated_sccs_preserve_internal_score_gaps( + tmp_path, pseudo +): + rank = _rank(tmp_path) + item_ids = ["0", "1", "2", "3"] + outcomes = [ + ("0", "1", 1.0), + ("0", "2", 1.0), + ("0", "3", 3.0), + ("3", "0", 1.0), + ("1", "2", 3.0), + ("2", "1", 1.0), + ("3", "1", 1.0), + ("3", "2", 1.0), + ] + scores = rank._fit_bt( + item_ids, + outcomes, + pseudo=pseudo, + max_iter=1000, + tol=1e-10, + ) + permuted_scores = rank._fit_bt( + ["0", "2", "3", "1"], + outcomes, + pseudo=pseudo, + max_iter=1000, + tol=1e-10, + ) + + assert all(np.isfinite(list(scores.values()))) + assert min(scores["0"], scores["3"]) > max(scores["1"], scores["2"]) + assert scores["0"] - scores["3"] == pytest.approx(np.log(3.0), abs=2e-9) + assert scores["1"] - scores["2"] == pytest.approx(np.log(3.0), abs=2e-9) + + expected_stratum_mean = 0.5 * (-np.log(pseudo) + np.log(8.0 / 3.0)) + top_mean = 0.5 * (scores["0"] + scores["3"]) + bottom_mean = 0.5 * (scores["1"] + scores["2"]) + assert top_mean == pytest.approx(expected_stratum_mean, abs=2e-9) + assert bottom_mean == pytest.approx(-expected_stratum_mean, abs=2e-9) + for item_id in item_ids: + assert permuted_scores[item_id] == pytest.approx( + scores[item_id], abs=2e-9 + ) + + +def test_tiny_pseudo_unequal_scc_sizes_converge_without_root_quantization( + tmp_path, +): + rank = _rank(tmp_path) + item_ids = ["0", "1", "2", "3"] + outcomes = [ + ("0", "1", 3.0), + ("1", "0", 1.0), + ("2", "0", 1.0), + ("3", "0", 1.0), + ("2", "1", 1.0), + ("1", "3", 1.0), + ("2", "3", 1.0), + ] + pseudo = 1e-18 + scores = rank._fit_bt( + item_ids, + outcomes, + pseudo=pseudo, + max_iter=1000, + tol=1e-6, + ) + observed_wins, fit_wins = rank._build_bt_win_matrices( + item_ids, outcomes, pseudo + ) + fitted = np.array([scores[item_id] for item_id in item_ids]) + + assert np.isfinite(fitted).all() + assert scores["2"] > max(scores["0"], scores["1"], scores["3"]) + assert rank._bt_scc_solution_is_certified( + fitted, + observed_wins, + fit_wins, + 1e-6, + ) + + +def test_tiny_pseudo_condensation_newton_resolves_sparse_scc_hierarchy( + tmp_path, +): + rng = np.random.default_rng(2277) + for _ in range(58): + n_items = int(rng.integers(2, 25)) + item_ids = [str(index) for index in range(n_items)] + edges = { + (index, int(rng.integers(index))) + for index in range(1, n_items) + } + density = float(rng.uniform(0.05, 0.35)) + for left in range(n_items): + for right in range(left + 1, n_items): + if rng.random() < density: + edges.add((left, right)) + + outcomes = [] + for left, right in sorted(edges): + state = int(rng.integers(3)) + if state == 0: + outcomes.append((str(left), str(right), 1.0)) + elif state == 1: + outcomes.append((str(right), str(left), 1.0)) + else: + outcomes.extend( + ( + (str(left), str(right), 3.0), + (str(right), str(left), 1.0), + ) + ) + pseudo = float(10.0 ** rng.uniform(-300.0, 3.0)) + + assert n_items == 16 + assert pseudo == pytest.approx(1.6510907983669315e-227) + rank = _rank(tmp_path) + scores = rank._fit_bt( + item_ids, + outcomes, + pseudo=pseudo, + max_iter=1000, + tol=1e-6, + ) + observed_wins, fit_wins = rank._build_bt_win_matrices( + item_ids, outcomes, pseudo + ) + fitted = np.array([scores[item_id] for item_id in item_ids]) + + assert np.isfinite(fitted).all() + assert rank._bt_scc_solution_is_certified( + fitted, + observed_wins, + fit_wins, + 1e-6, + ) + + +def test_observed_information_has_graph_laplacian_structure(tmp_path): + rank = _rank(tmp_path) + item_ids = ["a", "b", "c"] + outcomes = [ + ("a", "b", 3.0), + ("b", "a", 2.0), + ("b", "c", 4.0), + ("c", "b", 1.0), + ("c", "a", 2.0), + ("a", "c", 2.0), + ] + scores, n_ij, p_ij = _fit(rank, item_ids, outcomes) + q_ij = n_ij * p_ij * (1.0 - p_ij) + fisher = np.diag(q_ij.sum(axis=1)) - q_ij + + np.testing.assert_allclose(fisher, fisher.T, atol=1e-12) + np.testing.assert_allclose(fisher @ np.ones(len(item_ids)), 0.0, atol=1e-12) + assert np.all(fisher[np.triu_indices(len(item_ids), 1)] < 0) + eigenvalues = np.linalg.eigvalsh(fisher) + assert np.sum(eigenvalues < 1e-10) == 1 + + regularized_q = (n_ij + 0.1 * (n_ij > 0)) * p_ij * (1.0 - p_ij) + bread = np.diag(regularized_q.sum(axis=1)) - regularized_q + bread_inverse = np.linalg.pinv(bread, rcond=1e-12, hermitian=True) + expected_covariance = bread_inverse @ fisher @ bread_inverse + se = _standard_errors(rank, item_ids, scores, n_ij, p_ij) + np.testing.assert_allclose( + se, np.sqrt(np.diag(expected_covariance)), atol=1e-10 + ) + + +def test_unregularized_fit_rejects_ford_condition_failure(tmp_path): + rank = _rank(tmp_path) + outcomes = [ + ("a", "b"), + ("b", "a"), + ("c", "d"), + ("d", "c"), + ("a", "c"), + ] + + with pytest.raises(ValueError, match="directed win graph"): + rank._fit_bt( + ["a", "b", "c", "d"], + outcomes, + pseudo=0.0, + max_iter=20_000, + tol=1e-12, + ) + + +def test_abstentions_add_no_evidence_or_precision(tmp_path): + rank = _rank(tmp_path, insufficient_signal_policy="abstain") + base_outcomes = [("a", "b", 6.0), ("b", "a", 4.0)] + abstentions = [] + for _ in range(100): + _, decoded = rank._decode_pairwise_outcome( + "a", "b", "insufficient signal" + ) + abstentions.extend(decoded) + + base_scores, base_n, base_p = _fit(rank, ["a", "b"], base_outcomes) + augmented_scores, augmented_n, augmented_p = _fit( + rank, ["a", "b"], base_outcomes + abstentions + ) + base_se = _standard_errors(rank, ["a", "b"], base_scores, base_n, base_p) + augmented_se = _standard_errors( + rank, ["a", "b"], augmented_scores, augmented_n, augmented_p + ) + + assert augmented_scores == pytest.approx(base_scores, abs=1e-12) + np.testing.assert_allclose(augmented_n, base_n) + np.testing.assert_allclose(augmented_se, base_se, atol=1e-12) + + +def test_isolated_items_do_not_change_scores_or_precision(tmp_path): + rank = _rank(tmp_path) + outcomes = [("a", "b", 6.0), ("b", "a", 4.0)] + + base_ids = ["a", "b"] + base_scores, base_n, base_p = _fit(rank, base_ids, outcomes) + base_se = _standard_errors(rank, base_ids, base_scores, base_n, base_p) + + extended_ids = base_ids + [f"unused-{i}" for i in range(25)] + extended_scores, extended_n, extended_p = _fit( + rank, extended_ids, outcomes + ) + with pytest.warns(RuntimeWarning, match="comparison graph is disconnected"): + extended_se = _standard_errors( + rank, extended_ids, extended_scores, extended_n, extended_p + ) + + assert extended_scores["a"] - extended_scores["b"] == pytest.approx( + base_scores["a"] - base_scores["b"], abs=1e-12 + ) + np.testing.assert_allclose(extended_n[:2, :2], base_n) + np.testing.assert_allclose(extended_se[:2], base_se, atol=1e-12) + assert np.isnan(extended_se[2:]).all() + + +def test_zscores_are_normalized_within_comparison_components(tmp_path): + rank = _rank(tmp_path) + base = rank._component_zscores( + np.array([-1.0, 0.0, 1.0]), np.array([0, 0, 0]) + ) + extended = rank._component_zscores( + np.array([-1.0, 0.0, 1.0, -100.0, 100.0]), + np.array([0, 0, 0, 1, 1]), + ) + + np.testing.assert_allclose(extended[:3], base, atol=1e-12) + np.testing.assert_allclose(extended[3:], [-1.0, 1.0], atol=1e-12) + singleton = rank._component_zscores(np.array([0.0]), np.array([0])) + assert np.isnan(singleton[0]) + + +def test_invalid_fit_inputs_fail_loudly(tmp_path): + rank = _rank(tmp_path) + + with pytest.raises(ValueError, match="unique"): + _fit(rank, ["a", "a"], [("a", "a")]) + with pytest.raises(ValueError, match="non-negative"): + _fit(rank, ["a", "b"], [("a", "b", -1.0)]) + with pytest.raises(ValueError, match="present in item_ids"): + _fit(rank, ["a", "b"], [("a", "missing")]) + with pytest.raises(ValueError, match="pseudo"): + _fit(rank, ["a", "b"], [("a", "b")], pseudo=-0.1) + with pytest.raises(ValueError, match="learning_rate"): + _fit(rank, ["a", "b"], [("a", "b")], pseudo=0.0) + + +def test_slow_mm_fit_is_polished_before_inference(tmp_path): + rank = _rank(tmp_path) + outcomes = [("a", "b")] * 9 + [("b", "a")] + + scores = rank._fit_bt( + ["a", "b"], + outcomes, + pseudo=0.1, + max_iter=1, + tol=1e-15, + ) + expected_gap = np.log((9.0 + 0.05) / (1.0 + 0.05)) + assert scores["a"] - scores["b"] == pytest.approx( + expected_gap, abs=1e-10 + ) + + +def test_production_solver_converges_on_sparse_directed_chain(tmp_path): + rank = _rank(tmp_path) + item_ids = [str(index) for index in range(50)] + outcomes = [ + (str(index), str(index + 1)) for index in range(len(item_ids) - 1) + ] + + scores, n_ij, p_ij = rank._fit_bt( + item_ids, + outcomes, + pseudo=0.1, + max_iter=rank._MAX_ITER, + tol=rank._TOL, + return_info=True, + ) + score_values = np.array([scores[item_id] for item_id in item_ids]) + adjacent_gaps = score_values[:-1] - score_values[1:] + np.testing.assert_allclose(adjacent_gaps, np.log(21.0), atol=2e-6) + standard_errors = rank._bt_standard_errors( + score_values, + n_ij, + p_ij, + rcond=rank._SE_EIGEN_TOL, + regularization_strength=0.1, + ) + assert np.isfinite(standard_errors).all() + + +def test_live_rank_path_uses_fractional_draws(monkeypatch, tmp_path): + async def fake_get_all_responses(*, identifiers, **kwargs): + return pd.DataFrame( + { + "Identifier": identifiers, + "Response": ['{"quality": "draw"}'] * len(identifiers), + } + ) + + monkeypatch.setattr( + "gabriel.tasks.rank.get_all_responses", fake_get_all_responses + ) + rank = Rank( + RankConfig( + attributes={"quality": ""}, + save_dir=str(tmp_path), + file_name="rankings", + initial_rating_pass=False, + n_rounds=1, + matches_per_round=1, + ) + ) + + result = asyncio.run( + rank.run( + pd.DataFrame({"text": ["first", "second"]}), + column_name="text", + reset_files=True, + ) + ) + + np.testing.assert_allclose(result["quality_raw"], [0.0, 0.0], atol=1e-12) + np.testing.assert_allclose( + result["quality_se"], [1.0 / 1.1, 1.0 / 1.1], atol=1e-10 + ) + assert result["quality_component"].tolist() == [0, 0] + metadata = json.loads((tmp_path / "rankings_run_metadata.json").read_text()) + assert metadata["rank_estimator_version"] == 2 + assert metadata["insufficient_signal_policy"] == "tie" + assert metadata["learning_rate"] == 0.1 + assert metadata["last_completed_round"] == 0 + + diagnostics = pd.read_csv(tmp_path / "rankings_diagnostics.csv") + assert diagnostics.to_dict(orient="records") == [ + { + "attribute": "quality", + "circle_count": 0, + "square_count": 0, + "draw_count": 1, + "insufficient_signal_count": 0, + "invalid_count": 0, + "effective_comparison_weight": 1.0, + "comparison_components": 1, + "isolated_items": 0, + "finite_standard_errors": 2, + } + ] + + +def test_semantically_incomplete_response_is_not_committed( + monkeypatch, tmp_path +): + responses = ['{"quality": "draw"}', '{"quality": "draw", "novelty": "draw"}'] + calls = 0 + + async def fake_get_all_responses(*, identifiers, **kwargs): + nonlocal calls + response = responses[min(calls, len(responses) - 1)] + calls += 1 + return pd.DataFrame( + { + "Identifier": identifiers, + "Response": [response] * len(identifiers), + "Successful": [True] * len(identifiers), + } + ) + + monkeypatch.setattr( + "gabriel.tasks.rank.get_all_responses", fake_get_all_responses + ) + cfg = RankConfig( + attributes={"quality": "", "novelty": ""}, + save_dir=str(tmp_path), + initial_rating_pass=False, + n_rounds=1, + matches_per_round=1, + ) + data = pd.DataFrame({"id": ["a", "b"], "text": ["A", "B"]}) + + with pytest.raises(ValueError, match="semantically invalid"): + asyncio.run( + Rank(cfg).run( + data, + column_name="text", + id_column="id", + reset_files=True, + ) + ) + metadata = json.loads((tmp_path / "rankings_run_metadata.json").read_text()) + assert metadata["last_completed_round"] == -1 + assert not (tmp_path / "rankings_round0.csv").exists() + staged = pd.read_csv(tmp_path / ".rankings_round0.csv") + assert not staged["Successful"].astype(bool).any() + + result = asyncio.run( + Rank(cfg).run( + data, + column_name="text", + id_column="id", + reset_files=False, + ) + ) + assert calls == 2 + assert result[["quality_raw", "novelty_raw"]].notna().all().all() + + +def test_duplicate_and_singleton_inputs_are_handled_before_model_calls( + monkeypatch, tmp_path +): + calls = 0 + + async def fake_get_all_responses(**kwargs): + nonlocal calls + calls += 1 + raise AssertionError("no model call expected") + + monkeypatch.setattr( + "gabriel.tasks.rank.get_all_responses", fake_get_all_responses + ) + cfg = RankConfig( + attributes={"quality": ""}, + save_dir=str(tmp_path / "duplicate"), + initial_rating_pass=False, + n_rounds=1, + matches_per_round=1, + ) + rank = Rank(cfg) + + with pytest.raises(ValueError, match="unique identifier"): + asyncio.run( + rank.run( + pd.DataFrame({"id": [1, 1, 2], "text": ["a", "b", "c"]}), + column_name="text", + id_column="id", + reset_files=True, + ) + ) + + singleton_rank = Rank( + RankConfig( + attributes={"quality": ""}, + save_dir=str(tmp_path / "singleton"), + initial_rating_pass=False, + n_rounds=1, + matches_per_round=1, + ) + ) + singleton = asyncio.run( + singleton_rank.run( + pd.DataFrame({"id": [7], "text": ["only"]}), + column_name="text", + id_column="id", + reset_files=True, + ) + ) + + assert calls == 0 + assert singleton["quality_component"].tolist() == [0] + assert singleton[["quality", "quality_raw", "quality_se"]].isna().all().all() + + +def test_abstention_resume_tracks_numeric_ids_and_preserves_round_rows( + monkeypatch, tmp_path +): + prompt_counts = [] + + async def fake_get_all_responses(*, identifiers, **kwargs): + prompt_counts.append(len(identifiers)) + return pd.DataFrame( + { + "Identifier": identifiers, + "Response": [ + '{"quality": "insufficient signal"}' + ] + * len(identifiers), + } + ) + + monkeypatch.setattr( + "gabriel.tasks.rank.get_all_responses", fake_get_all_responses + ) + first_cfg = RankConfig( + attributes={"quality": ""}, + save_dir=str(tmp_path), + file_name="rankings", + initial_rating_pass=False, + insufficient_signal_policy="abstain", + n_rounds=1, + matches_per_round=1, + ) + first_data = pd.DataFrame({"id": [101, 202], "text": ["a", "b"]}) + first = asyncio.run( + Rank(first_cfg).run( + first_data, + column_name="text", + id_column="id", + reset_files=True, + ) + ) + original_round = pd.read_csv(tmp_path / "rankings_round0.csv") + original_rows = original_round.set_index("Identifier")[["Batch", "IdA", "IdB"]] + call_boundary = len(prompt_counts) + + second_cfg = RankConfig( + attributes={"quality": ""}, + save_dir=str(tmp_path), + file_name="rankings", + initial_rating_pass=False, + insufficient_signal_policy="abstain", + n_rounds=2, + matches_per_round=1, + ) + expanded = asyncio.run( + Rank(second_cfg).run( + pd.DataFrame( + {"id": [101, 202, 303], "text": ["a", "b", "c"]} + ), + column_name="text", + id_column="id", + reset_files=False, + ) + ) + + # Only the genuinely new numeric ID catches up; abstaining old pairs are + # still recognized as completed comparisons even though they add no wins. + assert prompt_counts[call_boundary:] == [1, 2] + updated_round = pd.read_csv(tmp_path / "rankings_round0.csv").set_index( + "Identifier" + ) + pd.testing.assert_frame_equal( + updated_round.loc[original_rows.index, ["Batch", "IdA", "IdB"]], + original_rows, + check_dtype=False, + ) + assert first[["quality", "quality_raw", "quality_se"]].isna().all().all() + assert expanded[["quality", "quality_raw", "quality_se"]].isna().all().all() + assert expanded["quality_component"].nunique() == 3 + + +@pytest.mark.parametrize( + "metadata", + [ + {}, + {"rank_estimator_version": 1, "insufficient_signal_policy": "tie", "learning_rate": 0.1}, + {"rank_estimator_version": 3, "insufficient_signal_policy": "tie", "learning_rate": 0.1}, + {"rank_estimator_version": 2, "insufficient_signal_policy": "abstain", "learning_rate": 0.1}, + {"rank_estimator_version": 2, "insufficient_signal_policy": "tie", "learning_rate": True}, + {"rank_estimator_version": 2, "insufficient_signal_policy": "tie", "learning_rate": 1e-16}, + ], +) +def test_resume_rejects_incompatible_final_artifacts_before_calls( + monkeypatch, tmp_path, metadata +): + calls = 0 + + async def fake_get_all_responses(**kwargs): + nonlocal calls + calls += 1 + raise AssertionError("incompatible cache must fail before model calls") + + monkeypatch.setattr( + "gabriel.tasks.rank.get_all_responses", fake_get_all_responses + ) + (tmp_path / "rankings_final.csv").write_text("text,quality\na,0\nb,0\n") + (tmp_path / "rankings_run_metadata.json").write_text(json.dumps(metadata)) + rank = Rank( + RankConfig( + attributes={"quality": ""}, + save_dir=str(tmp_path), + initial_rating_pass=False, + n_rounds=1, + matches_per_round=1, + ) + ) + + with pytest.raises(ValueError, match="incompatible"): + asyncio.run( + rank.run( + pd.DataFrame({"text": ["a", "b"]}), + column_name="text", + reset_files=False, + ) + ) + assert calls == 0 + + +def test_recursive_rank_rejects_disconnected_pruning_scores( + monkeypatch, tmp_path +): + async def fake_get_all_responses(*, identifiers, **kwargs): + return pd.DataFrame( + { + "Identifier": identifiers, + "Response": [ + '{"quality": "insufficient signal"}' + ] + * len(identifiers), + } + ) + + monkeypatch.setattr( + "gabriel.tasks.rank.get_all_responses", fake_get_all_responses + ) + rank = Rank( + RankConfig( + attributes={"quality": ""}, + save_dir=str(tmp_path), + initial_rating_pass=False, + insufficient_signal_policy="abstain", + recursive=True, + recursive_rate_first_round=False, + recursive_fraction=0.75, + recursive_min_remaining=1, + n_rounds=1, + matches_per_round=1, + ) + ) + + with pytest.raises(ValueError, match="disconnected comparison graph"): + asyncio.run( + rank.run( + pd.DataFrame({"text": ["a", "b", "c", "d"]}), + column_name="text", + reset_files=True, + ) + ) + + +def test_resume_rejects_changed_payload_for_a_persisted_id( + monkeypatch, tmp_path +): + calls = 0 + + async def fake_get_all_responses(*, identifiers, **kwargs): + nonlocal calls + calls += 1 + return pd.DataFrame( + { + "Identifier": identifiers, + "Response": ['{"quality": "draw"}'] * len(identifiers), + } + ) + + monkeypatch.setattr( + "gabriel.tasks.rank.get_all_responses", fake_get_all_responses + ) + cfg = RankConfig( + attributes={"quality": ""}, + save_dir=str(tmp_path), + initial_rating_pass=False, + n_rounds=1, + matches_per_round=1, + ) + asyncio.run( + Rank(cfg).run( + pd.DataFrame({"id": ["a", "b"], "text": ["old a", "old b"]}), + column_name="text", + id_column="id", + reset_files=True, + ) + ) + call_boundary = calls + + with pytest.raises(ValueError, match="payload changed"): + asyncio.run( + Rank(cfg).run( + pd.DataFrame( + {"id": ["a", "b"], "text": ["edited a", "old b"]} + ), + column_name="text", + id_column="id", + reset_files=False, + ) + ) + assert calls == call_boundary + + +def test_round_replay_preserves_string_identifiers_that_resemble_csv_values( + monkeypatch, tmp_path +): + prompt_counts = [] + + async def fake_get_all_responses(*, identifiers, **kwargs): + prompt_counts.append(len(identifiers)) + return pd.DataFrame( + { + "Identifier": identifiers, + "Response": ['{"quality": "draw"}'] * len(identifiers), + } + ) + + monkeypatch.setattr( + "gabriel.tasks.rank.get_all_responses", fake_get_all_responses + ) + data = pd.DataFrame( + { + "id": ["001", "002", "NA", "null"], + "text": ["a", "b", "c", "d"], + } + ) + first_cfg = RankConfig( + attributes={"quality": ""}, + save_dir=str(tmp_path), + initial_rating_pass=False, + n_rounds=1, + matches_per_round=1, + ) + first = asyncio.run( + Rank(first_cfg).run( + data, + column_name="text", + id_column="id", + reset_files=True, + ) + ) + boundary = len(prompt_counts) + second_cfg = RankConfig( + attributes={"quality": ""}, + save_dir=str(tmp_path), + initial_rating_pass=False, + n_rounds=2, + matches_per_round=1, + ) + second = asyncio.run( + Rank(second_cfg).run( + data, + column_name="text", + id_column="id", + reset_files=False, + ) + ) + + assert first["id"].tolist() == data["id"].tolist() + assert second["id"].tolist() == data["id"].tolist() + assert prompt_counts[boundary:] == [2] + round_zero = Rank._read_rank_checkpoint( + str(tmp_path / "rankings_round0.csv"), 1 + ) + assert set(round_zero["IdA"]) | set(round_zero["IdB"]) == set(data["id"]) + + +def test_reset_clears_stale_future_rounds_and_batch_state( + monkeypatch, tmp_path +): + async def fake_get_all_responses(*, identifiers, **kwargs): + return pd.DataFrame( + { + "Identifier": identifiers, + "Response": ['{"quality": "draw"}'] * len(identifiers), + } + ) + + monkeypatch.setattr( + "gabriel.tasks.rank.get_all_responses", fake_get_all_responses + ) + (tmp_path / "rankings_round4.csv").write_text("stale") + (tmp_path / "rankings_round4.csv.batch_state.json").write_text("{}") + result = asyncio.run( + Rank( + RankConfig( + attributes={"quality": ""}, + save_dir=str(tmp_path), + initial_rating_pass=False, + n_rounds=1, + matches_per_round=1, + ) + ).run( + pd.DataFrame({"text": ["a", "b"]}), + column_name="text", + reset_files=True, + ) + ) + + assert len(result) == 2 + assert not (tmp_path / "rankings_round4.csv").exists() + assert not (tmp_path / "rankings_round4.csv.batch_state.json").exists() + + +def test_failed_model_row_does_not_commit_and_is_retried(monkeypatch, tmp_path): + calls = 0 + + async def fake_get_all_responses(*, identifiers, **kwargs): + nonlocal calls + calls += 1 + if calls == 1: + return pd.DataFrame( + { + "Identifier": identifiers, + "Response": [""] * len(identifiers), + "Successful": [False] * len(identifiers), + } + ) + return pd.DataFrame( + { + "Identifier": identifiers, + "Response": ['{"quality": "circle"}'] * len(identifiers), + "Successful": [True] * len(identifiers), + } + ) + + monkeypatch.setattr( + "gabriel.tasks.rank.get_all_responses", fake_get_all_responses + ) + cfg = RankConfig( + attributes={"quality": ""}, + save_dir=str(tmp_path), + initial_rating_pass=False, + n_rounds=1, + matches_per_round=1, + ) + data = pd.DataFrame({"id": ["a", "b"], "text": ["A", "B"]}) + + with pytest.raises(ValueError, match="not committed"): + asyncio.run( + Rank(cfg).run( + data, + column_name="text", + id_column="id", + reset_files=True, + ) + ) + metadata = json.loads((tmp_path / "rankings_run_metadata.json").read_text()) + assert metadata["last_completed_round"] == -1 + + result = asyncio.run( + Rank(cfg).run( + data, + column_name="text", + id_column="id", + reset_files=False, + ) + ) + + assert calls == 2 + assert result["quality_raw"].notna().all() + metadata = json.loads((tmp_path / "rankings_run_metadata.json").read_text()) + assert metadata["last_completed_round"] == 0 + + +def test_committed_checkpoint_rejects_explicit_failure_or_blank_response(tmp_path): + checkpoint = pd.DataFrame( + { + "Identifier": ["response-1"], + "Response": [""], + "Successful": [False], + "Batch": [0], + "Pair": [0], + "IdA": ["a"], + "IdB": ["b"], + } + ) + path = tmp_path / "rankings_round0.csv" + checkpoint.to_csv(path, index=False) + + with pytest.raises(ValueError, match="failed or blank"): + Rank._read_rank_checkpoint(str(path), 1) + + +def test_round_marker_write_failure_promotes_without_repeating_paid_work( + monkeypatch, tmp_path +): + calls = 0 + + async def fake_get_all_responses(*, identifiers, **kwargs): + nonlocal calls + calls += 1 + return pd.DataFrame( + { + "Identifier": identifiers, + "Response": ['{"quality": "draw"}'] * len(identifiers), + } + ) + + rank_module = importlib.import_module("gabriel.tasks.rank") + real_update = rank_module.update_run_metadata + + def fail_marker(*args, **kwargs): + if kwargs.get("last_completed_round") == 0: + raise OSError("simulated durable-marker failure") + return real_update(*args, **kwargs) + + monkeypatch.setattr(rank_module, "get_all_responses", fake_get_all_responses) + monkeypatch.setattr(rank_module, "update_run_metadata", fail_marker) + cfg = RankConfig( + attributes={"quality": ""}, + save_dir=str(tmp_path), + initial_rating_pass=False, + n_rounds=1, + matches_per_round=1, + ) + data = pd.DataFrame({"id": ["a", "b"], "text": ["A", "B"]}) + + with pytest.raises(OSError, match="durable-marker"): + asyncio.run( + Rank(cfg).run( + data, + column_name="text", + id_column="id", + reset_files=True, + ) + ) + metadata = json.loads((tmp_path / "rankings_run_metadata.json").read_text()) + assert metadata["last_completed_round"] == -1 + assert (tmp_path / "rankings_round0.csv").exists() + + monkeypatch.setattr(rank_module, "update_run_metadata", real_update) + result = asyncio.run( + Rank(cfg).run( + data, + column_name="text", + id_column="id", + reset_files=False, + ) + ) + + assert calls == 1 + assert result["quality_raw"].notna().all() + metadata = json.loads((tmp_path / "rankings_run_metadata.json").read_text()) + assert metadata["last_completed_round"] == 0 + assert not (tmp_path / ".rankings_round0_plan.json").exists() + assert not (tmp_path / ".rankings_round0.csv").exists() + + +def test_interrupted_round_plan_survives_input_row_reordering( + monkeypatch, tmp_path +): + collector_calls = 0 + paid_judgments = 0 + prompt_snapshots = [] + + async def interrupted_collector(*, identifiers, prompts, save_path, **kwargs): + nonlocal collector_calls, paid_judgments + collector_calls += 1 + prompt_snapshots.append(list(prompts)) + if not Path(save_path).exists(): + paid_judgments += len(identifiers) + pd.DataFrame( + { + "Identifier": identifiers, + "Response": ['{"quality": "draw"}'] * len(identifiers), + "Successful": [True] * len(identifiers), + } + ).to_csv(save_path, index=False) + raise RuntimeError("simulated crash after durable paid responses") + return pd.read_csv(save_path) + + monkeypatch.setattr( + "gabriel.tasks.rank.get_all_responses", interrupted_collector + ) + cfg = RankConfig( + attributes={"quality": ""}, + save_dir=str(tmp_path), + initial_rating_pass=False, + n_rounds=1, + matches_per_round=1, + ) + original = pd.DataFrame( + {"id": ["a", "b", "c"], "text": ["A", "B", "C"]} + ) + with pytest.raises(RuntimeError, match="durable paid"): + asyncio.run( + Rank(cfg).run( + original, + column_name="text", + id_column="id", + reset_files=True, + ) + ) + + metadata_path = tmp_path / "rankings_run_metadata.json" + metadata_before_rejection = metadata_path.read_bytes() + calls_before_rejection = collector_calls + with pytest.raises(ValueError, match="different input set"): + asyncio.run( + Rank(cfg).run( + pd.concat( + [ + original, + pd.DataFrame({"id": ["g"], "text": ["old payload"]}), + ], + ignore_index=True, + ), + column_name="text", + id_column="id", + reset_files=False, + ) + ) + assert collector_calls == calls_before_rejection + assert metadata_path.read_bytes() == metadata_before_rejection + + result = asyncio.run( + Rank(cfg).run( + original.iloc[::-1].reset_index(drop=True), + column_name="text", + id_column="id", + reset_files=False, + ) + ) + + assert collector_calls == 2 + assert paid_judgments == len(prompt_snapshots[0]) + assert prompt_snapshots[1] == prompt_snapshots[0] + assert set(result["id"]) == {"a", "b", "c"} + assert not (tmp_path / ".rankings_round0_plan.json").exists() + assert not (tmp_path / ".rankings_round0.csv").exists() + + +def test_catchup_deduplicates_pair_deficits_and_uses_staging_file( + monkeypatch, tmp_path +): + calls = [] + + async def fake_get_all_responses(*, identifiers, save_path, **kwargs): + calls.append((len(identifiers), save_path)) + return pd.DataFrame( + { + "Identifier": identifiers, + "Response": ['{"quality": "draw"}'] * len(identifiers), + } + ) + + monkeypatch.setattr( + "gabriel.tasks.rank.get_all_responses", fake_get_all_responses + ) + cfg = RankConfig( + attributes={"quality": ""}, + save_dir=str(tmp_path), + initial_rating_pass=False, + n_rounds=1, + matches_per_round=1, + ) + first = Rank(cfg) + first._generate_pairs = lambda *, texts_by_id, **kwargs: [ + (("a", texts_by_id["a"]), ("b", texts_by_id["b"])), + (("c", texts_by_id["c"]), ("d", texts_by_id["d"])), + ] + asyncio.run( + first.run( + pd.DataFrame( + {"id": ["a", "b", "c", "d"], "text": ["A", "B", "C", "D"]} + ), + column_name="text", + id_column="id", + reset_files=True, + ) + ) + committed_before = (tmp_path / "rankings_round0.csv").read_bytes() + + result = asyncio.run( + Rank(cfg).run( + pd.DataFrame({"id": ["a", "c"], "text": ["A", "C"]}), + column_name="text", + id_column="id", + reset_files=False, + ) + ) + + assert calls[0][0] == 2 + assert calls[1][0] == 1 + assert calls[0][1].endswith("rankings_round0.csv") + assert calls[1][1].endswith(".rankings_catchup_round0.csv") + assert (tmp_path / "rankings_round0.csv").read_bytes() != committed_before + checkpoint = Rank._read_rank_checkpoint( + str(tmp_path / "rankings_round0.csv"), 1 + ) + assert len(checkpoint) == 3 + assert result["quality_raw"].notna().all() + assert result["quality_component"].nunique() == 1 + + +def test_interrupted_catchup_cannot_mutate_committed_round(monkeypatch, tmp_path): + phase = "initial" + + async def fake_get_all_responses(*, identifiers, save_path, **kwargs): + if phase == "catchup": + pd.DataFrame( + { + "Identifier": identifiers[:1], + "Response": ['{"quality": "draw"}'], + "Successful": [True], + } + ).to_csv(save_path, index=False) + raise RuntimeError("simulated interrupted catch-up") + return pd.DataFrame( + { + "Identifier": identifiers, + "Response": ['{"quality": "draw"}'] * len(identifiers), + } + ) + + monkeypatch.setattr( + "gabriel.tasks.rank.get_all_responses", fake_get_all_responses + ) + cfg = RankConfig( + attributes={"quality": ""}, + save_dir=str(tmp_path), + initial_rating_pass=False, + n_rounds=1, + matches_per_round=1, + ) + first = Rank(cfg) + first._generate_pairs = lambda *, texts_by_id, **kwargs: [ + (("a", texts_by_id["a"]), ("b", texts_by_id["b"])), + (("c", texts_by_id["c"]), ("d", texts_by_id["d"])), + ] + asyncio.run( + first.run( + pd.DataFrame( + {"id": ["a", "b", "c", "d"], "text": ["A", "B", "C", "D"]} + ), + column_name="text", + id_column="id", + reset_files=True, + ) + ) + committed_path = tmp_path / "rankings_round0.csv" + committed_before = committed_path.read_bytes() + phase = "catchup" + + with pytest.raises(RuntimeError, match="interrupted catch-up"): + asyncio.run( + Rank(cfg).run( + pd.DataFrame({"id": ["a", "c"], "text": ["A", "C"]}), + column_name="text", + id_column="id", + reset_files=False, + ) + ) + + assert committed_path.read_bytes() == committed_before + Rank._read_rank_checkpoint(str(committed_path), 1) + assert (tmp_path / ".rankings_catchup_round0.csv").exists() + + +def test_empty_view_preserves_existing_tournament_and_handles_empty_ids( + monkeypatch, tmp_path +): + calls = 0 + + async def fake_get_all_responses(*, identifiers, **kwargs): + nonlocal calls + calls += 1 + return pd.DataFrame( + { + "Identifier": identifiers, + "Response": ['{"quality": "draw"}'] * len(identifiers), + } + ) + + monkeypatch.setattr( + "gabriel.tasks.rank.get_all_responses", fake_get_all_responses + ) + cfg = RankConfig( + attributes={"quality": ""}, + save_dir=str(tmp_path), + initial_rating_pass=False, + n_rounds=1, + matches_per_round=1, + ) + data = pd.DataFrame({"id": ["a", "b"], "text": ["A", "B"]}) + asyncio.run( + Rank(cfg).run( + data, + column_name="text", + id_column="id", + reset_files=True, + ) + ) + paths = [ + tmp_path / "rankings_run_metadata.json", + tmp_path / "rankings_round0.csv", + tmp_path / "rankings_final.csv", + tmp_path / "rankings_diagnostics.csv", + ] + snapshots = {path: path.read_bytes() for path in paths} + + empty = asyncio.run( + Rank(cfg).run( + pd.DataFrame({"id": pd.Series(dtype=str), "text": pd.Series(dtype=str)}), + column_name="text", + id_column="id", + reset_files=False, + ) + ) + + assert empty.empty + assert calls == 1 + assert {path: path.read_bytes() for path in paths} == snapshots + replayed = asyncio.run( + Rank(cfg).run( + data, + column_name="text", + id_column="id", + reset_files=False, + ) + ) + assert calls == 1 + assert len(replayed) == 2 + + +def test_blank_identifier_is_dropped_before_any_model_call(monkeypatch, tmp_path): + calls = 0 + + async def fake_get_all_responses(**kwargs): + nonlocal calls + calls += 1 + raise AssertionError("a singleton should not call the model") + + monkeypatch.setattr( + "gabriel.tasks.rank.get_all_responses", fake_get_all_responses + ) + result = asyncio.run( + Rank( + RankConfig( + attributes={"quality": "", "novelty": ""}, + save_dir=str(tmp_path), + initial_rating_pass=False, + n_rounds=1, + matches_per_round=1, + ) + ).run( + pd.DataFrame({"id": [" ", "b"], "text": ["blank", "B"]}), + column_name="text", + id_column="id", + reset_files=True, + ) + ) + + assert calls == 0 + assert result["id"].tolist() == ["b"] + + +def test_local_media_content_change_invalidates_persisted_judgments( + monkeypatch, tmp_path +): + calls = 0 + + async def fake_get_all_responses(*, identifiers, **kwargs): + nonlocal calls + calls += 1 + return pd.DataFrame( + { + "Identifier": identifiers, + "Response": ['{"quality": "draw"}'] * len(identifiers), + } + ) + + monkeypatch.setattr( + "gabriel.tasks.rank.get_all_responses", fake_get_all_responses + ) + image_a = tmp_path / "a.img" + image_b = tmp_path / "b.img" + image_a.write_bytes(b"first image bytes") + image_b.write_bytes(b"second image bytes") + cfg = RankConfig( + attributes={"quality": ""}, + save_dir=str(tmp_path / "run"), + modality="image", + initial_rating_pass=False, + n_rounds=1, + matches_per_round=1, + ) + data = pd.DataFrame( + {"id": ["a", "b"], "image": [str(image_a), str(image_b)]} + ) + asyncio.run( + Rank(cfg).run( + data, + column_name="image", + id_column="id", + reset_files=True, + ) + ) + image_a.write_bytes(b"replacement bytes at the same path") + + with pytest.raises(ValueError, match="payload changed"): + asyncio.run( + Rank(cfg).run( + data, + column_name="image", + id_column="id", + reset_files=False, + ) + ) + assert calls == 1 + + +def test_recursive_empty_output_uses_documented_directory(tmp_path): + result = asyncio.run( + Rank( + RankConfig( + attributes={"quality": "", "novelty": ""}, + save_dir=str(tmp_path), + recursive=True, + recursive_rate_first_round=False, + initial_rating_pass=False, + ) + ).run( + pd.DataFrame({"text": pd.Series(dtype=str)}), + column_name="text", + reset_files=True, + ) + ) + + assert result.empty + assert {"quality", "novelty", "overall_rank", "exit_stage"}.issubset( + result.columns + ) + assert not any( + column.endswith(("_raw", "_se")) for column in result.columns + ) + assert (tmp_path / "rankings_recursive" / "recursive_final.csv").exists() + + +def test_recursive_empty_view_preserves_existing_final_and_empty_folder_is_new( + monkeypatch, tmp_path +): + calls = 0 + + async def fake_get_all_responses(*, identifiers, **kwargs): + nonlocal calls + calls += 1 + return pd.DataFrame( + { + "Identifier": identifiers, + "Response": ['{"quality": "draw"}'] * len(identifiers), + } + ) + + monkeypatch.setattr( + "gabriel.tasks.rank.get_all_responses", fake_get_all_responses + ) + recursive_folder = tmp_path / "rankings_recursive" + recursive_folder.mkdir() + cfg = RankConfig( + attributes={"quality": ""}, + save_dir=str(tmp_path), + recursive=True, + recursive_rate_first_round=False, + recursive_min_remaining=2, + recursive_final_round_multiplier=1, + initial_rating_pass=False, + n_rounds=1, + matches_per_round=1, + ) + asyncio.run( + Rank(cfg).run( + pd.DataFrame({"id": ["a", "b"], "text": ["A", "B"]}), + column_name="text", + id_column="id", + reset_files=False, + ) + ) + final_path = recursive_folder / "recursive_final.csv" + snapshot = final_path.read_bytes() + boundary = calls + + empty = asyncio.run( + Rank(cfg).run( + pd.DataFrame( + {"id": pd.Series(dtype=str), "text": pd.Series(dtype=str)} + ), + column_name="text", + id_column="id", + reset_files=False, + ) + ) + + assert empty.empty + assert calls == boundary + assert final_path.read_bytes() == snapshot + + +def test_measurement_fingerprint_guards_attributes_without_sidecars( + monkeypatch, tmp_path +): + prompt_batches = [] + + async def fake_get_all_responses(*, identifiers, prompts, **kwargs): + prompt_batches.append(list(prompts)) + return pd.DataFrame( + { + "Identifier": identifiers, + "Response": ['{"quality": "draw"}'] * len(identifiers), + } + ) + + monkeypatch.setattr( + "gabriel.tasks.rank.get_all_responses", fake_get_all_responses + ) + data = pd.DataFrame({"id": ["a", "b"], "text": ["A", "B"]}) + original_cfg = RankConfig( + attributes={"quality": "OLD_UNIQUE_DEFINITION"}, + save_dir=str(tmp_path), + initial_rating_pass=False, + n_rounds=1, + matches_per_round=1, + ) + asyncio.run( + Rank(original_cfg).run( + data, + column_name="text", + id_column="id", + reset_files=True, + ) + ) + for sidecar in (tmp_path / "attributes.json", tmp_path / "rankings_attrs.json"): + sidecar.unlink() + boundary = len(prompt_batches) + + incompatible_cfg = RankConfig( + attributes={"quality": "NEW_UNIQUE_DEFINITION"}, + save_dir=str(tmp_path), + initial_rating_pass=False, + n_rounds=2, + matches_per_round=1, + ) + with pytest.raises(ValueError, match="measurement_spec_fingerprint"): + asyncio.run( + Rank(incompatible_cfg).run( + data, + column_name="text", + id_column="id", + reset_files=False, + ) + ) + assert len(prompt_batches) == boundary + assert not (tmp_path / "attributes.json").exists() + assert not (tmp_path / "rankings_attrs.json").exists() + + compatible_cfg = RankConfig( + attributes={"quality": "OLD_UNIQUE_DEFINITION"}, + save_dir=str(tmp_path), + initial_rating_pass=False, + n_rounds=2, + matches_per_round=1, + ) + asyncio.run( + Rank(compatible_cfg).run( + data, + column_name="text", + id_column="id", + reset_files=False, + ) + ) + + assert len(prompt_batches) == boundary + 1 + assert all("OLD_UNIQUE_DEFINITION" in prompt for prompt in prompt_batches[-1]) + assert (tmp_path / "attributes.json").exists() + assert (tmp_path / "rankings_attrs.json").exists() + + +def test_recursive_orchestrator_guards_attributes_without_top_sidecars( + monkeypatch, tmp_path +): + calls = 0 + + async def fake_get_all_responses(*, identifiers, **kwargs): + nonlocal calls + calls += 1 + return pd.DataFrame( + { + "Identifier": identifiers, + "Response": ['{"quality": "draw"}'] * len(identifiers), + } + ) + + monkeypatch.setattr( + "gabriel.tasks.rank.get_all_responses", fake_get_all_responses + ) + data = pd.DataFrame({"id": ["a", "b"], "text": ["A", "B"]}) + original_cfg = RankConfig( + attributes={"quality": "OLD_RECURSIVE_DEFINITION"}, + save_dir=str(tmp_path), + recursive=True, + recursive_rate_first_round=False, + recursive_min_remaining=2, + recursive_final_round_multiplier=1, + initial_rating_pass=False, + n_rounds=1, + matches_per_round=1, + ) + asyncio.run( + Rank(original_cfg).run( + data, + column_name="text", + id_column="id", + reset_files=True, + ) + ) + for sidecar in (tmp_path / "attributes.json", tmp_path / "rankings_attrs.json"): + sidecar.unlink() + boundary = calls + + incompatible_cfg = RankConfig( + attributes={"replacement": "NEW_RECURSIVE_DEFINITION"}, + save_dir=str(tmp_path), + recursive=True, + recursive_rate_first_round=False, + recursive_min_remaining=2, + recursive_final_round_multiplier=1, + initial_rating_pass=False, + n_rounds=1, + matches_per_round=1, + ) + with pytest.raises(ValueError, match="measurement_spec_fingerprint"): + asyncio.run( + Rank(incompatible_cfg).run( + data, + column_name="text", + id_column="id", + reset_files=False, + ) + ) + + assert calls == boundary + assert not (tmp_path / "attributes.json").exists() + assert not (tmp_path / "rankings_attrs.json").exists() + + +def test_missing_persisted_input_fingerprints_fail_closed(monkeypatch, tmp_path): + calls = 0 + + async def fake_get_all_responses(*, identifiers, **kwargs): + nonlocal calls + calls += 1 + return pd.DataFrame( + { + "Identifier": identifiers, + "Response": ['{"quality": "draw"}'] * len(identifiers), + } + ) + + monkeypatch.setattr( + "gabriel.tasks.rank.get_all_responses", fake_get_all_responses + ) + cfg = RankConfig( + attributes={"quality": ""}, + save_dir=str(tmp_path), + initial_rating_pass=False, + n_rounds=1, + matches_per_round=1, + ) + original = pd.DataFrame({"id": ["a", "b"], "text": ["original A", "B"]}) + asyncio.run( + Rank(cfg).run( + original, + column_name="text", + id_column="id", + reset_files=True, + ) + ) + metadata_path = tmp_path / "rankings_run_metadata.json" + metadata = json.loads(metadata_path.read_text()) + metadata["input_fingerprints"] = {} + metadata_path.write_text(json.dumps(metadata)) + boundary = calls + + with pytest.raises(ValueError, match="lacks valid content fingerprints"): + asyncio.run( + Rank(cfg).run( + pd.DataFrame( + {"id": ["a", "b"], "text": ["changed A", "B"]} + ), + column_name="text", + id_column="id", + reset_files=False, + ) + ) + assert calls == boundary + + +def test_recursive_top_fingerprint_detects_local_media_change( + monkeypatch, tmp_path +): + calls = 0 + + async def fake_rate_responses(*, identifiers, **kwargs): + nonlocal calls + calls += len(identifiers) + return pd.DataFrame( + { + "Identifier": identifiers, + "Response": ['{"quality": 50}'] * len(identifiers), + } + ) + + monkeypatch.setattr( + "gabriel.tasks.rate.get_all_responses", fake_rate_responses + ) + image_a = tmp_path / "a.img" + image_b = tmp_path / "b.img" + image_a.write_bytes(b"image a version one") + image_b.write_bytes(b"image b") + cfg = RankConfig( + attributes={"quality": ""}, + save_dir=str(tmp_path / "run"), + modality="image", + recursive=True, + recursive_rate_first_round=True, + recursive_min_remaining=2, + initial_rating_pass=False, + n_rounds=1, + matches_per_round=1, + ) + data = pd.DataFrame( + {"id": ["a", "b"], "image": [str(image_a), str(image_b)]} + ) + asyncio.run( + Rank(cfg).run( + data, + column_name="image", + id_column="id", + reset_files=True, + ) + ) + boundary = calls + image_a.write_bytes(b"image a replacement at same path") + + with pytest.raises(ValueError, match="comparison payload changed"): + asyncio.run( + Rank(cfg).run( + data, + column_name="image", + id_column="id", + reset_files=False, + ) + ) + assert calls == boundary + + +def test_recursive_rate_semantics_are_in_top_fingerprint(monkeypatch, tmp_path): + models = [] + + async def fake_rate_responses(*, identifiers, model, **kwargs): + models.extend([model] * len(identifiers)) + return pd.DataFrame( + { + "Identifier": identifiers, + "Response": ['{"quality": 50}'] * len(identifiers), + } + ) + + monkeypatch.setattr( + "gabriel.tasks.rate.get_all_responses", fake_rate_responses + ) + data = pd.DataFrame({"id": ["a", "b"], "text": ["A", "B"]}) + first_cfg = RankConfig( + attributes={"quality": ""}, + save_dir=str(tmp_path), + recursive=True, + recursive_rate_first_round=True, + recursive_min_remaining=2, + initial_rating_pass=False, + rate_kwargs={"model": "rate-model-A"}, + n_rounds=1, + matches_per_round=1, + ) + asyncio.run( + Rank(first_cfg).run( + data, + column_name="text", + id_column="id", + reset_files=True, + ) + ) + boundary = len(models) + + changed_cfg = RankConfig( + attributes={"quality": ""}, + save_dir=str(tmp_path), + recursive=True, + recursive_rate_first_round=True, + recursive_min_remaining=2, + initial_rating_pass=False, + rate_kwargs={"model": "rate-model-B"}, + n_rounds=1, + matches_per_round=1, + ) + with pytest.raises(ValueError, match="measurement_spec_fingerprint"): + asyncio.run( + Rank(changed_cfg).run( + data, + column_name="text", + id_column="id", + reset_files=False, + ) + ) + + assert len(models) == boundary + assert set(models) == {"rate-model-A"} + + +def test_initial_rate_crash_still_persists_top_measurement_guard( + monkeypatch, tmp_path +): + models = [] + + async def fake_rate_responses(*, identifiers, model, **kwargs): + models.extend([model] * len(identifiers)) + return pd.DataFrame( + { + "Identifier": identifiers, + "Response": ['{"quality": 50}'] * len(identifiers), + } + ) + + monkeypatch.setattr( + "gabriel.tasks.rate.get_all_responses", fake_rate_responses + ) + first_cfg = RankConfig( + attributes={"quality": ""}, + save_dir=str(tmp_path), + initial_rating_pass=True, + rate_kwargs={"model": "rate-model-A"}, + n_rounds=1, + matches_per_round=1, + ) + first = Rank(first_cfg) + + def crash_after_rate(*args, **kwargs): + raise RuntimeError("simulated crash after initial Rate") + + monkeypatch.setattr(first, "_seed_ratings_from_rate", crash_after_rate) + data = pd.DataFrame({"id": ["a", "b"], "text": ["A", "B"]}) + with pytest.raises(RuntimeError, match="after initial Rate"): + asyncio.run( + first.run( + data, + column_name="text", + id_column="id", + reset_files=True, + ) + ) + metadata = json.loads((tmp_path / "rankings_run_metadata.json").read_text()) + assert metadata["last_completed_round"] == -1 + boundary = len(models) + + changed_cfg = RankConfig( + attributes={"quality": ""}, + save_dir=str(tmp_path), + initial_rating_pass=True, + rate_kwargs={"model": "rate-model-B"}, + n_rounds=1, + matches_per_round=1, + ) + with pytest.raises(ValueError, match="measurement_spec_fingerprint"): + asyncio.run( + Rank(changed_cfg).run( + data, + column_name="text", + id_column="id", + reset_files=False, + ) + ) + assert len(models) == boundary + + +def test_batch_state_only_round_is_guarded_before_resume_calls( + monkeypatch, tmp_path +): + models = [] + + async def interrupted_batch(*, save_path, model, **kwargs): + models.append(model) + Path(f"{save_path}.batch_state.json").write_text( + json.dumps({"batches": []}) + ) + raise RuntimeError("simulated crash after batch submission") + + monkeypatch.setattr( + "gabriel.tasks.rank.get_all_responses", interrupted_batch + ) + data = pd.DataFrame({"id": ["a", "b"], "text": ["A", "B"]}) + first_cfg = RankConfig( + attributes={"quality": ""}, + save_dir=str(tmp_path), + model="model-A", + initial_rating_pass=False, + n_rounds=1, + matches_per_round=1, + ) + with pytest.raises(RuntimeError, match="after batch submission"): + asyncio.run( + Rank(first_cfg).run( + data, + column_name="text", + id_column="id", + reset_files=True, + ) + ) + + changed_cfg = RankConfig( + attributes={"quality": ""}, + save_dir=str(tmp_path), + model="model-B", + initial_rating_pass=False, + n_rounds=1, + matches_per_round=1, + ) + with pytest.raises(ValueError, match="measurement_spec_fingerprint"): + asyncio.run( + Rank(changed_cfg).run( + data, + column_name="text", + id_column="id", + reset_files=False, + ) + ) + + with pytest.raises(ValueError, match="without a durable response"): + asyncio.run( + Rank(first_cfg).run( + data, + column_name="text", + id_column="id", + reset_files=False, + ) + ) + assert models == ["model-A"] + + +def test_submitted_batch_with_round_plan_reenters_recovery_collector( + monkeypatch, tmp_path +): + calls = 0 + + async def interrupted_batch(*, identifiers, save_path, **kwargs): + nonlocal calls + calls += 1 + if calls == 1: + pd.DataFrame(columns=["Identifier", "Response"]).to_csv( + save_path, index=False + ) + Path(f"{save_path}.batch_state.json").write_text( + json.dumps( + { + "batches": [ + { + "batch_id": "batch-paid-1", + "status": "submitted", + "total": len(identifiers), + } + ] + } + ) + ) + raise RuntimeError("simulated crash after batch submission") + return pd.DataFrame( + { + "Identifier": identifiers, + "Response": ['{"quality": "draw"}'] * len(identifiers), + "Successful": [True] * len(identifiers), + } + ) + + monkeypatch.setattr( + "gabriel.tasks.rank.get_all_responses", interrupted_batch + ) + cfg = RankConfig( + attributes={"quality": ""}, + save_dir=str(tmp_path), + initial_rating_pass=False, + n_rounds=1, + matches_per_round=1, + ) + data = pd.DataFrame({"id": ["a", "b"], "text": ["A", "B"]}) + with pytest.raises(RuntimeError, match="after batch submission"): + asyncio.run( + Rank(cfg).run( + data, + column_name="text", + id_column="id", + reset_files=True, + ) + ) + + result = asyncio.run( + Rank(cfg).run( + data, + column_name="text", + id_column="id", + reset_files=False, + ) + ) + assert calls == 2 + assert result["quality_raw"].notna().all() + assert not (tmp_path / ".rankings_round0.csv.batch_state.json").exists() + + +def test_resume_rejects_fewer_rounds_than_already_committed( + monkeypatch, tmp_path +): + calls = 0 + + async def fake_get_all_responses(*, identifiers, **kwargs): + nonlocal calls + calls += 1 + return pd.DataFrame( + { + "Identifier": identifiers, + "Response": ['{"quality": "draw"}'] * len(identifiers), + } + ) + + monkeypatch.setattr( + "gabriel.tasks.rank.get_all_responses", fake_get_all_responses + ) + data = pd.DataFrame({"id": ["a", "b"], "text": ["A", "B"]}) + asyncio.run( + Rank( + RankConfig( + attributes={"quality": ""}, + save_dir=str(tmp_path), + initial_rating_pass=False, + n_rounds=2, + matches_per_round=1, + ) + ).run( + data, + column_name="text", + id_column="id", + reset_files=True, + ) + ) + boundary = calls + + with pytest.raises(ValueError, match="fewer n_rounds"): + asyncio.run( + Rank( + RankConfig( + attributes={"quality": ""}, + save_dir=str(tmp_path), + initial_rating_pass=False, + n_rounds=1, + matches_per_round=1, + ) + ).run( + data, + column_name="text", + id_column="id", + reset_files=False, + ) + ) + assert calls == boundary + + +def test_resume_rejects_changed_runtime_judgment_settings_before_new_calls( + monkeypatch, tmp_path +): + calls = 0 + + async def fake_get_all_responses(*, identifiers, **kwargs): + nonlocal calls + calls += 1 + return pd.DataFrame( + { + "Identifier": identifiers, + "Response": ['{"quality": "draw"}'] * len(identifiers), + } + ) + + monkeypatch.setattr( + "gabriel.tasks.rank.get_all_responses", fake_get_all_responses + ) + data = pd.DataFrame({"id": ["a", "b"], "text": ["A", "B"]}) + first_cfg = RankConfig( + attributes={"quality": ""}, + save_dir=str(tmp_path), + initial_rating_pass=False, + n_rounds=1, + matches_per_round=1, + ) + asyncio.run( + Rank(first_cfg).run( + data, + column_name="text", + id_column="id", + reset_files=True, + web_search=False, + ) + ) + boundary = calls + + resumed_cfg = RankConfig( + attributes={"quality": ""}, + save_dir=str(tmp_path), + initial_rating_pass=False, + n_rounds=2, + matches_per_round=1, + ) + with pytest.raises(ValueError, match="measurement_spec_fingerprint"): + asyncio.run( + Rank(resumed_cfg).run( + data, + column_name="text", + id_column="id", + reset_files=False, + web_search=True, + ) + ) + assert calls == boundary + + task = Rank(first_cfg) + assert task._measurement_spec_fingerprint( + { + "image_detail": "low", + "api_key": "first-secret", + "skip_tail_fails": True, + "timeout": 1, + "return_raw": False, + "request_phase_callback": lambda *_: None, + } + ) == task._measurement_spec_fingerprint( + { + "image_detail": "low", + "api_key": "second-secret", + "skip_tail_fails": False, + "timeout": 120, + "return_raw": True, + "request_phase_callback": lambda *_: None, + } + ) + assert task._measurement_spec_fingerprint( + {"image_detail": "low"} + ) != task._measurement_spec_fingerprint({"image_detail": "high"}) + + +def test_custom_rank_judge_requires_explicit_version_before_calls(tmp_path): + calls = 0 + + def make_judge(label): + async def judge(*, identifiers, **kwargs): + nonlocal calls + calls += 1 + return pd.DataFrame( + { + "Identifier": identifiers, + "Response": [json.dumps({"quality": label})] + * len(identifiers), + } + ) + + return judge + + cfg = RankConfig( + attributes={"quality": ""}, + save_dir=str(tmp_path), + initial_rating_pass=False, + n_rounds=1, + matches_per_round=1, + ) + with pytest.raises(ValueError, match="judge_version is required"): + asyncio.run( + Rank(cfg).run( + pd.DataFrame({"id": ["a", "b"], "text": ["A", "B"]}), + column_name="text", + id_column="id", + reset_files=True, + get_all_responses_fn=make_judge("draw"), + ) + ) + assert calls == 0 + + +def test_recursive_rejected_superset_does_not_mutate_top_metadata( + monkeypatch, tmp_path +): + calls = 0 + + async def interrupted_collector(*, identifiers, **kwargs): + nonlocal calls + calls += 1 + raise RuntimeError("simulated recursive stage interruption") + + monkeypatch.setattr( + "gabriel.tasks.rank.get_all_responses", interrupted_collector + ) + cfg = RankConfig( + attributes={"quality": ""}, + save_dir=str(tmp_path), + initial_rating_pass=False, + recursive=True, + recursive_rate_first_round=False, + recursive_fraction=0.5, + recursive_min_remaining=1, + n_rounds=1, + matches_per_round=1, + ) + original = pd.DataFrame({"id": ["a", "b"], "text": ["A", "B"]}) + with pytest.raises(RuntimeError, match="recursive stage interruption"): + asyncio.run( + Rank(cfg).run( + original, + column_name="text", + id_column="id", + reset_files=True, + ) + ) + metadata_path = tmp_path / "rankings_run_metadata.json" + metadata_before = metadata_path.read_bytes() + call_boundary = calls + + with pytest.raises(ValueError, match="unfinished stage transaction"): + asyncio.run( + Rank(cfg).run( + pd.concat( + [ + original, + pd.DataFrame({"id": ["g"], "text": ["old payload"]}), + ], + ignore_index=True, + ), + column_name="text", + id_column="id", + reset_files=False, + ) + ) + assert calls == call_boundary + assert metadata_path.read_bytes() == metadata_before + + +def test_rank_forces_exact_tail_retries_for_rounds_and_catchup( + monkeypatch, tmp_path +): + tail_settings = [] + + async def fake_get_all_responses(*, identifiers, **kwargs): + tail_settings.append(kwargs.get("skip_tail_fails")) + return pd.DataFrame( + { + "Identifier": identifiers, + "Response": ['{"quality": "draw"}'] * len(identifiers), + } + ) + + monkeypatch.setattr( + "gabriel.tasks.rank.get_all_responses", fake_get_all_responses + ) + cfg = RankConfig( + attributes={"quality": ""}, + save_dir=str(tmp_path), + initial_rating_pass=False, + n_rounds=1, + matches_per_round=1, + ) + asyncio.run( + Rank(cfg).run( + pd.DataFrame({"id": ["a", "b"], "text": ["A", "B"]}), + column_name="text", + id_column="id", + reset_files=True, + skip_tail_fails=True, + ) + ) + asyncio.run( + Rank(cfg).run( + pd.DataFrame( + {"id": ["a", "b", "c"], "text": ["A", "B", "C"]} + ), + column_name="text", + id_column="id", + reset_files=False, + skip_tail_fails=True, + ) + ) + assert len(tail_settings) >= 2 + assert tail_settings == [False] * len(tail_settings) + + +@pytest.mark.parametrize("recursive", [False, True]) +def test_reset_refuses_to_orphan_durable_batch_state(tmp_path, recursive): + cfg = RankConfig( + attributes={"quality": ""}, + save_dir=str(tmp_path), + initial_rating_pass=False, + recursive=recursive, + recursive_rate_first_round=False, + n_rounds=1, + matches_per_round=1, + ) + if recursive: + state_path = ( + tmp_path + / "rankings_recursive" + / "stage1" + / ".rankings_round0.csv.batch_state.json" + ) + else: + state_path = tmp_path / ".rankings_round0.csv.batch_state.json" + state_path.parent.mkdir(parents=True, exist_ok=True) + state_path.write_text( + json.dumps( + { + "batches": [ + {"batch_id": "batch-paid", "status": "in_progress"} + ] + } + ) + ) + + with pytest.raises(ValueError, match="cannot reset.*Batch API state"): + asyncio.run( + Rank(cfg).run( + pd.DataFrame({"id": [], "text": []}), + column_name="text", + id_column="id", + reset_files=True, + ) + ) + assert state_path.exists() + + +def test_reset_refuses_active_initial_rate_batch_state(tmp_path): + state_path = ( + tmp_path + / "rankings_initial_rate" + / "rankings_initial_rate_raw_responses.csv.batch_state.json" + ) + state_path.parent.mkdir(parents=True) + state_path.write_text( + json.dumps( + {"batches": [{"batch_id": "paid-rate", "status": "in_progress"}]} + ) + ) + cfg = RankConfig( + attributes={"quality": ""}, + save_dir=str(tmp_path), + initial_rating_pass=True, + n_rounds=1, + matches_per_round=1, + ) + + with pytest.raises(ValueError, match="cannot reset.*Batch API state"): + asyncio.run( + Rank(cfg).run( + pd.DataFrame({"id": [], "text": []}), + column_name="text", + id_column="id", + reset_files=True, + ) + ) + assert state_path.exists() + + +def test_recursive_reset_clears_task_owned_stage_tree(tmp_path): + recursive_dir = tmp_path / "rankings_recursive" + sentinel = recursive_dir / "stage9" / "stale.txt" + sentinel.parent.mkdir(parents=True) + sentinel.write_text("stale") + cfg = RankConfig( + attributes={"quality": ""}, + save_dir=str(tmp_path), + initial_rating_pass=False, + recursive=True, + recursive_rate_first_round=False, + n_rounds=1, + matches_per_round=1, + ) + task = Rank(cfg) + + async def fake_recursive(df, column_name, **kwargs): + assert not sentinel.exists() + return df.copy() + + task._run_recursive = fake_recursive + result = asyncio.run( + task.run( + pd.DataFrame({"id": ["a"], "text": ["A"]}), + column_name="text", + id_column="id", + reset_files=True, + ) + ) + assert list(result["id"]) == ["a"] + assert not sentinel.exists() + + +def test_catchup_avoids_repeating_committed_survivor_pairs(monkeypatch, tmp_path): + prompt_counts = [] + + async def fake_get_all_responses(*, identifiers, **kwargs): + prompt_counts.append(len(identifiers)) + return pd.DataFrame( + { + "Identifier": identifiers, + "Response": ['{"quality": "draw"}'] * len(identifiers), + } + ) + + monkeypatch.setattr( + "gabriel.tasks.rank.get_all_responses", fake_get_all_responses + ) + cfg = RankConfig( + attributes={"quality": ""}, + save_dir=str(tmp_path), + initial_rating_pass=False, + n_rounds=1, + matches_per_round=2, + ) + first = Rank(cfg) + original_pairs = [ + ("a", "b"), + ("c", "d"), + ("a", "x"), + ("b", "z"), + ("c", "y"), + ("d", "w"), + ] + first._generate_pairs = lambda *, texts_by_id, **kwargs: [ + ((id_a, texts_by_id[id_a]), (id_b, texts_by_id[id_b])) + for id_a, id_b in original_pairs + ] + full_ids = ["a", "b", "c", "d", "x", "z", "y", "w"] + asyncio.run( + first.run( + pd.DataFrame({"id": full_ids, "text": [value.upper() for value in full_ids]}), + column_name="text", + id_column="id", + reset_files=True, + ) + ) + resumed = Rank(cfg) + resumed.rng.choice = lambda values: values[0] + survivors = ["a", "b", "c", "d"] + result = asyncio.run( + resumed.run( + pd.DataFrame( + {"id": survivors, "text": [value.upper() for value in survivors]} + ), + column_name="text", + id_column="id", + reset_files=False, + ) + ) + + assert prompt_counts == [6, 2] + checkpoint = Rank._read_rank_checkpoint( + str(tmp_path / "rankings_round0.csv"), 1 + ) + current_pairs = [ + tuple(sorted((id_a, id_b))) + for id_a, id_b in zip(checkpoint["IdA"], checkpoint["IdB"]) + if id_a in survivors and id_b in survivors + ] + assert len(current_pairs) == len(set(current_pairs)) == 4 + assert {("a", "b"), ("c", "d")}.issubset(current_pairs) + assert result["quality_component"].nunique() == 1 + + +def test_interrupted_catchup_resumes_persisted_plan_without_orphaning_rows( + monkeypatch, tmp_path +): + phase = "initial" + planned_ids = [] + paid_resume_calls = 0 + + async def fake_get_all_responses(*, identifiers, save_path, **kwargs): + nonlocal paid_resume_calls + if phase == "initial": + return pd.DataFrame( + { + "Identifier": identifiers, + "Response": ['{"quality": "draw"}'] * len(identifiers), + } + ) + if phase == "interrupt": + planned_ids[:] = list(identifiers) + pd.DataFrame( + { + "Identifier": identifiers[:1], + "Response": ['{"quality": "draw"}'], + "Successful": [True], + } + ).to_csv(save_path, index=False) + raise RuntimeError("simulated partial catch-up") + assert list(identifiers) == planned_ids + staged = pd.read_csv(save_path) + completed = set(staged["Identifier"].astype(str)) + missing = [identifier for identifier in identifiers if identifier not in completed] + paid_resume_calls += len(missing) + appended = pd.DataFrame( + { + "Identifier": missing, + "Response": ['{"quality": "draw"}'] * len(missing), + "Successful": [True] * len(missing), + } + ) + combined = pd.concat([staged, appended], ignore_index=True) + combined.to_csv(save_path, index=False) + return combined + + monkeypatch.setattr( + "gabriel.tasks.rank.get_all_responses", fake_get_all_responses + ) + cfg = RankConfig( + attributes={"quality": ""}, + save_dir=str(tmp_path), + initial_rating_pass=False, + n_rounds=1, + matches_per_round=1, + ) + initial = Rank(cfg) + initial._generate_pairs = lambda *, texts_by_id, **kwargs: [ + (("a", texts_by_id["a"]), ("b", texts_by_id["b"])) + ] + asyncio.run( + initial.run( + pd.DataFrame({"id": ["a", "b"], "text": ["A", "B"]}), + column_name="text", + id_column="id", + reset_files=True, + ) + ) + expanded = pd.DataFrame( + { + "id": ["a", "b", "c", "d", "e", "f"], + "text": ["A", "B", "C", "D", "E", "F"], + } + ) + phase = "interrupt" + with pytest.raises(RuntimeError, match="partial catch-up"): + asyncio.run( + Rank(cfg).run( + expanded, + column_name="text", + id_column="id", + reset_files=False, + ) + ) + assert len(planned_ids) == 2 + assert (tmp_path / ".rankings_catchup_round0_plan.json").exists() + + phase = "resume" + metadata_path = tmp_path / "rankings_run_metadata.json" + metadata_before_rejected_resume = metadata_path.read_bytes() + with pytest.raises(ValueError, match="different input set"): + asyncio.run( + Rank(cfg).run( + pd.concat( + [ + expanded, + pd.DataFrame({"id": ["g"], "text": ["G"]}), + ], + ignore_index=True, + ), + column_name="text", + id_column="id", + reset_files=False, + ) + ) + assert paid_resume_calls == 0 + assert metadata_path.read_bytes() == metadata_before_rejected_resume + + resumed = Rank(cfg) + resumed.rng.choice = lambda values: values[-1] + asyncio.run( + resumed.run( + expanded, + column_name="text", + id_column="id", + reset_files=False, + ) + ) + + assert paid_resume_calls == 1 + checkpoint = Rank._read_rank_checkpoint( + str(tmp_path / "rankings_round0.csv"), 1 + ) + assert set(planned_ids).issubset(set(checkpoint["Identifier"])) + assert not (tmp_path / ".rankings_catchup_round0_plan.json").exists() + assert not (tmp_path / ".rankings_catchup_round0.csv").exists() + + +def test_empty_catchup_batch_state_without_staging_fails_closed( + monkeypatch, tmp_path +): + phase = "initial" + calls = 0 + + async def fake_get_all_responses(*, identifiers, save_path, **kwargs): + nonlocal calls + calls += 1 + if phase == "initial": + return pd.DataFrame( + { + "Identifier": identifiers, + "Response": ['{"quality": "draw"}'] * len(identifiers), + } + ) + Path(f"{save_path}.batch_state.json").write_text( + json.dumps({"batches": []}) + ) + raise RuntimeError("simulated post-batch pre-staging crash") + + monkeypatch.setattr( + "gabriel.tasks.rank.get_all_responses", fake_get_all_responses + ) + cfg = RankConfig( + attributes={"quality": ""}, + save_dir=str(tmp_path), + initial_rating_pass=False, + n_rounds=1, + matches_per_round=1, + ) + asyncio.run( + Rank(cfg).run( + pd.DataFrame({"id": ["a", "b"], "text": ["A", "B"]}), + column_name="text", + id_column="id", + reset_files=True, + ) + ) + expanded = pd.DataFrame( + {"id": ["a", "b", "c", "d"], "text": ["A", "B", "C", "D"]} + ) + phase = "catchup" + with pytest.raises(RuntimeError, match="pre-staging"): + asyncio.run( + Rank(cfg).run( + expanded, + column_name="text", + id_column="id", + reset_files=False, + ) + ) + boundary = calls + + with pytest.raises(ValueError, match="without a durable response checkpoint"): + asyncio.run( + Rank(cfg).run( + expanded, + column_name="text", + id_column="id", + reset_files=False, + ) + ) + assert calls == boundary diff --git a/tests/test_reset_files.py b/tests/test_reset_files.py index 1064fc7..e1b8b9c 100644 --- a/tests/test_reset_files.py +++ b/tests/test_reset_files.py @@ -1,5 +1,10 @@ import asyncio +import json +from types import SimpleNamespace + import pandas as pd +import pytest + from gabriel.utils import openai_utils @@ -57,6 +62,176 @@ def test_resume_treats_string_success_values_as_completed(tmp_path): assert len(df) == 3 +def test_successful_retry_replaces_failed_checkpoint_row(tmp_path): + save_path = tmp_path / "out.csv" + pd.DataFrame( + { + "Identifier": ["retry-me"], + "Response": [""], + "Web Search Sources": ["[]"], + "Time Taken": [0.1], + "Input Tokens": [1], + "Reasoning Tokens": [0], + "Output Tokens": [0], + "Reasoning Effort": ["default"], + "Successful": [False], + "Error Log": ["previous failure"], + } + ).to_csv(save_path, index=False) + calls = [] + + async def responder(prompt: str, **kwargs): + calls.append(prompt) + return ["retried successfully"] + + result = asyncio.run( + openai_utils.get_all_responses( + prompts=["prompt"], + identifiers=["retry-me"], + save_path=str(save_path), + response_fn=responder, + reset_files=False, + skip_tail_fails=False, + n_parallels=1, + ramp_up_seconds=0, + dynamic_timeout=False, + manage_rate_limits=False, + status_report_interval=None, + verbose=False, + ) + ) + + assert calls == ["prompt"] + assert len(result) == 1 + assert bool(result.iloc[0]["Successful"]) + persisted = pd.read_csv(save_path) + assert len(persisted) == 1 + assert persisted.iloc[0]["Identifier"] == "retry-me" + assert bool(persisted.iloc[0]["Successful"]) + + +def test_batch_submission_intent_prevents_ambiguous_resubmission( + monkeypatch, tmp_path +): + monkeypatch.setenv("OPENAI_API_KEY", "test-key") + class FakeFiles: + def __init__(self): + self.count = 0 + + async def create(self, **kwargs): + self.count += 1 + return SimpleNamespace(id=f"file-{self.count}") + + class FakeBatches: + def __init__(self): + self.create_calls = 0 + + async def create(self, **kwargs): + self.create_calls += 1 + if self.create_calls == 2: + raise ConnectionError("ambiguous submission failure") + return SimpleNamespace(id="batch-1") + + client = SimpleNamespace(files=FakeFiles(), batches=FakeBatches()) + monkeypatch.setattr(openai_utils, "_get_client", lambda *args: client) + save_path = tmp_path / "batch.csv" + + with pytest.raises(ConnectionError, match="ambiguous"): + asyncio.run( + openai_utils.get_all_responses( + prompts=["first", "second"], + identifiers=["1", "2"], + save_path=str(save_path), + use_batch=True, + batch_wait_for_completion=False, + max_batch_requests=1, + print_example_prompt=False, + verbose=False, + ) + ) + + state_path = tmp_path / "batch.csv.batch_state.json" + state = json.loads(state_path.read_text()) + assert state["batches"][0]["batch_id"] == "batch-1" + assert state["batches"][1]["status"] == "submitting" + create_calls = client.batches.create_calls + with pytest.raises(RuntimeError, match="unresolved Batch API submission"): + asyncio.run( + openai_utils.get_all_responses( + prompts=["first", "second"], + identifiers=["1", "2"], + save_path=str(save_path), + use_batch=True, + batch_wait_for_completion=False, + max_batch_requests=1, + print_example_prompt=False, + verbose=False, + ) + ) + assert client.batches.create_calls == create_calls + + +def test_completed_batch_state_is_retained_until_rows_are_durable( + monkeypatch, tmp_path +): + monkeypatch.setenv("OPENAI_API_KEY", "test-key") + class FakeFiles: + async def content(self, output_file_id): + return json.dumps( + { + "custom_id": "1", + "response": { + "body": {"output_text": "paid result"}, + "usage": {}, + }, + } + ) + + class FakeBatches: + async def retrieve(self, batch_id): + return SimpleNamespace( + status="completed", + output_file_id="output-1", + error_file_id=None, + ) + + client = SimpleNamespace(files=FakeFiles(), batches=FakeBatches()) + monkeypatch.setattr(openai_utils, "_get_client", lambda *args: client) + save_path = tmp_path / "batch.csv" + state_path = tmp_path / "batch.csv.batch_state.json" + original_state = { + "batches": [ + { + "batch_id": "batch-1", + "input_file_id": "input-1", + "total": 1, + "status": "submitted", + } + ] + } + state_path.write_text(json.dumps(original_state)) + + def fail_csv_write(self, *args, **kwargs): + raise OSError("simulated durable CSV failure") + + monkeypatch.setattr(pd.DataFrame, "to_csv", fail_csv_write) + with pytest.raises(OSError, match="durable CSV"): + asyncio.run( + openai_utils.get_all_responses( + prompts=["first"], + identifiers=["1"], + save_path=str(save_path), + use_batch=True, + batch_wait_for_completion=True, + batch_poll_interval=0, + print_example_prompt=False, + verbose=False, + ) + ) + + assert json.loads(state_path.read_text()) == original_state + + def test_resume_skip_tail_fails_returns_checkpoint_without_retry(tmp_path, capsys): save_path = tmp_path / "out.csv" total_rows = 5_001