Upload modeling_phase2_tc.py with huggingface_hub

modeling_phase2_tc.py

@@ -0,0 +1,144 @@
+"""Token-classification head on the PHASE-2 bidirectional LFM2.5 MLM encoder
+(LiquidAI/mlm_phase2_bidir2_step140800, "theirs" bidir variant: full gated shortconv
+made symmetric + non-causal attention). Self-contained for trust_remote_code.
+
+The encoder section is inlined VERBATIM from the backbone's shipped
+`modeling_lfm2_bidir_theirs.py` (so the trained conv in_proj/out_proj/conv weights are
+used exactly as trained — unlike our v13 vendored patch which dropped in_proj/out_proj).
+On top sits our BIOES classifier head + class-weighted label-smoothed CE.
+"""
+from __future__ import annotations
+from typing import Optional
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers.modeling_outputs import TokenClassifierOutput
+from transformers.models.lfm2 import modeling_lfm2 as _lfm2_mod
+from transformers.models.lfm2.configuration_lfm2 import Lfm2Config
+from transformers.models.lfm2.modeling_lfm2 import (
+    Lfm2Attention, Lfm2Model, Lfm2PreTrainedModel, Lfm2ShortConv, apply_mask_to_padding_states,
+)
+
+
+# ---- bidirectional patches (verbatim from the backbone's modeling_lfm2_bidir_theirs) ----
+def _bidirectional_mask(config, input_embeds=None, attention_mask=None, cache_position=None,
+                        past_key_values=None, position_ids=None, **kwargs):
+    if input_embeds is None:
+        input_embeds = kwargs.get("inputs_embeds")
+    if config._attn_implementation == "flash_attention_2":
+        if attention_mask is not None and not attention_mask.all():
+            return attention_mask
+        return None
+    device = input_embeds.device; dtype = input_embeds.dtype
+    bsz, q_len = input_embeds.shape[:2]
+    past = past_key_values.get_seq_length() if past_key_values is not None else 0
+    kv_len = past + q_len
+    mask = torch.zeros((bsz, 1, q_len, kv_len), device=device, dtype=dtype)
+    if attention_mask is not None:
+        cur_len = attention_mask.size(-1)
+        key_pad_flags = (attention_mask == 0).to(device=device, dtype=torch.float32)
+        pad_vec = torch.zeros((bsz, kv_len), device=device, dtype=torch.float32)
+        if cur_len > 0:
+            pad_vec[:, past:past + cur_len] = key_pad_flags * -1e9
+        mask = mask + pad_vec.to(dtype)[:, None, None, :]
+    return mask
+
+
+def _noncausal_shortconv_forward(self, hidden_states, past_key_values=None, cache_position=None,
+                                 attention_mask=None, **kwargs):
+    x = apply_mask_to_padding_states(hidden_states, attention_mask)
+    BCx = self.in_proj(x).transpose(-1, -2)
+    B, C, x = BCx.chunk(3, dim=-2)
+    Bx = B * x
+    k = self.conv.weight.shape[-1]; pad = k // 2
+    conv_out = F.conv1d(Bx, weight=self.conv.weight, bias=self.conv.bias,
+                        stride=1, padding=pad, dilation=1, groups=Bx.shape[1])
+    if conv_out.shape[-1] > Bx.shape[-1]:
+        conv_out = conv_out[..., :Bx.shape[-1]]
+    elif conv_out.shape[-1] < Bx.shape[-1]:
+        conv_out = F.pad(conv_out, (0, Bx.shape[-1] - conv_out.shape[-1]))
+    y = C * conv_out
+    y = y.transpose(-1, -2).contiguous()
+    return self.out_proj(y)
+
+
+def _shortconv_forward(self, *args, **kwargs):
+    return self.slow_forward(*args, **kwargs)
+
+
+_PATCHED = False
+def _install_patches():
+    global _PATCHED
+    if _PATCHED:
+        return
+    _lfm2_mod.create_causal_mask = _bidirectional_mask
+    Lfm2ShortConv.slow_forward = _noncausal_shortconv_forward
+    Lfm2ShortConv.forward = _shortconv_forward
+    _PATCHED = True
+_install_patches()
+
+
+def _set_attention_noncausal(model):
+    for m in model.modules():
+        if isinstance(m, Lfm2Attention):
+            m.is_causal = False
+
+
+class Lfm2BidirectionalModel_theirs(Lfm2Model):
+    def __init__(self, config):
+        _install_patches()
+        super().__init__(config)
+        _set_attention_noncausal(self)
+
+
+# ---- token-classification head (BIOES) ----
+class Lfm2BidirP2ForTokenClassification(Lfm2PreTrainedModel):
+    config_class = Lfm2Config
+    base_model_prefix = "lfm2"
+
+    def __init__(self, config: Lfm2Config):
+        _install_patches()
+        config = type(config).from_dict({**config.to_dict(), "use_cache": False})
+        super().__init__(config)
+        self.num_labels = config.num_labels
+        self.lfm2 = Lfm2BidirectionalModel_theirs(config)
+        cd = getattr(config, "classifier_dropout", None)
+        self.dropout = nn.Dropout(cd if cd is not None else 0.1)
+        self.classifier = nn.Linear(config.hidden_size, config.num_labels)
+        self.register_buffer("class_weights", torch.ones(config.num_labels), persistent=True)
+        self.label_smoothing = float(getattr(config, "label_smoothing", 0.0))
+        self.post_init()
+
+    def set_class_weights(self, weights: torch.Tensor) -> None:
+        if weights.shape != (self.num_labels,):
+            raise ValueError(f"class_weights must be ({self.num_labels},), got {tuple(weights.shape)}")
+        self.class_weights = weights.to(device=self.class_weights.device, dtype=torch.float32)
+
+    def get_input_embeddings(self):
+        return self.lfm2.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.lfm2.embed_tokens = value
+
+    def forward(self, input_ids=None, attention_mask=None, position_ids=None,
+                inputs_embeds=None, labels=None, output_hidden_states=None,
+                output_attentions=None, return_dict=None, **kwargs) -> TokenClassifierOutput:
+        return_dict = True if return_dict is None else return_dict
+        outputs = self.lfm2(input_ids=input_ids, attention_mask=attention_mask,
+                            position_ids=position_ids, inputs_embeds=inputs_embeds,
+                            use_cache=False, output_attentions=output_attentions,
+                            output_hidden_states=output_hidden_states, return_dict=True)
+        hidden = self.dropout(outputs.last_hidden_state)
+        logits = self.classifier(hidden)
+        loss = None
+        if labels is not None:
+            cw = self.class_weights.float()
+            if not torch.isfinite(cw).all():
+                cw = torch.ones_like(cw)
+            loss = F.cross_entropy(logits.view(-1, self.num_labels).float(), labels.view(-1),
+                                   weight=cw, ignore_index=-100, label_smoothing=self.label_smoothing)
+        if not return_dict:
+            out = (logits,) + outputs[1:]
+            return ((loss,) + out) if loss is not None else out
+        return TokenClassifierOutput(loss=loss, logits=logits,
+                                     hidden_states=outputs.hidden_states, attentions=outputs.attentions)