| """MiniCPM5-2B-Jev: Native causal FlashAttention multi-branch decision model with lm_head letter readout.""" |
| from __future__ import annotations |
|
|
| import copy |
| from datetime import datetime |
| import json |
| import os |
| import re |
| import string |
| from pathlib import Path |
| from typing import Any |
|
|
| import numpy as np |
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| from peft import LoraConfig, PeftModel, get_peft_model |
| from transformers import AutoModelForCausalLM, AutoTokenizer |
|
|
| DEFAULT_BASE_MODEL = "openbmb/MiniCPM5-2B" |
|
|
| MAX_OPTIONS = 255 |
| MAX_STATE = 1152 |
| MAX_BRANCH = 1664 |
| MAX_PACKED = 2304 |
| SERVE_MAX_STATE = 2560 |
| SERVE_MAX_BRANCH = 3072 |
| SERVE_MAX_PACKED = 3584 |
|
|
| _SPECIAL_RE = re.compile(r"<\|([A-Za-z0-9_]+)\|>") |
| _TYPE_PREFIX_RE = re.compile(r"^\[(choice|noul|score)\]\s*", re.IGNORECASE) |
| _TOKENIZER_META_CACHE: dict[int, tuple[Any, dict[str, Any]]] = {} |
|
|
|
|
| class ContextOverflow(ValueError): |
| """Raised when a record exceeds state, branch, or packed token limits.""" |
|
|
|
|
| def load_tokenizer(name: str = DEFAULT_BASE_MODEL): |
| try: |
| return AutoTokenizer.from_pretrained(name, local_files_only=True) |
| except Exception: |
| return AutoTokenizer.from_pretrained(name) |
|
|
|
|
| def pad_id(tok) -> int: |
| return tok.pad_token_id if tok.pad_token_id is not None else 0 |
|
|
|
|
| def user_tokens(tok, text: str) -> list[int]: |
| """Tokenize user text while preventing control-token injection.""" |
| safe_text = _SPECIAL_RE.sub(r"<¦\1¦>", str(text or "")) |
| return tok(safe_text, add_special_tokens=False).input_ids |
|
|
|
|
| def tokenizer_decision_meta(tok) -> dict[str, Any]: |
| """Build and cache option letter token IDs (A..Z, AA..) and chat template wrappers.""" |
| key = id(tok) |
| if key in _TOKENIZER_META_CACHE: |
| return _TOKENIZER_META_CACHE[key][1] |
|
|
| upper = string.ascii_uppercase |
| candidates = list(upper) + [a + b for a in upper for b in upper] |
| ans_prefix = tok.encode("Answer: (", add_special_tokens=False) |
| open_paren = tok.encode("\n(", add_special_tokens=False) |
|
|
| label_names: list[str] = [] |
| label_ids: list[int] = [] |
| seen_ids: set[int] = set() |
| for cand in candidates: |
| enc = tok.encode(f"Answer: ({cand}", add_special_tokens=False) |
| if len(enc) == len(ans_prefix) + 1 and enc[: len(ans_prefix)] == ans_prefix: |
| tid = enc[-1] |
| if tid not in seen_ids: |
| seen_ids.add(tid) |
| label_names.append(cand) |
| label_ids.append(tid) |
| if len(label_names) == MAX_OPTIONS: |
| break |
|
|
| sentinel = "@@SYSTEMONE_USER_CONTENT@@" |
| msgs = [{"role": "user", "content": sentinel}] |
| try: |
| rendered = tok.apply_chat_template( |
| msgs, tokenize=False, add_generation_prompt=True, enable_thinking=False |
| ) |
| except Exception: |
| rendered = tok.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True) |
|
|
| head_str, tail_str = rendered.split(sentinel) |
| for opened, closed in (("<think>\n", "</think>\n\n"), ("<think>", "</think>\n\n")): |
| if tail_str.endswith(opened): |
| tail_str += closed |
|
|
| chat_head_ids = tok.encode(head_str + "Context:\n", add_special_tokens=False) |
| chat_tail_ids = tok.encode(tail_str + "Answer: (", add_special_tokens=False) |
|
|
| meta = { |
| "label_names": label_names, |
| "label_ids": label_ids, |
| "open_paren": open_paren, |
| "chat_head_ids": chat_head_ids, |
| "chat_tail_ids": chat_tail_ids, |
| } |
| _TOKENIZER_META_CACHE[key] = (tok, meta) |
| return meta |
|
|
|
|
| def _trim_state_tokens(state_tokens: list[int], keep_state: int) -> list[int]: |
| """Preserve both the header/rules (first 60%) and most recent events/claim (last 40%) when trimming long states.""" |
| if len(state_tokens) <= keep_state: |
| return state_tokens |
| head_n = int(keep_state * 0.6) |
| tail_n = max(0, keep_state - head_n) |
| if tail_n == 0: |
| return state_tokens[:keep_state] |
| return state_tokens[:head_n] + state_tokens[-tail_n:] |
|
|
|
|
| def encode_record( |
| tok, |
| rec: dict[str, Any], |
| max_state: int = MAX_STATE, |
| max_branch: int = MAX_BRANCH, |
| strict: bool = False, |
| ) -> dict[str, Any]: |
| """Encode a record into shared state tokens S and per-question branches. |
| Each question branch forms an independent causal sequence S + br_k (mathematically identical to block-causal masking, |
| while unlocking native hardware-fused is_causal=True FlashAttention on MPS/CUDA with zero 4D mask overhead). |
| """ |
| meta = tokenizer_decision_meta(tok) |
| label_names = meta["label_names"] |
| label_ids = meta["label_ids"] |
| open_paren = meta["open_paren"] |
| chat_head_ids = meta["chat_head_ids"] |
| chat_tail_ids = meta["chat_tail_ids"] |
|
|
| state_tokens = user_tokens(tok, rec["state"]) |
| if strict and len(chat_head_ids) + len(state_tokens) > max_state: |
| raise ContextOverflow( |
| f"state exceeds {max_state} tokens: {len(chat_head_ids) + len(state_tokens)}" |
| ) |
|
|
| keep_state = max(16, max_state - len(chat_head_ids)) |
| S = list(chat_head_ids) + _trim_state_tokens(state_tokens, keep_state) |
|
|
| branches: list[tuple[list[int], int, str]] = [] |
| max_br_len = 0 |
| for q in rec["questions"]: |
| raw_instr_full = str(q.get("instr", "")) |
| q_type = str(q.get("type") or "").lower() |
| if not q_type: |
| m = _TYPE_PREFIX_RE.match(raw_instr_full) |
| q_type = m.group(1).lower() if m else "choice" |
| raw_instr = _TYPE_PREFIX_RE.sub("", raw_instr_full).strip() |
| if q_type == "score": |
| q_head = user_tokens(tok, f"\n\nQuestion (ordered rating scale, lowest to highest): {raw_instr}\nOptions:") |
| elif q_type == "noul": |
| q_head = user_tokens(tok, f"\n\nQuestion (yes/no judgement): {raw_instr}\nOptions:") |
| else: |
| q_head = user_tokens(tok, f"\n\nQuestion: {raw_instr}\nOptions:") |
|
|
| opts = q["options"] |
| if len(opts) > len(label_names): |
| raise ContextOverflow(f"Too many options ({len(opts)} > {len(label_names)})") |
|
|
| opt_tokens: list[int] = [] |
| for j, opt_str in enumerate(opts): |
| clean_opt = str(opt_str).replace("\n", " ").strip() |
| if j < 26: |
| opt_tokens.extend(user_tokens(tok, f"\n({label_names[j]}) {clean_opt}")) |
| else: |
| opt_tokens.extend(open_paren + [label_ids[j]] + user_tokens(tok, f") {clean_opt}")) |
|
|
| br = q_head + opt_tokens + list(chat_tail_ids) |
| max_br_len = max(max_br_len, len(br)) |
| branches.append((br, len(opts), q_type)) |
|
|
| if len(S) + max_br_len > max_branch: |
| if strict: |
| raise ContextOverflow( |
| f"branch too long: {max_br_len} tokens with {len(S)}-token state (limit {max_branch})" |
| ) |
| allowed_state = max(64, max_branch - max_br_len - len(chat_head_ids)) |
| S = list(chat_head_ids) + _trim_state_tokens(state_tokens, allowed_state) |
|
|
| ids = list(S) |
| seg = [0] * len(S) |
| pos = list(range(len(S))) |
| decide_idx: list[int] = [] |
| opt_idx: list[list[int]] = [] |
| n_opts: list[int] = [] |
| q_types: list[str] = [] |
| causal_rows: list[list[int]] = [] |
|
|
| for k, (br, n_opt, q_type) in enumerate(branches, start=1): |
| base = len(ids) |
| p0 = len(S) |
| br_pos = list(range(p0, p0 + len(br))) |
| ids.extend(br) |
| seg.extend([k] * len(br)) |
| pos.extend(br_pos) |
| decide_idx.append(base + len(br) - 1) |
| opt_idx.append(list(range(n_opt))) |
| n_opts.append(n_opt) |
| q_types.append(q_type) |
| causal_rows.append(S + br) |
|
|
| return { |
| "ids": ids, |
| "seg": seg, |
| "pos": pos, |
| "state_ids": S, |
| "causal_rows": causal_rows, |
| "decide_idx": decide_idx, |
| "opt_idx": opt_idx, |
| "n_opts": n_opts, |
| "q_types": q_types, |
| "labels": [int(q.get("label", 0)) for q in rec["questions"]], |
| "target_probs": [q.get("target_probs") for q in rec["questions"]], |
| "state_truncated": len(chat_head_ids) + len(state_tokens) > len(S), |
| } |
|
|
|
|
| def record_fits( |
| rec: dict[str, Any], |
| tok, |
| max_state: int = MAX_STATE, |
| max_branch: int = MAX_BRANCH, |
| max_packed: int = MAX_PACKED, |
| ) -> bool: |
| try: |
| enc = encode_record(tok, rec, max_state=max_state, max_branch=max_branch, strict=True) |
| return len(enc["ids"]) <= max_packed |
| except Exception: |
| return False |
|
|
|
|
| class LetterReadoutHead(nn.Module): |
| """Native lm_head projection onto option letter tokens (A..Z, AA..) with PriDe prior debiasing and per-type temperature calibration.""" |
|
|
| MAX_K_TABLE = 32 |
|
|
| def __init__(self, letter_weight: torch.Tensor) -> None: |
| super().__init__() |
| n_max = int(letter_weight.shape[0]) |
| self.register_buffer("letter_weight", letter_weight.detach().clone().float()) |
| self.register_buffer("letter_bias", torch.zeros(n_max, dtype=torch.float32)) |
| self.register_buffer( |
| "choice_k_bias", |
| torch.zeros(self.MAX_K_TABLE + 1, self.MAX_K_TABLE, dtype=torch.float32), |
| ) |
| self.register_buffer( |
| "score_k_bias", |
| torch.zeros(self.MAX_K_TABLE + 1, self.MAX_K_TABLE, dtype=torch.float32), |
| ) |
| self.register_buffer("noul_bias", torch.zeros(2, dtype=torch.float32)) |
| self.register_buffer("score_bias", torch.zeros(11, dtype=torch.float32)) |
| self.register_buffer("type_temps", torch.ones(3, dtype=torch.float32)) |
| self.temperature = 1.0 |
| self.use_prior = True |
|
|
| def forward_raw(self, h_decide: torch.Tensor, n_opt: int) -> torch.Tensor: |
| return F.linear(h_decide, self.letter_weight[:n_opt]) |
|
|
| def forward(self, h_decide: torch.Tensor, n_opt: int, q_type: str = "choice") -> torch.Tensor: |
| z = F.linear(h_decide, self.letter_weight[:n_opt]) |
| if self.use_prior and not self.training: |
| if q_type == "noul" and n_opt == 2: |
| z = z - self.noul_bias |
| elif q_type == "score": |
| if ( |
| 2 <= n_opt <= self.MAX_K_TABLE |
| and torch.any(self.score_k_bias[n_opt, :n_opt] != 0) |
| ): |
| z = z - self.score_k_bias[n_opt, :n_opt] |
| elif n_opt == 11 and torch.any(self.score_bias != 0): |
| z = z - self.score_bias |
| else: |
| if ( |
| 2 <= n_opt <= self.MAX_K_TABLE |
| and torch.any(self.choice_k_bias[n_opt, :n_opt] != 0) |
| ): |
| z = z - self.choice_k_bias[n_opt, :n_opt] |
| else: |
| z = z - self.letter_bias[:n_opt] |
|
|
| if self.training: |
| return z |
|
|
| type_idx = 1 if q_type == "noul" else (2 if q_type == "score" else 0) |
| t_type = float(self.type_temps[type_idx].item()) |
| t = t_type if t_type != 1.0 else float(self.temperature) |
| return z if t == 1.0 else z / t |
|
|
|
|
| class MiniCPMSystemOne(nn.Module): |
| """MiniCPM5-2B backbone + All-Layer LoRA + native causal FlashAttention + LetterReadoutHead.""" |
|
|
| SHAPE_BUCKET = int(os.environ.get("MINICPM_S1_SHAPE_BUCKET", "64")) |
| MAX_ROWS_PER_PASS = int(os.environ.get("MINICPM_S1_MAX_ROWS", "8")) |
|
|
| def __init__( |
| self, |
| base_model: str = DEFAULT_BASE_MODEL, |
| tok=None, |
| device: str = "mps", |
| lora_r: int = 0, |
| head_dim: int = 256, |
| dtype: torch.dtype = torch.bfloat16, |
| adapter_dir: str | Path | None = None, |
| lora_targets: str = "all", |
| top_k_layers: int = 20, |
| grad_ckpt: bool = False, |
| ) -> None: |
| super().__init__() |
| self.base_model = base_model |
| self.tok = tok or load_tokenizer(base_model) |
| self.pad_id = pad_id(self.tok) |
| self.device = device |
| self.lora_r = lora_r |
| self.head_dim = head_dim |
| self.lora_targets = lora_targets |
| self.top_k_layers = top_k_layers |
|
|
| meta = tokenizer_decision_meta(self.tok) |
| label_ids = torch.tensor(meta["label_ids"], dtype=torch.long) |
|
|
| try: |
| full_lm = AutoModelForCausalLM.from_pretrained( |
| base_model, |
| dtype=dtype, |
| attn_implementation="sdpa", |
| local_files_only=True, |
| ) |
| except Exception: |
| full_lm = AutoModelForCausalLM.from_pretrained( |
| base_model, |
| dtype=dtype, |
| attn_implementation="sdpa", |
| ) |
| letter_weight = full_lm.lm_head.weight[label_ids].detach().clone().float() |
| base_lm = full_lm.model |
| del full_lm |
|
|
| for p in base_lm.parameters(): |
| p.requires_grad_(False) |
|
|
| has_adapter_files = ( |
| lora_r > 0 |
| and adapter_dir is not None |
| and (Path(adapter_dir) / "adapter_config.json").exists() |
| and (Path(adapter_dir) / "adapter_model.safetensors").exists() |
| ) |
| if has_adapter_files: |
| self.lm = PeftModel.from_pretrained(base_lm, str(adapter_dir), is_trainable=False) |
| elif lora_r > 0: |
| if grad_ckpt and hasattr(base_lm, "gradient_checkpointing_enable"): |
| base_lm.gradient_checkpointing_enable( |
| gradient_checkpointing_kwargs={"use_reentrant": False} |
| ) |
| if hasattr(base_lm, "enable_input_require_grads"): |
| base_lm.enable_input_require_grads() |
|
|
| target_map = { |
| "attn": ["q_proj", "k_proj", "v_proj", "o_proj"], |
| "all": ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"], |
| } |
| targets = target_map.get(lora_targets, target_map["all"]) |
| n_layers = len(base_lm.layers) |
| layers_to_transform = ( |
| list(range(max(0, n_layers - top_k_layers), n_layers)) |
| if 0 < top_k_layers < n_layers |
| else None |
| ) |
| try: |
| cfg = LoraConfig( |
| task_type="FEATURE_EXTRACTION", |
| r=lora_r, |
| lora_alpha=2 * lora_r, |
| lora_dropout=0.0, |
| target_modules=targets, |
| layers_to_transform=layers_to_transform, |
| layers_pattern="layers" if layers_to_transform else None, |
| ) |
| self.lm = get_peft_model(base_lm, cfg) |
| except Exception: |
| cfg = LoraConfig( |
| task_type="FEATURE_EXTRACTION", |
| r=lora_r, |
| lora_alpha=2 * lora_r, |
| lora_dropout=0.0, |
| target_modules=targets, |
| ) |
| self.lm = get_peft_model(base_lm, cfg) |
| if 0 < top_k_layers < n_layers: |
| for layer in base_lm.layers[:-top_k_layers]: |
| for p in layer.parameters(): |
| p.requires_grad_(False) |
| else: |
| self.lm = base_lm |
|
|
| self.head = LetterReadoutHead(letter_weight) |
| self.to(device) |
| self.head.float() |
|
|
| def _extract_causal_rows(self, enc: dict[str, Any]) -> list[list[int]]: |
| if "causal_rows" in enc: |
| return enc["causal_rows"] |
| |
| ids = enc["ids"] |
| seg = enc["seg"] |
| state_ids = [tid for tid, s in zip(ids, seg) if s == 0] |
| rows: list[list[int]] = [] |
| for k in range(1, len(enc["decide_idx"]) + 1): |
| br_ids = [tid for tid, s in zip(ids, seg) if s == k] |
| rows.append(state_ids + br_ids) |
| return rows |
|
|
| def _forward_causal_rows(self, flat_rows: list[list[int]]) -> torch.Tensor: |
| """Run a list of causal rows [S + br_k] using native unmasked causal SDPA (is_causal=True) |
| and return the hidden states at each row's terminal `Answer: (` decision token [N_rows, H]. |
| Right-padding after `len(row) - 1` cannot affect causal positions `< len(row)`. |
| """ |
| max_len = max(len(r) for r in flat_rows) |
| if max_len > 1536: |
| max_r = 1 |
| elif max_len > 896: |
| max_r = min(2, self.MAX_ROWS_PER_PASS) |
| else: |
| max_r = self.MAX_ROWS_PER_PASS |
|
|
| if len(flat_rows) > max_r: |
| chunks = [ |
| self._forward_causal_rows(flat_rows[i : i + max_r]) |
| for i in range(0, len(flat_rows), max_r) |
| ] |
| return torch.cat(chunks, dim=0) |
|
|
| L = max_len |
| if str(self.device) == "mps" and self.SHAPE_BUCKET > 1: |
| L = -(-L // self.SHAPE_BUCKET) * self.SHAPE_BUCKET |
| B = len(flat_rows) |
| padded = [r + [self.pad_id] * (L - len(r)) for r in flat_rows] |
| slot_idx = [len(r) - 1 for r in flat_rows] |
| ids = torch.tensor(padded, dtype=torch.long, device=self.device) |
| out = self.lm(input_ids=ids, use_cache=False) |
| h = out.last_hidden_state |
| batch_idx = torch.arange(B, device=self.device) |
| slot_ten = torch.tensor(slot_idx, dtype=torch.long, device=self.device) |
| res = h[batch_idx, slot_ten].float() |
| if str(self.device) == "mps" and max_len > 1536: |
| del out, h, ids |
| torch.mps.empty_cache() |
| return res |
|
|
| @torch.no_grad() |
| def _forward_record_eval(self, enc: dict[str, Any]) -> torch.Tensor: |
| """Fast inference for a single record: reuses the state prefix KV-cache when a record |
| has multiple questions and a long state (`len(state_ids) >= 256`), avoiding redundant state encoding. |
| """ |
| rows = self._extract_causal_rows(enc) |
| state_ids = enc.get("state_ids") |
| if state_ids is None: |
| ids = enc["ids"] |
| seg = enc["seg"] |
| state_ids = [tid for tid, s in zip(ids, seg) if s == 0] |
| L_s = len(state_ids) |
| if len(rows) <= 1 or L_s < 256: |
| return self._forward_causal_rows(rows) |
|
|
| s_ten = torch.tensor([state_ids], dtype=torch.long, device=self.device) |
| out_s = self.lm(input_ids=s_ten, use_cache=True) |
| pkv = out_s.past_key_values |
| h_list: list[torch.Tensor] = [] |
| for r in rows: |
| br = r[L_s:] |
| br_ten = torch.tensor([br], dtype=torch.long, device=self.device) |
| out_b = self.lm(input_ids=br_ten, past_key_values=pkv, use_cache=True) |
| h_list.append(out_b.last_hidden_state[0, -1].float()) |
| pkv.crop(-len(br)) |
| res = torch.stack(h_list, dim=0) |
| del out_s, pkv, s_ten |
| if str(self.device) == "mps": |
| torch.mps.empty_cache() |
| return res |
|
|
| def forward_batch(self, encs: list[dict[str, Any]]) -> list[list[torch.Tensor]]: |
| flat_rows: list[list[int]] = [] |
| counts: list[int] = [] |
| for e in encs: |
| rows = self._extract_causal_rows(e) |
| flat_rows.extend(rows) |
| counts.append(len(rows)) |
|
|
| h_decides = self._forward_causal_rows(flat_rows) |
| out: list[list[torch.Tensor]] = [] |
| offset = 0 |
| for e, c in zip(encs, counts): |
| n_opts = e.get("n_opts") or [len(oi) for oi in e["opt_idx"]] |
| q_types = e.get("q_types") or ["choice"] * len(n_opts) |
| h_sub = h_decides[offset : offset + c] |
| offset += c |
| out.append( |
| [ |
| self.head(h_sub[j], n_opt, q_type=qt) |
| for j, (n_opt, qt) in enumerate(zip(n_opts, q_types)) |
| ] |
| ) |
| return out |
|
|
| def forward(self, enc: dict[str, Any]) -> list[torch.Tensor]: |
| return self.forward_batch([enc])[0] |
|
|
| @torch.no_grad() |
| def raw_logits(self, enc: dict[str, Any]) -> list[torch.Tensor]: |
| h_decides = self._forward_record_eval(enc) |
| n_opts = enc.get("n_opts") or [len(oi) for oi in enc["opt_idx"]] |
| return [self.head.forward_raw(h_decides[j], n_opt).cpu() for j, n_opt in enumerate(n_opts)] |
|
|
| @torch.no_grad() |
| def raw_logits_batch(self, encs: list[dict[str, Any]]) -> list[list[torch.Tensor]]: |
| flat_rows: list[list[int]] = [] |
| counts: list[int] = [] |
| for e in encs: |
| rows = self._extract_causal_rows(e) |
| flat_rows.extend(rows) |
| counts.append(len(rows)) |
| h_decides = self._forward_causal_rows(flat_rows) |
| out: list[list[torch.Tensor]] = [] |
| offset = 0 |
| for e, c in zip(encs, counts): |
| n_opts = e.get("n_opts") or [len(oi) for oi in e["opt_idx"]] |
| h_sub = h_decides[offset : offset + c] |
| offset += c |
| out.append( |
| [self.head.forward_raw(h_sub[j], n_opt).cpu() for j, n_opt in enumerate(n_opts)] |
| ) |
| return out |
|
|
| @torch.no_grad() |
| def probs(self, enc: dict[str, Any]) -> list[torch.Tensor]: |
| return [F.softmax(z, dim=-1).cpu() for z in self.forward(enc)] |
|
|
| @torch.no_grad() |
| def logits_and_probs(self, enc: dict[str, Any]) -> list[tuple[torch.Tensor, torch.Tensor]]: |
| h_decides = self._forward_record_eval(enc) |
| n_opts = enc.get("n_opts") or [len(oi) for oi in enc["opt_idx"]] |
| q_types = enc.get("q_types") or ["choice"] * len(n_opts) |
| zs = [ |
| self.head(h_decides[j], n_opt, q_type=qt) |
| for j, (n_opt, qt) in enumerate(zip(n_opts, q_types)) |
| ] |
| return [(z.cpu(), F.softmax(z, dim=-1).cpu()) for z in zs] |
|
|
| @torch.no_grad() |
| def predict_record( |
| self, |
| rec_internal: dict[str, Any], |
| use_pride: bool = False, |
| max_state: int = SERVE_MAX_STATE, |
| max_branch: int = SERVE_MAX_BRANCH, |
| max_packed: int = SERVE_MAX_PACKED, |
| ) -> list[tuple[torch.Tensor, torch.Tensor]]: |
| """Predict (calibrated_logits, probs) for every question in rec_internal using fast native causal rows.""" |
| self.eval() |
| qs = rec_internal["questions"] |
| if not qs: |
| return [] |
|
|
| if not use_pride: |
| enc_full = encode_record( |
| self.tok, rec_internal, max_state=max_state, max_branch=max_branch, strict=False |
| ) |
| return self.logits_and_probs(enc_full) |
|
|
| |
| all_sub_qs: list[dict[str, Any]] = [] |
| q_specs: list[tuple[str, int, float, list[list[int]] | None]] = [] |
| for q in qs: |
| opts = list(q["options"]) |
| K = len(opts) |
| qt = str(q.get("type") or "choice").lower() |
| type_idx = 1 if qt == "noul" else (2 if qt == "score" else 0) |
| t_type = float(self.head.type_temps[type_idx].item()) |
| temp = t_type if t_type != 1.0 else float(self.head.temperature) |
|
|
| if qt == "score" or K <= 1 or K > 12: |
| all_sub_qs.append(q) |
| q_specs.append((qt, K, temp, None)) |
| else: |
| perms: list[list[int]] = [] |
| for s in range(K): |
| perm = [(j + s) % K for j in range(K)] |
| perms.append(perm) |
| q_shift = copy.deepcopy(q) |
| q_shift["options"] = [opts[idx] for idx in perm] |
| all_sub_qs.append(q_shift) |
| q_specs.append((qt, K, temp, perms)) |
|
|
| sub_enc = encode_record( |
| self.tok, |
| {**rec_internal, "questions": all_sub_qs}, |
| max_state=max_state, |
| max_branch=max_branch, |
| strict=False, |
| ) |
| h_decides = self._forward_record_eval(sub_enc) |
|
|
| results_pride: list[tuple[torch.Tensor, torch.Tensor]] = [] |
| offset = 0 |
| for qt, K, temp, perms in q_specs: |
| if perms is None: |
| z_cal = self.head(h_decides[offset], K, q_type=qt).cpu() |
| offset += 1 |
| results_pride.append((z_cal, F.softmax(z_cal, dim=-1))) |
| else: |
| unpermed_log_probs = np.zeros((K, K), dtype=np.float64) |
| for s, perm in enumerate(perms): |
| z_s = self.head.forward_raw(h_decides[offset + s], K).cpu().numpy().astype(np.float64) |
| z_c = z_s - np.max(z_s) |
| log_p = z_c - np.log(np.sum(np.exp(z_c))) |
| for pos_j, opt_idx in enumerate(perm): |
| unpermed_log_probs[s, opt_idx] = log_p[pos_j] |
| offset += K |
| debiased_logits = unpermed_log_probs.mean(axis=0) |
| debiased_logits = debiased_logits - debiased_logits.mean() |
| cal_logits = debiased_logits / max(temp, 1e-4) |
| z_ten = torch.tensor(cal_logits, dtype=torch.float32) |
| p_ten = F.softmax(z_ten, dim=-1) |
| results_pride.append((z_ten, p_ten)) |
|
|
| return results_pride |
|
|
| def trainable_parameters(self) -> list[nn.Parameter]: |
| return [p for p in self.parameters() if p.requires_grad] |
|
|
| def head_parameters(self) -> list[nn.Parameter]: |
| return [p for p in self.head.parameters() if p.requires_grad] |
|
|
| def lora_parameters(self) -> list[nn.Parameter]: |
| head_ids = {id(p) for p in self.head.parameters()} |
| return [p for p in self.parameters() if p.requires_grad and id(p) not in head_ids] |
|
|
| def save_checkpoint(self, out_dir: str | Path, extra_meta: dict[str, Any] | None = None) -> None: |
| out_path = Path(out_dir) |
| out_path.mkdir(parents=True, exist_ok=True) |
| if self.lora_r > 0 and hasattr(self.lm, "save_pretrained"): |
| self.lm.save_pretrained(str(out_path)) |
| else: |
| for stale in ("adapter_config.json", "adapter_model.safetensors"): |
| sp = out_path / stale |
| if sp.exists(): |
| sp.unlink() |
| type_temps_list = [round(float(x), 4) for x in self.head.type_temps.detach().cpu().tolist()] |
| payload = { |
| "state_dict": self.head.state_dict(), |
| "base_model": self.base_model, |
| "lora_r": self.lora_r, |
| "lora_targets": self.lora_targets, |
| "top_k_layers": self.top_k_layers, |
| "head_dim": self.head_dim, |
| "readout": "native_lm_head_letters_pride", |
| "temperature": float(self.head.temperature), |
| "type_temperatures": { |
| "choice": type_temps_list[0], |
| "noul": type_temps_list[1], |
| "score": type_temps_list[2], |
| }, |
| "meta": extra_meta or {}, |
| } |
| torch.save(payload, out_path / "head.pt") |
| with (out_path / "checkpoint_meta.json").open("w", encoding="utf-8") as f: |
| json.dump( |
| { |
| "base_model": self.base_model, |
| "lora_r": self.lora_r, |
| "lora_targets": self.lora_targets, |
| "top_k_layers": self.top_k_layers, |
| "head_dim": self.head_dim, |
| "readout": "native_lm_head_letters_pride", |
| "temperature": float(self.head.temperature), |
| "type_temperatures": { |
| "choice": type_temps_list[0], |
| "noul": type_temps_list[1], |
| "score": type_temps_list[2], |
| }, |
| **(extra_meta or {}), |
| }, |
| f, |
| indent=2, |
| ensure_ascii=False, |
| ) |
|
|
| @classmethod |
| def load_checkpoint( |
| cls, |
| ckpt_dir: str | Path, |
| device: str = "mps", |
| dtype: torch.dtype = torch.bfloat16, |
| ) -> tuple["MiniCPMSystemOne", Any]: |
| ckpt_path = Path(ckpt_dir) |
| head_pt = ckpt_path / "head.pt" |
| if head_pt.exists(): |
| head_data = torch.load(head_pt, map_location="cpu", weights_only=True) |
| else: |
| head_data = {} |
| base_model = head_data.get("base_model", DEFAULT_BASE_MODEL) |
| lora_r = int(head_data.get("lora_r", 0)) |
| lora_targets = str(head_data.get("lora_targets", "all")) |
| top_k_layers = int(head_data.get("top_k_layers", 0)) |
| head_dim = int(head_data.get("head_dim", 256)) |
| tok = load_tokenizer(base_model) |
| has_adapter = ( |
| lora_r > 0 |
| and (ckpt_path / "adapter_config.json").exists() |
| and (ckpt_path / "adapter_model.safetensors").exists() |
| ) |
| model = cls( |
| base_model=base_model, |
| tok=tok, |
| device=device, |
| lora_r=lora_r, |
| head_dim=head_dim, |
| dtype=dtype, |
| adapter_dir=ckpt_path if has_adapter else None, |
| lora_targets=lora_targets, |
| top_k_layers=top_k_layers, |
| grad_ckpt=False, |
| ) |
| if "state_dict" in head_data and "letter_weight" in head_data["state_dict"]: |
| model.head.load_state_dict(head_data["state_dict"], strict=False) |
| model.head.to(device) |
| model.head.temperature = float( |
| os.environ.get("MINICPM_S1_TEMPERATURE") or head_data.get("temperature", 1.0) |
| ) |
| if os.environ.get("MINICPM_S1_TEMPERATURE"): |
| t_override = float(os.environ["MINICPM_S1_TEMPERATURE"]) |
| model.head.type_temps.fill_(t_override) |
| model.eval() |
| return model, tok |
|
|
|
|
| SCORE_OPTIONS = [f"{i}: score {i} out of 10" for i in range(11)] |
| _LEVEL_PREFIX_RE = re.compile(r"^\s*-?\d+\s*:") |
| _NORM_TOKEN_RE = re.compile(r"[\s_\.\-/]+") |
| _MONTHS = "January|February|March|April|May|June|July|August|September|October|November|December" |
| _DATE_RE = re.compile( |
| rf"\b(?:{_MONTHS}) \d{{1,2}}, \d{{4}}\b|\b\d{{1,2}} (?:{_MONTHS}) \d{{4}}\b|\b\d{{4}}-\d{{2}}-\d{{2}}\b" |
| ) |
|
|
|
|
| def _norm_key(s: str) -> str: |
| return _NORM_TOKEN_RE.sub(" ", str(s or "").strip().lower()).strip() |
|
|
|
|
| def compute_date_facts(text: str, max_dates: int = 7) -> str: |
| """Deterministic date difference helper for states containing 2..max_dates absolute dates.""" |
| found: list[tuple[str, datetime]] = [] |
| for m in _DATE_RE.finditer(text): |
| raw = m.group(0) |
| d = None |
| for fmt in ("%B %d, %Y", "%d %B %Y", "%Y-%m-%d"): |
| try: |
| d = datetime.strptime(raw, fmt) |
| break |
| except ValueError: |
| continue |
| if d is not None and raw not in [r for r, _ in found]: |
| found.append((raw, d)) |
| if len(found) >= max_dates: |
| break |
| if len(found) < 2: |
| return "" |
| facts: list[str] = [] |
| for i in range(len(found)): |
| for j in range(i + 1, len(found)): |
| n = (found[j][1] - found[i][1]).days |
| if n == 0: |
| facts.append(f"{found[j][0]} is the same day as {found[i][0]}.") |
| else: |
| direction = "after" if n > 0 else "before" |
| facts.append( |
| f"{found[j][0]} is {abs(n)} day{'s' if abs(n) != 1 else ''} {direction} {found[i][0]}." |
| ) |
| return " ".join(facts[:10]) |
|
|
|
|
| def render_json_content(v: Any, indent: int = 0) -> str: |
| """Flatten str | dict | list into clean text, unpacking JSON-serialized string objects when present.""" |
| pad = " " * indent |
| if v is None: |
| return "" |
| if isinstance(v, str) and indent == 0: |
| s_strip = v.strip() |
| if (s_strip.startswith("{") and s_strip.endswith("}")) or ( |
| s_strip.startswith("[") and s_strip.endswith("]") |
| ): |
| try: |
| parsed = json.loads(s_strip) |
| if isinstance(parsed, (dict, list)): |
| return render_json_content(parsed, indent=0) |
| except Exception: |
| pass |
| return v |
| if isinstance(v, (str, int, float, bool)): |
| return str(v) |
| if isinstance(v, list): |
| return "\n".join(f"{pad}- {render_json_content(x, indent + 1).lstrip()}" for x in v) |
| if isinstance(v, dict): |
| return "\n".join( |
| f"{pad}{k}:\n{render_json_content(x, indent + 1)}" |
| if isinstance(x, (dict, list)) |
| else f"{pad}{k}: {render_json_content(x)}" |
| for k, x in v.items() |
| ) |
| return str(v) |
|
|
|
|
| def option_text(name: str, desc: Any) -> str: |
| """Render option name and optional description while suppressing redundant 'name: name' duplicates.""" |
| if desc is None or desc == "": |
| return str(name) |
| rendered = render_json_content(desc).strip() |
| if not rendered or _norm_key(rendered) == _norm_key(name): |
| return str(name) |
| return f"{name}: {rendered}" |
|
|
|
|
| def to_internal_record(rec: dict[str, Any], add_date_facts: bool = False) -> dict[str, Any]: |
| """Convert wire-format record (questions dict) into internal question list for encoding.""" |
| qs = [] |
| for name, q in rec["questions"].items(): |
| qtype = str(q["type"]).lower() |
| instr_str = render_json_content(q.get("instructions")) |
| if not instr_str and qtype == "noul": |
| instr_str = "Which answer fits the context?" |
| target_probs: list[float] | None = None |
|
|
| if qtype == "choice": |
| crit = q.get("criteria") or {} |
| if isinstance(crit, list): |
| crit = {str(x): None for x in crit} |
| keys = list(crit.keys()) |
| opts = [option_text(str(k), crit[k]) for k in keys] |
| lbl = q.get("label", q.get("answer", keys[0] if keys else 0)) |
| label = keys.index(lbl) if lbl in keys else int(lbl) |
| elif qtype == "noul": |
| crit = q.get("criteria") if isinstance(q.get("criteria"), dict) else {} |
| keys = ["false", "true"] |
| f_desc = crit.get("false", crit.get(False, "false, negative, does not hold")) |
| t_desc = crit.get("true", crit.get(True, "true, affirmative, holds")) |
| opts = [ |
| option_text("no", f_desc or "false, negative, does not hold"), |
| option_text("yes", t_desc or "true, affirmative, holds"), |
| ] |
| lbl = q.get("label", q.get("answer", False)) |
| if isinstance(lbl, bool): |
| label = int(lbl) |
| elif str(lbl).lower() in ("true", "yes", "1"): |
| label = 1 |
| else: |
| label = 0 |
| elif qtype == "score": |
| crit = q.get("criteria") |
| if isinstance(crit, dict): |
| crit = [crit[k] for k in sorted(crit, key=lambda x: float(x))] |
| if isinstance(crit, list) and crit: |
| opts = [] |
| for i, x in enumerate(crit): |
| rendered_x = render_json_content(x).strip() |
| if _LEVEL_PREFIX_RE.match(rendered_x): |
| opts.append(rendered_x) |
| else: |
| opts.append(f"{i}: {rendered_x}") |
| keys = [str(i) for i in range(len(opts))] |
| else: |
| opts = list(SCORE_OPTIONS) |
| keys = [str(i) for i in range(11)] |
| label = int(q.get("label", q.get("answer", 0))) |
| else: |
| raise ValueError(f"Unknown question type: {qtype}") |
|
|
| if isinstance(q.get("target_probs"), list) and len(q["target_probs"]) == len(opts): |
| target_probs = [float(p) for p in q["target_probs"]] |
| elif isinstance(q.get("target"), dict): |
| tgt_map = q["target"] |
| probs_raw = [float(tgt_map.get(k, 1.0 / len(keys))) for k in keys] |
| s_tot = sum(probs_raw) or 1.0 |
| target_probs = [p / s_tot for p in probs_raw] |
|
|
| qs.append( |
| { |
| "name": name, |
| "type": qtype, |
| "instr": f"[{qtype}] {instr_str}", |
| "options": opts, |
| "keys": keys, |
| "label": max(0, min(label, len(opts) - 1)), |
| "target_probs": target_probs, |
| } |
| ) |
| meta = rec.get("_meta") or {} |
| state_str = render_json_content(rec["state"]) |
| if add_date_facts and "date_facts:" not in state_str: |
| df = compute_date_facts(state_str) |
| if df: |
| state_str = f"{state_str}\n\ndate_facts: {df}" |
| return { |
| "id": rec.get("id") or meta.get("id", ""), |
| "source": rec.get("source") or meta.get("source", "custom"), |
| "state": state_str, |
| "questions": qs, |
| } |
|
|
|
|
