| """Inference-only formula excerpts from the source-bound canonical runtime.""" |
| from dataclasses import dataclass |
| import hashlib |
| import json |
| import re |
| from pathlib import Path |
| import torch |
| from torch import nn |
| MAX_TOKENS = 32768 |
| TRAIN_TOKENS = 768 |
| def file_sha256(path): |
| digest = hashlib.sha256() |
| with Path(path).open('rb') as stream: |
| for chunk in iter(lambda: stream.read(4 * 1024 * 1024), b''): |
| digest.update(chunk) |
| return digest.hexdigest() |
|
|
|
|
| def audit_loading_info(info, source_keys): |
| if not isinstance(info, dict) or not all(k in info for k in ('missing_keys', 'unexpected_keys', 'mismatched_keys')): |
| raise ValueError('missing loading audit fields') |
| unexpected = set(info['unexpected_keys']) |
| allowed = {'lm_head.weight'} if 'lm_head.weight' in source_keys else set() |
| if info['missing_keys'] or info['mismatched_keys'] or info.get('error_msgs') or unexpected - allowed: |
| raise ValueError(f'base loading audit failed: {info}') |
| return {'missing_keys': [], 'mismatched_keys': [], 'unexpected_keys': sorted(unexpected), |
| 'allowed_unexpected_keys': sorted(unexpected), |
| 'exception_reason': ('LM vocabulary head is deliberately excluded by AutoModel; source key confirmed' |
| if unexpected else None)} |
|
|
|
|
| def load_base(lock): |
| """Only this boundary touches pretrained assets; local and fail-closed.""" |
| from safetensors import safe_open |
| from transformers import AutoModel, AutoTokenizer |
| if lock.get('model') != 'ibm-granite/granite-4.0-350m': |
| raise ValueError('wrong locked model') |
| if not isinstance(lock.get('revision'), str) or not re.fullmatch('[0-9a-f]{40}', lock['revision']): |
| raise ValueError('invalid locked revision') |
| sha = lock.get('model_safetensors_sha256') |
| if not isinstance(sha, str) or not re.fullmatch('[0-9a-f]{64}', sha): |
| raise ValueError('invalid source hash') |
| local = Path(lock['local_path']) |
| if not local.is_absolute() or lock.get('max_context_tokens') != MAX_TOKENS: |
| raise ValueError('invalid local path or context lock') |
| source = local / 'model.safetensors' |
| actual = file_sha256(source) |
| if actual != sha: |
| raise ValueError('source model hash mismatch') |
| with safe_open(str(source), framework='pt', device='cpu') as handle: |
| source_keys = set(handle.keys()) |
| model, info = AutoModel.from_pretrained(str(local), local_files_only=True, |
| trust_remote_code=False, attn_implementation='sdpa', |
| dtype=torch.float32, output_loading_info=True) |
| audit = audit_loading_info(info, source_keys) |
| if model.config.model_type != 'granitemoehybrid' or model.config.hidden_size != 1024: |
| raise ValueError('unexpected base architecture (requires GraniteMoeHybrid hidden=1024)') |
| model.to(device='cpu', dtype=torch.float32).eval() |
| tokenizer = AutoTokenizer.from_pretrained(str(local), local_files_only=True, trust_remote_code=False) |
| assets = ('config.json', 'tokenizer.json', 'tokenizer_config.json', |
| 'special_tokens_map.json', 'vocab.json', 'merges.txt', 'added_tokens.json', 'tokenizer.model') |
| asset_hashes = {name: file_sha256(local/name) for name in assets if (local/name).is_file()} |
| return model, tokenizer, dict(audit, source_sha256=actual, source_file=str(source), asset_sha256=asset_hashes) |
|
|
|
|
| def architecture(model): |
| return {'schema_version': 1, 'backbone_class': type(model.backbone).__name__, |
| 'backbone_config': json.loads(model.backbone.config.to_json_string(use_diff=False)), |
| 'hidden_size': model.backbone.config.hidden_size, 'head_hidden': 128, |
| 'head_activation': 'GELU', 'head_outputs': 1, 'shared_option_head': True, |
| 'separate_none_head': True, 'rank': 8, 'alpha': 16, 'dropout': 0, |
| 'adapter_targets': model.adapter_targets, |
| 'option_readout': 'option_segment_last_token', 'none_readout': 'final_segment_last_token', |
| 'prompt_version': 'context_question_json_options_decision_v1', |
| 'tokenization': 'encode_each_segment_add_special_tokens_false', |
| 'max_input_tokens': MAX_TOKENS, 'train_max_tokens': TRAIN_TOKENS, |
| 'dtype': 'float32', 'device': 'cpu', 'attention': 'sdpa', 'use_cache': False} |
|
|
| def tensor_hash(items): |
| digest = hashlib.sha256() |
| for name, tensor in sorted(items): |
| value = tensor.detach().cpu().contiguous() |
| digest.update(name.encode()) |
| digest.update(str(value.dtype).encode()) |
| digest.update(str(tuple(value.shape)).encode()) |
| digest.update(value.reshape(-1).view(torch.uint8).numpy().tobytes()) |
| return digest.hexdigest() |
|
|
|
|
| def frozen_hash(model): |
| return tensor_hash([(n, p) for n, p in model.named_parameters() if not p.requires_grad] |
| + list(model.named_buffers())) |
|
|
|
|
| def trainable_state(model): |
| return {n: p.detach().cpu().contiguous() for n, p in model.named_parameters() if p.requires_grad} |
|
|
| class LoRALinear(nn.Module): |
| """Frozen linear + rank-eight update; zero B preserves initial backbone.""" |
| def __init__(self, base, rank=8, alpha=16): |
| super().__init__() |
| if not isinstance(base, nn.Linear): |
| raise ValueError('LoRA target must be nn.Linear') |
| self.base = base |
| self.base.requires_grad_(False) |
| self.rank, self.alpha = rank, alpha |
| self.A = nn.Parameter(torch.empty(rank, base.in_features, dtype=base.weight.dtype)) |
| self.B = nn.Parameter(torch.zeros(base.out_features, rank, dtype=base.weight.dtype)) |
| nn.init.kaiming_uniform_(self.A, a=5 ** 0.5) |
|
|
| def forward(self, x): |
| return self.base(x) + nn.functional.linear(nn.functional.linear(x, self.A), self.B) * (self.alpha / self.rank) |
|
|
|
|
| class OptionScorer(nn.Module): |
| def __init__(self, backbone): |
| super().__init__() |
| self.backbone = backbone |
| self.backbone.requires_grad_(False) |
| attention = [(i, layer.self_attn) for i, layer in enumerate(backbone.layers) |
| if hasattr(layer, 'self_attn') and layer.self_attn is not None] |
| if len(attention) < 4: |
| raise ValueError('backbone requires at least four attention layers') |
| self.adapter_targets = [] |
| for index, attn in attention[-4:]: |
| for name in ('q_proj', 'v_proj'): |
| setattr(attn, name, LoRALinear(getattr(attn, name))) |
| self.adapter_targets.append(f'layers.{index}.self_attn.{name}') |
| hidden = backbone.config.hidden_size |
| self.option_head = nn.Sequential(nn.Linear(hidden, 128), nn.GELU(), nn.Linear(128, 1)) |
| self.none_head = nn.Sequential(nn.Linear(hidden, 128), nn.GELU(), nn.Linear(128, 1)) |
| self.train(False) |
|
|
| def train(self, mode=True): |
| super().train(mode) |
| self.backbone.eval() |
| for module in self.backbone.modules(): |
| if isinstance(module, LoRALinear): |
| module.training = mode |
| return self |
|
|
| def forward(self, encoded): |
| ids = torch.tensor([encoded.input_ids], dtype=torch.long, device='cpu') |
| hidden = self.backbone(input_ids=ids, attention_mask=torch.ones_like(ids), |
| use_cache=False, return_dict=True).last_hidden_state[0] |
| options = self.option_head(hidden[encoded.option_positions]).squeeze(-1) |
| none = self.none_head(hidden[encoded.final_position]).reshape(1) |
| return torch.cat((options, none)) |
|
|
| def validate_request(request, *, labeled=False): |
| if not isinstance(request, dict): |
| raise ValueError('request must be an object') |
| for key in ('context', 'question'): |
| if not isinstance(request.get(key), str) or (key == 'question' and not request[key].strip()): |
| raise ValueError(f'{key} must be a string (question nonempty)') |
| options = request.get('options') |
| if not isinstance(options, list) or not 2 <= len(options) <= 16: |
| raise ValueError('options must contain 2..16 choices') |
| ids = [] |
| for option in options: |
| if not isinstance(option, dict): |
| raise ValueError('option must be an object') |
| for key in ('id', 'description'): |
| if not isinstance(option.get(key), str) or not option[key].strip(): |
| raise ValueError(f'option {key} must be nonempty string') |
| ids.append(option['id']) |
| if len(set(ids)) != len(ids): |
| raise ValueError('option ids must be unique') |
| if 'target' in request and request['target'] is not None and request['target'] not in ids: |
| raise ValueError('target must be option id or null') |
| if labeled: |
| for key in ('id', 'family'): |
| if not isinstance(request.get(key), str) or not request[key].strip(): |
| raise ValueError(f'{key} must be nonempty') |
| if 'target' not in request: |
| raise ValueError('labeled row requires target') |
| return request |
|
|
|
|
| @dataclass(frozen=True) |
| class EncodedRequest: |
| input_ids: list |
| option_positions: list |
| final_position: int |
| text: str |
|
|
