Upload folder using huggingface_hub

config.json

@@ -0,0 +1,65 @@
+{
+  "architectures": [
+    "Lfm2BidirForRuleMatching"
+  ],
+  "auto_map": {
+    "AutoModel": "modeling_lfm2_bidir_theirs.Lfm2BidirectionalModel_theirs",
+    "AutoModelForMaskedLM": "modeling_lfm2_bidir_theirs.Lfm2BidirForMaskedLM_theirs"
+  },
+  "block_auto_adjust_ff_dim": true,
+  "block_dim": 1024,
+  "block_ffn_dim_multiplier": 1.0,
+  "block_mlp_init_scale": 1.0,
+  "block_multiple_of": 256,
+  "block_norm_eps": 1e-05,
+  "block_out_init_scale": 1.0,
+  "block_use_swiglu": true,
+  "block_use_xavier_init": true,
+  "bos_token_id": 1,
+  "conv_L_cache": 3,
+  "conv_bias": false,
+  "conv_dim": 1024,
+  "conv_dim_out": 1024,
+  "conv_use_xavier_init": true,
+  "dtype": "float32",
+  "eos_token_id": 7,
+  "hidden_size": 1024,
+  "initializer_range": 0.02,
+  "intermediate_size": 6656,
+  "layer_types": [
+    "conv",
+    "conv",
+    "full_attention",
+    "conv",
+    "conv",
+    "full_attention",
+    "conv",
+    "conv",
+    "full_attention",
+    "conv",
+    "full_attention",
+    "conv",
+    "full_attention",
+    "conv",
+    "full_attention",
+    "conv"
+  ],
+  "max_position_embeddings": 128000,
+  "model_type": "lfm2",
+  "norm_eps": 1e-05,
+  "num_attention_heads": 16,
+  "num_heads": 16,
+  "num_hidden_layers": 16,
+  "num_key_value_heads": 8,
+  "pad_token_id": 0,
+  "rope_parameters": {
+    "rope_theta": 1000000.0,
+    "rope_type": "default"
+  },
+  "rule_proj_dim": 256,
+  "tie_word_embeddings": true,
+  "transformers_version": "5.1.0",
+  "use_cache": false,
+  "use_pos_enc": true,
+  "vocab_size": 65536
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:96c70cb285aa8515361b695cd18330a381a52122ebe91e1be9f224b6537ab4a9
+size 1420051844

modeling_lfm2_bidir_theirs.py

@@ -0,0 +1,236 @@
+"""
+Self-contained inference modeling for the Lfm2 bidirectional encoder MLM
+(BiEnc-preview shortconv variant, aka "bidirectional-2-exp" / mlm-bidir2).
+
+Designed to be SHIPPED ALONGSIDE THE CHECKPOINT via `trust_remote_code`:
+
+    >>> from transformers import AutoModelForMaskedLM, AutoTokenizer, AutoModel
+    >>> tok  = AutoTokenizer.from_pretrained("LiquidAI/mlm-bidir2", trust_remote_code=True)
+    >>> mlm  = AutoModelForMaskedLM.from_pretrained("LiquidAI/mlm-bidir2", trust_remote_code=True)
+    >>> body = AutoModel.from_pretrained("LiquidAI/mlm-bidir2", trust_remote_code=True)
+
+This file:
+  1. Installs the BiEnc-preview-style bidirectional patches:
+       * `create_causal_mask` -> non-causal padding-only additive mask (with an
+         FA2 path that returns the 2D padding mask)
+       * `Lfm2ShortConv.forward` -> full pipeline:
+           in_proj -> chunk(B,C,x) -> B*x -> conv1d(symmetric pad) -> C*conv_out -> out_proj
+         (UNLIKE the "bidirectional-1-exp" variant which used depthwise-only conv1d
+         on hidden_states.)
+  2. Exposes `Lfm2BidirectionalModel_theirs(Lfm2Model)` — the encoder base, with
+     `Lfm2Attention.is_causal = False`.
+  3. Exposes `Lfm2BidirForMaskedLM_theirs(Lfm2PreTrainedModel)` — adds an MLM
+     head tied to `embed_tokens.weight`.
+
+Compatible with `transformers >= 5.0`. AutoModelForMaskedLM dispatches via the
+`auto_map` in config.json. The forward signature absorbs kwargs to stay
+compatible across upstream signature drift between transformers minor versions.
+"""
+
+from typing import Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers.configuration_utils import PretrainedConfig
+from transformers.modeling_outputs import BaseModelOutput, MaskedLMOutput
+from transformers.modeling_utils import PreTrainedModel
+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,
+)
+
+
+# --------------------------------------------------------------------------- #
+#  Patch 1: bidirectional attention mask                                      #
+# --------------------------------------------------------------------------- #
+def _bidirectional_mask(
+    config,
+    input_embeds: torch.Tensor = None,
+    attention_mask: Optional[torch.Tensor] = None,
+    cache_position: Optional[torch.LongTensor] = None,
+    past_key_values=None,
+    position_ids: Optional[torch.LongTensor] = None,
+    **kwargs,
+) -> Optional[torch.Tensor]:
+    # transformers 5.x renamed input_embeds -> inputs_embeds; absorb both.
+    if input_embeds is None:
+        input_embeds = kwargs.get("inputs_embeds")
+
+    if config._attn_implementation == "flash_attention_2":
+        # FA2 only consumes the 2D padding mask to unpad; causality is
+        # controlled by Lfm2Attention.is_causal = False (set below).
+        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
+
+
+# --------------------------------------------------------------------------- #
+#  Patch 2: BiEnc-preview shortconv forward (full pipeline)                   #
+# --------------------------------------------------------------------------- #
+def _noncausal_shortconv_forward(
+    self,
+    hidden_states: torch.Tensor,
+    past_key_values=None,
+    cache_position=None,
+    attention_mask: Optional[torch.Tensor] = None,
+    **kwargs,
+) -> torch.Tensor:
+    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() -> None:
+    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) -> None:
+    for module in model.modules():
+        if isinstance(module, Lfm2Attention):
+            module.is_causal = False
+
+
+# --------------------------------------------------------------------------- #
+#  Base model — Lfm2 backbone, encoder-style                                  #
+# --------------------------------------------------------------------------- #
+class Lfm2BidirectionalModel_theirs(Lfm2Model):
+    """LFM2 backbone patched for encoder-style use:
+    full bidirectional attention + BiEnc-preview non-causal short-conv."""
+
+    def __init__(self, config):
+        _install_patches()
+        super().__init__(config)
+        _set_attention_noncausal(self)
+
+
+# --------------------------------------------------------------------------- #
+#  AutoModelForMaskedLM head                                                  #
+# --------------------------------------------------------------------------- #
+class Lfm2BidirForMaskedLM_theirs(Lfm2PreTrainedModel):
+    """Lfm2 bidirectional encoder + MLM head (tied to embed_tokens.weight)."""
+
+    config_class = Lfm2Config
+    base_model_prefix = "lfm2"
+    _tied_weights_keys = {"lm_head.weight": "lfm2.embed_tokens.weight"}
+
+    def __init__(self, config: Lfm2Config):
+        _install_patches()
+        # MLM never uses KV cache
+        config = type(config).from_dict({**config.to_dict(), "use_cache": False})
+        super().__init__(config)
+        self.lfm2 = Lfm2BidirectionalModel_theirs(config)
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        self.post_init()
+        # tie weights
+        self.lm_head.weight = self.lfm2.embed_tokens.weight
+
+    def get_input_embeddings(self):
+        return self.lfm2.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.lfm2.embed_tokens = value
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, new_embeddings):
+        self.lm_head = new_embeddings
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        output_hidden_states: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        **kwargs,
+    ) -> MaskedLMOutput:
+        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 = outputs.last_hidden_state
+        logits = self.lm_head(hidden)
+
+        loss = None
+        if labels is not None:
+            loss = F.cross_entropy(
+                logits.view(-1, self.config.vocab_size),
+                labels.view(-1),
+                ignore_index=-100,
+            )
+
+        if not return_dict:
+            out = (logits,) + outputs[1:]
+            return ((loss,) + out) if loss is not None else out
+        return MaskedLMOutput(
+            loss=loss,
+            logits=logits,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )

