764 lines
25 KiB
Python
764 lines
25 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
ONNX输入法模型推理脚本
|
||
|
||
使用ONNX Runtime进行推理，测量每个阶段的执行时长
|
||
|
||
使用方法:
|
||
python onnx_inference.py --context-encoder exported_models/context_encoder.onnx --decoder exported_models/decoder.onnx
|
||
|
||
交互模式: 分步询问输入
|
||
1. 上下文提示: 模型不掌握的专有词汇、姓名等（可为空）
|
||
2. 光标前文本: 光标前的连续文本
|
||
3. 光标后文本: 光标后的连续文本
|
||
4. 拼音: 当前输入的拼音
|
||
5. 槽位历史: 用户已确认的输入历史（如输入shanghai已确认"上"）
|
||
"""
|
||
|
||
import argparse
|
||
import os
|
||
import sys
|
||
import time
|
||
from pathlib import Path
|
||
from typing import List, Optional, Tuple
|
||
|
||
import numpy as np
|
||
import onnxruntime as ort
|
||
import torch
|
||
import torch.nn.functional as F
|
||
from modelscope import AutoTokenizer
|
||
|
||
from src.model.dataset import text_to_pinyin_ids
|
||
from src.model.query import QueryEngine
|
||
|
||
|
||
class ONNXInference:
|
||
"""ONNX输入法模型推理器"""
|
||
|
||
def __init__(
|
||
self,
|
||
context_encoder_path: str,
|
||
decoder_path: str,
|
||
vocab_size: int = 10019,
|
||
device: str = "cpu",
|
||
use_beam_search: bool = False,
|
||
beam_size: int = 5,
|
||
):
|
||
self.vocab_size = vocab_size
|
||
self.device = device
|
||
self.use_beam_search = use_beam_search
|
||
self.beam_size = beam_size
|
||
|
||
# 加载组件
|
||
print(f"正在加载上下文编码器: {context_encoder_path}")
|
||
load_start = time.perf_counter()
|
||
self.load_context_encoder(context_encoder_path)
|
||
self.context_encoder_load_time = (time.perf_counter() - load_start) * 1000
|
||
print(f" ✅ 上下文编码器加载完成 ({self.context_encoder_load_time:.2f}ms)")
|
||
|
||
print(f"正在加载解码器: {decoder_path}")
|
||
load_start = time.perf_counter()
|
||
self.load_decoder(decoder_path)
|
||
self.decoder_load_time = (time.perf_counter() - load_start) * 1000
|
||
print(f" ✅ 解码器加载完成 ({self.decoder_load_time:.2f}ms)")
|
||
|
||
# 加载tokenizer
|
||
print("正在加载tokenizer...")
|
||
load_start = time.perf_counter()
|
||
self.load_tokenizer()
|
||
self.tokenizer_load_time = (time.perf_counter() - load_start) * 1000
|
||
print(f" ✅ Tokenizer加载完成 ({self.tokenizer_load_time:.2f}ms)")
|
||
|
||
# 加载查询引擎
|
||
print("正在加载查询引擎...")
|
||
load_start = time.perf_counter()
|
||
self.load_query_engine()
|
||
self.query_engine_load_time = (time.perf_counter() - load_start) * 1000
|
||
print(f" ✅ 查询引擎加载完成 ({self.query_engine_load_time:.2f}ms)")
|
||
|
||
total_load_time = (
|
||
self.context_encoder_load_time
|
||
+ self.decoder_load_time
|
||
+ self.tokenizer_load_time
|
||
+ self.query_engine_load_time
|
||
)
|
||
print(f"\n✅ 推理器初始化完成 (设备: {device})")
|
||
print(f" 总加载时间: {total_load_time:.2f}ms")
|
||
|
||
# 尝试启用readline
|
||
try:
|
||
import readline
|
||
|
||
readline.set_completer_delims(" \t\n`~!@#$%^&*()-=+[{]}\\|;:'\",<>/?")
|
||
except ImportError:
|
||
pass
|
||
|
||
def load_context_encoder(self, model_path: str):
|
||
"""加载上下文编码器ONNX模型"""
|
||
providers = (
|
||
["CUDAExecutionProvider", "CPUExecutionProvider"]
|
||
if self.device == "cuda"
|
||
else ["CPUExecutionProvider"]
|
||
)
|
||
self.context_encoder_session = ort.InferenceSession(
|
||
model_path, providers=providers
|
||
)
|
||
|
||
self.context_input_names = [
|
||
inp.name for inp in self.context_encoder_session.get_inputs()
|
||
]
|
||
self.context_output_names = [
|
||
out.name for out in self.context_encoder_session.get_outputs()
|
||
]
|
||
|
||
def load_decoder(self, model_path: str):
|
||
"""加载解码器ONNX模型"""
|
||
providers = (
|
||
["CUDAExecutionProvider", "CPUExecutionProvider"]
|
||
if self.device == "cuda"
|
||
else ["CPUExecutionProvider"]
|
||
)
|
||
self.decoder_session = ort.InferenceSession(model_path, providers=providers)
|
||
|
||
self.decoder_input_names = [
|
||
inp.name for inp in self.decoder_session.get_inputs()
|
||
]
|
||
self.decoder_output_names = [
|
||
out.name for out in self.decoder_session.get_outputs()
|
||
]
|
||
|
||
def load_tokenizer(self):
|
||
"""加载tokenizer"""
|
||
try:
|
||
tokenizer_path = (
|
||
Path(__file__).parent / "src" / "model" / "assets" / "tokenizer"
|
||
)
|
||
self.tokenizer = AutoTokenizer.from_pretrained(str(tokenizer_path))
|
||
except Exception:
|
||
print(" ⚠️ 无法加载自定义tokenizer，使用bert-base-chinese")
|
||
self.tokenizer = AutoTokenizer.from_pretrained("bert-base-chinese")
|
||
|
||
def load_query_engine(self):
|
||
"""加载查询引擎"""
|
||
try:
|
||
self.query_engine = QueryEngine()
|
||
stats_path = (
|
||
Path(__file__).parent
|
||
/ "src"
|
||
/ "model"
|
||
/ "assets"
|
||
/ "pinyin_char_statistics.json"
|
||
)
|
||
if stats_path.exists():
|
||
self.query_engine.load(stats_path)
|
||
except Exception:
|
||
self.query_engine = None
|
||
|
||
def char_to_id(self, char: str, pinyin: Optional[str] = None) -> int:
|
||
"""将汉字转换为ID"""
|
||
if char == "//":
|
||
return 0
|
||
|
||
if self.query_engine is None:
|
||
return ord(char) if len(char) == 1 else 0
|
||
|
||
try:
|
||
if pinyin is not None:
|
||
info = self.query_engine.get_char_info_by_char_pinyin(char, pinyin)
|
||
if info:
|
||
return info.id
|
||
|
||
results = self.query_engine.query_by_char(char, limit=1)
|
||
if results:
|
||
return results[0][0]
|
||
return 0
|
||
except:
|
||
return 0
|
||
|
||
def id_to_char(self, id: int) -> str:
|
||
"""将ID转换为汉字"""
|
||
if id == 0:
|
||
return "//"
|
||
|
||
if self.query_engine is None:
|
||
return chr(id) if id < 0x110000 else f"[ID:{id}]"
|
||
|
||
try:
|
||
info = self.query_engine.query_by_id(id)
|
||
return info.char if info else f"[ID:{id}]"
|
||
except:
|
||
return f"[ID:{id}]"
|
||
|
||
def _clean_pinyin_input(self, pinyin: str) -> str:
|
||
"""清理拼音输入字符串"""
|
||
if not pinyin:
|
||
return ""
|
||
|
||
result = []
|
||
for c in pinyin:
|
||
is_valid = ("a" <= c <= "z") or ("A" <= c <= "Z") or c in ["`", "'", "-"]
|
||
if is_valid:
|
||
result.append(c.lower())
|
||
elif c == " ":
|
||
continue
|
||
elif c in ["\b", "\x7f", "\x08"]:
|
||
if result:
|
||
result.pop()
|
||
elif c == "\x1b":
|
||
result.clear()
|
||
return "".join(result)
|
||
|
||
def _safe_input(self, prompt: str, default: str = "") -> str:
|
||
"""安全的输入函数"""
|
||
try:
|
||
full_prompt = f"{prompt} [{default}]: " if default else f"{prompt}: "
|
||
result = input(full_prompt)
|
||
if not result and default:
|
||
return default
|
||
return result.strip()
|
||
except (EOFError, KeyboardInterrupt):
|
||
print()
|
||
return ""
|
||
except Exception as e:
|
||
print(f"\n⚠️ 输入错误: {e}")
|
||
return ""
|
||
|
||
def prepare_inputs(
|
||
self,
|
||
context_prompts: List[str],
|
||
text_before: str,
|
||
text_after: str,
|
||
pinyin: str,
|
||
slot_chars: List[str],
|
||
max_seq_len: int = 128,
|
||
) -> dict:
|
||
"""
|
||
准备模型输入
|
||
|
||
Returns:
|
||
dict: {
|
||
'preprocess_time': float, # 预处理时间(ms)
|
||
'input_ids': numpy array,
|
||
'attention_mask': numpy array,
|
||
'pinyin_ids': numpy array,
|
||
'history_slot_ids': numpy array,
|
||
}
|
||
"""
|
||
preprocess_start = time.perf_counter()
|
||
|
||
# 1. 构建tokenizer输入
|
||
context_text = "|".join(context_prompts) if context_prompts else ""
|
||
|
||
if context_text:
|
||
input_text = f"{context_text}|{text_before}"
|
||
else:
|
||
input_text = text_before
|
||
|
||
# 2. Tokenize
|
||
encoded = self.tokenizer(
|
||
input_text,
|
||
text_after,
|
||
max_length=max_seq_len,
|
||
padding="max_length",
|
||
truncation=True,
|
||
return_tensors="pt",
|
||
return_token_type_ids=True,
|
||
)
|
||
|
||
input_ids = encoded["input_ids"].numpy()
|
||
attention_mask = encoded["attention_mask"].numpy()
|
||
|
||
# 3. 处理拼音输入
|
||
cleaned_pinyin = self._clean_pinyin_input(pinyin)
|
||
pinyin_ids = text_to_pinyin_ids(cleaned_pinyin)
|
||
|
||
if len(pinyin_ids) < 24:
|
||
pinyin_ids.extend([0] * (24 - len(pinyin_ids)))
|
||
else:
|
||
pinyin_ids = pinyin_ids[:24]
|
||
|
||
pinyin_ids = np.array([pinyin_ids], dtype=np.int64)
|
||
|
||
# 4. 处理历史槽位
|
||
history_slot_ids = []
|
||
for char in slot_chars:
|
||
char_id = self.char_to_id(char)
|
||
history_slot_ids.append(char_id)
|
||
|
||
if len(history_slot_ids) < 8:
|
||
history_slot_ids.extend([0] * (8 - len(history_slot_ids)))
|
||
else:
|
||
history_slot_ids = history_slot_ids[:8]
|
||
|
||
history_slot_ids = np.array([history_slot_ids], dtype=np.int64)
|
||
|
||
preprocess_time = (time.perf_counter() - preprocess_start) * 1000
|
||
|
||
return {
|
||
"preprocess_time": preprocess_time,
|
||
"input_ids": input_ids,
|
||
"attention_mask": attention_mask,
|
||
"pinyin_ids": pinyin_ids,
|
||
"history_slot_ids": history_slot_ids,
|
||
}
|
||
|
||
def run_context_encoder(
|
||
self, input_ids: np.ndarray, pinyin_ids: np.ndarray, attention_mask: np.ndarray
|
||
) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
|
||
"""
|
||
运行上下文编码器
|
||
|
||
Returns:
|
||
context_H, pinyin_P, context_mask, pinyin_mask
|
||
"""
|
||
context_start = time.perf_counter()
|
||
|
||
inputs = {
|
||
"input_ids": input_ids,
|
||
"pinyin_ids": pinyin_ids,
|
||
"attention_mask": attention_mask,
|
||
}
|
||
|
||
outputs = self.context_encoder_session.run(self.context_output_names, inputs)
|
||
|
||
context_H, pinyin_P, context_mask, pinyin_mask = outputs
|
||
|
||
self.last_context_encoder_time = (time.perf_counter() - context_start) * 1000
|
||
|
||
return context_H, pinyin_P, context_mask, pinyin_mask
|
||
|
||
def run_decoder(
|
||
self,
|
||
context_H: np.ndarray,
|
||
pinyin_P: np.ndarray,
|
||
history_slot_ids: np.ndarray,
|
||
context_mask: np.ndarray,
|
||
pinyin_mask: np.ndarray,
|
||
) -> np.ndarray:
|
||
"""
|
||
运行解码器
|
||
|
||
Returns:
|
||
logits: [batch, vocab_size]
|
||
"""
|
||
decoder_start = time.perf_counter()
|
||
|
||
inputs = {
|
||
"context_H": context_H,
|
||
"pinyin_P": pinyin_P,
|
||
"history_slot_ids": history_slot_ids,
|
||
"context_mask": context_mask,
|
||
"pinyin_mask": pinyin_mask,
|
||
}
|
||
|
||
outputs = self.decoder_session.run(self.decoder_output_names, inputs)
|
||
|
||
self.last_decoder_time = (time.perf_counter() - decoder_start) * 1000
|
||
|
||
return outputs[0]
|
||
|
||
def predict(
|
||
self,
|
||
context_prompts: List[str],
|
||
text_before: str,
|
||
text_after: str,
|
||
pinyin: str,
|
||
slot_chars: List[str],
|
||
top_k: int = 20,
|
||
use_beam_search: bool = False,
|
||
beam_size: int = 5,
|
||
max_length: int = 10,
|
||
) -> Tuple[List[Tuple[str, float, int]], dict]:
|
||
"""
|
||
执行推理
|
||
|
||
Args:
|
||
context_prompts: 上下文提示
|
||
text_before: 光标前文本
|
||
text_after: 光标后文本
|
||
pinyin: 当前输入的拼音
|
||
slot_chars: 槽位内的汉字列表
|
||
top_k: 返回top-k个预测结果
|
||
use_beam_search: 是否使用束搜索
|
||
beam_size: 束大小
|
||
max_length: 最大生成长度
|
||
|
||
Returns:
|
||
(predictions, timing_info)
|
||
predictions: List[Tuple[char, prob, id]]
|
||
timing_info: 各阶段耗时字典
|
||
"""
|
||
total_start = time.perf_counter()
|
||
|
||
# 阶段1: 预处理
|
||
prep_start = time.perf_counter()
|
||
inputs = self.prepare_inputs(
|
||
context_prompts, text_before, text_after, pinyin, slot_chars
|
||
)
|
||
preprocess_time = inputs["preprocess_time"]
|
||
|
||
input_ids = inputs["input_ids"]
|
||
attention_mask = inputs["attention_mask"]
|
||
pinyin_ids = inputs["pinyin_ids"]
|
||
history_slot_ids = inputs["history_slot_ids"]
|
||
prep_time = (time.perf_counter() - prep_start) * 1000
|
||
|
||
# 阶段2: 上下文编码
|
||
context_start = time.perf_counter()
|
||
context_H, pinyin_P, context_mask, pinyin_mask = self.run_context_encoder(
|
||
input_ids, pinyin_ids, attention_mask
|
||
)
|
||
context_encoder_time = self.last_context_encoder_time
|
||
|
||
if use_beam_search:
|
||
# 阶段3: 束搜索解码
|
||
decode_start = time.perf_counter()
|
||
predictions, beam_decode_time = self._beam_search_decode(
|
||
context_H,
|
||
pinyin_P,
|
||
context_mask,
|
||
pinyin_mask,
|
||
beam_size,
|
||
max_length,
|
||
top_k,
|
||
)
|
||
decoder_time = beam_decode_time
|
||
else:
|
||
# 阶段3: 单步解码
|
||
decode_start = time.perf_counter()
|
||
logits = self.run_decoder(
|
||
context_H,
|
||
pinyin_P,
|
||
history_slot_ids,
|
||
context_mask,
|
||
pinyin_mask,
|
||
)
|
||
|
||
# 阶段4: 后处理
|
||
postprocess_start = time.perf_counter()
|
||
probs = self._softmax(logits)
|
||
top_indices, top_probs = self._topk(probs, top_k)
|
||
|
||
predictions = []
|
||
for i in range(top_k):
|
||
idx = int(top_indices[0, i])
|
||
prob = float(top_probs[0, i])
|
||
char = self.id_to_char(idx)
|
||
predictions.append((char, prob, idx))
|
||
|
||
postprocess_time = (time.perf_counter() - postprocess_start) * 1000
|
||
decoder_time = self.last_decoder_time
|
||
|
||
total_time = (time.perf_counter() - total_start) * 1000
|
||
|
||
timing_info = {
|
||
"预处理": prep_time,
|
||
"上下文编码": context_encoder_time,
|
||
"解码": decoder_time,
|
||
"后处理": postprocess_time if not use_beam_search else 0,
|
||
"总耗时": total_time,
|
||
}
|
||
|
||
return predictions, timing_info
|
||
|
||
def _softmax(self, logits: np.ndarray) -> np.ndarray:
|
||
"""计算softmax"""
|
||
exp_logits = np.exp(logits - np.max(logits, axis=-1, keepdims=True))
|
||
return exp_logits / np.sum(exp_logits, axis=-1, keepdims=True)
|
||
|
||
def _topk(self, probs: np.ndarray, k: int) -> Tuple[np.ndarray, np.ndarray]:
|
||
"""获取top-k"""
|
||
topk_indices = np.argsort(probs, axis=-1)[:, -k:][:, ::-1]
|
||
topk_probs = np.take_along_axis(probs, topk_indices, axis=-1)
|
||
return topk_indices, topk_probs
|
||
|
||
def _beam_search_decode(
|
||
self,
|
||
context_H: np.ndarray,
|
||
pinyin_P: np.ndarray,
|
||
context_mask: np.ndarray,
|
||
pinyin_mask: np.ndarray,
|
||
beam_size: int,
|
||
max_length: int,
|
||
top_k: int,
|
||
) -> Tuple[List[Tuple[str, float, int]], float]:
|
||
"""束搜索解码"""
|
||
beams = [([], 0.0)] # (序列, 对数概率)
|
||
|
||
for step in range(max_length):
|
||
new_beams = []
|
||
|
||
for seq, score in beams:
|
||
if len(seq) < 8:
|
||
history = seq + [0] * (8 - len(seq))
|
||
else:
|
||
history = seq[-8:]
|
||
|
||
history_tensor = np.array([history], dtype=np.int64)
|
||
|
||
logits = self.run_decoder(
|
||
context_H,
|
||
pinyin_P,
|
||
history_tensor,
|
||
context_mask,
|
||
pinyin_mask,
|
||
)
|
||
|
||
probs = self._softmax(logits)[0]
|
||
topk_indices = np.argsort(probs)[-beam_size:][::-1]
|
||
topk_probs = probs[topk_indices]
|
||
|
||
for idx, prob in zip(topk_indices, topk_probs):
|
||
new_seq = seq + [int(idx)]
|
||
new_score = score + np.log(prob + 1e-10)
|
||
new_beams.append((new_seq, new_score))
|
||
|
||
new_beams.sort(key=lambda x: x[1], reverse=True)
|
||
beams = new_beams[:beam_size]
|
||
|
||
all_ended = all(seq[-1] == 0 for seq, _ in beams if seq)
|
||
if all_ended:
|
||
break
|
||
|
||
# 返回top-k个候选
|
||
predictions = []
|
||
for seq, score in beams[:top_k]:
|
||
if seq:
|
||
char = self.id_to_char(seq[-1])
|
||
prob = np.exp(score / max(len(seq), 1))
|
||
else:
|
||
char = self.id_to_char(0)
|
||
prob = 0.0
|
||
predictions.append((char, prob, seq[-1] if seq else 0))
|
||
|
||
decode_time = self.last_decoder_time # 只记录最后一次解码的时间
|
||
|
||
return predictions, decode_time
|
||
|
||
def interactive_mode(self):
|
||
"""交互式推理模式"""
|
||
print("\n" + "=" * 60)
|
||
print("ONNX输入法模型推理 - 交互模式")
|
||
print("=" * 60)
|
||
|
||
encoding = sys.stdout.encoding or "unknown"
|
||
print(f"终端编码: {encoding}")
|
||
|
||
print("\n说明:")
|
||
print(" - 上下文提示: 模型不掌握的专有词汇、姓名等（可为空）")
|
||
print(" - 光标前文本: 光标前的连续文本")
|
||
print(" - 光标后文本: 光标后的连续文本")
|
||
print(" - 拼音: 当前输入的拼音")
|
||
print(" - 槽位历史: 用户已确认的输入历史")
|
||
if self.use_beam_search:
|
||
print(f" - 解码模式: 束搜索 (beam_size={self.beam_size})")
|
||
else:
|
||
print(" - 解码模式: 单步解码 (使用 --beam 启用束搜索)")
|
||
print("提示: 输入 'quit' 或 'exit' 或 'q' 可随时退出")
|
||
print("-" * 60)
|
||
|
||
while True:
|
||
try:
|
||
print("\n" + "=" * 60)
|
||
context_input = self._safe_input("第1步: 上下文提示（直接回车跳过）")
|
||
if context_input.lower() in ["quit", "exit", "q"]:
|
||
break
|
||
|
||
context_prompts = [
|
||
item.strip() for item in context_input.split(",") if item.strip()
|
||
]
|
||
|
||
print("\n" + "-" * 40)
|
||
text_before = self._safe_input("第2步: 光标前文本")
|
||
if text_before.lower() in ["quit", "exit", "q"]:
|
||
break
|
||
|
||
print("\n" + "-" * 40)
|
||
text_after = self._safe_input("第3步: 光标后文本")
|
||
if text_after.lower() in ["quit", "exit", "q"]:
|
||
break
|
||
|
||
print("\n" + "-" * 40)
|
||
pinyin = self._safe_input("第4步: 拼音输入")
|
||
if pinyin.lower() in ["quit", "exit", "q"]:
|
||
break
|
||
|
||
print("\n" + "-" * 40)
|
||
slot_input = self._safe_input("第5步: 槽位历史（直接回车表示无）")
|
||
if slot_input.lower() in ["quit", "exit", "q"]:
|
||
break
|
||
|
||
slot_chars = [
|
||
char.strip() for char in slot_input.split(",") if char.strip()
|
||
]
|
||
|
||
print("\n" + "=" * 60)
|
||
print("📝 输入汇总:")
|
||
print(f" 上下文提示: {context_prompts if context_prompts else '无'}")
|
||
print(f" 光标前文本: '{text_before}'")
|
||
print(f" 光标后文本: '{text_after}'")
|
||
print(f" 拼音: '{pinyin}'")
|
||
print(f" 槽位历史: {slot_chars if slot_chars else '无'}")
|
||
|
||
print("\n🔮 推理中...")
|
||
predictions, timing_info = self.predict(
|
||
context_prompts,
|
||
text_before,
|
||
text_after,
|
||
pinyin,
|
||
slot_chars,
|
||
top_k=20,
|
||
use_beam_search=self.use_beam_search,
|
||
beam_size=self.beam_size,
|
||
)
|
||
|
||
# 显示时间统计
|
||
print(f"\n⏱️ 执行时间统计:")
|
||
print("-" * 40)
|
||
for stage, duration in timing_info.items():
|
||
if duration > 0:
|
||
print(f" {stage:<12}: {duration:>8.2f} ms")
|
||
print("-" * 40)
|
||
|
||
# 显示结果
|
||
print("\n🏆 Top-20 预测结果:")
|
||
print("-" * 50)
|
||
for i, (char, prob, idx) in enumerate(predictions):
|
||
if char == "//":
|
||
print(f"{i + 1:2d}. {'//':<4} (结束符) - 概率: {prob:.4f}")
|
||
else:
|
||
print(
|
||
f"{i + 1:2d}. {char:<4} (ID: {idx:>5}) - 概率: {prob:.4f}"
|
||
)
|
||
|
||
# 显示拼音参考
|
||
if pinyin and self.query_engine:
|
||
print(f"\n📖 拼音 '{pinyin}' 的常见汉字:")
|
||
pinyin_results = self.query_engine.query_by_pinyin(pinyin, limit=10)
|
||
if pinyin_results:
|
||
for j, (pid, char, count) in enumerate(pinyin_results):
|
||
print(f" {char} (频次: {count:,})")
|
||
|
||
print("\n" + "-" * 40)
|
||
continue_input = (
|
||
self._safe_input("是否继续推理？(y/n)", "y").strip().lower()
|
||
)
|
||
if continue_input not in ["y", "yes", ""]:
|
||
break
|
||
|
||
except KeyboardInterrupt:
|
||
print("\n\n退出交互模式")
|
||
break
|
||
except Exception as e:
|
||
print(f"\n❌ 推理出错: {e}")
|
||
import traceback
|
||
|
||
traceback.print_exc()
|
||
|
||
|
||
def main():
|
||
parser = argparse.ArgumentParser(description="ONNX输入法模型推理")
|
||
parser.add_argument(
|
||
"--context-encoder",
|
||
"-c",
|
||
type=str,
|
||
required=True,
|
||
help="上下文编码器ONNX模型路径",
|
||
)
|
||
parser.add_argument(
|
||
"--decoder",
|
||
"-d",
|
||
type=str,
|
||
required=True,
|
||
help="解码器ONNX模型路径",
|
||
)
|
||
parser.add_argument(
|
||
"--vocab-size",
|
||
type=int,
|
||
default=10019,
|
||
help="词汇表大小 (默认: 10019)",
|
||
)
|
||
parser.add_argument(
|
||
"--device",
|
||
type=str,
|
||
default="cpu",
|
||
choices=["cpu", "cuda"],
|
||
help="推理设备 (默认: cpu)",
|
||
)
|
||
parser.add_argument(
|
||
"--interactive",
|
||
action="store_true",
|
||
default=True,
|
||
help="交互模式 (默认: True)",
|
||
)
|
||
parser.add_argument("--test", action="store_true", help="运行测试推理")
|
||
parser.add_argument(
|
||
"--beam",
|
||
action="store_true",
|
||
help="使用束搜索解码 (默认: 单步解码)",
|
||
)
|
||
parser.add_argument(
|
||
"--beam-size",
|
||
type=int,
|
||
default=5,
|
||
help="束大小 (默认: 5)",
|
||
)
|
||
|
||
args = parser.parse_args()
|
||
|
||
# 检查文件是否存在
|
||
if not os.path.exists(args.context_encoder):
|
||
print(f"❌ 错误: 上下文编码器文件不存在: {args.context_encoder}")
|
||
sys.exit(1)
|
||
if not os.path.exists(args.decoder):
|
||
print(f"❌ 错误: 解码器文件不存在: {args.decoder}")
|
||
sys.exit(1)
|
||
|
||
# 初始化推理器
|
||
inference = ONNXInference(
|
||
context_encoder_path=args.context_encoder,
|
||
decoder_path=args.decoder,
|
||
vocab_size=args.vocab_size,
|
||
device=args.device,
|
||
use_beam_search=args.beam,
|
||
beam_size=args.beam_size,
|
||
)
|
||
|
||
# 测试推理
|
||
if args.test:
|
||
print("\n🧪 运行测试推理...")
|
||
print("测试场景: 输入'shanghai'，已确认第一个字'上'，继续输入'tian'")
|
||
print("上下文提示: 张三、李四（模型不掌握的专有名词）")
|
||
|
||
predictions, timing_info = inference.predict(
|
||
context_prompts=["张三", "李四"],
|
||
text_before="今天天气",
|
||
text_after="很好",
|
||
pinyin="tian",
|
||
slot_chars=["上"],
|
||
use_beam_search=args.beam,
|
||
beam_size=args.beam_size,
|
||
)
|
||
|
||
print(f"\n⏱️ 执行时间统计:")
|
||
print("-" * 40)
|
||
for stage, duration in timing_info.items():
|
||
if duration > 0:
|
||
print(f" {stage:<12}: {duration:>8.2f} ms")
|
||
print("-" * 40)
|
||
|
||
print(f"\nTop-5 结果:")
|
||
for i, (char, prob, idx) in enumerate(predictions[:5]):
|
||
if char == "//":
|
||
print(f" {i + 1}. // (结束符) - 概率: {prob:.4f}")
|
||
else:
|
||
print(f" {i + 1}. {char} (ID: {idx}) - 概率: {prob:.4f}")
|
||
|
||
# 交互模式
|
||
if args.interactive:
|
||
inference.interactive_mode()
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|