| """Public document-state API and four-state answer semantics.""" |
|
|
| import copy |
| import json |
| import math |
| from pathlib import Path |
|
|
| from ._vendor.contract import _state_parts, decision, parse_questions, present |
| from ._vendor.prompts import boolean_block, label_block, listwise_block |
| from ._vendor.semantics import listed_probs, p_yes |
| from .artifacts import digest, sha256 |
|
|
| TASKS = ("boolean", "single", "ordered", "multilabel", "entity") |
|
|
|
|
| def branches(spec): |
| task, req = spec["task"], spec["request"] |
| if task == "boolean": |
| return [(boolean_block(req["question"]), 4, "boolean/state4")] |
| if task in ("single", "ordered"): |
| block, width = listwise_block( |
| req["question"], req["options"], ordered=task == "ordered", reserved=task == "single" |
| ) |
| return [(block, width, task + ("/choiceR" if task == "single" else "/choiceS"))] |
| if task == "entity": |
| return [ |
| (boolean_block(req["template"].replace("{entity}", c)), 4, "entity/state4") |
| for c in req["entities"] |
| ] |
| return [(*label_block(req["question"], c), "multilabel/state4") for c in req["labels"]] |
|
|
|
|
| def distributions(spec, rows, temperature): |
| if spec["task"] in ("boolean", "entity", "multilabel"): |
| values = [p_yes(r["letter_logits"], temperature) for r in rows] |
| return [[p, 1 - p] for p in values] |
| return [listed_probs(rows[0]["letter_logits"], len(spec["texts"]), temperature).tolist()] |
|
|
|
|
| def ordering_score(values): |
| """Product of per-unit top probabilities; not a calibrated joint probability.""" |
| if not values: |
| raise ValueError("At least one answer unit is required") |
| for p in values: |
| if len(p) < 2 or not all(math.isfinite(v) and v >= 0 for v in p) or abs(sum(p) - 1) > 1e-6: |
| raise ValueError("Invalid answer distribution") |
| return math.prod(max(p) for p in values) |
|
|
|
|
| class DocumentState: |
| def __init__(self, owner, data): |
| self._owner, self._data, self.closed = owner, data, False |
| self.image_hashes = {p["image"]: sha256(p["image"]) for p in data["parts"] if "image" in p} |
|
|
| @property |
| def prefix_tokens(self): |
| self._check() |
| return len(self._data["prefix_ids"]) |
|
|
| def _check(self): |
| if self.closed: |
| raise ValueError("Document state is closed") |
| if any(sha256(path) != value for path, value in self.image_hashes.items()): |
| raise ValueError("Document image changed after prefill") |
|
|
| def save(self, path): |
| """Save a source-bound replay recipe, never pickle executable cache objects.""" |
| self._check() |
| body = { |
| "format": "solomon-mlx-replay-v1", |
| "runtime": self._owner.identity["fingerprint"], |
| "parts": self._data["parts"], |
| "image_hashes": self.image_hashes, |
| "prefix_ids_sha256": digest(self._data["prefix_ids"]), |
| } |
| Path(path).write_text(json.dumps({**body, "sha256": digest(body)}, indent=2)) |
|
|
| def close(self): |
| with self._owner.engine.lock: |
| self._data.clear() |
| self.closed = True |
|
|
| def __enter__(self): |
| self._check() |
| return self |
|
|
| def __exit__(self, *args): |
| self.close() |
|
|
|
|
| class Solomon: |
| @classmethod |
| def load( |
| cls, |
| model_dir, |
| profile="quality", |
| *, |
| chunk_size=2048, |
| max_tokens=40960, |
| page_selector=None, |
| calibration=None, |
| ): |
| if profile != "quality": |
| raise ValueError("Only full BF16 quality is implemented; quantization is secondary") |
| from .engine import Engine |
|
|
| return cls( |
| Engine(model_dir, chunk_size=chunk_size, max_tokens=max_tokens), |
| page_selector=page_selector, |
| calibration=calibration, |
| ) |
|
|
| def __init__(self, engine, *, page_selector=None, calibration=None): |
| self.engine, self.identity, self.page_selector = engine, engine.identity, page_selector |
| self.temperatures = dict.fromkeys(TASKS, 1.0) |
| self.calibration_status = "uncalibrated" |
| if calibration is not None: |
| artifact = json.loads(Path(calibration).read_text()) |
| payload = {k: v for k, v in artifact.items() if k != "sha256"} |
| if ( |
| artifact.get("sha256") != digest(payload) |
| or artifact["runtime"] != self.identity["fingerprint"] |
| ): |
| raise ValueError("Calibration checksum or MLX runtime identity mismatch") |
| temps = artifact["temperatures"] |
| if set(temps) != set(TASKS) or any( |
| isinstance(v, bool) |
| or not isinstance(v, (int, float)) |
| or not math.isfinite(v) |
| or not 0 < v <= 20 |
| for v in temps.values() |
| ): |
| raise ValueError("Invalid temperatures") |
| self.temperatures, self.calibration_status = temps, "profile_fitted" |
|
|
| def prefill(self, document): |
| parts = copy.deepcopy(_state_parts(document)) |
| if not parts: |
| parts = [{"text": ""}] |
| for p in parts: |
| if not isinstance(p, dict) or set(p) not in ({"text"}, {"image"}): |
| raise ValueError("Each document part must contain only text or image") |
| if "text" in p and not isinstance(p["text"], str): |
| raise ValueError("Text parts must be strings") |
| if "image" in p: |
| p["image"] = str(Path(p["image"]).resolve(strict=True)) |
| hashes = {p["image"]: sha256(p["image"]) for p in parts if "image" in p} |
| state = DocumentState(self, self.engine.prefill(parts)) |
| if state.image_hashes != hashes: |
| state.close() |
| raise ValueError("Image changed while document was being prefilled") |
| return state |
|
|
| def replay(self, path): |
| body = json.loads(Path(path).read_text()) |
| expected = body.pop("sha256") |
| if ( |
| digest(body) != expected |
| or body["format"] != "solomon-mlx-replay-v1" |
| or body["runtime"] != self.identity["fingerprint"] |
| ): |
| raise ValueError("Replay checksum or runtime mismatch") |
| if any(sha256(p) != h for p, h in body["image_hashes"].items()): |
| raise ValueError("Replay image changed") |
| state = self.prefill(body["parts"]) |
| if digest(state._data["prefix_ids"]) != body["prefix_ids_sha256"]: |
| state.close() |
| raise ValueError("Replay tokenization differs") |
| return state |
|
|
| def _answer(self, state, spec, execution="cached"): |
| rows = [self.engine.ask(state._data, b, n, h, execution=execution) for b, n, h in branches(spec)] |
| dists = distributions(spec, rows, self.temperatures[spec["task"]]) |
| return { |
| **present(spec, dists), |
| "ordering_score": ordering_score(dists), |
| "temperature": self.temperatures[spec["task"]], |
| }, rows |
|
|
| def decide( |
| self, |
| *, |
| state, |
| questions, |
| evidence="support", |
| evidence_max_calls=64, |
| execution="cached", |
| diagnostics=False, |
| ): |
| if not isinstance(state, DocumentState) or state._owner is not self: |
| raise ValueError("State belongs to a different model instance") |
| if evidence not in ("none", "support", "sufficiency", "removal"): |
| raise ValueError("Invalid evidence level") |
| if type(evidence_max_calls) is not int or not 0 <= evidence_max_calls <= 512: |
| raise ValueError("Invalid evidence call budget") |
| specs = parse_questions(questions) |
| with self.engine.lock: |
| state._check() |
| answers, usage = {}, {"branches": 0, "input_tokens": 0, "evidence_calls": 0} |
| for spec in specs: |
| answer, rows = self._answer(state, spec, execution) |
| body = self._evidence(state, spec, answer, evidence, evidence_max_calls) |
| answer.update( |
| evidence=body["references"], evidence_status=body["status"], evidence_detail=body |
| ) |
| if diagnostics: |
| answer["branches"] = rows |
| answers[spec["id"]] = answer |
| usage["branches"] += len(rows) |
| usage["input_tokens"] += sum( |
| r["branch_tokens"] if execution == "cached" else r["prompt_tokens"] for r in rows |
| ) |
| usage["evidence_calls"] += body.get("calls", 0) |
| return { |
| "answers": answers, |
| "usage": usage, |
| "runtime": self.identity, |
| "calibration_status": self.calibration_status, |
| "answer_policy": "always_answers", |
| } |
|
|
| def _fresh(self, document, spec): |
| with self.prefill(document) as state: |
| answer, _ = self._answer(state, spec) |
| return decision(spec, answer) |
|
|
| def _evidence(self, state, spec, answer, level, budget): |
| from ._vendor import evidence_v3 as v3 |
| from ._vendor.evidence import image_pages, validate_pages, validate_spans |
| from ._vendor.retrieval import lexical_select, remove |
|
|
| body = { |
| "references": [], |
| "status": "not_requested", |
| "calls": 0, |
| "verification": "none", |
| "faithfulness_established": False, |
| } |
| if level == "none": |
| return body |
| parts, req = state._data["parts"], spec["request"] |
| task = spec["task"] |
| if task == "entity": |
| questions = [req["template"].replace("{entity}", c) for c in req["entities"]] |
| elif task == "multilabel": |
| questions = [req["question"] + " Label: " + c for c in req["labels"]] |
| else: |
| questions = [req["question"] + (" " + " ".join(req["options"]) if "options" in req else "")] |
| images = [p["image"] for p in parts if "image" in p] |
| needed = (1 if images else len(questions)) if level in ("sufficiency", "removal") else 0 |
| needed += int(level == "removal") |
| if needed > budget: |
| return {**body, "status": "budget_exhausted", "required_calls": needed} |
| baseline = decision(spec, answer) |
| if images: |
| if self.page_selector is None: |
| return {**body, "status": "unsupported_page_selector", "pages_available": len(images)} |
| pages = image_pages(images) |
| selector = self.page_selector |
| plan = None |
| if hasattr(selector, "plan"): |
| plan = selector.plan( |
| pages, questions, **({"task": task} if getattr(selector, "task_aware", False) else {}) |
| ) |
| if type(plan.get("calls")) is not int or plan["calls"] < 0: |
| raise ValueError("Invalid page selector call estimate") |
| if needed + plan["calls"] > budget: |
| return {**body, "status": "budget_exhausted", "required_calls": needed + plan["calls"]} |
| selection = ( |
| selector.execute(plan) if plan is not None else selector(copy.deepcopy(pages), questions) |
| ) |
| calls = selection.get("cost", {}).get("calls", 0) |
| if calls != (plan["calls"] if plan is not None else 0): |
| raise ValueError("Page selector exceeded its declared call budget") |
| refs = validate_pages(pages, selection["evidence"]) |
| body["calls"] = calls |
| selected = {r["page"] for r in refs} |
| page, remainder = 0, [] |
| for part in parts: |
| if "image" in part: |
| page += 1 |
| if page in selected: |
| continue |
| remainder.append(part) |
| subsets = [([{"image": r["path"]} for r in refs], spec)] |
| else: |
| text = "".join(p["text"] for p in parts) |
| selection = lexical_select(text, questions) |
| refs = validate_spans(text, selection["evidence"]) |
| structure = v3.Structure(text) |
| packages, subsets = [], [] |
| for i, q in enumerate(questions): |
| subject = req["entities"][i] if task == "entity" else None |
| package = v3.build(text, q, refs, subject=subject, structure=structure) |
| unit = copy.deepcopy(spec) |
| if "candidates" in spec: |
| candidate = spec["candidates"][i] |
| unit["candidates"] = [candidate] |
| unit["request"]["entities" if task == "entity" else "labels"] = [candidate] |
| subsets.append((package["text"], unit)) |
| packages.append({k: v for k, v in package.items() if k != "text"}) |
| body["packages"] = packages |
| remainder = remove(text, refs) |
| body.update( |
| references=refs, status="found" if refs else "no_support_found", verification="retrieval_only" |
| ) |
| if level in ("sufficiency", "removal"): |
| predictions = [self._fresh(doc, unit) for doc, unit in subsets] |
| assembled = ( |
| {k: v for d in predictions for k, v in d.items()} |
| if "candidates" in spec and not images |
| else predictions[0] |
| ) |
| body["evidence_only"] = {"prediction": assembled, "agrees_with_full": assembled == baseline} |
| body["calls"] += len(subsets) |
| body["verification"] = "fresh_source_reencoding" |
| if level == "removal": |
| removed = self._fresh(remainder, spec) |
| body["evidence_removed"] = {"prediction": removed, "agrees_with_full": removed == baseline} |
| body["calls"] += 1 |
| return body |
|
|
