| """Direct decision head on a small decoder (plan §6: 'same family with a direct decision head'). |
| |
| A causal LM (1-4B, LoRA-adapted) reads the serialized query + candidate list with symbols A, B, C, ... |
| and NONE. The decision is read from the next-token logits at the final position, restricted to the |
| symbol tokens. Single-pass; no autoregressive generation. Tokenization of every symbol is verified to be |
| exactly one token (plan §6: a decimal integer is not necessarily one token). |
| |
| prompt_format "chat" (Nov 8, 2026; post hoc arm) wraps the same prompt in the backbone's chat template. For |
| hybrid-thinking models (Qwen3) the thinking-off template ends in an empty "<think>\n\n</think>\n\n" block, so |
| training and thinking-off scoring read the symbol right after it. Thinking-on scoring (decoder_trainer. |
| predict_decoder_thinking) lets the model write its own reasoning inside that block first, then reads the symbol |
| logits at the same position after "</think>\n\n". |
| """ |
| from __future__ import annotations |
|
|
| from dataclasses import dataclass |
| from typing import Any |
|
|
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
| from ..data.serialization import NONE_SYMBOL, SYMBOLS, serialize_record |
|
|
|
|
| @dataclass |
| class DecoderHeadConfig: |
| backbone: str = "Qwen/Qwen2.5-1.5B-Instruct" |
| max_len: int = 2048 |
| serialization: str = "compact_v1" |
| lora_r: int = 16 |
| lora_alpha: int = 32 |
| lora_dropout: float = 0.05 |
| lora_targets: tuple[str, ...] = ("q_proj", "k_proj", "v_proj", "o_proj") |
| dtype: str = "bfloat16" |
| symbol_prefix: str = " " |
| gradient_checkpointing: bool = True |
| prompt_format: str = "raw" |
| max_think_tokens: int = 1024 |
|
|
| def to_dict(self) -> dict[str, Any]: |
| d = self.__dict__.copy() |
| d["lora_targets"] = list(self.lora_targets) |
| return d |
|
|
|
|
| PROMPT_HEADER = ( |
| "You are matching a query record to at most one catalog record.\n" |
| "Identity policy: {policy}\n" |
| "Query: {query}\n" |
| "Candidates:\n{candidates}\n" |
| "Answer with the single symbol of the matching candidate, or {none} if no candidate is the same entity.\n" |
| "Answer:" |
| ) |
|
|
|
|
| def render_prompt(query_fields: dict[str, Any], candidates: list[dict[str, Any]], policy: str, serialization: str, field_order: list[str] | None = None, compact_fields: list[str] | None = None) -> str: |
| q = serialize_record(query_fields, field_order, serialization, compact_fields) |
| lines = [f"{SYMBOLS[j]}. {serialize_record(c, field_order, serialization, compact_fields)}" for j, c in enumerate(candidates)] |
| return PROMPT_HEADER.format(policy=policy, query=q, candidates="\n".join(lines), none=NONE_SYMBOL) |
|
|
|
|
| def verify_symbol_tokens(tokenizer, k_max: int, prefix_candidates: tuple[str, ...] = (" ", "")) -> tuple[str, list[int], int]: |
| """Return (prefix, candidate_token_ids[k_max], none_token_id) such that every symbol is ONE token. |
| |
| Raises if no prefix variant makes all symbols single tokens. |
| """ |
| problems = {} |
| for prefix in prefix_candidates: |
| ids = [] |
| ok = True |
| for s in SYMBOLS[:k_max] + [NONE_SYMBOL]: |
| toks = tokenizer.encode(prefix + s, add_special_tokens=False) |
| if len(toks) != 1: |
| ok = False |
| problems[prefix + s] = toks |
| break |
| ids.append(toks[0]) |
| if ok and len(set(ids)) == len(ids): |
| return prefix, ids[:-1], ids[-1] |
| raise ValueError(f"no single-token symbol encoding for K={k_max}; offending: {problems}. Reduce K or change symbols.") |
|
|
|
|
| class DecoderDecisionHead(nn.Module): |
| def __init__(self, cfg: DecoderHeadConfig, k_max: int, use_lora: bool = True): |
| super().__init__() |
| from transformers import AutoModelForCausalLM, AutoTokenizer |
|
|
| self.cfg = cfg |
| self.tokenizer = AutoTokenizer.from_pretrained(cfg.backbone) |
| if self.tokenizer.pad_token is None: |
| self.tokenizer.pad_token = self.tokenizer.eos_token |
| self.tokenizer.padding_side = "left" |
| dtype = getattr(torch, cfg.dtype) |
| self.lm = AutoModelForCausalLM.from_pretrained(cfg.backbone, torch_dtype=dtype, attn_implementation="sdpa") |
| if cfg.gradient_checkpointing: |
| self.lm.gradient_checkpointing_enable() |
| self.lm.config.use_cache = False |
| if use_lora: |
| from peft import LoraConfig, get_peft_model |
|
|
| self.lm = get_peft_model(self.lm, LoraConfig(r=cfg.lora_r, lora_alpha=cfg.lora_alpha, lora_dropout=cfg.lora_dropout, target_modules=list(cfg.lora_targets), task_type="CAUSAL_LM")) |
| self.prefix, cand_ids, none_id = verify_symbol_tokens(self.tokenizer, k_max, (cfg.symbol_prefix, "" if cfg.symbol_prefix else " ")) |
| self.register_buffer("symbol_ids", torch.tensor(cand_ids + [none_id], dtype=torch.long), persistent=False) |
| self.k_max = k_max |
|
|
| def forward(self, input_ids: torch.Tensor, attention_mask: torch.Tensor, slot_mask: torch.Tensor) -> torch.Tensor: |
| """Return log-probs [B, k_max+1] over symbols (padded slots -inf). Left-padded inputs.""" |
| |
| out = self.lm(input_ids=input_ids, attention_mask=attention_mask, logits_to_keep=1) |
| last = out.logits[:, -1, :].float() |
| sym = last[:, self.symbol_ids] |
| full_mask = torch.cat([slot_mask, torch.ones(slot_mask.shape[0], 1, dtype=torch.bool, device=slot_mask.device)], 1) |
| sym = sym.masked_fill(~full_mask, float("-inf")) |
| return F.log_softmax(sym, dim=-1) |
|
|
| def format_prompts(self, prompts: list[str], thinking: bool = False) -> list[str]: |
| """Model input text for rendered prompts: unchanged for "raw"; chat template for "chat".""" |
| if self.cfg.prompt_format == "raw": |
| if thinking: |
| raise ValueError("thinking needs prompt_format: chat") |
| return prompts |
| if self.cfg.prompt_format != "chat": |
| raise ValueError(self.cfg.prompt_format) |
| suffix = "\nAnswer:" |
| return [self.tokenizer.apply_chat_template([{"role": "user", "content": p.removesuffix(suffix)}], tokenize=False, |
| add_generation_prompt=True, enable_thinking=thinking) for p in prompts] |
|
|
| def tokenize_prompts(self, prompts: list[str], formatted: bool = False, max_len: int | None = None) -> dict[str, torch.Tensor]: |
| texts = prompts if formatted else self.format_prompts(prompts) |
| enc = self.tokenizer(texts, return_tensors="pt", padding=True, truncation=False, add_special_tokens=self.cfg.prompt_format == "raw") |
| limit = max_len or self.cfg.max_len |
| if enc["input_ids"].shape[1] > limit: |
| |
| raise ValueError(f"prompt of {enc['input_ids'].shape[1]} tokens exceeds max_len {limit}") |
| return {"input_ids": enc["input_ids"], "attention_mask": enc["attention_mask"]} |
|
|
| def save_adapter(self, path: str) -> None: |
| self.lm.save_pretrained(path) |
| self.tokenizer.save_pretrained(path) |
|
|
