Token Classification
Transformers
ONNX
Safetensors
Polish
bert
polish
grammar-error-correction
gector
herbert
Instructions to use citan/plgec-herbert-large-v2 with libraries, inference providers, notebooks, and local apps. Follow these links to get started.
- Libraries
- Transformers
How to use citan/plgec-herbert-large-v2 with Transformers:
# Use a pipeline as a high-level helper from transformers import pipeline pipe = pipeline("token-classification", model="citan/plgec-herbert-large-v2")# pip install -U transformers accelerate # Load model directly from transformers import AutoTokenizer, AutoModelForTokenClassification tokenizer = AutoTokenizer.from_pretrained("citan/plgec-herbert-large-v2") model = AutoModelForTokenClassification.from_pretrained("citan/plgec-herbert-large-v2", device_map="auto") - Notebooks
- Google Colab
- Kaggle
Download inference_example.py from citan/plgec-herbert-large-v2: direct link, hf CLI and curl.
- Browser
- Download file 30.9 kB
-
https://huggingface.co/citan/plgec-herbert-large-v2/resolve/main/inference_example.py
- Command line
-
hf download hf://citan/plgec-herbert-large-v2/inference_example.py
-
curl -L -o inference_example.py https://huggingface.co/citan/plgec-herbert-large-v2/resolve/main/inference_example.py
30.9 kB
| #!/usr/bin/env python3 | |
| """Standalone inference example for citan/plgec-herbert-large-v2. | |
| Public dependencies only: onnxruntime, tokenizers, numpy (+ stdlib). | |
| Loads the fp32 ONNX model, tokenizer, label vocabulary and the | |
| consumer-side thresholds.yaml from a directory laid out like this | |
| repository (either the repository root or its onnx/ subdirectory), | |
| decodes sentences with the settings fixed at publication, applies the | |
| per-class filter, and prints one JSON object per sentence: | |
| {"sentence": ..., "corrected": ..., "edits": [{"token_idx", | |
| "char_span", "original", "replacement", "tag", "class", | |
| "confidence", "iteration"}, ...]} | |
| Decode semantics (fixed at publication; transcribed from the released | |
| checkpoint's threshold selection): keep_confidence_bias 0.25, | |
| min_error_probability 0.20, max_iters 1 β a single conservative pass. | |
| A word is edited only when 1 - min(1, P($KEEP) + 0.25) > 0.20 AND the | |
| top non-KEEP tag's probability > 0.20; confidence is the chosen tag's | |
| raw softmax probability under identity temperature (T = 1.0 β the v2 | |
| calibration block ships the explicit identity, no fitted constant). | |
| Edits are applied right-to-left on token indices; an insertion before | |
| the first token reports token_idx -1. | |
| Span caveat: an insertion after the FINAL token reports a character | |
| span whose end is one codepoint past the end of the text by design. | |
| $MERGE and other tags that do not render as single-token edits are | |
| skipped, exactly like the reference decoder. | |
| File map of this repository: vocab.json at the root is the GECToR | |
| LABEL vocabulary (6,757 tags, tag_to_id/id_to_tag) β onnx/vocab.json | |
| is the tokenizer's BPE word vocabulary; the two are distinguished by | |
| shape, this example resolves them by predicate. thresholds.yaml is the | |
| schema-v3 consumer-side filter (per-class bars, conf disabled by | |
| default with an optional floor, calibration block, `; -> ,` blocklist | |
| rule); thresholds.sha256 beside it carries the file's sha256 as | |
| published (sha256sum format) for local integrity checks. | |
| Phase 8 (plan 08-02): the thresholds file may additionally carry a | |
| schema-v2 calibration block (temperature T β applied as | |
| softmax(logits / T), argmax-preserving because T divides the logits | |
| PRE-softmax, at the same pre-decode point as the reference torch leg) | |
| and an exact-token blocklist (REPLACE-shaped `; -> ,`-style rules, | |
| dropped POST-decode in the same filter stage as the per-class bars β | |
| never re-decodes). Both apply only from schema_version 2; absent | |
| sections decode at the byte-identical v1 shape (T=1.0, no suppression). | |
| The calibration + blocklist code is DUPLICATED from | |
| publish/calibration.py β standalone-by-design (CONTEXT D-15); zero new | |
| dependencies. | |
| Phase 11 (plan 11-03, OQ-2): thresholds schema v3 adds an optional | |
| per-class confidence `floor` β an edit of a class whose block carries | |
| a floor survives only if its confidence clears BOTH the class bar AND | |
| the floor (ADDITIVE; an absent floor key is exact v2 behavior), | |
| dropped in the same post-decode filter stage as the bars and the | |
| blocklist (never re-decodes). tests/test_calibration_three_runtime_ | |
| parity.py is the drift guard across the three runtimes. | |
| No INT8 graph exists in this repository (no onnx_int8/ sibling): the | |
| v2 artifact is fp32 only. The v1 repository keeps its own INT8 | |
| evidence; it is not referenced here. | |
| """ | |
| from __future__ import annotations | |
| import json | |
| import math | |
| import re | |
| import sys | |
| from pathlib import Path | |
| import numpy as np | |
| import onnxruntime as ort | |
| from tokenizers import Tokenizer | |
| REPO_ID = "citan/plgec-herbert-large-v2" | |
| KEEP, DELETE = "$KEEP", "$DELETE" | |
| REPLACE_PREFIX, APPEND_PREFIX = "$REPLACE_", "$APPEND_" | |
| # Published decode triple (fixed at publication β the Stage-IV | |
| # threshold selection, combo 105: keep_confidence_bias 0.25 / | |
| # min_error_probability 0.20 / max_iters 1). | |
| KEEP_CONFIDENCE_BIAS = 0.25 | |
| MIN_ERROR_PROBABILITY = 0.20 | |
| MAX_ITERS = 1 | |
| DECODE_CHUNK = 32 # sentences per model call β padding stays reference-identical | |
| # Session options matching the reference decode (deterministic output). | |
| ORT_SPIN = {"session.intra_op.allow_spinning": "1", | |
| "session.intra_op.spin_duration_us": "1000", | |
| "session.intra_op.spin_backoff_max": "8"} | |
| # Canonical word tokenization: Polish/Latin word runs + single-char tokens. | |
| WORD_RUN = r"[A-Za-zΔ ΔΔΕΕΓ³ΕΕΊΕΌΔΔΔΕΕΓΕΕΉΕ»0-9]+" | |
| _TOKEN_RE = re.compile(WORD_RUN + r"|[^\s]") | |
| # Orthography swap pairs (attribution fallback when a tag carries no | |
| # class metadata in the shipped vocabulary). | |
| _TWO_CHAR_SWAPS = {"rz": "ΕΌ", "ch": "h"} | |
| _ONE_CHAR_SWAPS = {"ΕΌ": "rz", "Γ³": "u", "u": "Γ³", "Δ ": "Δ", "Δ": "Δ ", "h": "ch"} | |
| def tokenize_words(text): | |
| return _TOKEN_RE.findall(text) | |
| def token_char_spans(sentence): | |
| return [(m.start(), m.end()) for m in _TOKEN_RE.finditer(sentence)] | |
| def edit_char_span(sentence, edit, spans): | |
| """(start, end) codepoints an edit touches on the source text.""" | |
| idx = edit.get("token_idx") | |
| if idx is None or idx >= len(spans): | |
| return (0, 0) | |
| if idx < 0: # leading insertion | |
| return (0, spans[0][0]) if spans else (0, 0) | |
| start, end = spans[idx] | |
| if edit["tag"].startswith(APPEND_PREFIX): | |
| nxt = spans[idx + 1][0] if idx + 1 < len(spans) else len(sentence) | |
| return (end, max(end + 1, nxt)) if nxt <= end else (end, nxt) | |
| return (start, end) | |
| def _orth_swap_transform(token): | |
| result = token | |
| lowered = token.lower() | |
| sites, i = [], 0 | |
| while i < len(lowered): | |
| two = lowered[i:i + 2] | |
| if two in _TWO_CHAR_SWAPS: | |
| sites.append((i, 2, _TWO_CHAR_SWAPS[two])) | |
| i += 2 | |
| continue | |
| if lowered[i] in _ONE_CHAR_SWAPS: | |
| sites.append((i, 1, _ONE_CHAR_SWAPS[lowered[i]])) | |
| i += 1 | |
| for start, match_len, replacement in sorted(sites, reverse=True): | |
| span = result[start:start + match_len] | |
| repl = replacement | |
| if span[:1].isupper(): | |
| repl = repl[:1].upper() + repl[1:] | |
| result = result[:start] + repl + result[start + match_len:] | |
| return result | |
| def attribute_class(edit, tag_class_metadata): | |
| """Tag β class, metadata-first; orth-swap fallback for bare replaces.""" | |
| tag = edit.get("tag", "") | |
| if tag_class_metadata: | |
| classes = tag_class_metadata.get(tag) | |
| if classes: | |
| return classes[0] # pre-sorted at vocabulary build time | |
| if tag.startswith(REPLACE_PREFIX): | |
| original, replacement = edit["original"], edit["replacement"] | |
| if replacement != original and replacement == _orth_swap_transform( | |
| original): | |
| return "orth" | |
| return "infl" | |
| return "unknown" | |
| # --------------------------------------------------------------------------- | |
| # Consumer config (schema v2/v3 β plans 08-02 and 11-03; CONTEXT | |
| # D-12/D-13/D-15). | |
| # | |
| # DUPLICATED from publish/calibration.py β standalone-by-design (CONTEXT | |
| # D-15): the serve CLI imports the single definition site; the | |
| # standalone legs (this file and mcp/server.py) carry byte-identical | |
| # copies of this whole block (the parity tests and | |
| # tests/test_thresholds_yaml_schema_v2.py are the drift guards). No | |
| # pyyaml here β the public-deps contract keeps the parser hand-rolled. | |
| # --------------------------------------------------------------------------- | |
| SCHEMA_VERSION_CURRENT = 3 # Phase 8 (D-13): v2 added schema_version + | |
| # calibration + blocklist top-level sections; Phase 11 (plan 11-03, | |
| # OQ-2) bumps to 3 β the optional per-class confidence `floor` key | |
| # (ADDITIVE: an edit of the class must clear BOTH the class bar AND | |
| # the floor; an absent floor key is exact v2 behavior). | |
| def _yaml_unquote(value): | |
| """Strip one matching surrounding quote pair from a yaml scalar.""" | |
| if len(value) >= 2 and value[0] == value[-1] and value[0] in ("'", '"'): | |
| return value[1:-1] | |
| return value | |
| def _yaml_float(value, path, key): | |
| """Scalar -> float, rejecting bool spellings (the L-06 trap) and any | |
| non-numeric text with the fail-closed ValueError style.""" | |
| if value in ("true", "false"): | |
| raise ValueError(f"{path}: {key} {value!r} must be a number") | |
| try: | |
| return float(value) | |
| except ValueError: | |
| raise ValueError(f"{path}: {key} {value!r} must be a number") from None | |
| def load_thresholds(path): | |
| """Hand-rolled YAML reader for the shipped thresholds file (no pyyaml | |
| β the public-deps contract). | |
| Parses the full schema v2 plus the v3 per-class floor key: the | |
| per-class mapping {class: {min_error_probability: float, | |
| enabled: bool, floor: float in (0, 1] OPTIONAL β Phase 11, plan | |
| 11-03: the ADDITIVE confidence floor, an absent key is exact v2 | |
| behavior}} PLUS the top-level | |
| schema_version int scalar (a FUTURE version > 3 is refused loudly β | |
| an old runtime must never silently drop new semantics), the | |
| calibration block {temperature: float > 0, t_int8: null-or-float | |
| > 0} and the blocklist list-of-flat-mappings (string class/from/to/ | |
| reason, bool enabled, optional int >= 0 min_tokens_before guard). | |
| Returns the FULL document: per-class keys plus the special keys | |
| "schema_version", "calibration", "blocklist" whenever present in | |
| the file (mirroring the serve CLI's loader; the Corrector gates | |
| their APPLICATION on schema_version >= 2 β a v1 file's stray | |
| sections parse but never apply). Any other top-level scalar stays a | |
| hard ValueError; malformed v2 values fail closed, never a repaired | |
| default. | |
| """ | |
| root, current = {}, None | |
| section = "class" # active top-level section: class|calibration|blocklist | |
| schema_version, calibration, blocklist, rule = None, None, None, None | |
| for raw in Path(path).read_text(encoding="utf-8").splitlines(): | |
| line = raw.split("#", 1)[0].rstrip() | |
| if not line.strip(): | |
| continue | |
| indent = len(line) - len(line.lstrip()) | |
| stripped = line.strip() | |
| list_item = stripped.startswith("-") | |
| if list_item: | |
| stripped = stripped[1:].strip() | |
| key, _, value = stripped.partition(":") | |
| key, value = key.strip(), value.strip() | |
| if list_item: | |
| if (indent == 0 or section != "blocklist" | |
| or blocklist is None): | |
| raise ValueError(f"{path}: unexpected list item {raw!r}") | |
| rule = {} | |
| blocklist.append(rule) | |
| if not key: | |
| continue # bare '-'; payload on the following lines | |
| elif indent == 0: | |
| rule = None | |
| if key == "schema_version": | |
| if not value.isdigit(): | |
| raise ValueError(f"{path}: schema_version {value!r} " | |
| "is not an integer") | |
| if int(value) > SCHEMA_VERSION_CURRENT: | |
| raise ValueError( | |
| f"{path}: schema_version={int(value)} > current " | |
| f"runtime's {SCHEMA_VERSION_CURRENT}; upgrade " | |
| "the runtime") | |
| schema_version = int(value) | |
| continue | |
| if value: | |
| raise ValueError(f"{path}: top-level scalar {key!r} β " | |
| "expected class blocks") | |
| if key == "calibration": | |
| section, calibration, current = "calibration", {}, None | |
| elif key == "blocklist": | |
| section, blocklist, current = "blocklist", [], None | |
| else: | |
| section, current = "class", root.setdefault(key, {}) | |
| continue | |
| # indented content (or the payload of a "- key: value" row) | |
| if section == "class": | |
| if current is None or value == "": | |
| raise ValueError(f"{path}: unexpected line {raw!r}") | |
| if key == "floor": | |
| # Phase 11 (plan 11-03, OQ-2 additive): optional | |
| # per-class confidence floor, strictly in (0, 1]; | |
| # _yaml_float rejects the true/false bool spellings | |
| # (the L-06 trap) and any non-numeric text. | |
| floor = _yaml_float(value, path, "class floor") | |
| if not 0.0 < floor <= 1.0: | |
| raise ValueError(f"{path}: class floor {value!r} " | |
| "must be a number in (0, 1]") | |
| current["floor"] = floor | |
| elif value in ("true", "false"): | |
| current[key] = value == "true" | |
| else: | |
| current[key] = float(value) | |
| elif section == "calibration": | |
| if calibration is None: # unreachable: the section switch sets it | |
| raise ValueError(f"{path}: unexpected line {raw!r}") | |
| if key == "temperature": | |
| temperature = _yaml_float(value, path, | |
| "calibration temperature") | |
| if not math.isfinite(temperature) \ | |
| or not temperature > 0: | |
| # CR-01: 'not x > 0' rejects nan but NOT inf | |
| raise ValueError(f"{path}: calibration temperature " | |
| f"{value!r} must be a finite " | |
| "number > 0") | |
| calibration["temperature"] = temperature | |
| elif key == "t_int8": | |
| if value in ("", "null", "~"): | |
| calibration["t_int8"] = None # reserved slot (v2 ships null) | |
| else: | |
| t_int8 = _yaml_float(value, path, "calibration t_int8") | |
| if not math.isfinite(t_int8) or not t_int8 > 0: | |
| raise ValueError(f"{path}: calibration t_int8 " | |
| f"{value!r} must be a finite " | |
| "number > 0 or null") | |
| calibration["t_int8"] = t_int8 | |
| else: | |
| raise ValueError(f"{path}: unexpected calibration key " | |
| f"{key!r}") | |
| else: # section == "blocklist" | |
| if rule is None: | |
| raise ValueError(f"{path}: unexpected line {raw!r}") | |
| if key in ("class", "from", "to", "reason"): | |
| rule[key] = _yaml_unquote(value) | |
| elif key == "enabled": | |
| if value not in ("true", "false"): | |
| raise ValueError(f"{path}: blocklist rule enabled " | |
| f"{value!r} must be a bool") | |
| rule["enabled"] = value == "true" | |
| elif key == "min_tokens_before": | |
| if not value.isdigit(): # digits-only: rejects -1/true/text | |
| raise ValueError(f"{path}: blocklist rule " | |
| "min_tokens_before " | |
| f"{value!r} must be an integer >= 0") | |
| rule["min_tokens_before"] = int(value) | |
| else: | |
| raise ValueError(f"{path}: unexpected blocklist rule key " | |
| f"{key!r}") | |
| if calibration is not None and "temperature" not in calibration: | |
| raise ValueError(f"{path}: calibration block missing 'temperature'") | |
| for rule in blocklist or []: | |
| for required in ("class", "from", "to"): | |
| if required not in rule: | |
| raise ValueError(f"{path}: blocklist rule missing " | |
| f"{required!r}") | |
| document = dict(root) | |
| if schema_version is not None: | |
| document["schema_version"] = schema_version | |
| if calibration is not None: | |
| document["calibration"] = calibration | |
| if blocklist is not None: | |
| document["blocklist"] = blocklist | |
| return document | |
| def blocklist_match(edit: dict, rules) -> bool: | |
| """True if the edit is BLOCKED by an enabled rule (drop it). | |
| DUPLICATED from publish/calibration.py β standalone-by-design | |
| (CONTEXT D-15); keep byte-identical across the two standalone legs. | |
| Match semantics (CONTEXT D-12 + the match-discretion resolution): | |
| $REPLACE_-shaped tags only, exact string equality of | |
| (original, replacement) == (from, to), and the optional per-rule | |
| integer guard min_tokens_before β the rule fires only when the | |
| edit's token_idx >= the guard (absent guard = 0 = unconditional). | |
| $APPEND_/$DELETE never match (RESEARCH Pitfall 4); ';'/';;'/','- | |
| adjacent near-misses never match (exact tokens only). | |
| """ | |
| for rule in rules or []: | |
| if not isinstance(rule, dict): | |
| continue | |
| if not rule.get("enabled", True): | |
| continue # per-rule disable path (D-14, conf.enabled precedent) | |
| # exact-token REPLACE match β the only tag shape allowed | |
| if not str(edit.get("tag", "")).startswith("$REPLACE_"): | |
| continue | |
| if (edit.get("original") == rule.get("from") | |
| and edit.get("replacement") == rule.get("to")): | |
| guard = rule.get("min_tokens_before", 0) or 0 | |
| token_idx = edit.get("token_idx") | |
| if not isinstance(token_idx, int) or isinstance(token_idx, bool): | |
| token_idx = 0 # defensive: real REPLACE edits carry int idx | |
| if token_idx >= int(guard): | |
| return True | |
| return False | |
| def _first_json(candidates, predicate, what): | |
| for cand in candidates: | |
| if cand.is_file(): | |
| data = json.loads(cand.read_text(encoding="utf-8")) | |
| if predicate(data): | |
| return data | |
| raise SystemExit(f"FATAL: {what} not found under: " | |
| + ", ".join(str(c) for c in candidates)) | |
| class Corrector: | |
| """ONNX GECToR decoder over the published v2 checkpoint layout.""" | |
| def __init__(self, model_dir, thresholds=None): | |
| model_dir = Path(model_dir) | |
| # v2 ships fp32 only β no INT8 graph exists in this repository | |
| # (D-09); the candidates are the two fp32 layouts (onnx/ subdir | |
| # or a flat download). | |
| candidates = (model_dir / "onnx" / "model.onnx", | |
| model_dir / "model.onnx") | |
| onnx_path = next((p for p in candidates if p.is_file()), None) | |
| tok_path = next((p for p in (model_dir / "tokenizer.json", | |
| model_dir / "onnx" / "tokenizer.json") | |
| if p.is_file()), None) | |
| if onnx_path is None or tok_path is None: | |
| raise SystemExit( | |
| f"FATAL: {model_dir} is not a plgec-herbert-large-v2 " | |
| "download (need onnx/model.onnx and tokenizer.json)") | |
| vocab = _first_json( | |
| [model_dir / "vocab.json", model_dir / "onnx" / "vocab.json"], | |
| lambda d: isinstance(d, dict) and "tag_to_id" in d | |
| and "id_to_tag" in d, | |
| "GECToR label vocabulary (vocab.json with tag_to_id/id_to_tag)") | |
| self.tag_to_id = vocab["tag_to_id"] | |
| self.id_to_tag = {int(k): v for k, v in vocab["id_to_tag"].items()} | |
| self.keep_id = self.tag_to_id[KEEP] | |
| self.tag_class_metadata = vocab.get("tag_class_metadata") or None | |
| options = ort.SessionOptions() | |
| options.intra_op_num_threads = 8 | |
| options.execution_mode = ort.ExecutionMode.ORT_SEQUENTIAL | |
| options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL | |
| for key, value in ORT_SPIN.items(): | |
| options.add_session_config_entry(key, value) | |
| self.session = ort.InferenceSession(str(onnx_path), options, | |
| providers=["CPUExecutionProvider"]) | |
| self.tokenizer = Tokenizer.from_file(str(tok_path)) | |
| pad = self.tokenizer.token_to_id("<pad>") | |
| self.pad_id = 0 if pad is None else int(pad) | |
| self.thresholds = thresholds | |
| # -- consumer-side per-class filter (applied AFTER decode) ---------- | |
| def _drops(self, edit_class, confidence): | |
| if not self.thresholds: | |
| return False | |
| block = (self.thresholds.get(edit_class) | |
| or self.thresholds.get("default")) | |
| if not block: | |
| return False | |
| if not block.get("enabled", True): | |
| return True | |
| if confidence < float(block.get("min_error_probability", 0.0)): | |
| return True | |
| # Phase 11 (plan 11-03, OQ-2 additive): the optional per-class | |
| # confidence `floor` (schema v3) β the edit must ALSO clear it | |
| # (inclusive at the boundary), in the same post-decode filter | |
| # stage; never re-decodes; an absent key is exact v2 behavior. | |
| # Duplicated from the serve CLI's _threshold_drops semantics | |
| # (standalone-by-design, D-15) β the three-runtime parity test | |
| # is the drift guard. | |
| floor = block.get("floor") | |
| return floor is not None and confidence < float(floor) | |
| # -- Phase 8 v2 consumer config (plan 08-02) ------------------------ | |
| # DUPLICATED from publish/calibration.py semantics β standalone-by- | |
| # design (CONTEXT D-15); read dynamically per correct() call so | |
| # post-construction threshold mutation behaves like the per-class | |
| # bars. Both keys apply only from schema_version 2 (a v1 file's | |
| # stray sections parse but never apply β backward compat). | |
| def _calibration_temperature(self): | |
| """Decode temperature T: softmax(logits / T), argmax-preserving | |
| because T divides the LOGITS pre-softmax (CONTEXT D-01, never a | |
| post-softmax rescale). Absent calibration block / v1 file -> | |
| T=1.0 exact identity. The shipped v2 file carries the explicit | |
| identity temperature (T = 1.0).""" | |
| if not self.thresholds: | |
| return 1.0 | |
| if self.thresholds.get("schema_version", 1) < 2: | |
| return 1.0 | |
| block = self.thresholds.get("calibration") or {} | |
| return float(block.get("temperature", 1.0)) | |
| def _active_blocklist_rules(self): | |
| """Blocklist rules applying to this corrector (schema v2+ only).""" | |
| if not self.thresholds: | |
| return [] | |
| if self.thresholds.get("schema_version", 1) < 2: | |
| return [] | |
| rules = self.thresholds.get("blocklist") | |
| return list(rules) if isinstance(rules, list) else [] | |
| def _extract_edits(self, probs_row, positions, word_ids, words): | |
| edits = [] | |
| for pos in positions: | |
| position_probs = probs_row[pos] | |
| keep_conf = min(1.0, | |
| float(position_probs[self.keep_id]) | |
| + KEEP_CONFIDENCE_BIAS) | |
| if 1.0 - keep_conf <= MIN_ERROR_PROBABILITY: | |
| continue | |
| masked = position_probs.copy() | |
| masked[self.keep_id] = -1.0 | |
| top_id = int(np.argmax(masked)) | |
| top_p = float(position_probs[top_id]) | |
| if top_p <= MIN_ERROR_PROBABILITY: | |
| continue | |
| tag = self.id_to_tag[top_id] | |
| word_id = word_ids[pos] if pos < len(word_ids) else None | |
| if tag == DELETE: | |
| if word_id is None: | |
| continue | |
| edits.append({"token_idx": word_id, "original": words[word_id], | |
| "replacement": "", "tag": tag, | |
| "confidence": top_p}) | |
| elif tag.startswith(REPLACE_PREFIX): | |
| if word_id is None: | |
| continue | |
| edits.append({"token_idx": word_id, "original": words[word_id], | |
| "replacement": tag[len(REPLACE_PREFIX):], | |
| "tag": tag, "confidence": top_p}) | |
| elif tag.startswith(APPEND_PREFIX): | |
| appended = tag[len(APPEND_PREFIX):] | |
| edits.append({"token_idx": word_id if word_id is not None else -1, | |
| "original": "" if word_id is None else words[word_id], | |
| "replacement": appended, "tag": tag, | |
| "confidence": top_p}) | |
| return edits | |
| def _apply_edits(words, edits): | |
| """Apply one pass's edits; descending token order keeps indices valid.""" | |
| updated = list(words) | |
| for edit in sorted((e for e in edits if e["token_idx"] >= 0), | |
| key=lambda e: e["token_idx"], reverse=True): | |
| idx = edit["token_idx"] | |
| if edit["tag"] == DELETE: | |
| del updated[idx] | |
| elif edit["tag"].startswith(REPLACE_PREFIX): | |
| updated[idx] = edit["replacement"] | |
| elif edit["tag"].startswith(APPEND_PREFIX): | |
| updated.insert(idx + 1, edit["replacement"]) | |
| for edit in (e for e in edits if e["token_idx"] == -1): | |
| if edit["tag"].startswith(APPEND_PREFIX): | |
| updated.insert(0, edit["replacement"]) | |
| return updated | |
| def correct(corrector, texts): | |
| """[{sentence, corrected, edits}] for texts, reference-batched.""" | |
| originals = list(texts) | |
| current = list(texts) | |
| all_edits = [[] for _ in current] | |
| for _iteration in range(1, MAX_ITERS + 1): | |
| fired_any = False | |
| for c in range(0, len(current), DECODE_CHUNK): | |
| chunk = current[c:c + DECODE_CHUNK] | |
| words_per_sentence = [tokenize_words(s) for s in chunk] | |
| encodings = [corrector.tokenizer.encode(words, is_pretokenized=True) | |
| for words in words_per_sentence] | |
| id_rows = [enc.ids for enc in encodings] | |
| wid_rows = [enc.word_ids for enc in encodings] | |
| longest = max(len(ids) for ids in id_rows) | |
| input_ids = np.array( | |
| [ids + [corrector.pad_id] * (longest - len(ids)) | |
| for ids in id_rows], dtype=np.int64) | |
| attention = np.array( | |
| [[1] * len(ids) + [0] * (longest - len(ids)) | |
| for ids in id_rows], dtype=np.int64) | |
| logits = corrector.session.run(None, { | |
| "input_ids": input_ids, | |
| "attention_mask": attention, | |
| "token_type_ids": np.zeros_like(input_ids)})[0] | |
| # Softmax in float64 (accurate, deterministic across numpy | |
| # builds), cast to float32 to mirror the reference decoder's | |
| # float32 probability tensor: gate comparisons and the | |
| # emitted confidence both read the float32 values. Kernel | |
| # level difference vs the torch reference is bounded and | |
| # test-pinned (tests/test_public_example_parity.py). | |
| # Phase 8 (plan 08-02, CALIB-01): the consumer-side | |
| # temperature T enters HERE, pre-softmax on the shifted | |
| # logits β softmax((l - max) / T) == softmax(l / T), the | |
| # argmax-preserving form (D-01) at the same pre-decode | |
| # point as the torch leg; T = 1.0 divides exactly, so the | |
| # identity path stays byte-identical (the v2 file ships | |
| # the explicit identity). | |
| temperature = corrector._calibration_temperature() | |
| shifted = (logits - logits.max(axis=-1, keepdims=True) | |
| ).astype(np.float64) | |
| shifted /= temperature | |
| exp = np.exp(shifted) | |
| probs = (exp / exp.sum(axis=-1, keepdims=True) | |
| ).astype(np.float32) | |
| for row, (words, word_ids) in enumerate( | |
| zip(words_per_sentence, wid_rows)): | |
| seen, positions = set(), [] | |
| for pos, word_id in enumerate(word_ids): | |
| if word_id is not None and word_id not in seen: | |
| seen.add(word_id) | |
| positions.append(pos) | |
| edits = corrector._extract_edits( | |
| probs[row], [0] + positions, word_ids, words) | |
| if edits: | |
| fired_any = True | |
| for edit in edits: | |
| edit["iteration"] = _iteration | |
| all_edits[c + row].extend(edits) | |
| words_per_sentence[row] = _apply_edits(words, edits) | |
| current[c + row] = " ".join(words_per_sentence[row]) | |
| if not fired_any: | |
| break | |
| results = [] | |
| for source_text, final, edits in zip(originals, current, all_edits): | |
| words = tokenize_words(source_text) | |
| spans = token_char_spans(source_text) | |
| surviving = [] | |
| blocklist_rules = corrector._active_blocklist_rules() | |
| for edit in edits: | |
| edit_class = attribute_class(edit, corrector.tag_class_metadata) | |
| if corrector._drops(edit_class, edit["confidence"]): | |
| continue | |
| # Phase 8 (plan 08-02, SERVE-03): exact-token blocklist drop | |
| # in the same post-decode filter stage as the per-class bars | |
| # β never re-decodes, never reorders. | |
| if blocklist_match(edit, blocklist_rules): | |
| continue | |
| surviving.append({ | |
| "token_idx": edit["token_idx"], | |
| "char_span": list(edit_char_span(source_text, edit, spans)), | |
| "original": edit["original"], | |
| "replacement": edit["replacement"], | |
| "tag": edit["tag"], | |
| "class": edit_class, | |
| "confidence": edit["confidence"], | |
| "iteration": edit["iteration"], | |
| }) | |
| corrected = (" ".join(_apply_edits(words, surviving)) | |
| if surviving else source_text) | |
| results.append({"sentence": source_text, "corrected": corrected, | |
| "edits": surviving}) | |
| return results | |
| def main(argv): | |
| import argparse | |
| parser = argparse.ArgumentParser( | |
| description="Standalone inference for " + REPO_ID | |
| + " (public deps: onnxruntime, tokenizers, numpy).") | |
| parser.add_argument("sentences", nargs="+", | |
| help="sentence(s) to correct (quoted)") | |
| parser.add_argument("--model-dir", default=str(Path(__file__).parent), | |
| help="model download directory (default: this " | |
| "script's directory)") | |
| parser.add_argument("--thresholds", default=None, | |
| help="per-class thresholds YAML (default: " | |
| "thresholds.yaml in the model dir, if present; " | |
| "'none' disables consumer-side filtering)") | |
| args = parser.parse_args(argv) | |
| thresholds = args.thresholds | |
| if thresholds == "none": | |
| thresholds = None | |
| elif thresholds is None: | |
| candidate = Path(args.model_dir) / "thresholds.yaml" | |
| thresholds = str(candidate) if candidate.is_file() else None | |
| corrector = Corrector(args.model_dir, | |
| load_thresholds(thresholds) | |
| if thresholds else None) | |
| for record in correct(corrector, args.sentences): | |
| print(json.dumps(record, ensure_ascii=False)) | |
| if __name__ == "__main__": | |
| main(sys.argv[1:]) | |