| """Load and score the released value model. |
| |
| The inference API supports one unpadded sequence at a time and returns |
| pre-action values for prompt and response token IDs. |
| """ |
| from pathlib import Path |
| import json |
| from numbers import Integral |
|
|
| import torch |
| from torch import nn |
| from torch.nn import functional as F |
| from safetensors.torch import load_file |
|
|
| from configuration_gmpt import GMPTConfig |
| from modeling_backbone import GMPTModel, configure_attention |
|
|
|
|
| class ValueModel(nn.Module): |
| def __init__(self, config): |
| super().__init__() |
| self.config = config |
| self.backbone = GMPTModel(config) |
| self.value_head = nn.Linear(config.hidden_size, 1, dtype=torch.float32) |
|
|
| @property |
| def device(self): |
| return self.value_head.weight.device |
|
|
| @torch.no_grad() |
| def forward(self, input_ids): |
| """Return raw post-token values [1, L] for a single unpadded input.""" |
| if input_ids.ndim != 2 or input_ids.shape[0] != 1: |
| raise ValueError("Expected one unpadded sequence, shape [1, L].") |
| if input_ids.dtype != torch.long: |
| raise ValueError("input_ids must be torch.int64.") |
| n = input_ids.shape[1] |
| if not 0 < n <= self.config.max_position_embeddings: |
| raise ValueError("Input length is empty or exceeds the context; no automatic truncation.") |
| if self.config._attn_implementation == "eager" and n > 4096: |
| raise ValueError("The eager reference backend is limited to 4096 tokens; use FlashAttention for longer inputs.") |
| ids = input_ids.to(self.device) |
| positions = torch.arange(n, device=self.device, dtype=torch.long)[None, :] |
| with torch.autocast(device_type=self.device.type, enabled=False): |
| h = self.backbone(input_ids=ids, position_ids=positions).last_hidden_state |
| values = F.linear(h.float(), self.value_head.weight, self.value_head.bias).squeeze(-1) |
| return values |
|
|
| @torch.no_grad() |
| def score_tokens(self, prompt_ids, response_ids): |
| """Score original token IDs; continuation values precede each action.""" |
| for name, tokens in (("prompt_ids", prompt_ids), ("response_ids", response_ids)): |
| if not isinstance(tokens, (list, tuple)) or not tokens: |
| raise ValueError(f"{name} must be a nonempty list or tuple of integer token IDs.") |
| if any(isinstance(t, bool) or not isinstance(t, Integral) for t in tokens): |
| raise ValueError(f"{name} must contain integer token IDs; no implicit conversion.") |
| ids = torch.tensor([list(prompt_ids) + list(response_ids)], dtype=torch.long, device=self.device) |
| post = self(ids) |
| p, length = len(prompt_ids), len(response_ids) |
| pre = post[0, p - 1:p + length - 1] |
| if not torch.isfinite(post).all(): |
| raise RuntimeError("Non-finite value output.") |
| return { |
| "prompt_tokens": p, |
| "response_tokens": length, |
| "pre_action_values": pre.cpu().tolist(), |
| "prompt_value": pre[0].item(), |
| "mean_response_pre_action_value": pre.mean().item(), |
| "last_pre_action_value": pre[-1].item(), |
| "raw_value_after_complete_input": post[0, -1].item(), |
| } |
|
|
|
|
| def load_value_model(model_dir, *, device="cuda", attention_backend="eager", weights_path=None): |
| """Load every critic tensor, including its trained backbone and scalar head. |
| |
| FP32 storage is retained in the release. Inference uses BF16 backbone |
| parameters and the original FP32 scalar head, matching native compute. |
| """ |
| model_dir = Path(model_dir) |
| config = GMPTConfig(**json.loads((model_dir / "backbone_config.json").read_text())) |
| if attention_backend.startswith("flash_attention") and torch.device(device).type != "cuda": |
| raise ValueError("FlashAttention requires a supported CUDA GPU.") |
| configure_attention(attention_backend) |
| config._attn_implementation = attention_backend |
| weights_path = Path(weights_path) if weights_path is not None else model_dir / "model.safetensors" |
| state = load_file(str(weights_path), device="cpu") |
| head = state.get("value_head.weight") |
| if head is None or tuple(head.shape) != (1, config.hidden_size): |
| raise ValueError( |
| "These weights are not this value model: expected a " |
| f"value_head.weight of shape (1, {config.hidden_size})." |
| ) |
| with torch.device("meta"): |
| model = ValueModel(config) |
| model.load_state_dict(state, strict=True, assign=True) |
| del state |
| model.backbone.to(device=device, dtype=torch.bfloat16) |
| model.value_head.to(device=device, dtype=torch.float32) |
| return model.eval().requires_grad_(False) |
|
|
|
|
| _TOKENIZER_CACHE = {} |
|
|
|
|
| def _load_tokenizers(model_dir): |
| """Load and cache the tokenizer pair; repeated calls must not reparse it.""" |
| from transformers import AutoTokenizer |
| from tokenizers import Tokenizer |
|
|
| key = str(Path(model_dir).resolve()) |
| if key not in _TOKENIZER_CACHE: |
| _TOKENIZER_CACHE[key] = ( |
| AutoTokenizer.from_pretrained(model_dir, local_files_only=True), |
| Tokenizer.from_file(str(Path(model_dir) / "tokenizer.json")), |
| ) |
| return _TOKENIZER_CACHE[key] |
|
|
|
|
| def tokenize_pair(model_dir, prompt, response): |
| """Use the shipped template; append the continuation without synthetic EOS.""" |
| tokenizer, raw = _load_tokenizers(model_dir) |
| messages = [{"role": "user", "content": prompt}] |
| prompt_ids = tokenizer.apply_chat_template( |
| messages, tokenize=True, add_generation_prompt=True, return_dict=True, |
| )["input_ids"] |
| response_ids = tokenizer.encode(response, add_special_tokens=False) |
| formatted = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) |
| if prompt_ids != raw.encode(formatted, add_special_tokens=False).ids: |
| raise ValueError("Tokenizer IDs differ from the shipped tokenizer.json; do not silently migrate it.") |
| if response_ids != raw.encode(response, add_special_tokens=False).ids: |
| raise ValueError("Continuation token IDs differ from tokenizer.json.") |
| return prompt_ids, response_ids |
|
|
