"""HuggingFace PreTrainedModel wrapper for the dual-attention (DAT) decoder LM. Source-of-truth copy. scripts/convert_dat_to_hf.py copies it into a generated HF repository together with configuration_dat.py and the flattened model source (dat_config.py, dat_core.py, dat_symbols.py, dat_lm.py, transformer_core.py, transformer_components.py). Why every model module is imported directly here: with trust_remote_code on a local directory, transformers 4.46.3 only copies the entry module's *direct* relative imports into its dynamic-module cache (it does not recurse). Importing all flattened modules here forces every one of them to be copied, after which their own relative imports resolve because the files are co-located. """ import collections import torch from transformers import PreTrainedModel from transformers.modeling_outputs import BaseModelOutput, CausalLMOutput, MaskedLMOutput, SequenceClassifierOutput from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss # Force every flattened module to be copied into the dynamic-module cache. from .attention_config import SUPPORTED_SEQUENCE_BOUNDARY_POLICIES # noqa: F401 from .masks import build_decoder_attention_mask # noqa: F401 from .mlm import MaskClassifier # noqa: F401 from .transformer_components import FeedForward, PositionalEncoding # noqa: F401 from .transformer_core import SelfAttention # noqa: F401 from .dat_symbols import SymbolicAttentionRetriever # noqa: F401 from .dat_core import RelationalAttention # noqa: F401 from .dat_config import DatLMConfig from .dat_lm import DatDecoderLM from .configuration_dat import DatConfig, DAT_LM_FIELDS def _refresh_rope_buffers(model: DatDecoderLM, rope_theta: float) -> None: pe = model.position_encoder if hasattr(pe, "pe_type") and pe.pe_type == "rope": cos, sin = pe._precompute_rope_freqs(pe.embedding_dim, pe.max_len, rope_theta) pe.register_buffer("rope_cos", cos, persistent=False) pe.register_buffer("rope_sin", sin, persistent=False) pe._rope_uninitialized = True def _collect_shared_weight_keys(module: torch.nn.Module) -> dict[str, str]: # remove_duplicate=False is essential: by default named_parameters yields # each shared tensor only once, which would hide exactly the duplicates we # need to declare (and which state_dict still emits under every name). names_by_storage: dict[int, list[str]] = collections.defaultdict(list) for name, parameter in module.named_parameters(remove_duplicate=False): names_by_storage[id(parameter)].append(name) shared_keys: dict[str, str] = {} for names in names_by_storage.values(): if len(names) > 1: names = sorted(names) # Prefer token_embeddings.weight as the canonical source # (HF convention: output weights tied to input embeddings). source = names[0] for n in names: if n.endswith("token_embeddings.weight"): source = n break for target in names: if target != source: shared_keys[target] = source return shared_keys class _ExpandedTiedWeightsMixin: def get_expanded_tied_weights_keys(self, all_submodels: bool = False) -> dict: # HF's default returns {} when tie_word_embeddings=False, which would # discard the mlm_head.linear_out <-> lm_head tie for NextLat models # (tie_lm_head=False). Return the dynamically computed mapping directly # so all shared weights are properly tied on load. return dict(self._tied_weights_keys) class DatModel(_ExpandedTiedWeightsMixin, PreTrainedModel): config_class = DatConfig base_model_prefix = "" def __init__(self, config: DatConfig) -> None: super().__init__(config) dat_config = DatLMConfig(**{field: getattr(config, field) for field in DAT_LM_FIELDS}) self.model = DatDecoderLM(dat_config) self._tied_weights_keys = _collect_shared_weight_keys(self) self.post_init() def post_init(self) -> None: super().post_init() # Work around accelerate/transformers leaving persistent=False buffers uninitialized _refresh_rope_buffers(self.model, self.config.rope_theta) def get_input_embeddings(self) -> torch.nn.Module: return self.model.token_embeddings def set_input_embeddings(self, value: torch.nn.Module) -> None: self.model.token_embeddings = value def forward( self, input_ids: torch.Tensor, attention_mask: torch.Tensor | None = None, **kwargs, ) -> BaseModelOutput: _, hidden_states = self.model.encode_for_objective(input_ids, attention_mask=attention_mask) return BaseModelOutput(last_hidden_state=hidden_states) class DatForCausalLM(_ExpandedTiedWeightsMixin, PreTrainedModel): config_class = DatConfig base_model_prefix = "model" _tied_weights_keys = {"model.lm_head.weight": "model.token_embeddings.weight"} def __init__(self, config: DatConfig) -> None: super().__init__(config) dat_config = DatLMConfig(**{field: getattr(config, field) for field in DAT_LM_FIELDS}) self.model = DatDecoderLM(dat_config) self._tied_weights_keys = _collect_shared_weight_keys(self) self.post_init() def post_init(self) -> None: super().post_init() # Work around accelerate/transformers leaving persistent=False buffers uninitialized _refresh_rope_buffers(self.model, self.config.rope_theta) def get_input_embeddings(self) -> torch.nn.Module: return self.model.token_embeddings def set_input_embeddings(self, value: torch.nn.Module) -> None: self.model.token_embeddings = value def get_output_embeddings(self) -> torch.nn.Module: return self.model.lm_head def set_output_embeddings(self, new_embeddings: torch.nn.Module) -> None: self.model.lm_head = new_embeddings def forward( self, input_ids: torch.Tensor, attention_mask: torch.Tensor | None = None, labels: torch.Tensor | None = None, **kwargs, ) -> CausalLMOutput: logits, _ = self.model(input_ids, attention_mask=attention_mask) loss = None if labels is not None: shift_logits = logits[:, :-1, :].contiguous() shift_labels = labels[:, 1:].contiguous() loss = torch.nn.functional.cross_entropy( shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1), ignore_index=-100, ) return CausalLMOutput(loss=loss, logits=logits) class DatForMaskedLM(_ExpandedTiedWeightsMixin, PreTrainedModel): config_class = DatConfig base_model_prefix = "model" def __init__(self, config: DatConfig) -> None: super().__init__(config) dat_config = DatLMConfig(**{field: getattr(config, field) for field in DAT_LM_FIELDS}) self.model = DatDecoderLM(dat_config) self._tied_weights_keys = _collect_shared_weight_keys(self) self.post_init() def post_init(self) -> None: super().post_init() _refresh_rope_buffers(self.model, self.config.rope_theta) def get_input_embeddings(self) -> torch.nn.Module: return self.model.token_embeddings def set_input_embeddings(self, value: torch.nn.Module) -> None: self.model.token_embeddings = value def get_output_embeddings(self) -> torch.nn.Module: return self.model.mlm_head.linear_out def set_output_embeddings(self, new_embeddings: torch.nn.Module) -> None: self.model.mlm_head.linear_out = new_embeddings # Keep lm_head in sync with mlm_head.linear_out so that # resize_token_embeddings updates both tied heads. self.model.lm_head.weight = new_embeddings.weight def forward( self, input_ids: torch.Tensor, attention_mask: torch.Tensor | None = None, labels: torch.Tensor | None = None, **kwargs, ) -> MaskedLMOutput | tuple: return_dict = kwargs.pop("return_dict", None) if return_dict is None: return_dict = self.config.return_dict if kwargs: raise TypeError(f"Unexpected keyword argument(s): {list(kwargs)}") logits = self.model.forward_mlm(input_ids, attention_mask=attention_mask) loss = None if labels is not None: loss = CrossEntropyLoss(ignore_index=-100)( logits.view(-1, logits.size(-1)), labels.view(-1), ) if not return_dict: output = (logits,) return ((loss,) + output) if loss is not None else output return MaskedLMOutput(loss=loss, logits=logits) class DatForSequenceClassification(_ExpandedTiedWeightsMixin, PreTrainedModel): config_class = DatConfig base_model_prefix = "model" def __init__(self, config: DatConfig) -> None: super().__init__(config) self.num_labels = getattr(config, "num_labels", 2) dat_config = DatLMConfig(**{field: getattr(config, field) for field in DAT_LM_FIELDS}) self.model = DatDecoderLM(dat_config) self.score = torch.nn.Linear(config.hidden_dim, self.num_labels, bias=False) self._tied_weights_keys = _collect_shared_weight_keys(self) self.post_init() def post_init(self) -> None: super().post_init() _refresh_rope_buffers(self.model, self.config.rope_theta) def get_input_embeddings(self) -> torch.nn.Module: return self.model.token_embeddings def set_input_embeddings(self, value: torch.nn.Module) -> None: self.model.token_embeddings = value def forward( self, input_ids: torch.Tensor, attention_mask: torch.Tensor | None = None, labels: torch.Tensor | None = None, **kwargs, ) -> SequenceClassifierOutput: _, hidden_states = self.model.encode_for_objective(input_ids, attention_mask=attention_mask) logits = self.score(hidden_states) batch_size = input_ids.shape[0] if self.config.pad_token_id is None: last_non_pad_token = -1 else: non_pad_mask = (input_ids != self.config.pad_token_id).to(logits.device, torch.int32) if (non_pad_mask.sum(-1) == 0).any().item(): raise ValueError("Cannot pool sequence-classification logits for all-padding sequences") token_indices = torch.arange(input_ids.shape[-1], device=logits.device, dtype=torch.int32) last_non_pad_token = (token_indices * non_pad_mask).argmax(-1) pooled_logits = logits[torch.arange(batch_size, device=logits.device), last_non_pad_token] loss = None if labels is not None: if getattr(self.config, "problem_type", None) is None: if self.num_labels == 1: self.config.problem_type = "regression" elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int): self.config.problem_type = "single_label_classification" else: self.config.problem_type = "multi_label_classification" if self.config.problem_type == "regression": loss_fct = MSELoss() if self.num_labels == 1: loss = loss_fct(pooled_logits.squeeze(), labels.squeeze()) else: loss = loss_fct(pooled_logits, labels) elif self.config.problem_type == "single_label_classification": loss_fct = CrossEntropyLoss() loss = loss_fct(pooled_logits.view(-1, self.num_labels), labels.view(-1)) elif self.config.problem_type == "multi_label_classification": loss_fct = BCEWithLogitsLoss() loss = loss_fct(pooled_logits, labels) return SequenceClassifierOutput( loss=loss, logits=pooled_logits, hidden_states=hidden_states, )