| """ |
| TypoFormer: Standalone Python Inference Runner |
| Supports both Autocomplete (Causal LM) and Grammar Correction (GEC). |
| |
| Requirements: |
| pip install onnxruntime tokenizers numpy |
| """ |
|
|
| import json |
| from pathlib import Path |
| import numpy as np |
| import onnxruntime as ort |
| from tokenizers import Tokenizer, decoders |
|
|
| class TypoFormer: |
| def __init__(self, model_dir: str | None = None): |
| p = Path(__file__).resolve().parent if model_dir is None else Path(model_dir) |
| self.tokenizer = Tokenizer.from_file(str(p / "tokenizer.json")) |
| if self.tokenizer.decoder is None: |
| self.tokenizer.decoder = decoders.ByteLevel() |
| |
| |
| opts = ort.SessionOptions() |
| opts.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL |
| opts.intra_op_num_threads = 4 |
| |
| self.lm_session = ort.InferenceSession(str(p / "lm_stage2_int8.onnx"), opts) |
| self.gec_session = ort.InferenceSession(str(p / "gec_stage2_int8.onnx"), opts) |
| |
| |
| maps_path = p / "edit_maps.json" |
| if maps_path.exists(): |
| maps = json.loads(maps_path.read_text(encoding="utf-8")) |
| if "replace" in maps: |
| self.id_to_replace = {int(lab): int(tid) for tid, lab in maps["replace"].items()} |
| self.id_to_append = {int(lab): int(tid) for tid, lab in maps.get("append", {}).items()} |
| else: |
| self.id_to_replace = {int(k): int(v) for k, v in maps.get("id_to_replace", {}).items()} |
| self.id_to_append = {int(k): int(v) for k, v in maps.get("id_to_append", {}).items()} |
| else: |
| self.id_to_replace = {} |
| self.id_to_append = {} |
|
|
| def _prepare_inputs(self, text: str, is_causal: bool = False): |
| encoding = self.tokenizer.encode(text) |
| if is_causal: |
| token_ids = [2] + encoding.ids |
| tokens = ["<bos>"] + encoding.tokens |
| else: |
| token_ids = [2] + encoding.ids + [3] |
| tokens = ["<bos>"] + encoding.tokens + ["<eos>"] |
|
|
| |
| char_matrix = [] |
| for t in tokens: |
| raw = t.encode("utf-8", errors="replace")[:16] |
| b = [byte if 0 < byte < 128 else 127 for byte in raw] |
| b += [0] * (16 - len(b)) |
| char_matrix.append(b) |
| |
| token_tensor = np.array([token_ids], dtype=np.int64) |
| char_tensor = np.array([char_matrix], dtype=np.int64) |
| return token_tensor, char_tensor, tokens |
|
|
| def autocomplete(self, text: str, top_k: int = 3): |
| token_tensor, char_tensor, _ = self._prepare_inputs(text, is_causal=True) |
| outputs = self.lm_session.run(None, { |
| "token_ids": token_tensor, |
| "char_ids": char_tensor |
| }) |
| logits = outputs[0][0, -1] |
| exp = np.exp(logits - np.max(logits)) |
| probs = exp / np.sum(exp) |
| |
| top_indices = np.argsort(probs)[-top_k:][::-1] |
| results = [] |
| for idx in top_indices: |
| token = self.tokenizer.decode([idx]) |
| results.append((token, float(probs[idx]))) |
| return results |
|
|
| def correct_grammar(self, text: str, min_edit_prob: float = 0.60, edit_margin: float = 0.08): |
| token_tensor, char_tensor, tokens = self._prepare_inputs(text, is_causal=False) |
| outputs = self.gec_session.run(None, { |
| "token_ids": token_tensor, |
| "char_ids": char_tensor |
| }) |
| logits = outputs[0][0] |
| |
| |
| exp = np.exp(logits - np.max(logits, axis=-1, keepdims=True)) |
| probs = exp / np.sum(exp, axis=-1, keepdims=True) |
| actions = np.argmax(probs, axis=-1) |
| best_probs = np.max(probs, axis=-1) |
| keep_probs = probs[:, 0] |
| |
| |
| for i in range(len(actions)): |
| if actions[i] != 0: |
| if (best_probs[i] < min_edit_prob) or (best_probs[i] - keep_probs[i] < edit_margin): |
| actions[i] = 0 |
| |
| elif 0 < i < len(tokens) - 1: |
| is_subtoken = not tokens[i].startswith("Ġ") and tokens[i] not in ("<bos>", "<eos>") |
| if is_subtoken and best_probs[i] < 0.92: |
| actions[i] = 0 |
| |
| token_ids = token_tensor[0].tolist() |
| new_ids = [] |
| n_replace = len(self.id_to_replace) if self.id_to_replace else 4096 |
|
|
| for tid, act in zip(token_ids, actions): |
| if tid in (2, 3): |
| continue |
| if act == 0: |
| new_ids.append(tid) |
| elif act == 1: |
| continue |
| elif act == 2: |
| decoded = self.tokenizer.decode([tid]).capitalize() |
| rec = self.tokenizer.encode(decoded).ids |
| new_ids.append(rec[0] if rec else tid) |
| elif act == 3: |
| decoded = self.tokenizer.decode([tid]).lower() |
| rec = self.tokenizer.encode(decoded).ids |
| new_ids.append(rec[0] if rec else tid) |
| elif 4 <= act < 4 + n_replace: |
| repl_id = self.id_to_replace.get(act, tid) |
| new_ids.append(repl_id) |
| elif act >= 4 + n_replace: |
| new_ids.append(tid) |
| extra_id = self.id_to_append.get(act) |
| if extra_id is not None: |
| new_ids.append(extra_id) |
| else: |
| new_ids.append(tid) |
| return self.tokenizer.decode([i for i in new_ids if i > 3]).strip() |
|
|
|
|
| if __name__ == "__main__": |
| model = TypoFormer() |
| |
| |
| print("=== TypoFormer Live Inference Test ===") |
| print("\n--- 1. Autocomplete (Next-Word Generation) ---") |
| prefix = "Could you please confirm if the proposed implementation timeline aligns with your" |
| print(f"Prompt: {repr(prefix)}") |
| print("Predictions (Token, Prob):", model.autocomplete(prefix, top_k=3)) |
| |
| |
| print("\n--- 2. Keystroke Grammar & Punctuation Fix ---") |
| dirty = "We dont expect any downtime during the deployment phase." |
| print(f"Input: {repr(dirty)}") |
| print(f"Corrected: {repr(model.correct_grammar(dirty))}") |
|
|
|
|
