File size: 5,184 Bytes
45958cf | 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 | #!/usr/bin/env python3
"""OPUS-MT (Marian) translation on AX650, sentencepiece-only tokenizer."""
import json
import time
from pathlib import Path
import numpy as np
import sentencepiece as spm
DEFAULT_PREFIX = ">>cmn_Hans<< "
class Tokenizer:
def __init__(self, model_dir):
d = Path(model_dir)
self.piece_to_id = json.loads((d / "vocab.json").read_text(encoding="utf-8"))
self.id_to_piece = {v: k for k, v in self.piece_to_id.items()}
self.spm_src = spm.SentencePieceProcessor(model_file=str(d / "source.spm"))
self.spm_tgt = spm.SentencePieceProcessor(model_file=str(d / "target.spm"))
cfg = json.loads((d / "model_meta.json").read_text(encoding="utf-8"))
self.eos = int(cfg["eos_token_id"])
self.pad = int(cfg["pad_token_id"])
self.start = int(cfg["decoder_start_token_id"])
self.src_len = int(cfg["src_len"])
self.tgt_len = int(cfg["tgt_len"])
self.unk = self.piece_to_id.get("<unk>", 1)
self.special = {self.eos, self.pad, self.unk}
def encode(self, text):
code = None
rest = text
if text.startswith(">>"):
end = text.find("<<")
if end != -1:
code = text[: end + 2]
rest = text[end + 2 :]
pieces = ([code] if code else []) + self.spm_src.encode(rest, out_type=str)
ids = [self.piece_to_id.get(p, self.unk) for p in pieces]
truncated = len(ids) > self.src_len - 1
ids = ids[: self.src_len - 1] + [self.eos]
return ids, truncated
def decode(self, ids):
pieces = []
for i in ids:
i = int(i)
if i in self.special:
continue
p = self.id_to_piece.get(i)
if p is None or (p.startswith(">>") and p.endswith("<<")):
continue
pieces.append(p)
return self.spm_src.decode_pieces(pieces).replace("\u2581", " ").strip()
class Translator:
def __init__(self, direction, model_root=None, max_new_tokens=None):
if direction not in ("zh-en", "en-zh"):
raise ValueError("direction must be zh-en or en-zh")
self.direction = direction
root = Path(model_root) if model_root else Path(__file__).resolve().parents[1] / "models"
self.model_dir = Path(root) / direction
self.tok = Tokenizer(self.model_dir)
import axengine as axe
self.enc = axe.InferenceSession(str(self.model_dir / "encoder.axmodel"))
self.dec = axe.InferenceSession(str(self.model_dir / "decoder.axmodel"))
self.max_new_tokens = int(max_new_tokens or self.tok.tgt_len)
self.warmup()
def _prepare(self, text):
if self.direction == "en-zh" and not text.startswith(">>"):
text = DEFAULT_PREFIX + text
ids, truncated = self.tok.encode(text)
n = len(ids)
x = np.full((1, self.tok.src_len), self.tok.pad, dtype=np.int32)
x[0, :n] = ids
m = np.zeros((1, self.tok.src_len), dtype=np.float32)
m[0, :n] = 1.0
return x, m, truncated
def warmup(self):
for _ in range(2):
self.translate("你好" if self.direction == "zh-en" else "Hello")
def translate(self, text, return_stats=False):
if not text.strip():
return ("", {"src_tokens": 0, "out_tokens": 0, "enc_ms": 0.0,
"dec_ms": 0.0, "dec_step_ms": 0.0, "truncated": False}) if return_stats else ""
x, m, truncated = self._prepare(text)
t0 = time.perf_counter()
hidden = self.enc.run(None, {"input_ids": x, "attention_mask": m})[0]
t_enc = time.perf_counter() - t0
out = [self.tok.start]
dec_times = []
for _ in range(self.max_new_tokens):
t = len(out)
dec = np.full((1, self.tok.tgt_len), self.tok.pad, dtype=np.int32)
dec[0, self.tok.tgt_len - t :] = out
dm = np.zeros((1, self.tok.tgt_len), dtype=np.int32)
dm[0, self.tok.tgt_len - t :] = 1
pos = np.zeros((1, self.tok.tgt_len), dtype=np.int32)
pos[0, self.tok.tgt_len - t :] = np.arange(t, dtype=np.int32)
t0 = time.perf_counter()
logits = self.dec.run(None, {
"decoder_input_ids": dec,
"decoder_attention_mask": dm,
"position_ids": pos,
"encoder_hidden_states": hidden,
"encoder_attention_mask": m,
})[0][0, 0].astype(np.float32)
dec_times.append(time.perf_counter() - t0)
logits[self.tok.pad] = -np.inf
nxt = int(np.argmax(logits))
out.append(nxt)
if nxt == self.tok.eos:
break
result = self.tok.decode(out[1:])
if not return_stats:
return result
return result, {
"src_tokens": int(m.sum()),
"out_tokens": len(out) - 1,
"enc_ms": t_enc * 1000,
"dec_ms": float(np.sum(dec_times)) * 1000,
"dec_step_ms": float(np.mean(dec_times)) * 1000 if dec_times else 0.0,
"truncated": truncated,
}
|