File size: 6,554 Bytes
dbd7316 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 | """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)
|