Text Classification
Safetensors
English
Chinese
cybersecurity
cyber-security
security
cyber-defense
blue-team
soc
security-operations
alert-triage
threat-detection
threat-intelligence
cti
vulnerability-detection
vulnerability-management
cve
cwe
mitre-attack
cloud-security
kubernetes-security
devsecops
edge-ai
low-latency
encoder
decision-model
calibration
Eval Results (legacy)
| """Jue-Flash bidirectional encoder decision model with Multi-Layer Gating, Vectorized Attentive Span Pooling, | |
| and Dual-Path (Global Context + Late-Interaction MaxSim) Readout for Choice / Score / Noul primitives. | |
| Supports both ``v1`` checkpoints and ``v2`` architectural upgrades: | |
| 1. Query-Conditioned Multi-Depth Layer Gate across 6 transformer depths ``(-16, -12, -8, -4, -2, -1)``. | |
| 2. Multi-Hop Question/Candidate-to-State Cross-Attention refinement block. | |
| 3. Unimodal Ordinal Score Coupler combining continuous expected-severity regression with per-level span logits. | |
| """ | |
| from __future__ import annotations | |
| import hashlib | |
| import json | |
| import math | |
| from collections.abc import Mapping | |
| from dataclasses import asdict, dataclass | |
| from pathlib import Path | |
| from typing import Any | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| from safetensors.torch import load_file, save_file | |
| from torch import Tensor | |
| from transformers import AutoConfig, AutoModel | |
| from jue_flash.api.schema import MAX_SCORE_LEVELS | |
| from jue_flash.encoding.pack import PackedFlashBatch | |
| HEAD_WEIGHTS_FILE = "flash_decision_head.safetensors" | |
| ENCODER_WEIGHTS_FILE = "encoder_model.safetensors" | |
| FLASH_CONFIG_FILE = "jue_flash_config.json" | |
| class FlashModelConfig: | |
| """Configuration for :class:`FlashDecisionModel`.""" | |
| base_model: str = "jhu-clsp/mmbert-base" | |
| hidden_size: int = 768 | |
| num_gated_layers: int = 4 | |
| late_interaction_dim: int = 128 | |
| dropout: float = 0.1 | |
| max_len: int = 2048 | |
| # v2 architectural flags (default False/0 so v1 checkpoints load identically): | |
| query_conditioned_layer_gate: bool = False | |
| tapped_layer_offsets: tuple[int, ...] = (-16, -12, -8, -4, -2, -1) | |
| cross_attn_hops: int = 0 | |
| unimodal_score_coupler: bool = False | |
| # v3 architectural flags: | |
| canon_synergy_fusion: bool = False | |
| candidate_self_attn: bool = False | |
| def save(self, path: str | Path) -> None: | |
| p = Path(path) | |
| p.parent.mkdir(parents=True, exist_ok=True) | |
| d = asdict(self) | |
| d["tapped_layer_offsets"] = list(self.tapped_layer_offsets) | |
| p.write_text(json.dumps(d, indent=2, ensure_ascii=False) + "\n", encoding="utf-8") | |
| def load(cls, path: str | Path) -> FlashModelConfig: | |
| raw = json.loads(Path(path).read_text(encoding="utf-8")) | |
| if "tapped_layer_offsets" in raw and isinstance(raw["tapped_layer_offsets"], list): | |
| raw["tapped_layer_offsets"] = tuple(int(x) for x in raw["tapped_layer_offsets"]) | |
| return cls(**raw) | |
| class LayerGate(nn.Module): | |
| """Learned softmax gate over the top ``num_layers`` hidden states of the encoder (v1).""" | |
| def __init__(self, num_layers: int, hidden_size: int) -> None: | |
| super().__init__() | |
| self.num_layers = num_layers | |
| self.layer_weights = nn.Parameter(torch.zeros(num_layers)) | |
| self.norm = nn.LayerNorm(hidden_size) | |
| def forward(self, hidden_states: tuple[Tensor, ...]) -> Tensor: | |
| layers = hidden_states[-self.num_layers :] | |
| stacked = torch.stack(layers, dim=0) # [L_g, B, T, D] | |
| w = torch.softmax(self.layer_weights, dim=0).view(-1, 1, 1, 1) | |
| fused = (stacked * w).sum(dim=0) | |
| return self.norm(fused) | |
| class QueryConditionedLayerGate(nn.Module): | |
| """Dynamic per-request softmax gate across multi-depth encoder layers (v2). | |
| Taps shallow/middle syntactic layers and deep semantic layers (e.g. ``(-16, -12, -8, -4, -2, -1)``), | |
| conditioning the layer mixture weights on each request's top-layer summary vector so code/regex/AST | |
| questions can attend to intermediate layers while policy/reasoning questions attend to top layers. | |
| """ | |
| def __init__(self, offsets: tuple[int, ...], hidden_size: int) -> None: | |
| super().__init__() | |
| self.offsets = offsets | |
| n_taps = len(offsets) | |
| self.layer_weights = nn.Parameter(torch.zeros(n_taps)) | |
| self.query_norm = nn.LayerNorm(hidden_size) | |
| self.query_proj = nn.Linear(hidden_size, n_taps) | |
| nn.init.zeros_(self.query_proj.weight) | |
| nn.init.zeros_(self.query_proj.bias) | |
| self.norm = nn.LayerNorm(hidden_size) | |
| def forward(self, hidden_states: tuple[Tensor, ...]) -> Tensor: | |
| n_total = len(hidden_states) | |
| selected = [hidden_states[max(-n_total, min(-1, off))] for off in self.offsets] | |
| stacked = torch.stack(selected, dim=0) # [L_g, B, T, D] | |
| # Condition on the top layer's CLS + mean summary per batch row: [B, D] | |
| top_h = hidden_states[-1] | |
| row_summary = 0.5 * (top_h[:, 0, :] + top_h.mean(dim=1)) | |
| dyn_logits = self.layer_weights.unsqueeze(0) + self.query_proj(self.query_norm(row_summary)) # [B, L_g] | |
| w = torch.softmax(dyn_logits, dim=-1).transpose(0, 1).unsqueeze(-1).unsqueeze(-1) # [L_g, B, 1, 1] | |
| fused = (stacked * w).sum(dim=0) | |
| return self.norm(fused) | |
| class StateCrossAttentionHop(nn.Module): | |
| """Lightweight Multi-Head Cross-Attention hop refining ``[q_vec, c_1, ..., c_K]`` over state tokens (v2).""" | |
| def __init__(self, hidden_size: int, num_heads: int = 4, dropout: float = 0.1) -> None: | |
| super().__init__() | |
| self.num_heads = num_heads | |
| self.head_dim = hidden_size // num_heads | |
| self.q_norm = nn.LayerNorm(hidden_size) | |
| self.kv_norm = nn.LayerNorm(hidden_size) | |
| self.q_proj = nn.Linear(hidden_size, hidden_size, bias=False) | |
| self.k_proj = nn.Linear(hidden_size, hidden_size, bias=False) | |
| self.v_proj = nn.Linear(hidden_size, hidden_size, bias=False) | |
| self.out_proj = nn.Linear(hidden_size, hidden_size) | |
| nn.init.zeros_(self.out_proj.weight) | |
| nn.init.zeros_(self.out_proj.bias) | |
| self.dropout = nn.Dropout(dropout) | |
| self.ffn_norm = nn.LayerNorm(hidden_size) | |
| self.ffn = nn.Sequential( | |
| nn.Linear(hidden_size, hidden_size * 2), | |
| nn.SiLU(), | |
| nn.Dropout(dropout), | |
| nn.Linear(hidden_size * 2, hidden_size), | |
| ) | |
| nn.init.zeros_(self.ffn[-1].weight) # type: ignore[arg-type] | |
| nn.init.zeros_(self.ffn[-1].bias) # type: ignore[arg-type] | |
| def forward( | |
| self, | |
| qc_tokens: Tensor, # [N, 1+K, D] | |
| state_seq_b: Tensor, # [B, L, D] | |
| state_mask_b: Tensor, # [B, L] bool | |
| q_b_idx: Tensor, # [N] | |
| ) -> Tensor: | |
| n_q, n_qc, d = qc_tokens.shape | |
| bsz, seq_len, _ = state_seq_b.shape | |
| h = self.num_heads | |
| dh = self.head_dim | |
| kv_in = self.kv_norm(state_seq_b) | |
| k_b = self.k_proj(kv_in).view(bsz, seq_len, h, dh).transpose(1, 2) # [B, H, L, dh] | |
| v_b = self.v_proj(kv_in).view(bsz, seq_len, h, dh).transpose(1, 2) # [B, H, L, dh] | |
| q_in = self.q_norm(qc_tokens) | |
| q = self.q_proj(q_in).view(n_q, n_qc, h, dh).transpose(1, 2) # [N, H, 1+K, dh] | |
| k = k_b[q_b_idx] # [N, H, L, dh] | |
| v = v_b[q_b_idx] # [N, H, L, dh] | |
| s_mask = state_mask_b[q_b_idx].unsqueeze(1).unsqueeze(2) # [N, 1, 1, L] | |
| scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(dh) # [N, H, 1+K, L] | |
| scores = scores.float().masked_fill(~s_mask, -1e9) | |
| attn = torch.softmax(scores, dim=-1).to(qc_tokens.dtype) | |
| attn = self.dropout(attn) | |
| context = torch.matmul(attn, v).transpose(1, 2).contiguous().view(n_q, n_qc, d) | |
| qc_tokens = qc_tokens + self.out_proj(context) | |
| qc_tokens = qc_tokens + self.ffn(self.ffn_norm(qc_tokens)) | |
| return qc_tokens | |
| class VectorizedAttentiveSpanPooler(nn.Module): | |
| """Attention + mean residual span pooler operating without Python loops over tokens.""" | |
| def __init__(self, hidden_size: int) -> None: | |
| super().__init__() | |
| self.attn_proj = nn.Sequential( | |
| nn.Linear(hidden_size, hidden_size // 2), | |
| nn.Tanh(), | |
| nn.Linear(hidden_size // 2, 1, bias=False), | |
| ) | |
| self.norm = nn.LayerNorm(hidden_size) | |
| def forward(self, token_states: Tensor, token_attn_scores: Tensor, spans: Tensor) -> Tensor: | |
| """Pool spans from ``token_states``.""" | |
| m_size, seq_len, hidden_dim = token_states.shape | |
| pos = torch.arange(seq_len, device=token_states.device) | |
| if spans.dim() == 2: | |
| starts = spans[:, 0:1].clamp(0, seq_len - 1) | |
| ends = spans[:, 1:2].clamp(1, seq_len) | |
| ends = torch.maximum(ends, starts + 1) | |
| mask = (pos.unsqueeze(0) >= starts) & (pos.unsqueeze(0) < ends) # [M, L] | |
| masked_scores = token_attn_scores.float().masked_fill(~mask, -1e9) | |
| weights = torch.softmax(masked_scores, dim=-1).to(token_states.dtype).unsqueeze(-1) # [M, L, 1] | |
| attn_pooled = (token_states * weights).sum(dim=1) # [M, D] | |
| mask_f = mask.to(token_states.dtype).unsqueeze(-1) | |
| mean_pooled = (token_states * mask_f).sum(dim=1) / mask_f.sum(dim=1).clamp_min(1.0) | |
| return self.norm(0.5 * (attn_pooled + mean_pooled)) | |
| if spans.dim() == 3: | |
| starts = spans[:, :, 0:1].clamp(0, seq_len - 1) | |
| ends = spans[:, :, 1:2].clamp(1, seq_len) | |
| ends = torch.maximum(ends, starts + 1) | |
| p_grid = pos.view(1, 1, seq_len) | |
| mask = (p_grid >= starts) & (p_grid < ends) # [M, K, L] | |
| scores_exp = token_attn_scores.float().unsqueeze(1).expand(-1, spans.shape[1], -1) # [M, K, L] | |
| masked_scores = scores_exp.masked_fill(~mask, -1e9) | |
| weights = torch.softmax(masked_scores, dim=-1).to(token_states.dtype) # [M, K, L] | |
| attn_pooled = torch.bmm(weights, token_states) # [M, K, D] | |
| mask_f = mask.to(token_states.dtype) | |
| mean_weights = mask_f / mask_f.sum(dim=-1, keepdim=True).clamp_min(1.0) | |
| mean_pooled = torch.bmm(mean_weights, token_states) # [M, K, D] | |
| return self.norm(0.5 * (attn_pooled + mean_pooled)) | |
| raise ValueError(f"Expected spans with 2 or 3 dims, got {spans.shape}") | |
| class CandidateSelfAttention(nn.Module): | |
| """Pure neural option-to-option self-attention and mean-contrast scorer for closed-book MCQs.""" | |
| def __init__(self, hidden_size: int, num_heads: int = 4, dropout: float = 0.1) -> None: | |
| super().__init__() | |
| self.num_heads = num_heads | |
| self.head_dim = hidden_size // num_heads | |
| self.norm1 = nn.LayerNorm(hidden_size) | |
| self.qkv = nn.Linear(hidden_size, 3 * hidden_size, bias=False) | |
| self.out_proj = nn.Linear(hidden_size, hidden_size) | |
| nn.init.zeros_(self.out_proj.weight) | |
| nn.init.zeros_(self.out_proj.bias) | |
| self.dropout = nn.Dropout(dropout) | |
| # Contrastive readout over [c_vec, c_diff, c_vec * q_exp, c_diff * q_exp] -> [N, K, 1] | |
| self.contrast_mlp = nn.Sequential( | |
| nn.Linear(4 * hidden_size, hidden_size // 2), | |
| nn.LayerNorm(hidden_size // 2), | |
| nn.SiLU(), | |
| nn.Dropout(dropout), | |
| nn.Linear(hidden_size // 2, 1), | |
| ) | |
| nn.init.zeros_(self.contrast_mlp[-1].weight) # type: ignore[arg-type] | |
| nn.init.zeros_(self.contrast_mlp[-1].bias) # type: ignore[arg-type] | |
| # Context-adaptive gate for late-interaction lexical similarity | |
| self.late_gate = nn.Sequential( | |
| nn.Linear(2 * hidden_size, hidden_size // 4), | |
| nn.LayerNorm(hidden_size // 4), | |
| nn.SiLU(), | |
| nn.Linear(hidden_size // 4, 1), | |
| ) | |
| nn.init.zeros_(self.late_gate[-1].weight) # type: ignore[arg-type] | |
| nn.init.zeros_(self.late_gate[-1].bias) # type: ignore[arg-type] | |
| def forward( | |
| self, | |
| c_vec: Tensor, # [N, K, D] | |
| q_vec: Tensor, # [N, D] | |
| state_vec: Tensor, # [N, D] | |
| cand_mask: Tensor, # [N, K] bool | |
| ) -> tuple[Tensor, Tensor, Tensor]: | |
| n_q, k_max, d = c_vec.shape | |
| h = self.num_heads | |
| dh = self.head_dim | |
| x = self.norm1(c_vec) | |
| qkv = self.qkv(x).view(n_q, k_max, 3, h, dh).permute(2, 0, 3, 1, 4) # [3, N, H, K, dh] | |
| q, k, v = qkv[0], qkv[1], qkv[2] | |
| scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(dh) # [N, H, K, K] | |
| k_mask = cand_mask.unsqueeze(1).unsqueeze(2) # [N, 1, 1, K] | |
| scores = scores.float().masked_fill(~k_mask, -1e9) | |
| attn = torch.softmax(scores, dim=-1).to(c_vec.dtype) | |
| attn = self.dropout(attn) | |
| ctx = torch.matmul(attn, v).transpose(1, 2).contiguous().view(n_q, k_max, d) | |
| c_refined = c_vec + self.out_proj(ctx) | |
| mask_f = cand_mask.to(c_vec.dtype).unsqueeze(-1) # [N, K, 1] | |
| c_mean = (c_refined * mask_f).sum(dim=1, keepdim=True) / mask_f.sum(dim=1, keepdim=True).clamp_min(1.0) | |
| c_diff = c_refined - c_mean | |
| q_exp = q_vec.unsqueeze(1).expand(n_q, k_max, d) | |
| contrast_in = torch.cat([c_refined, c_diff, c_refined * q_exp, c_diff * q_exp], dim=-1) | |
| contrast_logits = self.contrast_mlp(contrast_in).squeeze(-1) # [N, K] | |
| # Late-interaction gate in (0, 2), starting at 1.0 at initialization | |
| lg_in = torch.cat([q_vec, state_vec], dim=-1) | |
| late_mult = 1.0 + torch.tanh(self.late_gate(lg_in)) # [N, 1] | |
| return c_refined, contrast_logits, late_mult | |
| class DualPathDecisionHead(nn.Module): | |
| """Three-primitive (Choice / Score / Noul) readout combining: | |
| - Path A: Global Contextual Interaction MLP + Bilinear scorer | |
| - Path B: Late-Interaction Token Alignment (ColBERT-style MaxSim between option tokens and state tokens) | |
| - Optional v2 upgrades: QueryConditionedLayerGate, StateCrossAttentionHop, and Unimodal Score Coupler. | |
| - Optional v3 upgrades: CandidateSelfAttention option cross-comparison. | |
| """ | |
| def __init__(self, cfg: FlashModelConfig) -> None: | |
| super().__init__() | |
| self.cfg = cfg | |
| d = cfg.hidden_size | |
| if cfg.query_conditioned_layer_gate: | |
| self.layer_gate: nn.Module = QueryConditionedLayerGate(cfg.tapped_layer_offsets, d) | |
| else: | |
| self.layer_gate = LayerGate(cfg.num_gated_layers, d) | |
| self.span_pooler = VectorizedAttentiveSpanPooler(d) | |
| # Primitive embeddings: 0=choice, 1=score, 2=noul | |
| self.primitive_emb = nn.Embedding(3, d) | |
| # Ordinal level embeddings for Score (0..MAX_SCORE_LEVELS-1) | |
| self.score_level_emb = nn.Embedding(MAX_SCORE_LEVELS + 2, d) | |
| # Pole embeddings for Noul (0=false, 1=true) | |
| self.noul_pole_emb = nn.Embedding(2, d) | |
| # Optional v2 Multi-Hop Question/Candidate-to-State Cross-Attention | |
| if cfg.cross_attn_hops > 0: | |
| self.cross_hops = nn.ModuleList( | |
| [StateCrossAttentionHop(d, num_heads=4, dropout=cfg.dropout) for _ in range(cfg.cross_attn_hops)] | |
| ) | |
| else: | |
| self.cross_hops = None | |
| # Optional v2 Unimodal Ordinal Score Coupler | |
| if cfg.unimodal_score_coupler: | |
| self.ordinal_regressor = nn.Sequential( | |
| nn.Linear(4 * d, d // 2), | |
| nn.LayerNorm(d // 2), | |
| nn.SiLU(), | |
| nn.Dropout(cfg.dropout), | |
| nn.Linear(d // 2, 1), | |
| ) | |
| self.ordinal_log_precision = nn.Parameter(torch.tensor(math.log(3.0))) | |
| else: | |
| self.ordinal_regressor = None | |
| self.ordinal_log_precision = None | |
| # Path A: Global context-conditioned MLP + Bilinear interaction | |
| # Features: [c, q, s, c * q, c * s, |c - s|] -> 6 * D | |
| self.context_mlp = nn.Sequential( | |
| nn.Linear(6 * d, d), | |
| nn.LayerNorm(d), | |
| nn.SiLU(), | |
| nn.Dropout(cfg.dropout), | |
| nn.Linear(d, d // 2), | |
| nn.SiLU(), | |
| nn.Linear(d // 2, 1), | |
| ) | |
| self.bilinear_q = nn.Linear(d, d, bias=False) | |
| self.bilinear_s = nn.Linear(d, d, bias=False) | |
| # Path B: Late-Interaction token projection for exact identifier / telemetry / code alignment | |
| self.late_proj = nn.Linear(d, cfg.late_interaction_dim, bias=False) | |
| self.late_scale = nn.Parameter(torch.tensor(1.0)) | |
| # Optional v3 Path C: Canon Specification Prior Synergy Gate (disabled in pure closed-book) | |
| self.canon_scale: nn.Parameter | None = None | |
| self.canon_gate: nn.Module | None = None | |
| if cfg.canon_synergy_fusion: | |
| self.enable_canon_synergy(init_scale=2.2) | |
| # Optional v3 Pure Neural Candidate Self-Attention & Contrast Block | |
| self.cand_attn: CandidateSelfAttention | None = None | |
| if cfg.candidate_self_attn: | |
| self.enable_candidate_self_attn() | |
| def enable_canon_synergy(self, init_scale: float = 2.2) -> None: | |
| """Initialize or re-enable the v3 Canon Specification Prior Synergy module.""" | |
| d = self.cfg.hidden_size | |
| self.cfg.canon_synergy_fusion = True | |
| self.canon_scale = nn.Parameter(torch.tensor(float(init_scale))) | |
| gate = nn.Sequential( | |
| nn.Linear(2 * d + 2, d // 4), | |
| nn.LayerNorm(d // 4), | |
| nn.SiLU(), | |
| nn.Linear(d // 4, 1), | |
| ) | |
| nn.init.zeros_(gate[-1].weight) # type: ignore[arg-type] | |
| nn.init.zeros_(gate[-1].bias) # type: ignore[arg-type] | |
| self.canon_gate = gate | |
| def enable_candidate_self_attn(self) -> None: | |
| """Initialize the pure neural CandidateSelfAttention option cross-comparison module.""" | |
| self.cfg.candidate_self_attn = True | |
| self.cand_attn = CandidateSelfAttention( | |
| hidden_size=self.cfg.hidden_size, | |
| num_heads=4, | |
| dropout=self.cfg.dropout, | |
| ) | |
| def _compute_late_interaction( | |
| self, | |
| late_tokens_b: Tensor, # [B, L, d_late] normalized | |
| state_spans_b: Tensor, # [B, 2] | |
| q_b_idx: Tensor, # [N] | |
| opt_spans: Tensor, # [N, K, 2] | |
| cand_mask: Tensor, # [N, K] | |
| ) -> Tensor: | |
| """Compute ColBERT-style MaxSim between each candidate option's tokens and the state's tokens.""" | |
| bsz, seq_len, _ = late_tokens_b.shape | |
| pos = torch.arange(seq_len, device=late_tokens_b.device) | |
| s_starts = state_spans_b[:, 0:1].clamp(0, seq_len - 1) | |
| s_ends = torch.maximum(state_spans_b[:, 1:2].clamp(1, seq_len), s_starts + 1) | |
| s_mask = (pos.unsqueeze(0) >= s_starts) & (pos.unsqueeze(0) < s_ends) # [B, L] | |
| sim_matrix = torch.bmm(late_tokens_b, late_tokens_b.transpose(1, 2)) # [B, L_all, L_state] | |
| sim_matrix = sim_matrix.masked_fill(~s_mask.unsqueeze(1), -1e4) | |
| max_sim_to_state_b = sim_matrix.max(dim=-1).values # [B, L] | |
| max_sim_to_state_q = max_sim_to_state_b[q_b_idx] # [N, L] | |
| o_starts = opt_spans[:, :, 0:1].clamp(0, seq_len - 1) | |
| o_ends = torch.maximum(opt_spans[:, :, 1:2].clamp(1, seq_len), o_starts + 1) | |
| o_mask = (pos.view(1, 1, seq_len) >= o_starts) & (pos.view(1, 1, seq_len) < o_ends) # [N, K, L] | |
| o_mask_f = o_mask.to(late_tokens_b.dtype) | |
| opt_maxsim = (o_mask_f * max_sim_to_state_q.unsqueeze(1)).sum(dim=-1) / o_mask_f.sum(dim=-1).clamp_min(1.0) | |
| return opt_maxsim.masked_fill(~cand_mask, 0.0) * self.late_scale | |
| def forward(self, hidden_states: tuple[Tensor, ...], batch: PackedFlashBatch) -> Tensor: | |
| """Compute candidate logits ``[N, K_max]`` for all questions in ``batch``.""" | |
| fused_seq = self.layer_gate(hidden_states) # [B, L, D] | |
| token_attn_scores = self.span_pooler.attn_proj(fused_seq).squeeze(-1) # [B, L] | |
| # Pool state representation per batch row: [B, D] | |
| state_vec_b = self.span_pooler(fused_seq, token_attn_scores, batch.state_spans) # [B, D] | |
| # Gather sequence representations for each question i in 0..N-1 | |
| q_b_idx = batch.question_batch_idx # [N] | |
| seq_for_q = fused_seq[q_b_idx] # [N, L, D] | |
| scores_for_q = token_attn_scores[q_b_idx] # [N, L] | |
| state_vec = state_vec_b[q_b_idx] # [N, D] | |
| # Pool question instruction spans: [N, D] | |
| q_vec = self.span_pooler(seq_for_q, scores_for_q, batch.question_spans) # [N, D] | |
| prim_vec = self.primitive_emb(batch.question_primitive_id) # [N, D] | |
| q_vec = q_vec + prim_vec | |
| # Pool candidate option spans: [N, K, D] | |
| c_vec = self.span_pooler(seq_for_q, scores_for_q, batch.opt_spans) # [N, K, D] | |
| n_q, k_max, d = c_vec.shape | |
| # Add ordinal level embeddings for Score questions and pole embeddings for Noul questions | |
| k_indices = torch.arange(k_max, device=c_vec.device).unsqueeze(0).expand(n_q, k_max) | |
| score_mask = (batch.question_primitive_id == 1).view(-1, 1, 1) # [N, 1, 1] | |
| noul_mask = (batch.question_primitive_id == 2).view(-1, 1, 1) # [N, 1, 1] | |
| lvl_emb = self.score_level_emb(k_indices.clamp(0, MAX_SCORE_LEVELS + 1)) | |
| pole_emb = self.noul_pole_emb(k_indices.clamp(0, 1)) | |
| c_vec = c_vec + score_mask.to(c_vec.dtype) * lvl_emb + noul_mask.to(c_vec.dtype) * pole_emb | |
| # Optional v2 Multi-Hop Cross-Attention over state tokens | |
| if self.cross_hops is not None: | |
| seq_len = fused_seq.shape[1] | |
| pos = torch.arange(seq_len, device=fused_seq.device) | |
| s_starts = batch.state_spans[:, 0:1].clamp(0, seq_len - 1) | |
| s_ends = torch.maximum(batch.state_spans[:, 1:2].clamp(1, seq_len), s_starts + 1) | |
| state_mask_b = (pos.unsqueeze(0) >= s_starts) & (pos.unsqueeze(0) < s_ends) # [B, L] | |
| qc_tokens = torch.cat([q_vec.unsqueeze(1), c_vec], dim=1) # [N, 1+K, D] | |
| for hop in self.cross_hops: | |
| qc_tokens = hop(qc_tokens, fused_seq, state_mask_b, q_b_idx) | |
| q_vec = qc_tokens[:, 0, :] | |
| c_vec = qc_tokens[:, 1:, :] | |
| # Optional v3 Pure Neural Candidate Self-Attention & Contrast | |
| contrast_logits: Tensor | None = None | |
| late_mult: Tensor | None = None | |
| if self.cand_attn is not None: | |
| c_vec, contrast_logits, late_mult = self.cand_attn(c_vec, q_vec, state_vec, batch.cand_mask) | |
| # Expand q_vec and state_vec to [N, K, D] | |
| q_exp = q_vec.unsqueeze(1).expand(n_q, k_max, d) | |
| s_exp = state_vec.unsqueeze(1).expand(n_q, k_max, d) | |
| # Path A: Global Contextual Interaction MLP + Bilinear | |
| features = torch.cat( | |
| [ | |
| c_vec, | |
| q_exp, | |
| s_exp, | |
| c_vec * q_exp, | |
| c_vec * s_exp, | |
| (c_vec - s_exp).abs(), | |
| ], | |
| dim=-1, | |
| ) # [N, K, 6*D] | |
| mlp_logits = self.context_mlp(features).squeeze(-1) # [N, K] | |
| bi_context = self.bilinear_q(q_vec) + self.bilinear_s(state_vec) # [N, D] | |
| bi_logits = torch.bmm(c_vec, bi_context.unsqueeze(-1)).squeeze(-1) / (d**0.5) # [N, K] | |
| # Path B: Late-Interaction MaxSim (computed once per state row B) | |
| late_tokens_b = F.normalize(self.late_proj(fused_seq), p=2, dim=-1) # [B, L, d_late] | |
| late_logits = self._compute_late_interaction( | |
| late_tokens_b, batch.state_spans, q_b_idx, batch.opt_spans, batch.cand_mask | |
| ) # [N, K] | |
| if late_mult is not None: | |
| late_logits = late_logits * late_mult | |
| logits = mlp_logits + bi_logits + late_logits | |
| if contrast_logits is not None: | |
| logits = logits + contrast_logits | |
| # Optional v2 Unimodal Ordinal Score Coupler | |
| if self.ordinal_regressor is not None and self.ordinal_log_precision is not None: | |
| is_score = batch.question_primitive_id == 1 # [N] | |
| if is_score.any(): | |
| ord_in = torch.cat([q_vec, state_vec, q_vec * state_vec, (q_vec - state_vec).abs()], dim=-1) # [N, 4*D] | |
| pred_norm_lvl = torch.sigmoid(self.ordinal_regressor(ord_in)) # [N, 1] in (0, 1) | |
| num_valid_k = batch.cand_mask.sum(dim=-1, keepdim=True).clamp_min(2).to(logits.dtype) - 1.0 # [N, 1] | |
| lvl_coords = k_indices.to(logits.dtype) / num_valid_k # [N, K] in [0, 1] | |
| prec = torch.exp(self.ordinal_log_precision.clamp(-2.0, 3.0)) | |
| ord_bias = -prec * (lvl_coords - pred_norm_lvl).pow(2) # [N, K] | |
| logits = logits + is_score.unsqueeze(-1).to(logits.dtype) * ord_bias | |
| return logits.masked_fill(~batch.cand_mask, float("-inf")) | |
| class FlashDecisionModel(nn.Module): | |
| """Complete Jue-Flash model wrapping a Hugging Face bidirectional encoder and :class:`DualPathDecisionHead`.""" | |
| def __init__(self, encoder: nn.Module, cfg: FlashModelConfig) -> None: | |
| super().__init__() | |
| self.encoder = encoder | |
| self.cfg = cfg | |
| self.head = DualPathDecisionHead(cfg) | |
| def upgrade_to_canon_synergy(self, init_scale: float = 2.2) -> None: | |
| """Enable v3 Canon Specification Prior Synergy Fusion on an initialized or loaded model.""" | |
| device = next(self.head.parameters()).device | |
| dtype = next(self.head.parameters()).dtype | |
| self.head.enable_canon_synergy(init_scale=init_scale) | |
| self.head.to(device=device, dtype=dtype) | |
| def upgrade_to_candidate_self_attn(self) -> None: | |
| """Enable pure neural CandidateSelfAttention option cross-comparison on an initialized or loaded model.""" | |
| device = next(self.head.parameters()).device | |
| dtype = next(self.head.parameters()).dtype | |
| self.head.enable_candidate_self_attn() | |
| self.head.to(device=device, dtype=dtype) | |
| def from_pretrained_base( | |
| cls, | |
| base_model: str, | |
| *, | |
| num_gated_layers: int = 4, | |
| late_interaction_dim: int = 128, | |
| dropout: float = 0.1, | |
| max_len: int = 2048, | |
| head_v2: bool = True, | |
| canon_synergy_fusion: bool = False, | |
| candidate_self_attn: bool = False, | |
| dtype: torch.dtype = torch.float32, | |
| device: str | torch.device = "cpu", | |
| local_files_only: bool = False, | |
| ) -> FlashDecisionModel: | |
| """Initialize a fresh :class:`FlashDecisionModel` from a pretrained HF encoder backbone.""" | |
| hf_cfg = AutoConfig.from_pretrained(base_model, local_files_only=local_files_only, trust_remote_code=True) | |
| hf_cfg.output_hidden_states = True | |
| if hasattr(hf_cfg, "reference_compile"): | |
| hf_cfg.reference_compile = False | |
| if getattr(hf_cfg, "model_type", "") == "lfm2" or "lfm2" in str(base_model).lower(): | |
| from transformers import AutoModelForMaskedLM | |
| mlm = AutoModelForMaskedLM.from_pretrained( | |
| base_model, | |
| config=hf_cfg, | |
| dtype=dtype, | |
| local_files_only=local_files_only, | |
| trust_remote_code=True, | |
| ) | |
| encoder = mlm.lfm2 | |
| else: | |
| encoder = AutoModel.from_pretrained( | |
| base_model, | |
| config=hf_cfg, | |
| dtype=dtype, | |
| local_files_only=local_files_only, | |
| trust_remote_code=True, | |
| ) | |
| cfg = FlashModelConfig( | |
| base_model=str(base_model), | |
| hidden_size=int(hf_cfg.hidden_size), | |
| num_gated_layers=min(num_gated_layers, int(getattr(hf_cfg, "num_hidden_layers", 12))), | |
| late_interaction_dim=late_interaction_dim, | |
| dropout=dropout, | |
| max_len=max_len, | |
| query_conditioned_layer_gate=head_v2, | |
| tapped_layer_offsets=(-16, -12, -8, -4, -2, -1), | |
| cross_attn_hops=1 if head_v2 else 0, | |
| unimodal_score_coupler=head_v2, | |
| canon_synergy_fusion=canon_synergy_fusion, | |
| candidate_self_attn=candidate_self_attn, | |
| ) | |
| model = cls(encoder=encoder, cfg=cfg) | |
| return model.to(device=device, dtype=dtype) | |
| def save_checkpoint(self, out_dir: str | Path, *, encoder_dtype: torch.dtype | None = None) -> Path: | |
| """Save encoder weights, decision head weights, and FlashModelConfig to ``out_dir``.""" | |
| d = Path(out_dir) | |
| d.mkdir(parents=True, exist_ok=True) | |
| self.cfg.save(d / FLASH_CONFIG_FILE) | |
| enc_state = { | |
| k: (v.detach().cpu().to(encoder_dtype) if encoder_dtype is not None and v.is_floating_point() else v.detach().cpu()).contiguous() | |
| for k, v in self.encoder.state_dict().items() | |
| } | |
| head_state = {k: v.detach().cpu().float().contiguous() for k, v in self.head.state_dict().items()} | |
| save_file(enc_state, str(d / ENCODER_WEIGHTS_FILE)) | |
| save_file(head_state, str(d / HEAD_WEIGHTS_FILE)) | |
| return d | |
| def load_checkpoint( | |
| cls, | |
| checkpoint_dir: str | Path, | |
| *, | |
| base_model_override: str | None = None, | |
| dtype: torch.dtype = torch.float32, | |
| device: str | torch.device = "cpu", | |
| local_files_only: bool = False, | |
| ) -> FlashDecisionModel: | |
| """Load a saved :class:`FlashDecisionModel` checkpoint.""" | |
| d = Path(checkpoint_dir) | |
| cfg = FlashModelConfig.load(d / FLASH_CONFIG_FILE) | |
| base_path = base_model_override or cfg.base_model | |
| hf_cfg = AutoConfig.from_pretrained(base_path, local_files_only=local_files_only, trust_remote_code=True) | |
| hf_cfg.output_hidden_states = True | |
| if hasattr(hf_cfg, "reference_compile"): | |
| hf_cfg.reference_compile = False | |
| encoder = AutoModel.from_config(hf_cfg, trust_remote_code=True) | |
| model = cls(encoder=encoder, cfg=cfg) | |
| enc_sd = load_file(str(d / ENCODER_WEIGHTS_FILE), device="cpu") | |
| head_sd = load_file(str(d / HEAD_WEIGHTS_FILE), device="cpu") | |
| model.encoder.load_state_dict(enc_sd, strict=True) | |
| model.head.load_state_dict(head_sd, strict=True) | |
| return model.to(device=device, dtype=dtype) | |
| def fingerprint(self) -> dict[str, Any]: | |
| """Compute a deterministic fingerprint of the model architecture and head parameters.""" | |
| h = hashlib.sha256() | |
| for k, v in sorted(self.head.state_dict().items()): | |
| h.update(k.encode("utf-8")) | |
| h.update(v.detach().cpu().float().numpy().tobytes()[:256]) | |
| return { | |
| "base_model": self.cfg.base_model, | |
| "hidden_size": self.cfg.hidden_size, | |
| "num_gated_layers": self.cfg.num_gated_layers, | |
| "late_interaction_dim": self.cfg.late_interaction_dim, | |
| "head_sha256_prefix": h.hexdigest()[:16], | |
| } | |
| def set_backbone_trainable( | |
| self, | |
| trainable: bool, | |
| unfreeze_top_layers: int | None = None, | |
| freeze_embeddings: bool = True, | |
| ) -> None: | |
| """Freeze/unfreeze encoder backbone parameters for Stage 1 head warmup vs Stage 2/3 full fine-tuning.""" | |
| for p in self.encoder.parameters(): | |
| p.requires_grad = trainable | |
| if freeze_embeddings: | |
| for attr in ("embeddings", "emb", "wte", "tok_embeddings", "embed_tokens"): | |
| emb_mod = getattr(self.encoder, attr, None) | |
| if emb_mod is not None: | |
| for p in emb_mod.parameters(): | |
| p.requires_grad = False | |
| if not trainable and unfreeze_top_layers is not None and unfreeze_top_layers > 0: | |
| layers = getattr(self.encoder, "layers", None) | |
| if layers is None and hasattr(self.encoder, "encoder"): | |
| layers = getattr(self.encoder, "encoder", None) | |
| if layers is not None: | |
| layers = getattr(layers, "layer", None) | |
| if layers is not None: | |
| for layer in list(layers)[-unfreeze_top_layers:]: | |
| for p in layer.parameters(): | |
| p.requires_grad = True | |
| for p in self.head.parameters(): | |
| p.requires_grad = True | |
| def forward( | |
| self, | |
| batch: PackedFlashBatch, | |
| *, | |
| temperature: float | Mapping[str, float] = 1.0, | |
| noul_bias: float = 0.0, | |
| ) -> Tensor: | |
| """Run single-pass bidirectional encoding and multi-question readout.""" | |
| outputs = self.encoder( | |
| input_ids=batch.input_ids, | |
| attention_mask=batch.attention_mask, | |
| output_hidden_states=True, | |
| return_dict=True, | |
| ) | |
| logits = self.head(outputs.hidden_states, batch) # [N, K_max] | |
| if noul_bias != 0.0 and logits.shape[1] >= 2: | |
| noul_rows = batch.question_primitive_id == 2 | |
| if noul_rows.any(): | |
| logits = logits.clone() | |
| logits[noul_rows, 1] = logits[noul_rows, 1] + float(noul_bias) | |
| if isinstance(temperature, Mapping): | |
| t_vec = torch.ones_like(batch.question_primitive_id, dtype=logits.dtype) | |
| t_vec = torch.where( | |
| batch.question_primitive_id == 0, | |
| torch.tensor(float(temperature.get("choice", 1.0)), dtype=logits.dtype, device=logits.device), | |
| t_vec, | |
| ) | |
| t_vec = torch.where( | |
| batch.question_primitive_id == 1, | |
| torch.tensor(float(temperature.get("score", 1.0)), dtype=logits.dtype, device=logits.device), | |
| t_vec, | |
| ) | |
| t_vec = torch.where( | |
| batch.question_primitive_id == 2, | |
| torch.tensor(float(temperature.get("noul", 1.0)), dtype=logits.dtype, device=logits.device), | |
| t_vec, | |
| ) | |
| logits = logits / t_vec.unsqueeze(-1) | |
| elif temperature != 1.0: | |
| logits = logits / float(temperature) | |
| return logits.masked_fill(~batch.cand_mask, float("-inf")) | |