overturn-trainer 0.1.0__tar.gz

overturn_trainer-0.1.0/.gitignore

@@ -0,0 +1,3 @@
+.venv/
+dist/
+__pycache__/

overturn_trainer-0.1.0/PKG-INFO

@@ -0,0 +1,38 @@
+Metadata-Version: 2.5
+Name: overturn-trainer
+Version: 0.1.0
+Summary: Train an open-weight model on your Overturn corrections in your own cloud, then register it for the gate
+Project-URL: Homepage, https://github.com/overrule-ai/overturn
+Project-URL: Documentation, https://github.com/overrule-ai/overturn/blob/main/docs/owned-models.md
+Author: Overturn contributors
+License-Expression: Apache-2.0
+Keywords: agents,fine-tuning,llm,lora,overturn
+Classifier: License :: OSI Approved :: Apache Software License
+Classifier: Programming Language :: Python :: 3
+Classifier: Typing :: Typed
+Requires-Python: >=3.10
+Requires-Dist: httpx>=0.27
+Provides-Extra: train
+Requires-Dist: accelerate>=0.33; extra == 'train'
+Requires-Dist: datasets>=2.20; extra == 'train'
+Requires-Dist: peft>=0.12; extra == 'train'
+Requires-Dist: torch>=2.3; extra == 'train'
+Requires-Dist: transformers>=4.44; extra == 'train'
+Requires-Dist: trl>=0.12; extra == 'train'
+Description-Content-Type: text/markdown
+
+# overturn-trainer
+
+Train an open-weight model on your [Overturn](https://overturn.run) corrections **in your own cloud**, then register it so Overturn gates it against your suite, shadows it on live traffic, and lets a human promote it.
+
+```bash
+pip install 'overturn-trainer[train]'
+export OVERTURN_API_KEY=ov_live_…
+
+overturn-trainer run --base meta-llama/Llama-3.1-8B-Instruct --agent ticket-triage --out ./triage-lora
+# serve ./triage-lora/adapter on vLLM (--enable-lora --lora-modules triage-lora=./triage-lora/adapter)
+# point the Overturn API at it: LOCAL_MODEL_BASE_URL=http://vllm:8000/v1
+overturn-trainer run … --register local/triage-lora     # or: overturn models register --base … --served-as local/triage-lora
+```
+
+Nothing but the export leaves Overturn; nothing but the registration goes back. `--backend dry` runs the whole pipeline without the ML stack (export, dataset, manifest, registration) — what CI tests.

overturn_trainer-0.1.0/README.md

@@ -0,0 +1,15 @@
+# overturn-trainer
+
+Train an open-weight model on your [Overturn](https://overturn.run) corrections **in your own cloud**, then register it so Overturn gates it against your suite, shadows it on live traffic, and lets a human promote it.
+
+```bash
+pip install 'overturn-trainer[train]'
+export OVERTURN_API_KEY=ov_live_…
+
+overturn-trainer run --base meta-llama/Llama-3.1-8B-Instruct --agent ticket-triage --out ./triage-lora
+# serve ./triage-lora/adapter on vLLM (--enable-lora --lora-modules triage-lora=./triage-lora/adapter)
+# point the Overturn API at it: LOCAL_MODEL_BASE_URL=http://vllm:8000/v1
+overturn-trainer run … --register local/triage-lora     # or: overturn models register --base … --served-as local/triage-lora
+```
+
+Nothing but the export leaves Overturn; nothing but the registration goes back. `--backend dry` runs the whole pipeline without the ML stack (export, dataset, manifest, registration) — what CI tests.

overturn_trainer-0.1.0/overturn_trainer/__init__.py

@@ -0,0 +1,7 @@
+"""overturn-trainer: fine-tune an open-weight model on your Overturn corrections, in your own cloud."""
+
+from .api import OverturnApi
+from .train import TrainResult, train
+
+__all__ = ["OverturnApi", "TrainResult", "train"]
+__version__ = "0.1.0"

overturn_trainer-0.1.0/overturn_trainer/api.py

@@ -0,0 +1,51 @@
+"""The two Overturn calls the trainer needs: pull the export, register the result."""
+
+from __future__ import annotations
+
+import json
+import os
+from typing import Any
+
+import httpx
+
+
+class OverturnApi:
+    def __init__(self, api_key: str | None = None, base_url: str | None = None) -> None:
+        self.api_key = api_key or os.environ.get("OVERTURN_API_KEY") or ""
+        self.base_url = (
+            base_url or os.environ.get("OVERTURN_BASE_URL") or "https://api.overturn.run"
+        ).rstrip("/")
+        if not self.api_key:
+            raise SystemExit("OVERTURN_API_KEY is not set (Settings → API keys)")
+        self._client = httpx.Client(timeout=120)
+
+    def _headers(self) -> dict[str, str]:
+        return {"authorization": f"Bearer {self.api_key}", "user-agent": "overturn-trainer/0.1.0"}
+
+    def export(self, fmt: str, **params: Any) -> list[dict[str, Any]]:
+        """`GET /v1/export?format=sft|dpo` as rows (one JSON object per JSONL line)."""
+        res = self._client.get(
+            f"{self.base_url}/v1/export",
+            params={"format": fmt, **{k: v for k, v in params.items() if v is not None}},
+            headers=self._headers(),
+        )
+        if not res.is_success:
+            raise SystemExit(f"export failed: {res.status_code} {res.text[:200]}")
+        return [json.loads(line) for line in res.text.splitlines() if line.strip()]
+
+    def export_stats(self, fmt: str, **params: Any) -> dict[str, Any]:
+        res = self._client.get(
+            f"{self.base_url}/v1/export/stats",
+            params={"format": fmt, **{k: v for k, v in params.items() if v is not None}},
+            headers=self._headers(),
+        )
+        if not res.is_success:
+            raise SystemExit(f"export stats failed: {res.status_code} {res.text[:200]}")
+        return res.json()
+
+    def register(self, body: dict[str, Any]) -> dict[str, Any]:
+        """`POST /v1/models/register`: the gate replays the suite against `servedAs` at once."""
+        res = self._client.post(f"{self.base_url}/v1/models/register", json=body, headers=self._headers())
+        if not res.is_success:
+            raise SystemExit(f"register failed: {res.status_code} {res.text[:300]}")
+        return res.json()

overturn_trainer-0.1.0/overturn_trainer/cli.py

@@ -0,0 +1,116 @@
+"""`overturn-trainer run`: export → fine-tune (LoRA) → save → register with Overturn for the gate."""
+
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+from typing import Any, cast
+
+from .api import OverturnApi
+from .train import train
+
+
+def build_parser() -> argparse.ArgumentParser:
+    p = argparse.ArgumentParser(prog="overturn-trainer", description=__doc__)
+    sub = p.add_subparsers(dest="cmd", required=True)
+    run = sub.add_parser(
+        "run", help="Export, train a LoRA, save it, and (optionally) register it for the gate"
+    )
+    run.add_argument(
+        "--base", required=True, help="Hugging Face model id, e.g. meta-llama/Llama-3.1-8B-Instruct"
+    )
+    run.add_argument(
+        "--out", default="./overturn-model", help="Output directory (adapter/, merged/, manifest)"
+    )
+    run.add_argument("--method", choices=["sft", "dpo"], default="sft")
+    run.add_argument("--epochs", type=int, default=3)
+    run.add_argument("--lora-r", type=int, default=16)
+    run.add_argument("--learning-rate", type=float, default=2e-4)
+    run.add_argument("--max-seq-length", type=int, default=4096)
+    run.add_argument("--merge", action="store_true", help="Also save merged full weights (merged/)")
+    run.add_argument("--agent", help="Only this agent's cases")
+    run.add_argument("--tag", help="Only cases with this tag")
+    run.add_argument("--since", help="Cases since (ISO date or 30d)")
+    run.add_argument("--min-rows", type=int, default=50, help="Refuse to train on fewer rows than this")
+    run.add_argument("--backend", choices=["trl", "dry"], default="trl")
+    run.add_argument(
+        "--register",
+        metavar="PROVIDER/MODEL",
+        help="Where the weights will answer from (e.g. local/triage-lora); registers the model so the gate runs",
+    )
+    run.add_argument("--name", help="Model name in Overturn")
+    run.add_argument("--api-key", help="Overturn API key (default: OVERTURN_API_KEY)")
+    run.add_argument("--base-url", help="Overturn API URL (default: OVERTURN_BASE_URL or hosted)")
+    return p
+
+
+def cmd_run(a: argparse.Namespace) -> int:
+    api = OverturnApi(a.api_key, a.base_url)
+    params: dict[str, Any] = {"agent": a.agent, "tag": a.tag, "since": a.since, "includeApproved": "true"}
+    stats = api.export_stats(a.method, **params)
+    rows_n = int(stats.get("rows", 0))
+    if rows_n < a.min_rows:
+        print(
+            f"{rows_n} {a.method} rows is under --min-rows {a.min_rows}. Until the corrections pile up, "
+            "put them in the prompt instead: overturn replay --examples 5",
+            file=sys.stderr,
+        )
+        return 2
+    rows = api.export(a.method, **params)
+    print(f"exported {len(rows)} {a.method} rows ({json.dumps(stats.get('byKind', {}))})")
+    result = train(
+        rows,
+        base_model=a.base,
+        output_dir=a.out,
+        method=a.method,
+        epochs=a.epochs,
+        lora_r=a.lora_r,
+        learning_rate=a.learning_rate,
+        max_seq_length=a.max_seq_length,
+        merge=a.merge,
+        backend=a.backend,
+    )
+    print(f"trained {result.base_model} ({result.backend}) in {result.seconds:.0f}s → {result.output_dir}")
+    if a.register:
+        filt: dict[str, Any] = {
+            k: v for k, v in {"agent": a.agent, "tag": a.tag, "since": a.since}.items() if v
+        }
+        m = api.register(
+            {
+                "name": a.name,
+                "baseModel": a.base,
+                "servedAs": a.register,
+                "method": a.method,
+                "epochs": a.epochs,
+                "filter": filt,
+                "datasetStats": stats,
+            }
+        )
+        gate: dict[str, Any] = m.get("gate") or {}
+        raw_reasons: Any = gate.get("reasons") or []
+        reasons: list[str] = [str(r) for r in cast(list[Any], raw_reasons)]
+        verdict = (
+            f", gate {'passed' if gate.get('pass') else 'failed'} ({'; '.join(reasons) or 'no regressions'})"
+            if gate
+            else ""
+        )
+        print(f"registered {m.get('name')} as {m.get('servedAs')}: status {m.get('status')}{verdict}")
+        return 0 if m.get("status") in ("passed", "trained", "gating") else 1
+    print(
+        "serve the adapter on any OpenAI-compatible server, then: "
+        f"overturn models register --base {a.base} --served-as local/<name>"
+        + (f" --agent {a.agent}" if a.agent else "")
+    )
+    return 0
+
+
+def main(argv: list[str] | None = None) -> int:
+    a = build_parser().parse_args(argv)
+    if a.cmd == "run":
+        return cmd_run(a)
+    return 2
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

overturn_trainer-0.1.0/overturn_trainer/train.py

@@ -0,0 +1,69 @@
+"""LoRA fine-tuning on Overturn's export. The `trl` backend needs the `[train]` extra and a GPU
+worth the name; the `dry` backend writes the dataset and a manifest and is what the tests run."""
+
+from __future__ import annotations
+
+import json
+import time
+from dataclasses import asdict, dataclass
+from pathlib import Path
+from typing import Any
+
+
+@dataclass
+class TrainResult:
+    backend: str
+    base_model: str
+    method: str
+    rows: int
+    epochs: int
+    lora_r: int
+    output_dir: str
+    seconds: float
+    merged: bool
+
+    def manifest(self) -> dict[str, Any]:
+        return asdict(self)
+
+
+def write_dataset(rows: list[dict[str, Any]], output_dir: Path, method: str) -> Path:
+    output_dir.mkdir(parents=True, exist_ok=True)
+    path = output_dir / f"train-{method}.jsonl"
+    with path.open("w", encoding="utf-8") as f:
+        for r in rows:
+            f.write(json.dumps(r, ensure_ascii=False) + "\n")
+    return path
+
+
+def train(
+    rows: list[dict[str, Any]],
+    *,
+    base_model: str,
+    output_dir: str | Path,
+    method: str = "sft",
+    epochs: int = 3,
+    lora_r: int = 16,
+    learning_rate: float = 2e-4,
+    max_seq_length: int = 4096,
+    merge: bool = False,
+    backend: str = "trl",
+) -> TrainResult:
+    """Fine-tune `base_model` on `rows` (Overturn's sft/dpo export) and save the adapter to `output_dir`."""
+    out = Path(output_dir)
+    t0 = time.time()
+    dataset = write_dataset(rows, out, method)
+    if backend == "dry":
+        result = TrainResult(
+            "dry", base_model, method, len(rows), epochs, lora_r, str(out), time.time() - t0, False
+        )
+    elif backend == "trl":
+        from .trl_backend import train_trl
+
+        train_trl(dataset, out, base_model, method, epochs, lora_r, learning_rate, max_seq_length, merge)
+        result = TrainResult(
+            "trl", base_model, method, len(rows), epochs, lora_r, str(out), time.time() - t0, merge
+        )
+    else:
+        raise SystemExit(f"unknown backend {backend!r} (trl | dry)")
+    (out / "overturn-manifest.json").write_text(json.dumps(result.manifest(), indent=2))
+    return result

overturn_trainer-0.1.0/overturn_trainer/trl_backend.py

@@ -0,0 +1,88 @@
+# pyright: basic
+"""The TRL/PEFT backend: LoRA SFT or DPO on the export. Imported lazily; needs `overturn-trainer[train]`
+and a GPU worth the name. Untyped third-party stack, hence the relaxed checker for this file only."""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any
+
+
+def train_trl(  # pragma: no cover - needs the ML stack and a GPU; exercised by hand, not in CI
+    dataset: Path,
+    out: Path,
+    base_model: str,
+    method: str,
+    epochs: int,
+    lora_r: int,
+    learning_rate: float,
+    max_seq_length: int,
+    merge: bool,
+) -> None:
+    try:
+        import torch  # type: ignore[import-not-found]
+        from datasets import load_dataset  # type: ignore[import-not-found]
+        from peft import LoraConfig  # type: ignore[import-not-found]
+        from transformers import AutoModelForCausalLM, AutoTokenizer  # type: ignore[import-not-found]
+    except ImportError as err:
+        raise SystemExit("the trl backend needs: pip install 'overturn-trainer[train]'") from err
+
+    tokenizer: Any = AutoTokenizer.from_pretrained(base_model)
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+    model: Any = AutoModelForCausalLM.from_pretrained(
+        base_model,
+        torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
+        device_map="auto",
+    )
+    peft_config: Any = LoraConfig(
+        r=lora_r,
+        lora_alpha=lora_r * 2,
+        lora_dropout=0.05,
+        target_modules="all-linear",
+        task_type="CAUSAL_LM",
+    )
+    data: Any = load_dataset("json", data_files=str(dataset), split="train")
+    if method == "dpo":
+        from trl import DPOConfig, DPOTrainer  # type: ignore[import-not-found]
+
+        trainer: Any = DPOTrainer(
+            model=model,
+            args=DPOConfig(
+                output_dir=str(out / "checkpoints"),
+                num_train_epochs=epochs,
+                learning_rate=learning_rate,
+                beta=0.1,
+                logging_steps=10,
+                bf16=torch.cuda.is_available(),
+                report_to=[],
+            ),
+            train_dataset=data,
+            processing_class=tokenizer,
+            peft_config=peft_config,
+        )
+    else:
+        from trl import SFTConfig, SFTTrainer  # type: ignore[import-not-found]
+
+        trainer = SFTTrainer(
+            model=model,
+            args=SFTConfig(
+                output_dir=str(out / "checkpoints"),
+                num_train_epochs=epochs,
+                learning_rate=learning_rate,
+                max_length=max_seq_length,
+                logging_steps=10,
+                bf16=torch.cuda.is_available(),
+                report_to=[],
+            ),
+            train_dataset=data,
+            processing_class=tokenizer,
+            peft_config=peft_config,
+        )
+    trainer.train()
+    trainer.save_model(str(out / "adapter"))
+    tokenizer.save_pretrained(str(out / "adapter"))
+    if merge:
+        merged: Any = trainer.model.merge_and_unload()
+        merged.save_pretrained(str(out / "merged"))
+        tokenizer.save_pretrained(str(out / "merged"))

overturn_trainer-0.1.0/package.json

@@ -0,0 +1,12 @@
+{
+  "name": "@overturn-ai/trainer",
+  "version": "0.1.0",
+  "private": true,
+  "description": "Turbo wrapper for the trainer (published to PyPI as `overturn-trainer`)",
+  "scripts": {
+    "setup": "uv sync --group dev",
+    "typecheck": "uv run --no-sync pyright",
+    "lint": "uv run --no-sync ruff check . && uv run --no-sync ruff format --check .",
+    "test": "uv run --no-sync pytest -q"
+  }
+}

overturn_trainer-0.1.0/pyproject.toml

@@ -0,0 +1,47 @@
+[project]
+name = "overturn-trainer"
+version = "0.1.0"
+description = "Train an open-weight model on your Overturn corrections in your own cloud, then register it for the gate"
+readme = "README.md"
+license = "Apache-2.0"
+requires-python = ">=3.10"
+authors = [{ name = "Overturn contributors" }]
+dependencies = ["httpx>=0.27"]
+keywords = ["llm", "fine-tuning", "lora", "agents", "overturn"]
+classifiers = [
+  "Programming Language :: Python :: 3",
+  "License :: OSI Approved :: Apache Software License",
+  "Typing :: Typed",
+]
+
+[project.optional-dependencies]
+# the actual training stack; the CLI runs without it in --backend dry (export + manifest + register)
+train = ["torch>=2.3", "transformers>=4.44", "peft>=0.12", "trl>=0.12", "datasets>=2.20", "accelerate>=0.33"]
+
+[project.scripts]
+overturn-trainer = "overturn_trainer.cli:main"
+
+[project.urls]
+Homepage = "https://github.com/overrule-ai/overturn"
+Documentation = "https://github.com/overrule-ai/overturn/blob/main/docs/owned-models.md"
+
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[tool.hatch.build.targets.wheel]
+packages = ["overturn_trainer"]
+
+[dependency-groups]
+dev = ["pytest>=8", "pyright>=1.1.400", "ruff>=0.6"]
+
+[tool.pyright]
+venvPath = "."
+venv = ".venv"
+include = ["overturn_trainer", "tests"]
+typeCheckingMode = "strict"
+reportMissingImports = false
+
+[tool.ruff]
+line-length = 110
+target-version = "py310"

overturn_trainer-0.1.0/tests/test_cli.py

@@ -0,0 +1,155 @@
+"""The trainer against a fake Overturn API: export in, manifest out, register with the gate's answer."""
+
+from __future__ import annotations
+
+import json
+import threading
+from http.server import BaseHTTPRequestHandler, HTTPServer
+from pathlib import Path
+from typing import Any
+
+import pytest
+
+from overturn_trainer.cli import main
+
+ROWS = [
+    {
+        "messages": [
+            {"role": "system", "content": "triage"},
+            {"role": "user", "content": f"ticket {i}"},
+            {"role": "assistant", "content": '{"priority":"P1"}'},
+        ]
+    }
+    for i in range(6)
+]
+
+
+CALLS: list[dict[str, Any]] = []
+
+
+class Handler(BaseHTTPRequestHandler):
+    def _json(self, status: int, body: Any) -> None:
+        data = json.dumps(body).encode()
+        self.send_response(status)
+        self.send_header("content-type", "application/json")
+        self.send_header("content-length", str(len(data)))
+        self.end_headers()
+        self.wfile.write(data)
+
+    def do_GET(self) -> None:
+        if self.path.startswith("/v1/export/stats"):
+            return self._json(200, {"rows": len(ROWS), "byKind": {"approve": 6}, "format": "sft"})
+        if self.path.startswith("/v1/export"):
+            CALLS.append({"export": self.path, "auth": self.headers.get("authorization")})
+            text = "\n".join(json.dumps(r) for r in ROWS) + "\n"
+            data = text.encode()
+            self.send_response(200)
+            self.send_header("content-type", "application/jsonl")
+            self.send_header("content-length", str(len(data)))
+            self.end_headers()
+            self.wfile.write(data)
+            return
+        self._json(404, {"message": "no"})
+
+    def do_POST(self) -> None:
+        length = int(self.headers.get("content-length") or 0)
+        body = json.loads(self.rfile.read(length) or b"{}")
+        if self.path == "/v1/models/register":
+            CALLS.append({"register": body})
+            return self._json(
+                201,
+                {
+                    "name": body.get("name") or "ext-1",
+                    "servedAs": body["servedAs"],
+                    "status": "passed",
+                    "gate": {"pass": True, "reasons": []},
+                },
+            )
+        self._json(404, {"message": "no"})
+
+    def log_message(self, format: str, *args: Any) -> None:
+        return
+
+
+@pytest.fixture()
+def server() -> Any:
+    srv = HTTPServer(("127.0.0.1", 0), Handler)
+    t = threading.Thread(target=srv.serve_forever, daemon=True)
+    t.start()
+    CALLS.clear()
+    yield f"http://127.0.0.1:{srv.server_port}"
+    srv.shutdown()
+
+
+def test_run_dry_exports_writes_manifest_and_registers(server: str, tmp_path: Path) -> None:
+    out = tmp_path / "model"
+    code = main(
+        [
+            "run",
+            "--base",
+            "Qwen/Qwen2.5-7B-Instruct",
+            "--out",
+            str(out),
+            "--backend",
+            "dry",
+            "--agent",
+            "ticket-triage",
+            "--min-rows",
+            "5",
+            "--register",
+            "local/triage-lora",
+            "--name",
+            "mine",
+            "--api-key",
+            "ov_test",
+            "--base-url",
+            server,
+        ]
+    )
+    assert code == 0
+    assert (out / "train-sft.jsonl").read_text().count("\n") == 6
+    manifest = json.loads((out / "overturn-manifest.json").read_text())
+    assert manifest == {
+        "backend": "dry",
+        "base_model": "Qwen/Qwen2.5-7B-Instruct",
+        "method": "sft",
+        "rows": 6,
+        "epochs": 3,
+        "lora_r": 16,
+        "output_dir": str(out),
+        "seconds": manifest["seconds"],
+        "merged": False,
+    }
+    exp = next(c for c in CALLS if "export" in c)
+    assert "format=sft" in exp["export"] and "agent=ticket-triage" in exp["export"]
+    assert exp["auth"] == "Bearer ov_test"
+    reg = next(c for c in CALLS if "register" in c)["register"]
+    assert reg == {
+        "name": "mine",
+        "baseModel": "Qwen/Qwen2.5-7B-Instruct",
+        "servedAs": "local/triage-lora",
+        "method": "sft",
+        "epochs": 3,
+        "filter": {"agent": "ticket-triage"},
+        "datasetStats": {"rows": 6, "byKind": {"approve": 6}, "format": "sft"},
+    }
+
+
+def test_refuses_under_min_rows(server: str, tmp_path: Path) -> None:
+    code = main(
+        [
+            "run",
+            "--base",
+            "x/y",
+            "--out",
+            str(tmp_path),
+            "--backend",
+            "dry",
+            "--api-key",
+            "k",
+            "--base-url",
+            server,
+        ]
+    )
+    assert code == 2
+    assert not (tmp_path / "overturn-manifest.json").exists()