ruling 0.2.0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
ruling/cli.py ADDED
@@ -0,0 +1,397 @@
1
+ """Command line: serve the API, ask one question, evaluate, calibrate, or build datasets."""
2
+
3
+ from __future__ import annotations
4
+
5
+ import argparse
6
+ import dataclasses
7
+ import json
8
+ import logging
9
+ from pathlib import Path
10
+
11
+ import numpy as np
12
+ import uvicorn
13
+
14
+ from ruling import corpus, states
15
+ from ruling.against import collect_live
16
+ from ruling.build import build_every, build_sst5, build_typesafe, build_wanli, write
17
+ from ruling.calibration import Calibration
18
+ from ruling.config import Settings, configure_memory
19
+ from ruling.engine import Engine
20
+ from ruling.evals import collect, fit, load_dataset, predict, report, write_json, write_predictions
21
+ from ruling.questions import Choice, Noul, Score, SystemOneRequest
22
+ from ruling.hosted import OpenRouterEngine
23
+ from ruling.registry import build_engine
24
+ from ruling.server import app_from_env
25
+
26
+
27
+ # The training settings a retraining run may set; the LoRA layout always follows the adapter it starts from.
28
+ RETRAIN_OPTIONS = (("steps", int), ("batch-size", int), ("learning-rate", float), ("warmup", int), ("eval-every", int),
29
+ ("valid-fraction", float), ("seed", int))
30
+
31
+
32
+ def bundled_evaluations() -> list[Path]:
33
+ """The labeled evaluation datasets shipped in `datasets/`, excluding row-selection manifests."""
34
+ folder = Path(__file__).resolve().parents[1] / "datasets"
35
+ return sorted(p for p in folder.glob("*.jsonl") if not p.name.endswith("-selection.jsonl"))
36
+
37
+
38
+ def main(argv: list[str] | None = None) -> None:
39
+ parser = argparse.ArgumentParser(prog="ruling")
40
+ commands = parser.add_subparsers(dest="command", required=True)
41
+
42
+ commands.add_parser("serve", help="run the HTTP API")
43
+
44
+ ask = commands.add_parser("ask", help="answer questions once, without a server")
45
+ ask.add_argument("state", help="state text, or @path to read a file")
46
+ ask.add_argument("--choice", action="append", default=[], metavar="ID=Q|a,b,c",
47
+ help="choice question: id=instructions|key1,key2,...")
48
+ ask.add_argument("--score", action="append", default=[], metavar="ID=Q|lo,mid,hi",
49
+ help="score question: id=instructions|level1,level2,...")
50
+ ask.add_argument("--noul", action="append", default=[], metavar="ID=STATEMENT")
51
+ _engine_overrides(ask)
52
+
53
+ for name, help_text in (("eval", "score a labeled dataset and report metrics"),
54
+ ("calibrate", "fit per-type temperatures on a labeled dataset")):
55
+ sub = commands.add_parser(name, help=help_text)
56
+ sub.add_argument("dataset", type=Path)
57
+ sub.add_argument("--output", type=Path, help="write JSON here instead of stdout")
58
+ _engine_overrides(sub)
59
+ sub.add_argument("--think-tokens", type=int,
60
+ help="let a decoder model reason up to this many tokens before each answer (slow; for teachers)")
61
+ if name == "eval":
62
+ sub.add_argument("--predictions", type=Path, help="also write each question's distribution as JSONL")
63
+ sub.add_argument("--against", choices=["typesafe"],
64
+ help="also send every record to the hosted TypeSafe API (TYPESAFE_API_KEY) and report agreement")
65
+ sub.add_argument("--typesafe-model", help="model alias for --against typesafe; the service default when omitted")
66
+
67
+ predicted = commands.add_parser("predict", help="answer every question of every row, labels ignored, as JSONL")
68
+ predicted.add_argument("rows", type=Path, help="records or hosted-platform rows, .jsonl or a .json array")
69
+ predicted.add_argument("--output", type=Path, required=True,
70
+ help="one line per row, in order; a run that stops resumes after the last line written")
71
+ _engine_overrides(predicted)
72
+
73
+ build = commands.add_parser("build-dataset", help="download a pinned public dataset into ruling format")
74
+ build.add_argument("name", choices=["wanli", "sst5", "typesafe", "every", "dev"])
75
+ build.add_argument("--output", type=Path, required=True)
76
+ build.add_argument("--limit", type=int, default=300, help="rows to sample (sst5), or rows per source (dev)")
77
+
78
+ data = commands.add_parser("build-corpus", help="build training records from open data, program state, or synthetic worlds")
79
+ data.add_argument("kind", choices=["public", "states", "worlds", "questions", "answers"])
80
+ data.add_argument("--output", type=Path, required=True)
81
+ data.add_argument("--count", type=int, help="records per public source, or total states, worlds, or question sets")
82
+ data.add_argument("--seed", type=int, default=0)
83
+ data.add_argument("--catalog", type=Path, default=Path(__file__).resolve().parents[1] / "datasets" / "worlds.json")
84
+ data.add_argument("--writer", help="writer and judge model for worlds; defaults to RULING_MODEL")
85
+ data.add_argument("--batch-size", type=int, default=32)
86
+ data.add_argument("--min-probability", type=float, default=0.05,
87
+ help="drop a world record when the judge gives any labeled answer less than this")
88
+ data.add_argument("--temperature", type=float, default=0.9)
89
+ data.add_argument("--states", type=Path, nargs="*", default=[], help="record files whose states get questions")
90
+ data.add_argument("--questions", type=Path, help="question records to answer (answers)")
91
+ data.add_argument("--teacher", choices=["model", "typesafe"], default="model",
92
+ help="answer with the --writer model (local, or openrouter:<model>) or the hosted TypeSafe API")
93
+ data.add_argument("--questions-per-state", type=int, default=6)
94
+ data.add_argument("--max-options", type=int, default=8)
95
+ data.add_argument("--rotations", type=int, default=3, help="option orderings the teacher averages")
96
+ data.add_argument("--think-tokens", type=int, help="let the teacher reason up to this many tokens before each answer")
97
+ data.add_argument("--holdout", type=Path, nargs="*",
98
+ default=bundled_evaluations(),
99
+ help="evaluation datasets whose states are never distilled")
100
+
101
+ train = commands.add_parser("train", help="LoRA fine-tune the answer readout on training records")
102
+ train.add_argument("--data", type=Path, nargs="+", required=True)
103
+ train.add_argument("--output", type=Path, required=True)
104
+ train.add_argument("--holdout", type=Path, nargs="*",
105
+ default=bundled_evaluations(),
106
+ help="evaluation datasets whose states must not appear in training data")
107
+ train.add_argument("--model", help="base model; defaults to RULING_MODEL")
108
+ for name, kind in (("valid-fraction", float), ("steps", int), ("batch-size", int), ("learning-rate", float), ("max-grad-norm", float),
109
+ ("warmup", int), ("lora-layers", int), ("lora-rank", int), ("lora-scale", float),
110
+ ("lora-dropout", float), ("max-tokens", int), ("eval-every", int), ("valid-examples", int), ("seed", int),
111
+ ("label-smoothing", float), ("brier-weight", float), ("ranking-weight", float), ("ranking-margin", float),
112
+ ("max-image-tokens", int)):
113
+ train.add_argument(f"--{name}", type=kind)
114
+ train.add_argument("--no-grad-checkpoint", dest="grad_checkpoint", action="store_false",
115
+ help="keep all activations; faster but needs far more memory")
116
+ train.add_argument("--resume", action="store_true",
117
+ help="continue from the checkpoint in --output; only --steps may differ from the original run")
118
+ train.add_argument("--from-adapter", type=Path,
119
+ help="start a fresh run from this adapter's weights; its LoRA settings must match")
120
+
121
+ export = commands.add_parser("export-adapter", help="write a ruling adapter in PEFT layout for transformers")
122
+ export.add_argument("adapter", type=Path, help="a LoRA directory from `ruling train`")
123
+ export.add_argument("--output", type=Path, required=True)
124
+ export.add_argument("--base-model", required=True, help="Hugging Face id of the unquantized base PEFT should load it on")
125
+
126
+ imported = commands.add_parser("import-adapter", help="write a PEFT LoRA adapter in the layout ruling serves")
127
+ imported.add_argument("peft", type=Path, help="a directory holding adapter_model.safetensors and adapter_config.json")
128
+ imported.add_argument("--output", type=Path, required=True)
129
+
130
+ decisions = commands.add_parser("export-decisions",
131
+ help="write the served decisions that have outcomes as training records")
132
+ decisions.add_argument("--ledger", type=Path, required=True, help="the server's RULING_LEDGER file")
133
+ decisions.add_argument("--output", type=Path, required=True)
134
+
135
+ again = commands.add_parser("retrain", help="train on the ledger's resolved decisions and promote the result "
136
+ "only if it beats the served adapter on held-out data")
137
+ again.add_argument("--ledger", type=Path, required=True, help="the server's RULING_LEDGER file")
138
+ again.add_argument("--adapters", type=Path, required=True,
139
+ help="versions v1, v2, ... and `current`, the one to serve: RULING_ADAPTER=<this>/current")
140
+ again.add_argument("--holdout", type=Path, nargs="*", default=bundled_evaluations(),
141
+ help="labeled sets both versions are scored on; their states are never trained on")
142
+ again.add_argument("--model", help="base model; defaults to RULING_MODEL")
143
+ again.add_argument("--min-records", type=int, default=100,
144
+ help="do nothing until this many decisions have outcomes")
145
+ for name, kind in RETRAIN_OPTIONS:
146
+ again.add_argument(f"--{name}", type=kind)
147
+
148
+ encoder = commands.add_parser("train-encoder", help="fully fine-tune a decision encoder on training records")
149
+ encoder.add_argument("--data", type=Path, nargs="+", required=True)
150
+ encoder.add_argument("--output", type=Path, required=True)
151
+ encoder.add_argument("--base", required=True, help="a ModernBERT-architecture encoder repository")
152
+ encoder.add_argument("--revision", required=True, help="the base repository commit to pin")
153
+ encoder.add_argument("--holdout", type=Path, nargs="*",
154
+ default=bundled_evaluations(),
155
+ help="evaluation datasets whose states must not appear in training data")
156
+ for name, kind in (("valid-fraction", float), ("steps", int), ("batch-tokens", int), ("learning-rate", float),
157
+ ("head-learning-rate", float), ("weight-decay", float), ("warmup", int), ("max-tokens", int),
158
+ ("eval-every", int), ("valid-records", int), ("seed", int)):
159
+ encoder.add_argument(f"--{name}", type=kind)
160
+ encoder.add_argument("--no-grad-checkpoint", dest="grad_checkpoint", action="store_false",
161
+ help="store all activations; faster but uses far more memory")
162
+
163
+ args = parser.parse_args(argv)
164
+ logging.basicConfig(level=logging.INFO, format="%(asctime)s %(name)s %(message)s")
165
+ configure_memory(Settings.from_env())
166
+
167
+ if args.command == "serve":
168
+ settings = Settings.from_env()
169
+ uvicorn.run(app_from_env(), host=settings.host, port=settings.port)
170
+ return
171
+ if args.command == "build-corpus":
172
+ if args.kind != "answers" and args.count is None:
173
+ parser.error(f"build-corpus {args.kind} needs --count")
174
+ if args.kind == "answers" and args.questions is None:
175
+ parser.error("build-corpus answers needs --questions")
176
+ if args.kind == "questions":
177
+ from ruling.distill import write_questions
178
+ from ruling.train import _state_key
179
+ held = {_state_key(r.state) for path in args.holdout for r in load_dataset(path)}
180
+ pool = {}
181
+ for path in args.states:
182
+ for record in load_dataset(path):
183
+ key = _state_key(record.state)
184
+ if key not in held:
185
+ pool.setdefault(key, (record.family or path.stem, record.state))
186
+ chosen = [list(pool.values())[i] for i in np.random.default_rng(args.seed).permutation(len(pool))[:args.count]]
187
+ settings = dataclasses.replace(Settings.from_env(), model=args.writer or Settings.from_env().model, rotations=1,
188
+ prior_debias=False)
189
+ writer = build_engine(settings)
190
+ if isinstance(writer, OpenRouterEngine):
191
+ generate = lambda prompts: writer.write(prompts, max_tokens=1200)
192
+ else:
193
+ from ruling.distill import local_generator
194
+ generate = local_generator(writer, args.seed, args.temperature)
195
+ _stream(args.output, write_questions(chosen, generate, args.questions_per_state, args.batch_size, args.max_options),
196
+ append=False)
197
+ _report_spend(writer)
198
+ return
199
+ if args.kind == "answers":
200
+ from ruling.distill import unanswered
201
+ records = load_dataset(args.questions)
202
+ pending = unanswered(records, args.output)
203
+ print(f"{len(records) - len(pending)} already answered, {len(pending)} to go")
204
+ if args.teacher == "typesafe":
205
+ from ruling.against import live_answers
206
+ answered = (pending[i].model_copy(update={"references": {qid: {"target": t} for qid, t in answers.items()}})
207
+ for i, answers in live_answers(pending))
208
+ else:
209
+ from ruling.distill import answer_locally
210
+ teacher = build_engine(dataclasses.replace(Settings.from_env(), model=args.writer or Settings.from_env().model,
211
+ rotations=args.rotations, prior_debias=False))
212
+ if args.think_tokens:
213
+ from ruling.thinking import ThinkingEngine
214
+ teacher = ThinkingEngine(teacher, args.think_tokens)
215
+ answered = answer_locally(pending, teacher)
216
+ _stream(args.output, ([r] for r in answered), append=True)
217
+ if args.teacher == "model":
218
+ _report_spend(teacher)
219
+ return
220
+ if args.kind == "public":
221
+ records = corpus.build(args.count, args.seed)
222
+ elif args.kind == "states":
223
+ records = states.generate(args.count, args.seed)
224
+ else:
225
+ from ruling.worlds import Catalog, synthesize
226
+ model = args.writer or Settings.from_env().model
227
+ engine = Engine(model, Calibration(), Settings.from_env().max_input_tokens, rotations=1, prior_debias=False)
228
+ written = 0
229
+ with args.output.open("w") as out:
230
+ for accepted in synthesize(Catalog.load(args.catalog), engine, args.count, args.seed, args.batch_size,
231
+ args.min_probability, args.temperature):
232
+ out.writelines(r.model_dump_json() + "\n" for r in accepted)
233
+ out.flush()
234
+ written += len(accepted)
235
+ print(f"wrote {written} records to {args.output}")
236
+ return
237
+ write([r.model_dump(mode="json") for r in records], args.output)
238
+ print(f"wrote {len(records)} records to {args.output}")
239
+ return
240
+ if args.command == "train":
241
+ from ruling.train import Trainer, TrainConfig
242
+ options = {name.replace("-", "_"): getattr(args, name.replace("-", "_"))
243
+ for name in ("valid-fraction", "steps", "batch-size", "learning-rate", "max-grad-norm", "warmup", "lora-layers", "lora-rank",
244
+ "lora-scale", "lora-dropout", "max-tokens", "eval-every", "valid-examples", "seed",
245
+ "label-smoothing", "brier-weight", "ranking-weight", "ranking-margin", "max-image-tokens")}
246
+ config = TrainConfig(model=args.model or Settings.from_env().model, data=args.data, output=args.output,
247
+ holdout=args.holdout, grad_checkpoint=args.grad_checkpoint, from_adapter=args.from_adapter,
248
+ **{k: v for k, v in options.items() if v is not None})
249
+ Trainer(config, resume=args.resume).run()
250
+ return
251
+ if args.command == "retrain":
252
+ from ruling.ledger import DecisionLedger
253
+ from ruling.retrain import retrain
254
+ settings = Settings.from_env()
255
+ options = {name.replace("-", "_"): getattr(args, name.replace("-", "_")) for name, _ in RETRAIN_OPTIONS}
256
+ retrain(DecisionLedger(args.ledger), args.adapters, args.model or settings.model, args.holdout, settings,
257
+ args.min_records, {k: v for k, v in options.items() if v is not None})
258
+ return
259
+ if args.command == "export-adapter":
260
+ from ruling.peft import export_adapter
261
+ export_adapter(args.adapter, args.output, args.base_model)
262
+ print(f"wrote {args.output}")
263
+ return
264
+ if args.command == "import-adapter":
265
+ from ruling.peft import import_adapter
266
+ import_adapter(args.peft, args.output)
267
+ print(f"wrote {args.output}")
268
+ return
269
+ if args.command == "export-decisions":
270
+ from ruling.ledger import DecisionLedger
271
+ records = list(DecisionLedger(args.ledger).records())
272
+ args.output.write_text("".join(r.model_dump_json(exclude_none=True) + "\n" for r in records))
273
+ print(f"wrote {len(records)} records to {args.output}")
274
+ return
275
+ if args.command == "train-encoder":
276
+ from ruling.decision_train import DecisionTrainConfig, DecisionTrainer
277
+ names = ("valid_fraction", "steps", "batch_tokens", "learning_rate", "head_learning_rate", "weight_decay", "warmup",
278
+ "max_tokens", "eval_every", "valid_records", "seed")
279
+ config = DecisionTrainConfig(base=args.base, revision=args.revision, data=args.data, holdout=args.holdout,
280
+ output=args.output, grad_checkpoint=args.grad_checkpoint,
281
+ **{n: getattr(args, n) for n in names if getattr(args, n) is not None})
282
+ DecisionTrainer(config).run()
283
+ return
284
+ if args.command == "build-dataset":
285
+ builders = {"wanli": build_wanli, "sst5": lambda: build_sst5(args.limit),
286
+ "typesafe": build_typesafe, "every": build_every,
287
+ "dev": lambda: [r.model_dump(mode="json") for r in
288
+ corpus.build(args.limit, seed=1, split="dev") + states.generate(args.limit, seed=10_001)]}
289
+ records = builders[args.name]()
290
+ write(records, args.output)
291
+ print(f"wrote {len(records)} records to {args.output}")
292
+ return
293
+
294
+ engine = build_engine(_settings(args))
295
+ if getattr(args, "think_tokens", None):
296
+ from ruling.thinking import ThinkingEngine
297
+ engine = ThinkingEngine(engine, args.think_tokens)
298
+ if args.command == "ask":
299
+ state = Path(args.state[1:]).read_text() if args.state.startswith("@") else args.state
300
+ try:
301
+ state = json.loads(state)
302
+ except json.JSONDecodeError:
303
+ pass
304
+ questions = {}
305
+ for spec in args.choice:
306
+ qid, instructions, items = _parse_spec(spec)
307
+ questions[qid] = Choice(instructions=instructions, criteria={k: None for k in items})
308
+ for spec in args.score:
309
+ qid, instructions, items = _parse_spec(spec)
310
+ questions[qid] = Score(instructions=instructions, criteria=items)
311
+ for spec in args.noul:
312
+ qid, _, statement = spec.partition("=")
313
+ questions[qid] = Noul(instructions=statement)
314
+ if not questions:
315
+ parser.error("ask needs at least one --choice, --score, or --noul")
316
+ response = engine.evaluate(SystemOneRequest(state=state, questions=questions))
317
+ print(response.model_dump_json(indent=2))
318
+ return
319
+
320
+ if args.command == "predict":
321
+ records = load_dataset(args.rows)
322
+ done = sum(1 for line in args.output.open() if line.strip()) if args.output.exists() else 0
323
+ with args.output.open("a") as out:
324
+ for line in predict(engine, records[done:]):
325
+ out.write(json.dumps(line) + "\n")
326
+ out.flush()
327
+ print(f"wrote {len(records) - done} rows to {args.output} ({done} were already there)")
328
+ return
329
+
330
+ records = load_dataset(args.dataset)
331
+ live = collect_live(records, model=args.typesafe_model) if getattr(args, "against", None) else None
332
+ observations = collect(engine, records, live)
333
+ header = {"model": engine.model_id, "revision": engine.revision, "adapter": engine.adapter, "rotations": engine.rotations,
334
+ "prior_debias": engine.prior_debias, "dataset": str(args.dataset)}
335
+ _report_spend(engine)
336
+ if args.command == "eval":
337
+ data = {**header, **report(observations, engine.calibration)}
338
+ if args.predictions:
339
+ write_predictions(observations, engine.calibration, args.predictions)
340
+ else:
341
+ calibration = fit(observations, engine.provenance)
342
+ data = {**header, "temperatures": calibration.temperatures,
343
+ "before": report(observations, engine.calibration)["question_types"],
344
+ "after": report(observations, calibration)["question_types"]}
345
+ if args.output:
346
+ calibration.save(args.output)
347
+ print(f"wrote {args.output}")
348
+ print(json.dumps({k: data[k] for k in ("temperatures", "before", "after")}, indent=2))
349
+ return
350
+ if args.output:
351
+ write_json(args.output, data)
352
+ print(f"wrote {args.output}")
353
+ else:
354
+ print(json.dumps(data, indent=2))
355
+
356
+
357
+ def _report_spend(engine) -> None:
358
+ if isinstance(engine, OpenRouterEngine):
359
+ print(f"OpenRouter spend this run: ${engine.budget.spent_usd:.4f}")
360
+
361
+
362
+ def _stream(output: Path, batches, append: bool) -> None:
363
+ """Write record batches as they arrive, so an interrupted run keeps everything finished so far."""
364
+ written = 0
365
+ with output.open("a" if append else "w") as out:
366
+ for records in batches:
367
+ out.writelines(r.model_dump_json() + "\n" for r in records)
368
+ out.flush()
369
+ written += len(records)
370
+ print(f"wrote {written} records to {output}")
371
+
372
+
373
+ def _engine_overrides(sub: argparse.ArgumentParser) -> None:
374
+ sub.add_argument("--model", help="override RULING_MODEL")
375
+ sub.add_argument("--rotations", type=int, help="override RULING_ROTATIONS")
376
+ sub.add_argument("--prior-debias", choices=["on", "off"], help="override RULING_PRIOR_DEBIAS")
377
+ sub.add_argument("--calibration", type=Path, help="override RULING_CALIBRATION")
378
+ sub.add_argument("--adapter", type=Path, help="override RULING_ADAPTER, a LoRA directory from `ruling train`")
379
+
380
+
381
+ def _settings(args: argparse.Namespace) -> Settings:
382
+ overrides = {
383
+ "model": args.model,
384
+ "rotations": args.rotations,
385
+ "prior_debias": None if args.prior_debias is None else args.prior_debias == "on",
386
+ "calibration_path": args.calibration,
387
+ "adapter_path": args.adapter,
388
+ }
389
+ return dataclasses.replace(Settings.from_env(), **{k: v for k, v in overrides.items() if v is not None})
390
+
391
+
392
+ def _parse_spec(spec: str) -> tuple[str, str, list[str]]:
393
+ qid, _, rest = spec.partition("=")
394
+ instructions, _, items = rest.partition("|")
395
+ if not (qid and instructions and items):
396
+ raise SystemExit(f"bad question spec {spec!r}; expected id=instructions|item1,item2,...")
397
+ return qid, instructions, [item.strip() for item in items.split(",") if item.strip()]
ruling/codes.py ADDED
@@ -0,0 +1,30 @@
1
+ """Single-token answer codes, verified against the tokenizer at load time."""
2
+
3
+ from __future__ import annotations
4
+
5
+ from dataclasses import dataclass
6
+
7
+ CANDIDATES = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789"
8
+
9
+
10
+ @dataclass(frozen=True)
11
+ class AnswerCodes:
12
+ codes: list[str]
13
+ token_ids: list[int]
14
+
15
+ @classmethod
16
+ def from_tokenizer(cls, tokenizer) -> "AnswerCodes":
17
+ codes: list[str] = []
18
+ ids: list[int] = []
19
+ for code in CANDIDATES:
20
+ tokens = tokenizer.encode(code, add_special_tokens=False)
21
+ if len(tokens) == 1:
22
+ codes.append(code)
23
+ ids.append(tokens[0])
24
+ if len(codes) < 2:
25
+ raise ValueError("tokenizer yields fewer than two single-token answer codes")
26
+ return cls(codes=codes, token_ids=ids)
27
+
28
+ @property
29
+ def max_options(self) -> int:
30
+ return len(self.codes)
ruling/config.py ADDED
@@ -0,0 +1,120 @@
1
+ """Runtime settings, read from the environment."""
2
+
3
+ from __future__ import annotations
4
+
5
+ import json
6
+ import os
7
+ from dataclasses import dataclass, field
8
+ from pathlib import Path
9
+
10
+ import mlx.core as mx
11
+
12
+ from ruling.against import API_KEY_ENV, BASE_URL_ENV, DEFAULT_BASE_URL
13
+ from ruling.hosted import TOP_LOGPROBS
14
+
15
+ DEFAULT_MODEL = "mlx-community/Qwen3.5-4B-4bit"
16
+
17
+
18
+ def configure_memory(settings: "Settings") -> None:
19
+ """Keep MLX well inside physical memory.
20
+
21
+ MLX's default ceiling is 1.5 times the GPU's recommended working set, 60 GB
22
+ on a 48 GB machine, so it swaps rather than stopping. Batches of varying
23
+ length also leave differently sized freed buffers in its cache, which by
24
+ default may grow to that ceiling: 33 GB after twenty training steps of a
25
+ 150M-parameter encoder, which got the process killed once another model
26
+ was loaded alongside it. The memory limit makes MLX evaluate in smaller
27
+ pieces and raise if a job cannot fit; the cache limit returns freed buffers.
28
+ """
29
+ info = mx.device_info()
30
+ if "memory_size" not in info:
31
+ # The CPU backend on Linux reports no unified memory, and with an
32
+ # OpenAI-backend model MLX holds no weights here: nothing to budget.
33
+ return
34
+ physical = info["memory_size"]
35
+ mx.set_memory_limit(int(physical * settings.memory_fraction))
36
+ mx.set_cache_limit(int(settings.cache_limit_gb * 2**30))
37
+
38
+
39
+ def _flag(value: str) -> bool:
40
+ if value.lower() in ("1", "true", "yes", "on"):
41
+ return True
42
+ if value.lower() in ("0", "false", "no", "off"):
43
+ return False
44
+ raise ValueError(f"expected a boolean, got {value!r}")
45
+
46
+
47
+ @dataclass(frozen=True)
48
+ class Settings:
49
+ model: str = DEFAULT_MODEL
50
+ calibration_path: Path | None = None
51
+ adapter_path: Path | None = None
52
+ max_input_tokens: int = 65_536
53
+ max_branch_tokens: int = 32_768
54
+ max_image_tokens: int = 1024
55
+ rotations: int = 3
56
+ prior_debias: bool = False
57
+ api_key: str | None = None
58
+ cache_limit_gb: float = 4.0
59
+ memory_fraction: float = 0.5
60
+ cascade_to: str | None = None
61
+ cascade_threshold: float | None = None
62
+ failover_to: str | None = None
63
+ openai_base_url: str = ""
64
+ openai_api_key: str | None = None
65
+ openai_top_logprobs: int = TOP_LOGPROBS
66
+ openai_concurrency: int = 16
67
+ openai_extra_body: dict = field(default_factory=dict)
68
+ vllm_base_url: str = ""
69
+ typesafe_base_url: str = DEFAULT_BASE_URL
70
+ typesafe_api_key: str | None = None
71
+ vllm_api_key: str | None = None
72
+ openrouter_providers: tuple[str, ...] = ()
73
+ openrouter_budget_usd: float = 1.0
74
+ models_path: Path | None = None
75
+ ledger_path: Path | None = None
76
+ ledger_retain_days: float | None = None
77
+ max_pending: int = 64
78
+ host: str = "127.0.0.1"
79
+ port: int = 8010
80
+
81
+ @classmethod
82
+ def from_env(cls) -> "Settings":
83
+ env = os.environ
84
+ calibration = env.get("RULING_CALIBRATION")
85
+ adapter = env.get("RULING_ADAPTER")
86
+ models = env.get("RULING_MODELS")
87
+ ledger = env.get("RULING_LEDGER")
88
+ return cls(
89
+ model=env.get("RULING_MODEL", DEFAULT_MODEL),
90
+ calibration_path=Path(calibration) if calibration else None,
91
+ adapter_path=Path(adapter) if adapter else None,
92
+ max_input_tokens=int(env.get("RULING_MAX_INPUT_TOKENS", cls.max_input_tokens)),
93
+ max_branch_tokens=int(env.get("RULING_MAX_BRANCH_TOKENS", cls.max_branch_tokens)),
94
+ max_image_tokens=int(env.get("RULING_MAX_IMAGE_TOKENS", cls.max_image_tokens)),
95
+ rotations=int(env.get("RULING_ROTATIONS", cls.rotations)),
96
+ prior_debias=_flag(env.get("RULING_PRIOR_DEBIAS", str(cls.prior_debias))),
97
+ api_key=env.get("RULING_API_KEY") or None,
98
+ cache_limit_gb=float(env.get("RULING_CACHE_LIMIT_GB", cls.cache_limit_gb)),
99
+ memory_fraction=float(env.get("RULING_MEMORY_FRACTION", cls.memory_fraction)),
100
+ cascade_to=env.get("RULING_CASCADE_TO") or None,
101
+ cascade_threshold=float(env["RULING_CASCADE_THRESHOLD"]) if env.get("RULING_CASCADE_THRESHOLD") else None,
102
+ failover_to=env.get("RULING_FAILOVER_TO") or None,
103
+ openai_base_url=env.get("RULING_OPENAI_BASE_URL", cls.openai_base_url),
104
+ openai_api_key=env.get("RULING_OPENAI_API_KEY") or None,
105
+ openai_top_logprobs=int(env.get("RULING_OPENAI_TOP_LOGPROBS", cls.openai_top_logprobs)),
106
+ openai_concurrency=int(env.get("RULING_OPENAI_CONCURRENCY", cls.openai_concurrency)),
107
+ openai_extra_body=json.loads(env.get("RULING_OPENAI_EXTRA_BODY", "{}")),
108
+ vllm_base_url=env.get("RULING_VLLM_BASE_URL", cls.vllm_base_url),
109
+ typesafe_base_url=env.get(BASE_URL_ENV, cls.typesafe_base_url),
110
+ typesafe_api_key=env.get(API_KEY_ENV) or None,
111
+ vllm_api_key=env.get("RULING_VLLM_API_KEY") or None,
112
+ openrouter_providers=tuple(p for p in env.get("RULING_OPENROUTER_PROVIDERS", "").split(",") if p),
113
+ openrouter_budget_usd=float(env.get("RULING_OPENROUTER_BUDGET_USD", cls.openrouter_budget_usd)),
114
+ models_path=Path(models) if models else None,
115
+ ledger_path=Path(ledger) if ledger else None,
116
+ ledger_retain_days=float(env["RULING_LEDGER_RETAIN_DAYS"]) if env.get("RULING_LEDGER_RETAIN_DAYS") else None,
117
+ max_pending=int(env.get("RULING_MAX_PENDING", cls.max_pending)),
118
+ host=env.get("RULING_HOST", cls.host),
119
+ port=int(env.get("RULING_PORT", cls.port)),
120
+ )