tokenizer_config.json

@@ -0,0 +1,11 @@
+{
+  "backend": "tokenizers",
+  "bos_token": "<|startoftext|>",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|im_end|>",
+  "is_local": false,
+  "mask_token": "<|mask|>",
+  "model_max_length": 1000000000000000019884624838656,
+  "pad_token": "<|pad|>",
+  "tokenizer_class": "TokenizersBackend"
+}

train_bizlint_v02.py

@@ -0,0 +1,237 @@
+#!/usr/bin/env python3
+"""bizlint v02 — GLiNER-style rule matching on the LFM2.5 bidirectional encoder.
+
+No fixed classes: each policy RULE in the input is a label. Rule representations
+are mean-pooled from the same forward pass; every text token is scored against
+every rule via projected dot-product; sigmoid per (token, rule). Neutral is the
+default (no rule above threshold), and new rule types need no retraining.
+
+    score[t, r] = <P_tok(h_t), P_rule(mean(h[rule_r tokens]))> / sqrt(d) + b
+
+Input:  "Policy:\n- <rule 1>\n- <rule 2>\n\nText:\n<doc>"
+Labels: (T_text_tokens x R) binary matrix from span<->rule-idx supervision.
+Eval:   span-level F1 (span + correct rule) at sigmoid > 0.5.
+"""
+import argparse
+import json
+import os
+
+import numpy as np
+import torch
+import torch.nn as nn
+from transformers import AutoTokenizer, Trainer, TrainingArguments
+from transformers.models.lfm2.configuration_lfm2 import Lfm2Config
+from transformers.models.lfm2.modeling_lfm2 import Lfm2PreTrainedModel
+
+import modeling_lfm2_bidir_theirs as bidir
+
+PROJ_D = 256
+
+
+class Lfm2BidirForRuleMatching(Lfm2PreTrainedModel):
+    config_class = Lfm2Config
+    base_model_prefix = "lfm2"
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.lfm2 = bidir.Lfm2BidirectionalModel_theirs(config)
+        d = getattr(config, "rule_proj_dim", PROJ_D)
+        self.tok_proj = nn.Linear(config.hidden_size, d)
+        self.rule_proj = nn.Linear(config.hidden_size, d)
+        self.score_bias = nn.Parameter(torch.tensor(-2.0))  # start conservative
+        self.post_init()
+
+    def forward(self, input_ids=None, attention_mask=None, rule_pool=None,
+                labels=None, label_mask=None, **kw):
+        # rule_pool: (B, R, T) normalized pooling weights over rule token ranges
+        h = self.lfm2(input_ids=input_ids, attention_mask=attention_mask,
+                      use_cache=False, return_dict=True).last_hidden_state  # (B,T,H)
+        rule_rep = torch.bmm(rule_pool, h)                                  # (B,R,H)
+        tp = self.tok_proj(h)                                               # (B,T,d)
+        rp = self.rule_proj(rule_rep)                                       # (B,R,d)
+        scores = torch.einsum("btd,brd->btr", tp, rp) / (tp.shape[-1] ** 0.5) + self.score_bias
+        loss = None
+        if labels is not None:
+            m = label_mask.bool()
+            if m.any():
+                pos_w = torch.tensor(8.0, device=scores.device)
+                loss = nn.functional.binary_cross_entropy_with_logits(
+                    scores[m], labels[m], pos_weight=pos_w)
+            else:
+                loss = scores.sum() * 0.0
+        return {"loss": loss, "logits": scores}
+
+
+def build_prompt(policies):
+    return "Policy:\n" + "\n".join(f"- {p}" for p in policies) + "\n\nText:\n"
+
+
+def encode_row(row, tok, max_len):
+    pols = row["policies"] if row["policies"] else ["(none)"]
+    prefix = build_prompt(pols)
+    full = prefix + row["text"]
+    enc = tok(full, truncation=True, max_length=max_len, return_offsets_mapping=True)
+    off = enc["offset_mapping"]
+    t0 = len(prefix)
+
+    # char ranges of each rule inside the prefix
+    ranges = []
+    pos = len("Policy:\n")
+    for ptxt in pols:
+        start = pos + 2                      # after "- "
+        ranges.append((start, start + len(ptxt)))
+        pos = start + len(ptxt) + 1          # + newline
+
+    T = len(off)
+    R = len(pols)
+    # rule token-pooling sets
+    pool = np.zeros((R, T), dtype=np.float32)
+    for ri, (rs, re_) in enumerate(ranges):
+        idxs = [i for i, (a, b) in enumerate(off) if a < re_ and b > rs and a != b]
+        for i in idxs:
+            pool[ri, i] = 1.0 / max(len(idxs), 1)
+
+    # labels over text tokens
+    spans = [(s + t0, e + t0, ri) for s, e, ri in row["spans"]]
+    labels = np.zeros((T, R), dtype=np.float32)
+    lmask = np.zeros((T, R), dtype=np.float32)
+    for i, (a, b) in enumerate(off):
+        if b <= t0 or a == b:
+            continue
+        lmask[i, :] = 1.0
+        for (s, e, ri) in spans:
+            if a < e and b > s and 0 <= ri < R:
+                labels[i, ri] = 1.0
+    return {"input_ids": enc["input_ids"], "attention_mask": enc["attention_mask"],
+            "pool": pool, "labels": labels, "label_mask": lmask}
+
+
+def decode_rule_spans(score_row, lmask_row, thr=0.5):
+    """(T,R) sigmoid scores -> set of (start,end,rule) spans over masked tokens."""
+    T, R = score_row.shape
+    spans = set()
+    for r in range(R):
+        cur = None
+        for i in range(T + 1):
+            on = i < T and lmask_row[i, r] > 0 and score_row[i, r] > thr
+            if on:
+                cur = (cur[0], i + 1, r) if cur else (i, i + 1, r)
+            else:
+                if cur: spans.add(cur)
+                cur = None
+    return spans
+
+
+def gold_rule_spans(labels_row, lmask_row):
+    return decode_rule_spans(labels_row, lmask_row, thr=0.5)
+
+
+def compute_metrics(eval_pred):
+    (scores, labels, lmask) = eval_pred.predictions if isinstance(eval_pred.predictions, tuple) else (eval_pred.predictions, None, None)
+    # Trainer passes logits only; labels via label_ids is our labels tensor
+    scores = 1 / (1 + np.exp(-scores))
+    labels, lmask = eval_pred.label_ids
+    tp = fp = fn = 0
+    for s_row, l_row, m_row in zip(scores, labels, lmask):
+        ps = decode_rule_spans(s_row, m_row)
+        gs = gold_rule_spans(l_row, m_row)
+        tp += len(ps & gs); fp += len(ps - gs); fn += len(gs - ps)
+    p = tp / max(tp + fp, 1); r = tp / max(tp + fn, 1)
+    return {"span_precision": p, "span_recall": r,
+            "span_f1": 2 * p * r / max(p + r, 1e-9)}
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--data", required=True)
+    ap.add_argument("--base", default="LiquidAI/LFM2.5-Encoder-350M")
+    ap.add_argument("--out", default="bizlint_v02_ckpt")
+    ap.add_argument("--max-len", type=int, default=320)
+    ap.add_argument("--epochs", type=float, default=4)
+    ap.add_argument("--bsz", type=int, default=48)
+    ap.add_argument("--lr", type=float, default=3e-5)
+    ap.add_argument("--max-steps", type=int, default=-1)
+    args = ap.parse_args()
+
+    tok = AutoTokenizer.from_pretrained(args.base, trust_remote_code=True)
+    rows = [json.loads(l) for l in open(args.data)]
+    ds = {"train": [], "val": []}
+    for r in rows:
+        if r["split"] in ds:
+            ds[r["split"]].append(encode_row(r, tok, args.max_len))
+    print(f"train {len(ds['train'])}  val {len(ds['val'])}")
+
+    cfg = Lfm2Config.from_pretrained(args.base)
+    cfg.rule_proj_dim = PROJ_D
+    model, info = Lfm2BidirForRuleMatching.from_pretrained(
+        args.base, config=cfg, torch_dtype=torch.float32, output_loading_info=True)
+    missing = [k for k in info["missing_keys"]
+               if not (k.startswith("tok_proj") or k.startswith("rule_proj") or k.startswith("score_bias"))]
+    assert not missing, f"body weights missing: {missing[:5]}"
+    print("load gate OK — fresh:", sorted(info["missing_keys"]))
+
+    class Collator:
+        def __call__(self, feats):
+            B = len(feats)
+            T = max(len(f["input_ids"]) for f in feats)
+            R = max(f["pool"].shape[0] for f in feats)
+            pad = tok.pad_token_id or 0
+            ii = np.full((B, T), pad, dtype=np.int64)
+            am = np.zeros((B, T), dtype=np.int64)
+            pool = np.zeros((B, R, T), dtype=np.float32)
+            lab = np.zeros((B, T, R), dtype=np.float32)
+            lm = np.zeros((B, T, R), dtype=np.float32)
+            for i, f in enumerate(feats):
+                n = len(f["input_ids"]); r = f["pool"].shape[0]
+                ii[i, :n] = f["input_ids"]; am[i, :n] = f["attention_mask"]
+                pool[i, :r, :n] = f["pool"]
+                lab[i, :n, :r] = f["labels"]; lm[i, :n, :r] = f["label_mask"]
+            return {"input_ids": torch.from_numpy(ii), "attention_mask": torch.from_numpy(am),
+                    "rule_pool": torch.from_numpy(pool),
+                    "labels": torch.from_numpy(lab), "label_mask": torch.from_numpy(lm)}
+
+    class RuleTrainer(Trainer):
+        def prediction_step(self, model, inputs, prediction_loss_only, ignore_keys=None):
+            with torch.no_grad():
+                out = model(**{k: v.to(model.device) for k, v in inputs.items()})
+            return (out["loss"].detach() if out["loss"] is not None else None,
+                    out["logits"].detach().float().cpu(),
+                    (inputs["labels"].cpu(), inputs["label_mask"].cpu()))
+
+    targs = TrainingArguments(
+        output_dir=args.out,
+        num_train_epochs=args.epochs, max_steps=args.max_steps,
+        per_device_train_batch_size=args.bsz, per_device_eval_batch_size=args.bsz,
+        learning_rate=args.lr, warmup_ratio=0.06, weight_decay=0.01,
+        lr_scheduler_type="cosine", bf16=True,
+        logging_steps=20, eval_strategy="epoch", save_strategy="no",
+        report_to=[], remove_unused_columns=False, seed=0,
+        eval_do_concat_batches=False,
+    )
+
+    def cm(eval_pred):
+        # eval_do_concat_batches=False: predictions/label_ids are LISTS of batches
+        tp = fp = fn = 0
+        for scores, (labels, lmask) in zip(eval_pred.predictions, eval_pred.label_ids):
+            sc = 1 / (1 + np.exp(-np.asarray(scores)))
+            for s_row, l_row, m_row in zip(sc, np.asarray(labels), np.asarray(lmask)):
+                ps = decode_rule_spans(s_row, m_row)
+                gs = gold_rule_spans(l_row, m_row)
+                tp += len(ps & gs); fp += len(ps - gs); fn += len(gs - ps)
+        p = tp / max(tp + fp, 1); r = tp / max(tp + fn, 1)
+        return {"span_precision": p, "span_recall": r, "span_f1": 2 * p * r / max(p + r, 1e-9)}
+
+    trainer = RuleTrainer(model=model, args=targs, train_dataset=ds["train"],
+                          eval_dataset=ds["val"], data_collator=Collator(),
+                          compute_metrics=cm)
+    trainer.train()
+    m = trainer.evaluate()
+    print("FINAL_VAL:", json.dumps({k: round(v, 4) for k, v in m.items() if isinstance(v, float)}))
+
+    model.save_pretrained(os.path.join(args.out, "final"))
+    tok.save_pretrained(os.path.join(args.out, "final"))
+    print("SAVED", os.path.join(args.out, "final"))
+
+
+if __name__ == "__main__":
+    main()