No description
Find a file
2026-08-05 17:10:24 -04:00
examples coherence checking 2026-08-05 14:52:09 -04:00
src/cai_tournament Update src/cai_tournament/coherence.py 2026-08-05 17:10:24 -04:00
tests coherence checking 2026-08-05 14:52:09 -04:00
.gitignore initslop 2026-08-05 13:19:34 -04:00
.python-version initslop 2026-08-05 13:19:34 -04:00
pyproject.toml lora support for example 2026-08-05 13:49:51 -04:00
README.md coherence checking 2026-08-05 14:52:09 -04:00
uv.lock coherence checking 2026-08-05 14:52:09 -04:00

cai-tournament

Tournament-style Constitutional AI reward scoring for RL training. Rollouts for the same prompt are paired off, a judge model picks a winner per pair against a randomly sampled constitutional principle, and winners advance through a bracket. Every win (including byes) is one point; a rollout's reward is its wins divided by the number of rounds in its bracket. Half of all rollouts get at least one point, which gives a dense reward signal.

The core package has no dependency on any RL environment library. Integrations live in cai_tournament.adapters and are thin: all an adapter does is flatten its library's generations into plain-text Rollout objects, call TournamentScorer.score(), and write the rewards back.

Installation

pip install .                    # core only
pip install ".[verifiers]"       # + the verifiers adapter
pip install ".[trl]"             # + the TRL adapter (GRPOTrainer)

Core usage

import asyncio

from cai_tournament import OpenAIJudge, Rollout, TournamentScorer, load_constitution

judge = OpenAIJudge(
    model="my-judge-model",
    base_url="http://localhost:8000/v1",  # any OpenAI-compatible endpoint
)
scorer = TournamentScorer(
    principles=load_constitution("const.txt"),
    judge=judge,
    max_concurrent_judges=16,
    max_concurrent_tournaments=4,
)

rollouts = [
    Rollout(example_id=0, prompt="...", completion="...", system_prompt="..."),
    ...
]

result = asyncio.run(scorer.score(rollouts))
result.rewards   # float per rollout, aligned with the input list
result.wins      # raw win counts
result.history   # example_id -> list of per-match records (round, principle, winner, rationale)

A rollout group with a single response gets reward 1.0 without any judging. If the judge times out, errors, or returns an unparseable verdict, that match falls back to a random winner and the reason is recorded in the rationale, so a flaky judge never aborts a training run.

An auxiliary coherence checker can veto individual responses independently of the bracket: pass any CoherenceChecker (for example OpenAICoherenceChecker, pointed at the judge model) to TournamentScorer, and every completion is checked on its own, concurrently with the tournaments. Completions judged incoherent get reward 0.0 whatever their bracket result. Checker failures (timeouts, API errors, unparseable verdicts) count as equivalent to "coherent", so a flaky checker never silently zeroes rewards. Per-rollout verdicts are in result.coherence (empty when no checker is configured).

from cai_tournament import OpenAICoherenceChecker, TournamentScorer

scorer = TournamentScorer(
    principles=load_constitution("const.txt"),
    judge=judge,
    coherence_checker=OpenAICoherenceChecker(
        model="my-judge-model",
        base_url="http://localhost:8000/v1",
    ),
)

verifiers adapter

from cai_tournament.adapters.verifiers_v0 import load_environment

env = load_environment(
    jsonl_path="data.jsonl",           # or dataset_name=... for HuggingFace
    constitution_path="const.txt",
    judge_model="my-judge-model",
    judge_base_url="http://localhost:8000/v1",
)

The adapter wraps a vf.SingleTurnEnv: it lets the base env generate unscored, runs tournaments over the rollouts of each example, and overwrites results.reward. Per-rollout match records are stored in state["reward_breakdown"]["tournament"].

TRL adapter (GRPO)

For TRL's GRPOTrainer, the adapter plugs in as a custom reward function: the trainer hands it the whole generation batch (each prompt repeated num_generations times), and completions sharing an example_id play one bracket. The callable is synchronous from the trainer's point of view (TRL versions differ in whether they await async reward functions); judge calls still run concurrently inside each call. See examples/train_grpo.py for a full script; the minimal wiring is:

from trl import GRPOConfig, GRPOTrainer

from cai_tournament import (
    OpenAICoherenceChecker,
    OpenAIJudge,
    TournamentScorer,
    load_constitution,
)
from cai_tournament.adapters.trl import TournamentReward

scorer = TournamentScorer(
    principles=load_constitution("const.txt"),
    judge=OpenAIJudge(model="my-judge-model", base_url="http://localhost:8000/v1"),
    coherence_checker=OpenAICoherenceChecker(
        model="my-judge-model", base_url="http://localhost:8000/v1",
    ),  # optional: zeroes reward for incoherent responses
)

trainer = GRPOTrainer(
    model="Qwen/Qwen2.5-0.5B-Instruct",
    args=GRPOConfig(num_generations=8, ...),  # num_generations = bracket size
    reward_funcs=TournamentReward(scorer),
    train_dataset=dataset,  # "prompt" column; optional "example_id" column
)
trainer.train()

If the dataset has an example_id column the trainer passes it through and it determines the brackets; otherwise completions with identical prompts are grouped. The reward also reports tournament/frac_rewarded, tournament/num_tournaments, tournament/judge_fallbacks, and tournament/frac_incoherent (when a coherence checker is configured) through the trainer's log_metric hook, and a tournament_wins column through log_extra.

Writing an adapter for another RL env library

The contract is three steps:

  1. Let the base library generate rollouts without scoring them.
  2. Build one Rollout per generation: example_id (which prompt it came from), prompt, completion, and optionally system_prompt, all as plain text.
  3. await scorer.score(rollouts), assign result.rewards back to the generations in the same order, and optionally stash result.history / result.wins wherever the library keeps per-rollout metadata.

If the default OpenAIJudge does not fit (a different API, a local model, a mock for tests), implement the Judge protocol instead. Similarly, to use a custom coherence checking backend, implement CoherenceChecker:

from cai_tournament import JudgeDecision

class MyJudge:
    async def judge(self, prompt, response_a, response_b, principle, system_prompt=""):
        ...
        return JudgeDecision(winner="A", rationale="...")

and pass it to TournamentScorer(principles=..., judge=MyJudge()). For a custom coherence checker, pass it as coherence_checker=MyChecker().

Package layout

  • cai_tournament.typesRollout, JudgeDecision, MatchResult, TournamentResult
  • cai_tournament.constitution — constitution file loading and principle sampling
  • cai_tournament.coherence — coherence prompt building, verdict parsing, CoherenceChecker protocol, OpenAICoherenceChecker
  • cai_tournament.judge — judge prompt building, response parsing, Judge protocol, OpenAIJudge
  • cai_tournament.tournament — bracket logic and TournamentScorer
  • cai_tournament.data — JSONL loading with category labels, balanced batch sampling
  • cai_tournament.messages — flattening chat-style messages to plain text
  • cai_tournament.adapters.verifiers — the verifiers integration
  • cai_tournament.adapters.trl — the TRL GRPOTrainer integration

Development

uv sync          # installs the package (editable) plus pytest
uv run pytest    # runs the test suite
uv build         # builds sdist and wheel