| """gemmaExtract: typed field extraction in one forward pass (no text generation) on a Gemma 4 backbone.""" |
| import math |
| import torch |
| import torch.nn as nn |
|
|
| TYPES = {"text": 0, "list": 1, "number": 2, "date": 3, "time": 4, "duration": 5, "enum": 6} |
| OPS = ["none", "sum", "diff", "avg", "count", "div", "mul"] |
| ORDERED = {"diff", "div"} |
|
|
|
|
| def num(t, en=False): |
| """Chuoi -> float hoac None. en=True: ',' hang nghin, '.' thap phan (DROP, TAT-QA). Mac dinh (hoa don VN/ID, kieu hon hop): |
| 2 loai dau -> dau cuoi la thap phan; 1 loai dau -> thap phan tru khi theo sau dung 3 chu so; '(12)' -> -12. |
| Giong data/common.py parse_num.""" |
| import re |
| m = re.search(r"-?\d[\d.,]*\d|-?\d", t or "") |
| if not m: |
| return None |
| x = m.group() |
| neg = x.startswith("-") or ("(" in t and ")" in t) |
| x = x.lstrip("-") |
| if en: |
| x = x.replace(",", "") |
| if x.count(".") > 1: |
| return None |
| else: |
| seps = [c for c in x if c in ".,"] |
| if seps: |
| last = max(x.rfind("."), x.rfind(",")) |
| dec = len(set(seps)) == 2 or (len(seps) == 1 and len(x) - last - 1 != 3) |
| x = re.sub(r"[.,]", "", x[:last]) + "." + x[last + 1:] if dec else re.sub(r"[.,]", "", x) |
| try: |
| v = float(x) |
| except ValueError: |
| return None |
| return -v if neg else v |
|
|
|
|
| MONTHS = {m: i + 1 for i, m in enumerate(["jan", "feb", "mar", "apr", "may", "jun", "jul", "aug", "sep", "oct", "nov", "dec"])} |
|
|
|
|
| def date(t, dayfirst=True): |
| """Chuoi ngay -> 'YYYY-MM-DD' hoac None (ngay truoc thang mac dinh, nhu hoa don VN/MY). Giong data/common.py parse_date.""" |
| import datetime, re |
| m = re.search(r"(\d{4})[-/.](\d{1,2})[-/.](\d{1,2})", t or "") |
| vn = re.search(r"(\d{1,2})\s*tháng\s*(\d{1,2})\s*,?\s*năm\s*(\d{4})", t or "", re.I) |
| if m: |
| y, mo, d = map(int, m.groups()) |
| elif vn: |
| d, mo, y = map(int, vn.groups()) |
| else: |
| m = re.search(r"(\d{1,2})[-/.](\d{1,2})[-/.](\d{4}|\d{2})", t or "") |
| m2 = re.search(r"(\d{1,2})[-/ .]*([A-Za-z]{3})[A-Za-z]*[-/ .,]*(\d{4}|\d{2})", t or "") |
| m3 = re.search(r"([A-Za-z]{3})[A-Za-z]*[-/ .]*(\d{1,2})[-/ .,]+(\d{4}|\d{2})", t or "") |
| if m: |
| d, mo, y = map(int, m.groups()) |
| if not dayfirst: |
| d, mo = mo, d |
| elif m2 and m2.group(2).lower() in MONTHS: |
| d, mo, y = int(m2.group(1)), MONTHS[m2.group(2).lower()], int(m2.group(3)) |
| elif m3 and m3.group(1).lower() in MONTHS: |
| mo, d, y = MONTHS[m3.group(1).lower()], int(m3.group(2)), int(m3.group(3)) |
| else: |
| return None |
| y += 2000 if y < 100 else 0 |
| try: |
| return datetime.date(y, mo, d).isoformat() |
| except ValueError: |
| return None |
|
|
|
|
| def number_style(text): |
| """Tu nhan kieu so cua van ban: True = EN (1,234.5), False = kieu cham hang nghin (1.234,5 VN/ID). |
| Dem '1,234' + '5.56' so voi '1.234' + '5,56'; ngang nhau -> False (num() kieu hon hop van doc dung phan lon).""" |
| import re |
| en = len(re.findall(r"\d,\d{3}(?!\d)", text)) + len(re.findall(r"\d\.\d{1,2}(?!\d)", text)) |
| vn = len(re.findall(r"\d\.\d{3}(?!\d)", text)) + len(re.findall(r"\d,\d{1,2}(?!\d)", text)) |
| return en > vn |
|
|
|
|
| def snap(text, a, b): |
| """Noi span ra tron con so: tokenizer Gemma tach tung chu so nen dau ra hay cat vun ('1951' -> '1' + '51'; dev DROP/TAT-QA: |
| 152/280 loi so phai tinh la vi vay). Gom ca dau ngoac am '(46)'.""" |
| dig = lambda i: 0 <= i < len(text) and text[i].isdigit() |
| while a > 0 and (dig(a - 1) or (text[a - 1] in ".," and dig(a - 2))): |
| a -= 1 |
| while b < len(text) and (dig(b) or (text[b] in ".," and dig(b + 1))): |
| b += 1 |
| if a > 0 and b < len(text) and text[a - 1] == "(" and text[b] == ")": |
| a, b = a - 1, b + 1 |
| return a, b |
|
|
|
|
| def trim(text, a, b): |
| """Bo khoang trang hai dau span (offset token Gemma kem dau cach dung truoc).""" |
| while a < b and text[a].isspace(): |
| a += 1 |
| while b > a and text[b - 1].isspace(): |
| b -= 1 |
| return a, b |
|
|
|
|
| def merge(spans): |
| out = [] |
| for a, b in sorted(spans): |
| if out and a <= out[-1][1]: |
| out[-1] = (out[-1][0], max(out[-1][1], b)) |
| else: |
| out.append((a, b)) |
| return out |
|
|
|
|
| def compute(op, vals, reverse=False): |
| if op == "count": |
| return float(len(vals)) |
| if not vals: |
| return None |
| if op == "sum": |
| return sum(vals) |
| if op == "avg": |
| return sum(vals) / len(vals) |
| if op == "mul": |
| out = 1.0 |
| for v in vals: |
| out *= v |
| return out |
| if len(vals) < 2: |
| return None |
| a, b = (vals[1], vals[0]) if reverse else (vals[0], vals[1]) |
| if op == "diff": |
| return a - b |
| return a / b if b else None |
|
|
|
|
| SPAN_TYPES = {"text", "number", "date", "time", "duration"} |
| MAX_SPAN_TOK = 64 |
|
|
|
|
| def _ids(tok, text): |
| return tok(text, add_special_tokens=False)["input_ids"] |
|
|
|
|
| def encode(tok, rec, max_len=1024): |
| """-> dict: ids, state (a, b) vi tri token state, off (offset ky tu cua token state), fields [(ten, kieu, (a, b) span instruction)].""" |
| schema, fields = _ids(tok, "\n\nSCHEMA FIELDS:\n"), [] |
| for i, (name, f) in enumerate(rec["fields"].items()): |
| schema += _ids(tok, "\nFIELD %d\nID: %s\nTYPE: %s\nINSTRUCTION: " % (i + 1, name, f["type"])) |
| a = len(schema) |
| schema += _ids(tok, f.get("instructions") or name) |
| fields.append((name, f["type"], (a, len(schema)))) |
| if f.get("options"): |
| schema += _ids(tok, "\nOPTIONS: " + " | ".join(map(str, f["options"]))) |
| schema += _ids(tok, "\nEND FIELD\n") |
| prefix = [tok.bos_token_id] + _ids(tok, "STATE:\n") |
| t = tok(rec["state"], add_special_tokens=False, return_offsets_mapping=True) |
| room = max_len - len(prefix) - len(schema) |
| if room < 16: |
| raise ValueError("schema qua dai: %d token" % len(schema)) |
| sid, soff = t["input_ids"][:room], [tuple(o) for o in t["offset_mapping"][:room]] |
| base = len(prefix) + len(sid) |
| return {"ids": prefix + sid + schema, "state": (len(prefix), base), "off": soff, |
| "fields": [(n, ty, (a + base, b + base)) for n, ty, (a, b) in fields]} |
|
|
|
|
| def labels(enc, rec): |
| """Nhan train cho moi field: null, cap (start, end) token chap nhan duoc (moi lan xuat hien), nhan B/I/O cho list. |
| Field co dap an nhung moi span bi cat do max_len -> bo khoi loss (None).""" |
| off, out = enc["off"], [] |
| for name, ty, _ in enc["fields"]: |
| ans = rec["answers"].get(name) |
| if ans is None: |
| out.append(None); continue |
| pairs, tags, opnd = [], [0] * len(off), [0] * len(off) |
| der = None if ans["null"] else ans.get("derivation") |
| if der and der.get("op") in OPS and der["op"] != "none": |
| spans_txt = [] |
| for s, e in sorted(der["operands"]): |
| toks = [i for i, (a, b) in enumerate(off) if a < e and b > s] |
| for i in toks: |
| opnd[i] = 1 |
| spans_txt.append(rec["state"][s:e]) |
| if not any(opnd): |
| out.append(None); continue |
| vals, order = [num(t, en=True) for t in spans_txt], None |
| if der["op"] in ORDERED and len(vals) == 2 and None not in vals and isinstance(ans.get("value"), (int, float)): |
| for rev in (False, True): |
| v = compute(der["op"], vals, rev) |
| if v is not None and abs(v - ans["value"]) <= 1e-6 + 1e-3 * abs(ans["value"]): |
| order = int(rev); break |
| out.append({"null": False, "pairs": [], "tags": tags, "op": OPS.index(der["op"]), "opnd": opnd, "order": order}) |
| continue |
| for s, e in ([] if ans["null"] else ans["spans"]): |
| toks = [i for i, (a, b) in enumerate(off) if a < e and b > s] |
| if toks: |
| pairs.append((toks[0], toks[-1])) |
| tags[toks[0]] = 1 |
| for i in toks[1:]: |
| tags[i] = 2 |
| if not ans["null"] and not pairs: |
| out.append(None); continue |
| out.append({"null": ans["null"], "pairs": pairs, "tags": tags, "op": 0 if ty == "number" and not ans["null"] else None, |
| "opnd": None, "order": None}) |
| return out |
|
|
|
|
| class Routing(nn.Module): |
| """Query field cross-attention vao memory + FFN (EvidenceRoutingLayer cua Clef).""" |
|
|
| def __init__(self, w, heads, ff): |
| super().__init__() |
| self.qn, self.mn, self.fn = nn.LayerNorm(w), nn.LayerNorm(w), nn.LayerNorm(w) |
| self.att = nn.MultiheadAttention(w, heads, batch_first=True) |
| self.ff = nn.Sequential(nn.Linear(w, ff), nn.GELU(), nn.Linear(ff, w)) |
|
|
| def forward(self, q, mem): |
| m = self.mn(mem) |
| q = q + self.att(self.qn(q), m, m, need_weights=False)[0] |
| return q + self.ff(self.fn(q)) |
|
|
|
|
| class S1Head(nn.Module): |
| def __init__(self, hidden, width=1024, routing=2, layers=4, heads=16, ff=4096): |
| super().__init__() |
| self.w = width |
| self.norm = nn.LayerNorm(hidden) |
| self.mem = nn.Linear(hidden, width, bias=False) |
| self.qproj = nn.Linear(hidden, width, bias=False) |
| self.gproj = nn.Linear(hidden, width, bias=False) |
| self.type_emb = nn.Embedding(len(TYPES), width) |
| self.route = nn.ModuleList([Routing(width, heads, ff) for _ in range(routing)]) |
| self.joint = nn.ModuleList([nn.TransformerDecoderLayer(width, heads, ff, activation="gelu", batch_first=True, norm_first=True) |
| for _ in range(layers)]) |
| self.fnorm = nn.LayerNorm(width) |
| self.null = nn.Linear(width, 1) |
| self.ks, self.ke, self.kt = (nn.Linear(hidden, width, bias=False) for _ in range(3)) |
| self.qs, self.qe = nn.Linear(width, width), nn.Linear(width, width) |
| self.qt = nn.Linear(width, 3 * width) |
| self.op = nn.Linear(width, len(OPS)) |
| self.kn, self.qn = nn.Linear(hidden, width, bias=False), nn.Linear(width, width) |
| self.order = nn.Linear(width, 2) |
|
|
| def forward(self, hidden, enc): |
| """hidden: (L, H) cua 1 record. -> null (F,), start/end (F, S), tag (F, S, 3).""" |
| h = self.norm(hidden.float()) |
| mem = self.mem(h).unsqueeze(0) |
| q = torch.stack([h[a:b].mean(0) for _, _, (a, b) in enc["fields"]]) |
| ty = torch.tensor([TYPES[t] for _, t, _ in enc["fields"]], device=h.device) |
| q = (self.qproj(q) + self.gproj(h[-1]) + self.type_emb(ty)).unsqueeze(0) |
| for layer in self.route: |
| q = layer(q, mem) |
| for layer in self.joint: |
| q = layer(q, mem) |
| f = self.fnorm(q[0]) |
| a, b = enc["state"] |
| hs = h[a:b] |
| sc = 1 / math.sqrt(self.w) |
| start = self.qs(f) @ self.ks(hs).T * sc |
| end = self.qe(f) @ self.ke(hs).T * sc |
| tag = torch.einsum("fcw,sw->fsc", self.qt(f).view(len(f), 3, self.w), self.kt(hs)) * sc |
| opnd = self.qn(f) @ self.kn(hs).T * sc |
| return self.null(f).squeeze(-1), start, end, tag, self.op(f), opnd, self.order(f) |
|
|
|
|
| def loss(out, labs, types): |
| null, start, end, tag, op, opnd, order = out |
| terms = [] |
| for i, (lab, ty) in enumerate(zip(labs, types)): |
| if lab is None: |
| continue |
| terms.append(nn.functional.binary_cross_entropy_with_logits(null[i], torch.tensor(float(lab["null"]), device=null.device))) |
| if ty == "list": |
| terms.append(nn.functional.cross_entropy(tag[i], torch.tensor(lab["tags"], device=tag.device))) |
| elif not lab["null"] and lab["pairs"]: |
| ls, le = start[i].log_softmax(-1), end[i].log_softmax(-1) |
| terms.append(-torch.logsumexp(torch.stack([ls[s] + le[e] for s, e in lab["pairs"]]), 0)) |
| if lab.get("op") is not None: |
| terms.append(nn.functional.cross_entropy(op[i:i + 1], torch.tensor([lab["op"]], device=op.device))) |
| if lab.get("opnd") is not None: |
| terms.append(nn.functional.binary_cross_entropy_with_logits(opnd[i], torch.tensor(lab["opnd"], dtype=torch.float, device=opnd.device))) |
| if lab.get("order") is not None: |
| terms.append(nn.functional.cross_entropy(order[i:i + 1], torch.tensor([lab["order"]], device=order.device))) |
| return torch.stack(terms).sum() / max(1, len(terms)) if terms else None |
|
|
|
|
| @torch.no_grad() |
| def decode(out, enc, state_text, null_threshold=0.5, en=None): |
| """-> {ten: {null, p_null, guess, confidence, spans}}. guess tra ca khi null (nhu Clairic).""" |
| null, start, end, tag, op, opnd, order = out |
| off, res = enc["off"], {} |
| if en is None: |
| en = number_style(state_text) |
| for i, (name, ty, _) in enumerate(enc["fields"]): |
| p_null = torch.sigmoid(null[i]).item() |
| if ty == "list": |
| tg = tag[i].argmax(-1).tolist() |
| pr = tag[i].softmax(-1) |
| spans, j = [], 0 |
| while j < len(tg): |
| if tg[j] == 1 or (tg[j] == 2 and (j == 0 or tg[j - 1] == 0)): |
| k = j + 1 |
| while k < len(tg) and tg[k] == 2: |
| k += 1 |
| spans.append((*trim(state_text, off[j][0], off[k - 1][1]), pr[j, tg[j]].item())) |
| j = k |
| else: |
| j += 1 |
| vals = list(dict.fromkeys(state_text[s:e].strip() for s, e, _ in spans if state_text[s:e].strip())) |
| conf = min((p for *_, p in spans), default=1 - p_null) |
| res[name] = {"null": p_null > null_threshold or not vals, "p_null": p_null, "guess": vals, "confidence": conf, |
| "spans": [[s, e] for s, e, _ in spans]} |
| continue |
| ps, pe = start[i].softmax(-1), end[i].softmax(-1) |
| S = len(off) |
| band = torch.triu(torch.ones(S, S, device=ps.device)) - torch.triu(torch.ones(S, S, device=ps.device), diagonal=MAX_SPAN_TOK) |
| joint = ps[:, None] * pe[None, :] * band |
| best = int(joint.argmax()) |
| s, e = divmod(best, S) |
| cs, ce = trim(state_text, off[s][0], off[e][1]) |
| r = {"null": p_null > null_threshold, "p_null": p_null, "guess": state_text[cs:ce].strip(), |
| "confidence": joint.view(-1)[best].item() * (1 - p_null), "spans": [[cs, ce]]} |
| if ty == "number": |
| o = OPS[int(op[i].argmax())] |
| r["op"] = o |
| if o != "none": |
| tg = (opnd[i] > 0).tolist() |
| spans, j = [], 0 |
| while j < len(tg): |
| if tg[j]: |
| k = j |
| while k + 1 < len(tg) and tg[k + 1]: |
| k += 1 |
| spans.append((off[j][0], off[k][1])); j = k + 1 |
| else: |
| j += 1 |
| spans = merge(snap(state_text, a, b) for a, b in spans) |
| vals = [v for v in (num(state_text[a:b], en) for a, b in spans) if v is not None] |
| rev = bool(order[i].argmax()) |
| v = compute(o, vals, rev) |
| r.update({"value": v, "guess": ("%g" % v) if v is not None else r["guess"], "operands": [[a, b] for a, b in spans], |
| "reverse": rev}) |
| else: |
| cs, ce = snap(state_text, cs, ce) |
| r.update({"guess": state_text[cs:ce].strip(), "spans": [[cs, ce]]}) |
| r["value"] = num(r["guess"], en) |
| elif ty == "date": |
| r["value"] = date(r["guess"]) |
| res[name] = r |
| return res |
|
|
|
|
| |
| from transformers.masking_utils import create_masks_for_generate |
|
|
|
|
| def enable_noncausal(wrapper, sliding=False): |
| cfg = wrapper.config.get_text_config() if hasattr(wrapper.config, "get_text_config") else wrapper.config |
| assert getattr(cfg, "_attn_implementation", "sdpa") == "sdpa", "noncausal can sdpa" |
| W = getattr(cfg, "sliding_window", None) |
| if sliding and not W: |
| raise ValueError("noncausal sliding: model khong co sliding_window") |
|
|
| def pre(module, args, kwargs): |
| if args or kwargs.get("past_key_values") is not None or kwargs.get("use_cache"): |
| raise ValueError("noncausal: chi goi keyword, khong KV cache") |
| ids, emb = kwargs.get("input_ids"), kwargs.get("inputs_embeds") |
| B, L = (ids.shape if ids is not None else emb.shape[:2]) |
| dev = ids.device if ids is not None else emb.device |
| pad = kwargs.get("attention_mask") |
| if pad is None: |
| pad = torch.ones((B, L), dtype=torch.long, device=dev) |
| if isinstance(pad, dict): |
| return args, kwargs |
| like = torch.empty((B, L, 1), dtype=next(module.parameters()).dtype, device=dev) |
| pos = kwargs.get("position_ids") |
| if pos is None: |
| pos = torch.arange(L, device=dev).unsqueeze(0).expand(B, -1) |
| m = create_masks_for_generate(config=cfg, inputs_embeds=like, attention_mask=pad, past_key_values=None, position_ids=pos) |
| if not isinstance(m, dict): |
| m = {"full_attention": m} |
| m = dict(m) |
| m["full_attention"] = pad[:, None, None, :].bool() |
| if sliding: |
| i = torch.arange(L, device=dev) |
| |
| m["sliding_attention"] = (((i[:, None] - i[None, :]).abs() < W) & pad[:, None, None, :].bool()) | torch.eye(L, dtype=torch.bool, device=dev) |
| kwargs["attention_mask"] = m |
| return args, kwargs |
|
|
| wrapper.register_forward_pre_hook(pre, with_kwargs=True) |
|
|
|
|
| |
| import json as _json, os as _os, re as _re |
|
|
| _VI = _re.compile(r"[ăâđêôơưạảấầẩẫậắằẳẵặẹẻẽếềểễệỉịọỏốồổỗộớờởỡợụủứừửữựỳỵỷỹ]", _re.I) |
| SUPPORTED = {"text", "list", "number", "date"} |
|
|
|
|
| def detect_language(text): |
| """'vi' neu co dau tieng Viet; 'en' neu chu cai gan nhu toan ASCII; con lai 'default' (nguong null cao hon: |
| do tren XQuAD 7 ngon ngu khac, nguong 0.4 tra null nham 3-8%, 0.8 con 0.3-3%).""" |
| if len(_VI.findall(text)) > 3: |
| return "vi" |
| letters = [c for c in text if c.isalpha()] |
| if letters and sum(c.isascii() for c in letters) / len(letters) > 0.98: |
| return "en" |
| return "default" |
|
|
|
|
| class GemmaExtract: |
| """Typed field extraction in one forward pass, no text generation. |
| |
| from gemma_extract import GemmaExtract |
| m = GemmaExtract.from_pretrained("beyoru/gemmaExtract") |
| m.extract("Invoice No: INV-0042 ... Total: 1,200.00", |
| {"invoice_no": {"type": "text", "instructions": "Invoice number"}, |
| "total": {"type": "number", "instructions": "Total amount due"}}) |
| """ |
|
|
| def __init__(self, path, device="cuda", dtype=torch.bfloat16): |
| from transformers import AutoModelForCausalLM, AutoTokenizer |
| self.cfg = _json.load(open(_os.path.join(path, "extract_config.json"))) |
| self.tok = AutoTokenizer.from_pretrained(path) |
| top = AutoModelForCausalLM.from_pretrained(path, dtype=dtype, attn_implementation="sdpa").to(device) |
| self.enc = top.model |
| enable_noncausal(self.enc, sliding=False) |
| self.enc.eval() |
| h = self.cfg["head"] |
| self.head = S1Head(self.cfg["hidden"], width=h["width"], routing=h["routing"], layers=h["layers"], heads=h["heads"], ff=h["ff"]).to(device) |
| from safetensors.torch import load_file |
| self.head.load_state_dict(load_file(_os.path.join(path, "head.safetensors"))) |
| self.head.eval() |
| self.device, self.dtype = device, dtype |
|
|
| @classmethod |
| def from_pretrained(cls, repo_or_path, **kw): |
| if not _os.path.isdir(repo_or_path): |
| from huggingface_hub import snapshot_download |
| repo_or_path = snapshot_download(repo_or_path) |
| return cls(repo_or_path, **kw) |
|
|
| def _passes(self, fields): |
| names, m = list(fields), self.cfg["max_fields_per_pass"] |
| return [{n: fields[n] for n in names[i:i + m]} for i in range(0, len(names), m)] |
|
|
| @torch.no_grad() |
| def extract(self, state, fields, null_threshold=None, number_style=None): |
| """state: document text. fields: {name: {"type": "text|list|number|date", "instructions": "...", "null_threshold": float}}. |
| null_threshold: default per language (auto-detected). number_style: "en" (1,234.5) / "vi" (1.234,5) / None = auto. |
| Returns {"answers": {name: {value, guess, null, p_null, confidence, spans, ...}}, "usage": {...}}.""" |
| for k, v in fields.items(): |
| if v.get("type") not in SUPPORTED: |
| raise ValueError("field %s: type must be one of %s" % (k, sorted(SUPPORTED))) |
| lang = detect_language(state) |
| thr_all = self.cfg["null_threshold"].get(lang, self.cfg["null_threshold"]["default"]) if null_threshold is None else null_threshold |
| en = {"en": True, "vi": False}.get(number_style) |
| answers, usage = {}, {"state_tokens": 0, "passes": 0} |
| queue = self._passes({k: {"type": v["type"], "instructions": v.get("instructions") or k} for k, v in fields.items()}) |
| for part in queue: |
| try: |
| e = encode(self.tok, {"state": state, "fields": part}, self.cfg["max_length"]) |
| except ValueError: |
| if len(part) > 1: |
| queue += [{k: v} for k, v in part.items()] |
| continue |
| raise |
| with torch.autocast("cuda" if "cuda" in str(self.device) else "cpu", dtype=self.dtype): |
| h = self.enc(input_ids=torch.tensor([e["ids"]], device=self.device)).last_hidden_state[0] |
| out = self.head(h, e) |
| for k, r in decode(out, e, state, 0.5, en=en).items(): |
| t = fields[k].get("null_threshold", thr_all) |
| r["null"] = bool(r["p_null"] > t) or (isinstance(r["guess"], list) and not r["guess"]) |
| r["spans"] = [{"start": a, "end": b, "text": state[a:b]} for a, b in r.get("spans", [])] |
| r["value"] = None if r["null"] else r.get("value", r["guess"]) |
| answers[k] = r |
| usage["state_tokens"] = max(usage["state_tokens"], e["state"][1] - e["state"][0]) |
| usage["passes"] += 1 |
| usage["truncated"] = usage["state_tokens"] < len(self.tok(state, add_special_tokens=False)["input_ids"]) |
| return {"answers": answers, "usage": usage, "language": lang} |
|
|
