| |
| """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, |
| } |
|
|
