| """Dazo Hugging Face wrapper. |
| |
| Dazo keeps backward-compatible experimental paths, plus a compact shared-marker path used by |
| Dazo-15M. The compact path runs the encoder exactly once per question over query + all candidate |
| markers + evidence, gathers the contextualized marker states, and lets a tiny recurrent core refine |
| them. This preserves typed decisions while avoiding one backbone pass per option. |
| """ |
| from __future__ import annotations |
|
|
| import re |
| from typing import Optional |
|
|
| import torch |
| import torch.nn as nn |
| from transformers import AutoConfig, AutoModel, PreTrainedModel |
|
|
| from .configuration_dazo import DazoConfig |
| from .core import DazoCore, masked_softmax |
|
|
|
|
| class DazoForDecision(PreTrainedModel): |
| config_class = DazoConfig |
| base_model_prefix = "dazo" |
| all_tied_weights_keys = {} |
|
|
| def __init__(self, config: DazoConfig): |
| super().__init__(config) |
| backbone_cfg = self._resolve_backbone_config(config) |
| self.backbone = AutoModel.from_config(backbone_cfg) |
| hidden = int(self.backbone.config.hidden_size) |
| self.core = DazoCore( |
| context_dim=hidden, |
| option_dim=hidden, |
| latent_dim=config.latent_dim, |
| n_evidence_slots=config.n_evidence_slots, |
| n_hypothesis_slots=config.n_hypothesis_slots, |
| n_critic_slots=config.n_critic_slots, |
| n_control_slots=config.n_control_slots, |
| layers_per_loop=config.layers_per_loop, |
| heads=config.num_attention_heads, |
| dropout=config.dropout, |
| max_loop_embeddings=config.max_loop_embeddings, |
| max_rank=config.max_rank, |
| ) |
|
|
| if config.query_conditioning: |
| heads = self._valid_heads(hidden, config.num_attention_heads) |
| self.query_norm = nn.LayerNorm(hidden) |
| self.query_option_norm = nn.LayerNorm(hidden) |
| self.query_cross_attn = nn.MultiheadAttention(hidden, heads, dropout=config.dropout, batch_first=True) |
| self.query_fuse = nn.Sequential( |
| nn.LayerNorm(hidden * 4), nn.Linear(hidden * 4, hidden), nn.GELU(), nn.Linear(hidden, hidden) |
| ) |
| self.query_fuse_norm = nn.LayerNorm(hidden) |
|
|
| if config.joint_candidate_encoding or getattr(config, "shared_joint_encoding", False): |
| self.marker_scorer = nn.Sequential( |
| nn.LayerNorm(hidden), nn.Linear(hidden, hidden), nn.GELU(), nn.Linear(hidden, 1) |
| ) |
|
|
| if config.base_compatibility: |
| heads = self._valid_heads(hidden, config.num_attention_heads) |
| self.base_context_norm = nn.LayerNorm(hidden) |
| self.base_option_norm = nn.LayerNorm(hidden) |
| self.base_cross_attn = nn.MultiheadAttention(hidden, heads, dropout=config.dropout, batch_first=True) |
| if config.joint_decision_head: |
| decision_dim = int(config.decision_dim) |
| decision_heads = self._valid_heads(decision_dim, config.num_attention_heads) |
| self.decision_pair_proj = nn.Sequential(nn.LayerNorm(hidden * 4), nn.Linear(hidden * 4, decision_dim), nn.GELU()) |
| self.decision_global_proj = nn.Sequential(nn.LayerNorm(hidden), nn.Linear(hidden, decision_dim)) |
| layer = nn.TransformerEncoderLayer( |
| d_model=decision_dim, nhead=decision_heads, dim_feedforward=decision_dim * 4, |
| dropout=config.dropout, activation="gelu", batch_first=True, norm_first=True, |
| ) |
| self.decision_head = nn.TransformerEncoder(layer, num_layers=max(1, int(config.decision_head_layers)), enable_nested_tensor=False) |
| self.decision_score = nn.Sequential(nn.LayerNorm(decision_dim), nn.Linear(decision_dim, 1)) |
| else: |
| self.base_pair_score = nn.Sequential( |
| nn.LayerNorm(hidden * 4), nn.Linear(hidden * 4, hidden), nn.GELU(), nn.Linear(hidden, 1) |
| ) |
|
|
| if config.freeze_backbone: |
| self.freeze_backbone() |
|
|
| @staticmethod |
| def _valid_heads(hidden: int, requested: int) -> int: |
| heads = min(requested, hidden) |
| while heads > 1 and hidden % heads != 0: |
| heads -= 1 |
| return heads |
|
|
| @staticmethod |
| def _resolve_backbone_config(config: DazoConfig): |
| saved = getattr(config, "backbone_config", None) |
| if saved: |
| saved = dict(saved) |
| model_type = saved.pop("model_type") |
| return AutoConfig.for_model(model_type, **saved) |
| backbone_cfg = AutoConfig.from_pretrained(config.backbone_name) |
| config.backbone_config = backbone_cfg.to_dict() |
| return backbone_cfg |
|
|
| @classmethod |
| def from_backbone_pretrained(cls, config: DazoConfig) -> "DazoForDecision": |
| backbone_cfg = AutoConfig.from_pretrained(config.backbone_name) |
| config.backbone_config = backbone_cfg.to_dict() |
| model = cls(config) |
| model.backbone = AutoModel.from_pretrained(config.backbone_name, config=backbone_cfg) |
| if config.freeze_backbone: |
| model.freeze_backbone() |
| if getattr(config, "unfreeze_last_n_layers", 0) > 0: |
| model.unfreeze_last_backbone_layers(config.unfreeze_last_n_layers) |
| return model |
|
|
| def freeze_backbone(self): |
| for p in self.backbone.parameters(): |
| p.requires_grad = False |
| return self |
|
|
| def unfreeze_backbone(self): |
| for p in self.backbone.parameters(): |
| p.requires_grad = True |
| return self |
|
|
| def unfreeze_last_backbone_layers(self, n: int): |
| n = int(n) |
| if n <= 0: |
| return self |
| layer_ids = set() |
| for name, _p in self.backbone.named_parameters(): |
| m = re.search(r"(?:^|\.)layers\.(\d+)\.", name) or re.search(r"(?:^|\.)layer\.(\d+)\.", name) |
| if m: |
| layer_ids.add(int(m.group(1))) |
| if not layer_ids: |
| raise RuntimeError("Could not identify transformer layers in backbone parameter names") |
| selected = set(sorted(layer_ids)[-n:]) |
| count = 0 |
| for name, p in self.backbone.named_parameters(): |
| m = re.search(r"(?:^|\.)layers\.(\d+)\.", name) or re.search(r"(?:^|\.)layer\.(\d+)\.", name) |
| if m and int(m.group(1)) in selected: |
| p.requires_grad = True |
| count += p.numel() |
| print(f"unfrozen_backbone_layers={sorted(selected)} trainable_backbone_params={count}") |
| return self |
|
|
| @staticmethod |
| def _mean_pool(hidden: torch.Tensor, mask: torch.Tensor) -> torch.Tensor: |
| w = mask.to(hidden.dtype).unsqueeze(-1) |
| return (hidden * w).sum(1) / w.sum(1).clamp_min(1.0) |
|
|
| def _encode_options(self, option_input_ids: torch.Tensor, option_attention_mask: torch.Tensor) -> torch.Tensor: |
| bsz, nopt, length = option_input_ids.shape |
| flat_ids = option_input_ids.reshape(bsz * nopt, length) |
| flat_mask = option_attention_mask.reshape(bsz * nopt, length) |
| out = self.backbone(input_ids=flat_ids, attention_mask=flat_mask).last_hidden_state |
| pooled = self._mean_pool(out, flat_mask) |
| return pooled.reshape(bsz, nopt, -1) |
|
|
| def _encode_joint_candidates(self, joint_input_ids, joint_attention_mask, marker_positions): |
| bsz, nopt, length = joint_input_ids.shape |
| flat_ids = joint_input_ids.reshape(bsz * nopt, length) |
| flat_mask = joint_attention_mask.reshape(bsz * nopt, length) |
| hidden = self.backbone(input_ids=flat_ids, attention_mask=flat_mask).last_hidden_state |
| pos = marker_positions.reshape(-1).long().clamp(0, length - 1) |
| row = torch.arange(hidden.size(0), device=hidden.device) |
| return hidden[row, pos].reshape(bsz, nopt, -1) |
|
|
| def _encode_shared(self, shared_input_ids, shared_attention_mask, shared_marker_positions): |
| hidden = self.backbone(input_ids=shared_input_ids, attention_mask=shared_attention_mask).last_hidden_state |
| pos = shared_marker_positions.long().clamp(0, hidden.size(1) - 1) |
| row = torch.arange(hidden.size(0), device=hidden.device)[:, None].expand_as(pos) |
| markers = hidden[row, pos] |
| return hidden, markers |
|
|
| def _condition_options_on_query(self, option_hidden, query_hidden, query_attention_mask): |
| options = self.query_option_norm(option_hidden) |
| query = self.query_norm(query_hidden) |
| attended, _ = self.query_cross_attn(options, query, query, key_padding_mask=~query_attention_mask.bool(), need_weights=False) |
| pair = torch.cat([options, attended, options * attended, (options - attended).abs()], dim=-1) |
| return self.query_fuse_norm(option_hidden + self.query_fuse(pair)) |
|
|
| def _joint_decision_logits(self, pair, option_mask, query_hidden, query_attention_mask, context_hidden, attention_mask): |
| option_tokens = self.decision_pair_proj(pair) |
| global_hidden = self._mean_pool(query_hidden, query_attention_mask) if query_hidden is not None else self._mean_pool(context_hidden, attention_mask) |
| global_token = self.decision_global_proj(global_hidden).unsqueeze(1) |
| tokens = torch.cat([global_token, option_tokens], dim=1) |
| valid = torch.cat([torch.ones(option_mask.size(0), 1, dtype=torch.bool, device=option_mask.device), option_mask.bool()], dim=1) |
| contextualized = self.decision_head(tokens, src_key_padding_mask=~valid) |
| return self.decision_score(contextualized[:, 1:]).squeeze(-1).masked_fill(~option_mask.bool(), -1e4) |
|
|
| def _base_compatibility_logits(self, context_hidden, attention_mask, option_hidden, option_mask, query_hidden=None, query_attention_mask=None): |
| options = self.base_option_norm(option_hidden) |
| context = self.base_context_norm(context_hidden) |
| attended, _ = self.base_cross_attn(options, context, context, key_padding_mask=~attention_mask.bool(), need_weights=False) |
| pair = torch.cat([options, attended, options * attended, (options - attended).abs()], dim=-1) |
| if self.config.joint_decision_head: |
| return self._joint_decision_logits(pair, option_mask, query_hidden, query_attention_mask, context_hidden, attention_mask) |
| return self.base_pair_score(pair).squeeze(-1).masked_fill(~option_mask.bool(), -1e4) |
|
|
| def _apply_base_compatibility(self, out, base_logits, option_mask): |
| scale = float(self.config.recurrent_logit_scale) |
| per_step_logits = base_logits[:, None, :] + scale * out.per_step_logits |
| step_mask = option_mask[:, None, :].bool().expand_as(per_step_logits) |
| per_step_probs = masked_softmax(per_step_logits, step_mask) |
| selected = out.selected_step.long().clamp_min(1) - 1 |
| batch = torch.arange(base_logits.size(0), device=base_logits.device) |
| out.per_step_logits = per_step_logits |
| out.per_step_probs = per_step_probs |
| out.logits = per_step_logits[batch, selected] |
| out.probs = per_step_probs[batch, selected] |
| out.energy = -torch.logsumexp(out.logits.float().masked_fill(~option_mask.bool(), -1e4), dim=-1) |
| return out |
|
|
| def forward( |
| self, |
| input_ids: torch.Tensor, |
| attention_mask: torch.Tensor, |
| option_input_ids: torch.Tensor, |
| option_attention_mask: torch.Tensor, |
| option_mask: torch.Tensor, |
| query_input_ids: Optional[torch.Tensor] = None, |
| query_attention_mask: Optional[torch.Tensor] = None, |
| joint_input_ids: Optional[torch.Tensor] = None, |
| joint_attention_mask: Optional[torch.Tensor] = None, |
| marker_positions: Optional[torch.Tensor] = None, |
| shared_input_ids: Optional[torch.Tensor] = None, |
| shared_attention_mask: Optional[torch.Tensor] = None, |
| shared_marker_positions: Optional[torch.Tensor] = None, |
| task_type: Optional[torch.Tensor] = None, |
| rank_ids: Optional[torch.Tensor] = None, |
| max_steps: Optional[int] = None, |
| min_steps: Optional[int] = None, |
| adaptive: bool = False, |
| halt_threshold: Optional[float] = None, |
| convergence_threshold: Optional[float] = None, |
| correctness_threshold: Optional[float] = None, |
| **unused, |
| ): |
| marker_logits = None |
| query_hidden = None |
|
|
| if getattr(self.config, "shared_joint_encoding", False): |
| if shared_input_ids is None or shared_attention_mask is None or shared_marker_positions is None: |
| raise ValueError("shared_joint_encoding=True requires shared_input_ids/shared_attention_mask/shared_marker_positions") |
| context, options = self._encode_shared(shared_input_ids, shared_attention_mask, shared_marker_positions) |
| attention_mask = shared_attention_mask |
| marker_logits = self.marker_scorer(options).squeeze(-1).masked_fill(~option_mask.bool(), -1e4) |
| reasoning_context = context |
| reasoning_mask = shared_attention_mask.bool() |
| else: |
| context = self.backbone(input_ids=input_ids, attention_mask=attention_mask).last_hidden_state |
| if self.config.joint_candidate_encoding: |
| if joint_input_ids is None or joint_attention_mask is None or marker_positions is None: |
| raise ValueError("joint_candidate_encoding=True requires joint_input_ids/joint_attention_mask/marker_positions") |
| options = self._encode_joint_candidates(joint_input_ids, joint_attention_mask, marker_positions) |
| marker_logits = self.marker_scorer(options).squeeze(-1).masked_fill(~option_mask.bool(), -1e4) |
| else: |
| options = self._encode_options(option_input_ids, option_attention_mask) |
| reasoning_context = context |
| reasoning_mask = attention_mask.bool() |
| if self.config.query_conditioning: |
| if query_input_ids is None or query_attention_mask is None: |
| raise ValueError("query_conditioning=True requires query_input_ids and query_attention_mask") |
| query_hidden = self.backbone(input_ids=query_input_ids, attention_mask=query_attention_mask).last_hidden_state |
| if not self.config.joint_candidate_encoding: |
| options = self._condition_options_on_query(options, query_hidden, query_attention_mask) |
| reasoning_context = torch.cat([query_hidden, context], dim=1) |
| reasoning_mask = torch.cat([query_attention_mask.bool(), attention_mask.bool()], dim=1) |
|
|
| out = self.core( |
| context_hidden=reasoning_context, |
| context_mask=reasoning_mask, |
| option_hidden=options, |
| option_mask=option_mask.bool(), |
| task_type=task_type, |
| rank_ids=rank_ids, |
| max_steps=max_steps or self.config.max_steps, |
| min_steps=min_steps or self.config.min_steps, |
| adaptive=adaptive, |
| halt_threshold=halt_threshold if halt_threshold is not None else self.config.halt_threshold, |
| convergence_threshold=convergence_threshold if convergence_threshold is not None else self.config.convergence_threshold, |
| correctness_threshold=correctness_threshold if correctness_threshold is not None else self.config.correctness_threshold, |
| ) |
|
|
| if marker_logits is not None: |
| out = self._apply_base_compatibility(out, marker_logits, option_mask) |
| elif self.config.base_compatibility: |
| base_logits = self._base_compatibility_logits(context, attention_mask, options, option_mask, query_hidden, query_attention_mask) |
| out = self._apply_base_compatibility(out, base_logits, option_mask) |
| return out.to_dict() |
|
|
