vela3-flash / span_decode_v2.py
adaamko's picture
Vela 3.0
b89343c
Raw History Blame Contribute Delete
13.7 kB
"""Word-aware span decoding for the schema encoder's sigmoid span head (PII long-doc fix, 26 Sep).
The v1 decoder (span_data.decode) thresholds every token x label independently and turns each above-threshold
run into a span. With sub-word tokens (the Vela tokenizer splits digits one by one) a single below-threshold
piece fragments one entity into several false positives ('ois', '0', 'example'), and every label is decoded
separately, so overlapping labels each cost a false positive.
decode_v2(probs, offs, names, text, thr, ...) -> [dict(start, end, label)] (code-point offsets)
1. label resolution: each token takes its argmax label if that probability > thr, else O
(overlapping labels resolved by max score; upstream Vela PII argmax-BIO likewise gives one label per token).
2. word units from the text (GLiNER2 word splitter, see _UNIT): URL | run of letters/digits/@._-+ | any other
single non-space char; every CJK / kana / hangul / full-width character is its own unit (character level).
A token belongs to the unit that holds its first non-space character.
3. within a unit, one label: the label with the largest summed probability over the unit's labelled tokens
wins and the other labelled tokens of the unit are relabelled to it.
4. fill: unlabelled tokens BETWEEN two tokens of the same label inside one unit take that label;
extension (extend="unit", default): to the outermost letter/digit tokens of the unit, through inner
punctuation ('ois' -> 'lois123@x.example.org'), never onto a bare edge '[' / ']' / '.';
extend="alnum": only through directly adjacent letter/digit tokens; extend=False: off.
5. contiguous tokens of one label form a span (never joins different labels: upstream decode_entities rule);
whitespace is trimmed (upstream), and so are enclosing bracket / quote characters at the span edges.
6. optional: bridge=True joins two same-label spans separated by one unit whose probability for that
label is > thr * bridge_frac; min_chars={label: n} drops shorter spans.
Upstream semantics kept: exact code-point spans, whitespace trim, different types never joined. Not
reproducible from a sigmoid head: the B- tag that splits two adjacent same-type entities (it has no B/I).
"""
import re
import numpy as np
_CJK = ("⺀-⿿ -〿぀-ヿ㄀-ㇿ㐀-䶿一-鿿ꀀ-꓏"
"가-힯豈-﫿︰-﹏＀-￯")
# Word units = GLiNER2's splitters (gliner2/processing/word_splitter.py, Apache-2.0): WhitespaceTokenSplitter's URL
# rule, CharLevelSplitter's [A-Za-z0-9@._\-+]+ runs (e-mails, IPs, phones, IDs stay one word), widened to Unicode
# letters/digits (Straße, título), and every CJK / kana / hangul / full-width character on its own (char level).
_UNIT = re.compile(rf"(?:https?://|www\.)[^\s\]\)\}}>「」『』\"'“”{_CJK}]+"
rf"|(?:(?![{_CJK}])[^\W_]|[@._\-+])+"
rf"|\S")
_ALNUM = re.compile(r"[^\W_]")
EDGE = set("[](){}<>「」『』【】《》〈〉\"'“”‘’«»‹›")
def units(text):
"""char index -> unit id (-1 for whitespace)."""
u = np.full(len(text) + 1, -1, dtype=np.int64)
for k, m in enumerate(_UNIT.finditer(text)):
u[m.start():m.end()] = k
return u
def _first_nonspace(text, a, b):
while a < b and text[a].isspace():
a += 1
return a
def trim(text, s, e):
while s < e and (text[s].isspace() or text[s] in EDGE):
s += 1
while e > s and (text[e - 1].isspace() or text[e - 1] in EDGE):
e -= 1
return s, e
def decode_v2(probs, offs, names, text, thr, fill=True, extend="unit", unit_vote=True, trim_edges=True,
bridge=False, bridge_frac=0.5, min_chars=None, u=None):
probs = np.asarray(probs, dtype=np.float32); offs = np.asarray(offs)
T = len(offs)
if T == 0:
return []
valid = offs[:, 1] > offs[:, 0]
k = probs.argmax(1); pk = probs[np.arange(T), k]
lab = np.where(valid & (pk > thr), k, -1)
u = units(text) if u is None else u
tu = np.full(T, -1, dtype=np.int64)
alnum = np.zeros(T, dtype=bool)
for t in range(T):
if not valid[t]:
continue
a, b = int(offs[t, 0]), int(offs[t, 1])
f = _first_nonspace(text, a, b)
tu[t] = u[f] if f < b and f < len(text) else -1
alnum[t] = bool(_ALNUM.search(text[a:b]))
# group tokens by unit (tokens are in text order, units contiguous)
if unit_vote or fill or extend:
t = 0
while t < T:
if tu[t] < 0:
t += 1; continue
t2 = t
while t2 + 1 < T and (tu[t2 + 1] == tu[t] or not valid[t2 + 1]):
t2 += 1
idx = [i for i in range(t, t2 + 1) if valid[i]]
L = [i for i in idx if lab[i] >= 0]
if L:
if unit_vote and len({int(lab[i]) for i in L}) > 1:
sc = {}
for i in L:
sc[int(lab[i])] = sc.get(int(lab[i]), 0.0) + float(pk[i])
w = max(sc, key=sc.get)
for i in L:
lab[i] = w
if fill:
for c in set(int(lab[i]) for i in L):
pos = [i for i in L if lab[i] == c]
for i in idx:
if pos[0] < i < pos[-1] and lab[i] < 0:
lab[i] = c
if extend:
for c in set(int(lab[i]) for i in idx if lab[i] >= 0):
pos = [idx.index(i) for i in idx if lab[i] == c]
if extend == "unit": # to the outermost letter/digit tokens of the unit, through inner punctuation
an = [j for j, i in enumerate(idx) if alnum[i]]
lo, hi = min([pos[0]] + an[:1]), max([pos[-1]] + an[-1:])
if all(lab[idx[j]] < 0 or lab[idx[j]] == c for j in range(lo, hi + 1)):
for j in range(lo, hi + 1):
lab[idx[j]] = c
else: # "alnum": only through adjacent letter/digit tokens
j = pos[0] - 1
while j >= 0 and lab[idx[j]] < 0 and alnum[idx[j]]:
lab[idx[j]] = c; j -= 1
j = pos[-1] + 1
while j < len(idx) and lab[idx[j]] < 0 and alnum[idx[j]]:
lab[idx[j]] = c; j += 1
t = t2 + 1
# contiguous runs of one label (skipping zero-width tokens)
spans = []
cur = None
for t in range(T):
if not valid[t]:
continue
c = int(lab[t])
if cur is not None and c == cur[0]:
cur[2] = int(offs[t, 1]); cur[3].append(t); continue
if cur is not None:
spans.append(cur)
cur = [c, int(offs[t, 0]), int(offs[t, 1]), [t]] if c >= 0 else None
if cur is not None:
spans.append(cur)
if bridge and len(spans) > 1:
out = [spans[0]]
for sp in spans[1:]:
pv = out[-1]
if sp[0] == pv[0]:
gap = [t for t in range(pv[3][-1] + 1, sp[3][0]) if valid[t]]
gu = {int(tu[t]) for t in gap if tu[t] >= 0}
if gap and len(gu) <= 1 and all(probs[t, sp[0]] > thr * bridge_frac for t in gap if tu[t] >= 0):
pv[2] = sp[2]; pv[3] = pv[3] + gap + sp[3]; continue
out.append(sp)
spans = out
res = []
for c, s, e, _ in spans:
s, e = trim(text, s, e) if trim_edges else trim_ws(text, s, e)
if e <= s:
continue
if min_chars and e - s < min_chars.get(names[c], 0):
continue
res.append(dict(start=s, end=e, label=names[c]))
return res
def trim_ws(text, s, e):
while s < e and text[s].isspace():
s += 1
while e > s and text[e - 1].isspace():
e -= 1
return s, e
def decode_v1_trimmed(probs, offs, names, text, thr):
"""The scorer's current rule (span_data.decode + whitespace trim), for before/after tables."""
from span_data import decode
out = []
for s in decode(probs, offs, names, thr):
a, b = trim_ws(text, s["start"], s["end"])
if b > a:
out.append(dict(start=a, end=b, label=s["label"]))
return out
# ---------------------------------------------------------------------------------------------------------------
# row-level API (rows = dicts with text, spans [{start,end,label}], names, offs, probs; see dump_span_probs_v2.py)
# ---------------------------------------------------------------------------------------------------------------
GRID = [round(x, 2) for x in np.arange(0.05, 0.96, 0.05)]
DECODERS = {
"v1": lambda r, thr: decode_v1_trimmed(r["probs"], r["offs"], r["names"], r["text"], thr),
"v2": lambda r, thr: decode_v2(r["probs"], r["offs"], r["names"], r["text"], thr),
"v2_noext": lambda r, thr: decode_v2(r["probs"], r["offs"], r["names"], r["text"], thr, extend=False),
"v2_alnum": lambda r, thr: decode_v2(r["probs"], r["offs"], r["names"], r["text"], thr, extend="alnum"),
"v2_notrim": lambda r, thr: decode_v2(r["probs"], r["offs"], r["names"], r["text"], thr, trim_edges=False),
"v2_bridge": lambda r, thr: decode_v2(r["probs"], r["offs"], r["names"], r["text"], thr, bridge=True),
"v2_argmax_only": lambda r, thr: decode_v2(r["probs"], r["offs"], r["names"], r["text"], thr, fill=False,
extend=False, unit_vote=False, trim_edges=False),
}
def decode_rows(rows, thr, decoder="v2"):
f = DECODERS[decoder] if isinstance(decoder, str) else decoder
return [f(r, thr) for r in rows]
def _psd(span_data_path):
"""Load the upstream semantic-router span_data.py (src/training/model_classifier/pii_model_fine_tuning_lora/span_data.py)
from span_data_path. Eval helper only (dev threshold picks); inference never calls it."""
import importlib.util
from pathlib import Path
if span_data_path is None:
raise ValueError("pii_metrics needs span_data_path: the path of the upstream span_data.py")
spec = importlib.util.spec_from_file_location("pii_span_data_up", Path(span_data_path))
m = importlib.util.module_from_spec(spec); spec.loader.exec_module(m)
return m
PII_TYPES = ["AGE", "CREDIT_CARD", "DATE_TIME", "DOMAIN_NAME", "EMAIL_ADDRESS", "GPE", "IBAN_CODE", "IP_ADDRESS", "NRP",
"ORGANIZATION", "PERSON", "PHONE_NUMBER", "STREET_ADDRESS", "TITLE", "US_DRIVER_LICENSE", "US_SSN", "ZIP_CODE"]
def pii_metrics(rows, preds, span_data_path):
"""Upstream Vela PII exact-entity metric (span_data.entity_metrics, whitespace-trimmed gold); span_data_path: see _psd."""
P = _psd(span_data_path)
recs = [dict(full_text=r["text"], spans=[dict(start_position=s["start"], end_position=s["end"], entity_type=s["label"],
entity_value=r["text"][s["start"]:s["end"]]) for s in r["spans"]]) for r in rows]
return P.entity_metrics(recs, [{(s["label"], s["start"], s["end"]) for s in p} for p in preds], PII_TYPES)
def toxic_score(rows, preds):
from span_data import char_f1
return float(np.mean([char_f1(p, r["spans"]) for r, p in zip(rows, preds)]))
def halu_metrics(rows, preds):
from train_schema_span import ragtruth_metrics_spans
return ragtruth_metrics_spans([dict(text=r["text"], spans=r["spans"]) for r in rows], preds)
def family_objective(family, span_data_path=None):
"""Dev objective per span label family (what the threshold is chosen to maximise); span_data_path: PII only (see _psd)."""
if family == "pii":
return lambda rows, preds: pii_metrics(rows, preds, span_data_path)["micro"]["f1"]
if family == "toxic":
return toxic_score
if family == "halu":
return lambda rows, preds: halu_metrics(rows, preds)["example_f1"]
raise KeyError(family)
def pick_threshold(dev_sets, family, decoder="v2", grid=GRID, span_data_path=None):
"""Choose one threshold per label family on DEV rows only.
dev_sets: list of row lists; the objective is the unweighted mean over the sets (for PII: dev short, dev 4096,
dev 8192 -> upstream's length-macro selection, so the dense long rows cannot dominate). Returns (thr, curve)."""
obj = family_objective(family, span_data_path)
curve = {t: float(np.mean([obj(rows, decode_rows(rows, t, decoder)) for rows in dev_sets])) for t in grid}
best = max(curve.values())
thr = min(t for t, v in curve.items() if v >= best - 1e-12) # ties -> lowest threshold
return thr, curve
if __name__ == "__main__": # self-test on the fragment cases seen in run C
text = "Le formulaire indique [pavel_200015@summit15.example.org] dans adresse courriel."
names = ["EMAIL_ADDRESS", "PERSON"]
import unicodedata # noqa
# fake token offsets: split at every non-alnum boundary
offs = [(m.start(), m.end()) for m in re.finditer(r"\s?\w+|\s?[^\w\s]", text)]
p = np.zeros((len(offs), 2), dtype=np.float32)
s0 = text.index("pavel"); e0 = text.index("]")
for i, (a, b) in enumerate(offs):
if a >= s0 - 1 and b <= e0 + 1:
p[i, 0] = 0.9 if i % 3 else 0.2 # fragmented
p[[i for i, (a, b) in enumerate(offs) if text[a:b].strip() == "["][0], 0] = 0.8 # bracket above threshold
out = decode_v2(p, np.array(offs), names, text, 0.5)
assert [(o["label"], text[o["start"]:o["end"]]) for o in out] == [("EMAIL_ADDRESS", "pavel_200015@summit15.example.org")], out
t2 = "表格在日期一栏填写了「2015-06-06」。"
print(units(t2)[:14], "ok")