Update model card and license

config.json

@@ -3,8 +3,8 @@
     "Lfm2BidirForRuleMatching"
   ],
   "auto_map": {
-    "AutoModel": "modeling_bizlint_rule_matching.Lfm2BidirForRuleMatching",
-    "AutoModelForMaskedLM": "modeling_lfm2_bidirectional.Lfm2BidirForMaskedLM"
+    "AutoModel": "modeling_lfm2_bidir_theirs.Lfm2BidirectionalModel_theirs",
+    "AutoModelForMaskedLM": "modeling_lfm2_bidir_theirs.Lfm2BidirForMaskedLM_theirs"
   },
   "block_auto_adjust_ff_dim": true,
   "block_dim": 1024,
@@ -62,4 +62,4 @@
   "use_cache": false,
   "use_pos_enc": true,
   "vocab_size": 65536
-}
+}

modeling_bizlint_rule_matching.py

@@ -1,55 +0,0 @@
-"""Self-contained inference class for the bizlint GLiNER-style rule-matching
-head — mirrors the training-time class so trust_remote_code loading restores
-every checkpoint weight (backbone under `lfm2.*` + tok_proj/rule_proj/
-score_bias head).
-
-    score[t, r] = <P_tok(h_t), P_rule(mean(h[rule_r tokens]))> / sqrt(d) + b
-
-Input format: "Policy:\n- <rule 1>\n- <rule 2>\n\nText:\n<doc>"
-`rule_pool` is a (B, R, T) matrix of normalized pooling weights over each
-rule's token range; sigmoid(score) > 0.5 flags a (token, rule) match.
-"""
-import torch
-import torch.nn as nn
-from transformers.models.lfm2.configuration_lfm2 import Lfm2Config
-from transformers.models.lfm2.modeling_lfm2 import Lfm2PreTrainedModel
-
-from .modeling_lfm2_bidirectional import Lfm2BidirectionalModel
-
-PROJ_D = 256
-
-
-class Lfm2BidirForRuleMatching(Lfm2PreTrainedModel):
-    config_class = Lfm2Config
-    base_model_prefix = "lfm2"
-
-    def __init__(self, config):
-        super().__init__(config)
-        self.lfm2 = Lfm2BidirectionalModel(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))
-        self.post_init()
-
-    def forward(self, input_ids=None, attention_mask=None, rule_pool=None,
-                labels=None, label_mask=None, **kw):
-        h = self.lfm2(input_ids=input_ids, attention_mask=attention_mask,
-                      use_cache=False, return_dict=True).last_hidden_state
-        if rule_pool is None:
-            return {"loss": None, "logits": None, "last_hidden_state": h}
-        rule_rep = torch.bmm(rule_pool, h)
-        tp = self.tok_proj(h)
-        rp = self.rule_proj(rule_rep)
-        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():
-                loss = nn.functional.binary_cross_entropy_with_logits(
-                    scores[m], labels[m],
-                    pos_weight=torch.tensor(8.0, device=scores.device))
-            else:
-                loss = scores.sum() * 0.0
-        return {"loss": loss, "logits": scores, "last_hidden_state": h}

modeling_lfm2_bidirectional.py → modeling_lfm2_bidir_theirs.py

