RESMP-DEV/LFM2.5-Encoder-230M-Code-MXFP4-GPTQ

1"""LFM2 backbone with bidirectional attention + non-causal short-conv.2 3Wired into the HF repo via `auto_map` in config.json so that4 5AutoModel.from_pretrained(repo, trust_remote_code=True)6AutoModelForMaskedLM.from_pretrained(repo, trust_remote_code=True)7 8both return a model with the encoder-style patches already applied.9 10Supports `attn_implementation` in {"eager", "sdpa", "flash_attention_2"}:11 12eager/sdpa consume a 4D additive pad-only mask and reproduce the exact13training-time behavior; flash_attention_2 receives the 2D padding mask (or14None) and runs the kernel non-causally via `Lfm2Attention.is_causal = False`,15yielding outputs equivalent to the unpadded forward.16"""17 18from typing import Optional19 20import torch21import torch.nn as nn22import torch.nn.functional as F23from transformers.configuration_utils import PretrainedConfig24from transformers.modeling_outputs import BaseModelOutput, MaskedLMOutput25from transformers.modeling_utils import PreTrainedModel26from transformers.models.lfm2 import modeling_lfm2 as _lfm2_mod27from transformers.models.lfm2.configuration_lfm2 import Lfm2Config28from transformers.models.lfm2.modeling_lfm2 import (29    Lfm2Attention,30    Lfm2Model,31    Lfm2PreTrainedModel,32    Lfm2ShortConv,33    apply_mask_to_padding_states,34)35 36 37def _bidirectional_mask(38    config,39    input_embeds: torch.Tensor = None,40    attention_mask: Optional[torch.Tensor] = None,41    cache_position: Optional[torch.LongTensor] = None,42    past_key_values=None,43    position_ids: Optional[torch.LongTensor] = None,44    **kwargs,45) -> Optional[torch.Tensor]:46    # transformers has renamed the embeds kwarg across versions47    # (input_embeds <-> inputs_embeds); accept either to stay forward-compatible.48    if input_embeds is None:49        input_embeds = kwargs.get("inputs_embeds")50 51    if config._attn_implementation == "flash_attention_2":52        # FA2 only uses the 2D padding mask to unpad sequences; causality is53        # controlled by `Lfm2Attention.is_causal` (set to False below).54        if attention_mask is not None and not attention_mask.all():55            return attention_mask56        return None57 58    device = input_embeds.device59    dtype = input_embeds.dtype60    bsz, q_len = input_embeds.shape[:2]61    past = past_key_values.get_seq_length() if past_key_values is not None else 062    kv_len = past + q_len63 64    mask = torch.zeros((bsz, 1, q_len, kv_len), device=device, dtype=dtype)65    if attention_mask is not None:66        cur_len = attention_mask.size(-1)67        key_pad_flags = (attention_mask == 0).to(device=device, dtype=torch.float32)68        pad_vec = torch.zeros((bsz, kv_len), device=device, dtype=torch.float32)69        if cur_len > 0:70            pad_vec[:, past:past + cur_len] = key_pad_flags * -1e971        mask = mask + pad_vec.to(dtype)[:, None, None, :]72    return mask73 74 75def _noncausal_shortconv_forward(76    self,77    hidden_states: torch.Tensor,78    past_key_values=None,79    cache_position=None,80    attention_mask: Optional[torch.Tensor] = None,81    **kwargs,82) -> torch.Tensor:83    x = apply_mask_to_padding_states(hidden_states, attention_mask)84 85    BCx = self.in_proj(x).transpose(-1, -2)86    B, C, x = BCx.chunk(3, dim=-2)87    Bx = B * x88 89    k = self.conv.weight.shape[-1]90    pad = k // 291    conv_out = F.conv1d(92        Bx, weight=self.conv.weight, bias=self.conv.bias,93        stride=1, padding=pad, dilation=1, groups=Bx.shape[1],94    )95    if conv_out.shape[-1] > Bx.shape[-1]:96        conv_out = conv_out[..., :Bx.shape[-1]]97    elif conv_out.shape[-1] < Bx.shape[-1]:98        conv_out = F.pad(conv_out, (0, Bx.shape[-1] - conv_out.shape[-1]))99 100    y = C * conv_out101    y = y.transpose(-1, -2).contiguous()102    return self.out_proj(y)103 104 105def _shortconv_forward(self, *args, **kwargs):106    return self.slow_forward(*args, **kwargs)107 108 109_PATCHED = False110 111 112def _install_patches() -> None:113    global _PATCHED114    if _PATCHED:115        return116    _lfm2_mod.create_causal_mask = _bidirectional_mask117    Lfm2ShortConv.slow_forward = _noncausal_shortconv_forward118    Lfm2ShortConv.forward = _shortconv_forward119    _PATCHED = True120 121 122_install_patches()123 124 125def _set_attention_noncausal(model) -> None:126    for module in model.modules():127        if isinstance(module, Lfm2Attention):128            module.is_causal = False129 130 131class Lfm2BidirectionalModel(Lfm2Model):132    """LFM2 patched for encoder-style use: 133    full bidirectional attention + non-causal short-conv."""134 135    def __init__(self, config):136        _install_patches()137        super().__init__(config)138        _set_attention_noncausal(self)139 140 141class Lfm2BidirectionalForMaskedLM(Lfm2PreTrainedModel):142    """LFM2 bidirectional encoder with a tied masked-LM head."""143 144    config_class = Lfm2Config145    base_model_prefix = "lfm2"146    _tied_weights_keys = {"lm_head.weight": "lfm2.embed_tokens.weight"}147 148    def __init__(self, config: Lfm2Config):149        _install_patches()150        config = type(config).from_dict({**config.to_dict(), "use_cache": False})151        super().__init__(config)152        self.lfm2 = Lfm2BidirectionalModel(config)153        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)154        self.post_init()155        self.lm_head.weight = self.lfm2.embed_tokens.weight156 157    def get_input_embeddings(self):158        return self.lfm2.embed_tokens159 160    def set_input_embeddings(self, value):161        self.lfm2.embed_tokens = value162 163    def get_output_embeddings(self):164        return self.lm_head165 166    def set_output_embeddings(self, new_embeddings):167        self.lm_head = new_embeddings168 169    def forward(170        self,171        input_ids: Optional[torch.LongTensor] = None,172        attention_mask: Optional[torch.Tensor] = None,173        position_ids: Optional[torch.LongTensor] = None,174        inputs_embeds: Optional[torch.FloatTensor] = None,175        labels: Optional[torch.LongTensor] = None,176        output_hidden_states: Optional[bool] = None,177        output_attentions: Optional[bool] = None,178        return_dict: Optional[bool] = None,179        **kwargs,180    ) -> MaskedLMOutput:181        return_dict = True if return_dict is None else return_dict182        outputs = self.lfm2(183            input_ids=input_ids,184            attention_mask=attention_mask,185            position_ids=position_ids,186            inputs_embeds=inputs_embeds,187            use_cache=False,188            output_attentions=output_attentions,189            output_hidden_states=output_hidden_states,190            return_dict=True,191        )192        hidden = outputs.last_hidden_state193        logits = self.lm_head(hidden)194 195        loss = None196        if labels is not None:197            loss = F.cross_entropy(198                logits.view(-1, self.config.vocab_size),199                labels.view(-1),200                ignore_index=-100,201            )202 203        if not return_dict:204            out = (logits,) + outputs[1:]205            return ((loss,) + out) if loss is not None else out206        return MaskedLMOutput(207            loss=loss,208            logits=logits,209            hidden_states=outputs.hidden_states,210            attentions=outputs.attentions,211        )212