File size: 5,749 Bytes
b03c74b 30d39c9 b03c74b 217c325 b03c74b 30d39c9 b03c74b 30d39c9 b03c74b 217c325 30d39c9 217c325 b03c74b 30d39c9 b03c74b 30d39c9 b03c74b 30d39c9 b03c74b 30d39c9 b03c74b 30d39c9 b03c74b 30d39c9 b03c74b 30d39c9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 | """Portable Unicode char_wb preprocessing and sparse ONNX inference."""
import hashlib
import json
import re
from collections import Counter
from pathlib import Path
import numpy as np
ROOT = Path(__file__).resolve().parent
def sha256(path):
with open(path, "rb") as stream:
return hashlib.file_digest(stream, "sha256").hexdigest()
def char_wb(text, ngram_range):
# Match sklearn's _char_wb_ngrams, including short-word handling.
text = re.sub(r"\s\s+", " ", text.lower())
for word in text.split():
word = " " + word + " "
for n in range(ngram_range[0], ngram_range[1] + 1):
offset = 0
yield word[offset:offset + n]
while offset + n < len(word):
offset += 1
yield word[offset:offset + n]
if offset == 0:
break
WORD_RE = re.compile(r"(?u)\b\w\w+\b")
def word_ngrams(text, ngram_range):
tokens = WORD_RE.findall(text.lower())
for n in range(ngram_range[0], ngram_range[1] + 1):
for i in range(len(tokens) - n + 1):
yield " ".join(tokens[i:i + n])
class LanguageIdentifier:
def __init__(self, model_dir=None, batch_size=256):
import onnxruntime as ort
if model_dir is not None:
self.model_dir = Path(model_dir)
else:
self.model_dir = ROOT / "model" if (ROOT / "model" / "META.json").is_file() else ROOT
config_path = self.model_dir / "config.json"
self.config = json.loads(config_path.read_text(encoding="utf-8")) if config_path.is_file() else None
self.meta = json.loads((self.model_dir / "META.json").read_text(encoding="utf-8"))
for name in ("model.onnx", "vectorizer.json"):
if name in self.meta.get("sha256", {}):
digest = self.meta["sha256"][name]
if sha256(self.model_dir / name) != digest:
raise ValueError(f"LID artifact checksum mismatch: {name}")
cfg = json.loads((self.model_dir / "vectorizer.json").read_text(encoding="utf-8"))
self.char_vocab = cfg.get("char_vocabulary", cfg.get("vocabulary", {}))
self.char_ngram_range = tuple(cfg.get("char_ngram_range", cfg.get("ngram_range", (2, 5))))
self.word_vocab = cfg.get("word_vocabulary", {})
self.word_ngram_range = tuple(cfg.get("word_ngram_range", (1, 2)))
self.classes = np.asarray(self.meta["classes"])
if self.config is not None:
expected = {"classes": self.classes.tolist(), "model_file": "model.onnx",
"vectorizer_file": "vectorizer.json", "metadata_file": "META.json"}
if any(self.config.get(k) != v for k, v in expected.items()):
raise ValueError("LID config does not match model artifacts")
if batch_size < 1:
raise ValueError("batch_size must be positive")
self.batch_size = batch_size
options = ort.SessionOptions()
options.intra_op_num_threads = 1
self.session = ort.InferenceSession(str(self.model_dir / "model.onnx"), options,
providers=["CPUExecutionProvider"])
def features(self, texts):
char_rows = []
word_rows = []
for text in texts:
if not isinstance(text, str):
raise TypeError("LID input must contain strings")
c_counts = Counter(self.char_vocab[g] for g in char_wb(text, self.char_ngram_range)
if g in self.char_vocab)
char_rows.append(c_counts)
if self.word_vocab:
w_counts = Counter(self.word_vocab[w] for w in word_ngrams(text, self.word_ngram_range)
if w in self.word_vocab)
word_rows.append(w_counts)
w_c = max(1, max(map(len, char_rows), default=0))
c_ids = np.zeros((len(char_rows), w_c), dtype=np.int64)
c_counts = np.zeros((len(char_rows), w_c), dtype=np.float32)
for i, row in enumerate(char_rows):
if row:
c_ids[i, :len(row)] = list(row.keys())
c_counts[i, :len(row)] = list(row.values())
if self.word_vocab:
w_w = max(1, max(map(len, word_rows), default=0))
w_ids = np.zeros((len(word_rows), w_w), dtype=np.int64)
w_counts = np.zeros((len(word_rows), w_w), dtype=np.float32)
for i, row in enumerate(word_rows):
if row:
w_ids[i, :len(row)] = list(row.keys())
w_counts[i, :len(row)] = list(row.values())
return {"char_ids": c_ids, "char_counts": c_counts, "word_ids": w_ids, "word_counts": w_counts}
return {"char_ids": c_ids, "char_counts": c_counts}
def predict_proba(self, texts):
texts = list(texts)
if not texts:
return np.empty((0, len(self.classes)), dtype=np.float32)
results = []
for start in range(0, len(texts), self.batch_size):
batch_texts = texts[start:start + self.batch_size]
results.append(self.session.run(["probabilities"], self.features(batch_texts))[0])
return np.concatenate(results) if results else np.empty((0, len(self.classes)), dtype=np.float32)
def predict(self, texts):
texts = list(texts)
if not texts:
return np.empty(0, dtype=str)
return self.classes[self.predict_proba(texts).argmax(axis=1)]
def describe(self):
return {"name": self.meta["id"], "version": self.meta["version"],
"sha256": self.meta["sha256"], "backend": "onnxruntime",
"decision_rule": "argmax probabilities; classes ordered as in model passport"}
|