| """Local browser decision inference. No downloads, teacher or game runtime.""" |
| from contextlib import nullcontext |
| import json,math,os |
| from pathlib import Path |
| from schema import validate_request |
|
|
| def unique_object(pairs): |
| obj = {} |
| for key, value in pairs: |
| if key in obj: |
| raise ValueError(f"JSON contains duplicate key: {key}") |
| obj[key] = value |
| return obj |
|
|
| def reject_nonfinite(value): |
| raise ValueError(f"JSON disallows non-finite number: {value}") |
|
|
| def read_json(path): |
| return json.loads(Path(path).read_text(encoding="utf-8"), object_pairs_hook=unique_object, |
| parse_constant=reject_nonfinite) |
|
|
| def nonempty_text(value): |
| return isinstance(value, str) and bool(value.strip()) |
|
|
| def prepare_examples(payload, tokenizer, max_length): |
| """Segment encoding, candidate text and EOS follow the upstream training tokenization exactly.""" |
| states = validate_request(payload) |
| if type(max_length) is not int or max_length <= 0: |
| raise ValueError("max_length must be a positive integer") |
| if type(tokenizer.eos_token_id) is not int or tokenizer.eos_token_id < 0: |
| raise ValueError("Checkpoint tokenizer must have a valid eos_token_id") |
| examples = [] |
| for row in states: |
| for qid, q in row["questions"].items(): |
| typ = q["type"] |
| ids = list(q["criteria"]) |
| texts = [f"{key}: {q['criteria'][key]}" for key in ids] |
| segments = [f"State:\n{row['state']}\n", |
| f"Question type: {typ}\nQuestion:\n{q['instructions']}\n"] |
| prefix = sum([tokenizer.encode(t, add_special_tokens=False) for t in segments], []) |
| leaves = [prefix + tokenizer.encode(f"Candidate:\n{t}\nDecision:", add_special_tokens=False) |
| + [tokenizer.eos_token_id] for t in texts] |
| largest = max(map(len, leaves)) |
| if largest > max_length: |
| raise ValueError(f"{row['id']}:{qid} candidate path has {largest} tokens, exceeding max_length={max_length}; input was not truncated") |
| examples.append({"id": f"{row['id']}:{qid}", "state_id": row["id"], "qid": qid, |
| "type": typ, "candidate_ids": ids, "candidate_texts": texts, |
| "leaf_tokens": leaves}) |
| return examples |
|
|
| def complete_question_batches(examples, batch_questions=0): |
| if type(batch_questions) is not int or batch_questions < 0: |
| raise ValueError("batch_questions must be nonnegative; zero batches all questions together") |
| size = batch_questions or len(examples) |
| if not examples: |
| return [] |
| return [examples[i:i + size] for i in range(0, len(examples), size)] |
|
|
| def answer_from_probabilities(example, probabilities): |
| ids = example["candidate_ids"] |
| if len(probabilities) != len(ids) or not all(math.isfinite(p) and 0 <= p <= 1 for p in probabilities): |
| raise ValueError("Model produced invalid probabilities") |
| if abs(math.fsum(probabilities) - 1.0) > 1e-5: |
| raise ValueError("Model probabilities do not sum to one") |
| best = max(range(len(ids)), key=probabilities.__getitem__) |
| result = {"type": example["type"], "probabilities": dict(zip(ids, probabilities))} |
| result.update(choice=ids[best], value=ids[best]) |
| return result |
|
|
| def local_checkpoint_files(checkpoint_dir): |
| root = Path(checkpoint_dir).expanduser().resolve(strict=True) |
| if not root.is_dir(): |
| raise ValueError("Checkpoint must be a local directory") |
| paths = {"run_config": root / "config.json", "body_config": root / "backbone_config", |
| "tokenizer": root / "tokenizer", "weights": root / "best.safetensors"} |
| for label, path in paths.items(): |
| if not path.exists(): |
| raise ValueError(f"Checkpoint missing {label}: {path.name}") |
| if not paths["run_config"].is_file() or not paths["weights"].is_file(): |
| raise ValueError("config.json and best.safetensors must be files") |
| if not paths["body_config"].is_dir() or not paths["tokenizer"].is_dir(): |
| raise ValueError("backbone_config and tokenizer must be directories") |
| return root, paths |
|
|
| class DecisionPredictor: |
| """Persistent local inference: load weights once, evaluate complete candidate sets per question.""" |
|
|
| def __init__(self, checkpoint_dir, max_length=None, device_name="mps", |
| disable_native_triton=False, precision="fp32"): |
| if precision != "fp32": |
| raise ValueError("precision must be fp32") |
| root, paths = local_checkpoint_files(checkpoint_dir) |
| run_config = read_json(paths["run_config"]) |
| if not isinstance(run_config, dict) or run_config.get("set_head") not in {"none", "attention"}: |
| raise ValueError("Checkpoint config has no valid set_head") |
| |
| |
| os.environ["HF_HUB_OFFLINE"] = "1" |
| os.environ["TRANSFORMERS_OFFLINE"] = "1" |
| os.environ["HF_HUB_DISABLE_TELEMETRY"] = "1" |
| import torch |
| from safetensors.torch import load_file |
| from transformers import AutoConfig, AutoModel, AutoTokenizer |
| |
| device = torch.device(device_name) |
| if device.type == "mps": |
| if not torch.backends.mps.is_available(): |
| raise ValueError("MPS is unavailable; native Apple Silicon Python is required") |
| elif device.type != "cpu": |
| raise ValueError("This package supports mps and cpu") |
| if precision != "fp32": |
| raise ValueError("This package uses FP32 inference") |
|
|
| tokenizer = AutoTokenizer.from_pretrained(str(paths["tokenizer"]), local_files_only=True, |
| trust_remote_code=False) |
| if tokenizer.pad_token_id is None: |
| tokenizer.pad_token = tokenizer.eos_token |
| body_config = AutoConfig.from_pretrained(str(paths["body_config"]), local_files_only=True, |
| trust_remote_code=False) |
| body_config.use_cache = False |
| limit = run_config.get("max_length", 512) if max_length is None else max_length |
| if type(limit) is not int or limit <= 0: |
| raise ValueError("max-length must be a positive integer") |
| context_limit = getattr(body_config, "max_position_embeddings", None) |
| if isinstance(context_limit, int) and limit > context_limit: |
| raise ValueError("max-length exceeds the backbone context length") |
| |
| |
| body = AutoModel.from_config(body_config, attn_implementation="sdpa", trust_remote_code=False).float() |
| from model import DecisionModel |
| model = DecisionModel(body, run_config["set_head"]) |
| weights = load_file(str(paths["weights"]), device="cpu") |
| model.load_state_dict(weights, strict=True) |
| del weights |
| model.to(device=device, dtype=torch.float32) |
| model.eval() |
| self.model = model |
| self.tokenizer = tokenizer |
| self.root = root |
| self.run_config = run_config |
| self.limit = limit |
| self.device = device |
| self.precision = precision |
| self.disable_native_triton = disable_native_triton |
| self.inference_calls = 0 |
| self._torch = torch |
|
|
| def predict(self, payload, batch_questions=0, temperature=1.0): |
| states = validate_request(payload) |
| if not isinstance(temperature, (int, float)) or isinstance(temperature, bool) or not math.isfinite(temperature) or temperature <= 0: |
| raise ValueError("temperature must be finite and positive") |
| torch = self._torch |
| model, tokenizer = self.model, self.tokenizer |
| root, run_config, limit = self.root, self.run_config, self.limit |
| device, precision = self.device, self.precision |
| disable_native_triton = self.disable_native_triton |
| examples = prepare_examples(payload, tokenizer, limit) |
| batches = complete_question_batches(examples, batch_questions) |
| self.inference_calls += 1 |
| model.eval() |
| outputs = {state["id"]: {"id": state["id"], "answers": {}} for state in states} |
| with torch.inference_mode(): |
| for batch in batches: |
| context = nullcontext() |
| with context: |
| logits, _ = model(batch, tokenizer.pad_token_id) |
| for example, values in zip(batch, logits): |
| k = len(example["candidate_ids"]) |
| scores = values[:k].float() |
| if not torch.isfinite(scores).all(): |
| raise ValueError("Model produced non-finite logits; no partial prediction returned") |
| probabilities = (scores / temperature).softmax(-1).cpu().tolist() |
| outputs[example["state_id"]]["answers"][example["qid"]] = answer_from_probabilities(example, probabilities) |
| return { |
| "schema_version": "nanojev-browser-inference-v1", |
| "checkpoint": {"id": "NanoJev-Web/browser-head-v5", "base_model": run_config.get("model"), |
| "base_revision": run_config.get("resolved_model_revision"), "set_head": run_config["set_head"]}, |
| "temperature": {"value": float(temperature), "fitted_by_this_command": False, |
| "note": "Applies the supplied scalar; the default of 1 does not imply calibration."}, |
| "execution": {"device": str(device), "parameter_storage": "float32", "precision": precision, |
| "forward_autocast": "bfloat16" if precision == "bf16" else "disabled", |
| "states": len(states), "questions": len(examples), |
| "candidate_paths": sum(len(ex["leaf_tokens"]) for ex in examples), |
| "input_tokens": sum(len(path) for ex in examples for path in ex["leaf_tokens"]), |
| "padded_input_tokens": sum( |
| max(len(path) for ex in batch for path in ex["leaf_tokens"]) |
| * sum(len(ex["leaf_tokens"]) for ex in batch) for batch in batches), |
| "output_tokens": 0, |
| "forward_passes": len(batches), "batch_questions_limit": batch_questions or "all", |
| "autoregressive_decode_steps": 0, "prefix_sharing": False, |
| "max_length": limit, "disable_native_triton": disable_native_triton, |
| "network_model_calls": 0, "persistent_model_load_count": 1, |
| "inference_call_index": self.inference_calls}, |
| "states": list(outputs.values()), |
| } |
|
|
