File size: 7,442 Bytes
e79feac | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 | """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 = " " # token form tried first: " A" (leading space), then "A"
gradient_checkpointing: bool = True
prompt_format: str = "raw" # raw: completion prompt ending in "Answer:" | chat: chat template (see module doc)
max_think_tokens: int = 1024 # thinking-on scoring: reasoning budget before "</think>" is forced
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" # last position = decision position
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."""
# logits_to_keep=1: only the decision position's vocabulary logits (same values; avoids [B, L, V] at long L)
out = self.lm(input_ids=input_ids, attention_mask=attention_mask, logits_to_keep=1)
last = out.logits[:, -1, :].float() # [B, V]
sym = last[:, self.symbol_ids] # [B, k_max+1]
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:
# right truncation would cut the decision position (the last token); never truncate silently
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)
|