@@ -1,21 +1,32 @@
-"""LFM2 backbone with bidirectional attention + non-causal short-conv.
-
-Wired into the HF repo via `auto_map` in config.json so that
-
-AutoModel.from_pretrained(repo, trust_remote_code=True)
-AutoModelForMaskedLM.from_pretrained(repo, trust_remote_code=True)
-
-both return a model with the encoder-style patches already applied.
-
-Supports `attn_implementation` in {"eager", "sdpa", "flash_attention_2"}:
-
-eager/sdpa consume a 4D additive pad-only mask and reproduce the exact
-training-time behavior; flash_attention_2 receives the 2D padding mask (or
-None) and runs the kernel non-causally via `Lfm2Attention.is_causal = False`,
-yielding outputs equivalent to the unpadded forward.
+"""
+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.
 """
 
-import math
 from typing import Optional
 
 import torch
@@ -35,6 +46,9 @@ from transformers.models.lfm2.modeling_lfm2 import (
 )
 
 
+# --------------------------------------------------------------------------- #
+#  Patch 1: bidirectional attention mask                                      #
+# --------------------------------------------------------------------------- #
 def _bidirectional_mask(
     config,
     input_embeds: torch.Tensor = None,
@@ -44,14 +58,13 @@ def _bidirectional_mask(
     position_ids: Optional[torch.LongTensor] = None,
     **kwargs,
 ) -> Optional[torch.Tensor]:
-    # transformers has renamed the embeds kwarg across versions
-    # (input_embeds <-> inputs_embeds); accept either to stay forward-compatible.
+    # 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 uses the 2D padding mask to unpad sequences; causality is
-        # controlled by `Lfm2Attention.is_causal` (set to False below).
+        # 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
@@ -73,6 +86,9 @@ def _bidirectional_mask(
     return mask
 
 
+# --------------------------------------------------------------------------- #
+#  Patch 2: BiEnc-preview shortconv forward (full pipeline)                   #
+# --------------------------------------------------------------------------- #
 def _noncausal_shortconv_forward(
     self,
     hidden_states: torch.Tensor,
@@ -129,9 +145,12 @@ def _set_attention_noncausal(model) -> None:
             module.is_causal = False
 
 
-class Lfm2BidirectionalModel(Lfm2Model):
-    """LFM2 patched for encoder-style use: 
-    full bidirectional attention + non-causal short-conv."""
+# --------------------------------------------------------------------------- #
+#  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()
@@ -139,8 +158,11 @@ class Lfm2BidirectionalModel(Lfm2Model):
         _set_attention_noncausal(self)
 
 
-class Lfm2BidirectionalForMaskedLM(Lfm2PreTrainedModel):
-    """LFM2 bidirectional encoder with a tied masked-LM head."""
+# --------------------------------------------------------------------------- #
+#  AutoModelForMaskedLM head                                                  #
+# --------------------------------------------------------------------------- #
+class Lfm2BidirForMaskedLM_theirs(Lfm2PreTrainedModel):
+    """Lfm2 bidirectional encoder + MLM head (tied to embed_tokens.weight)."""
 
     config_class = Lfm2Config
     base_model_prefix = "lfm2"
@@ -148,11 +170,13 @@ class Lfm2BidirectionalForMaskedLM(Lfm2PreTrainedModel):
 
     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(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):
@@ -210,102 +234,3 @@ class Lfm2BidirectionalForMaskedLM(Lfm2PreTrainedModel):
             hidden_states=outputs.hidden_states,
             attentions=outputs.attentions,
         )
-
-
-class Lfm2BidirForSequenceRouting(Lfm2PreTrainedModel):
-    """Zero-shot prompt router built on the bidirectional LFM2 encoder."""
-
-    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.lfm2 = Lfm2BidirectionalModel(config)
-        proj_dim = getattr(config, "rule_proj_dim", 256)
-        self.tok_proj = nn.Linear(config.hidden_size, proj_dim)
-        self.rule_proj = nn.Linear(config.hidden_size, proj_dim)
-        self.score_bias = nn.Parameter(torch.tensor(0.0))
-        self.logit_scale = nn.Parameter(torch.tensor(1.0))
-        self.post_init()
-
-    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: Optional[torch.LongTensor] = None,
-        attention_mask: Optional[torch.Tensor] = None,
-        text_pool: Optional[torch.Tensor] = None,
-        category_pool: Optional[torch.Tensor] = None,
-        **kwargs,
-    ):
-        outputs = self.lfm2(
-            input_ids=input_ids,
-            attention_mask=attention_mask,
-            use_cache=False,
-            return_dict=True,
-        )
-        hidden = outputs.last_hidden_state
-        text_rep = torch.bmm(text_pool, hidden).squeeze(1)
-        category_rep = torch.bmm(category_pool, hidden)
-        query = F.normalize(self.tok_proj(text_rep), dim=-1)
-        categories = F.normalize(self.rule_proj(category_rep), dim=-1)
-        scale = torch.clamp(self.logit_scale.exp(), max=30.0)
-        logits = torch.einsum("bd,brd->br", query, categories) * scale + self.score_bias
-        return {"logits": logits}
-
-    @staticmethod
-    def _prefix(routes):
-        body = "\n".join(f"- {route}" for route in routes) if routes else "- (none)"
-        return f"Categories:\n{body}\n\nText:\n"
-
-    @staticmethod
-    def _category_ranges(routes):
-        ranges = []
-        pos = len("Categories:\n")
-        for route in routes:
-            start = pos + 2
-            end = start + len(route)
-            ranges.append((start, end))
-            pos = end + 1
-        return ranges
-
-    @torch.no_grad()
-    def route(self, text, routes, tokenizer, threshold=None):
-        prefix = self._prefix(routes)
-        full_text = prefix + text
-        enc = tokenizer(full_text, return_offsets_mapping=True, return_tensors="pt")
-        offsets = enc.pop("offset_mapping")[0].tolist()
-        enc = {k: v.to(self.device) for k, v in enc.items()}
-
-        text_start = len(prefix)
-        text_idxs = [
-            i for i, (start, end) in enumerate(offsets)
-            if end > text_start and start != end
-        ]
-        text_pool = torch.zeros(1, 1, len(offsets), device=self.device)
-        if text_idxs:
-            text_pool[0, 0, text_idxs] = 1 / len(text_idxs)
-
-        category_pool = torch.zeros(1, len(routes), len(offsets), device=self.device)
-        for route_idx, (start, end) in enumerate(self._category_ranges(routes)):
-            token_idxs = [
-                i for i, (tok_start, tok_end) in enumerate(offsets)
-                if tok_start < end and tok_end > start and tok_start != tok_end
-            ]
-            if token_idxs:
-                category_pool[0, route_idx, token_idxs] = 1 / len(token_idxs)
-
-        logits = self(**enc, text_pool=text_pool, category_pool=category_pool)["logits"][0]
-        probs = logits.softmax(dim=-1).detach().cpu()
-        results = [
-            {"route": route, "score": float(prob)}
-            for route, prob in zip(routes, probs)
-            if threshold is None or prob >= threshold
-        ]
-        return sorted(results, key=lambda item: item["score"], reverse=True)

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()