231 lines
8.1 KiB
Python
231 lines
8.1 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
ONNX模型推理示例
|
||
|
||
展示如何使用导出的两个ONNX模型进行推理，
|
||
包括束搜索（beam search）算法。
|
||
"""
|
||
|
||
import numpy as np
|
||
import onnxruntime as ort
|
||
import torch
|
||
import torch.nn.functional as F
|
||
from typing import List, Tuple
|
||
|
||
|
||
class ONNXInference:
|
||
"""ONNX模型推理器"""
|
||
|
||
def __init__(self, context_encoder_path, decoder_path):
|
||
"""
|
||
初始化ONNX推理器
|
||
|
||
Args:
|
||
context_encoder_path: 上下文编码器ONNX模型路径
|
||
decoder_path: 解码器ONNX模型路径
|
||
"""
|
||
# 创建ONNX Runtime会话
|
||
self.context_encoder_session = ort.InferenceSession(
|
||
context_encoder_path,
|
||
providers=['CPUExecutionProvider'] # 或 'CUDAExecutionProvider'
|
||
)
|
||
self.decoder_session = ort.InferenceSession(
|
||
decoder_path,
|
||
providers=['CPUExecutionProvider']
|
||
)
|
||
|
||
# 获取输入输出名称
|
||
self.context_input_names = [input.name for input in self.context_encoder_session.get_inputs()]
|
||
self.context_output_names = [output.name for output in self.context_encoder_session.get_outputs()]
|
||
self.decoder_input_names = [input.name for input in self.decoder_session.get_inputs()]
|
||
self.decoder_output_names = [output.name for output in self.decoder_session.get_outputs()]
|
||
|
||
print(f"上下文编码器输入: {self.context_input_names}")
|
||
print(f"上下文编码器输出: {self.context_output_names}")
|
||
print(f"解码器输入: {self.decoder_input_names}")
|
||
print(f"解码器输出: {self.decoder_output_names}")
|
||
|
||
def prepare_inputs(self, text_before, text_after, pinyin, slot_chars, tokenizer, query_engine, max_seq_len=128):
|
||
"""
|
||
准备模型输入（与原始推理脚本保持一致）
|
||
|
||
注意: 这里需要实现文本到token的转换
|
||
为了简化示例，假设已经实现了相关函数
|
||
"""
|
||
# 这里应该调用实际的预处理函数
|
||
# 返回: input_ids, pinyin_ids, attention_mask, history_slot_ids
|
||
raise NotImplementedError("请实现实际的输入预处理")
|
||
|
||
def run_context_encoder(self, input_ids, pinyin_ids, attention_mask):
|
||
"""
|
||
运行上下文编码器
|
||
|
||
Args:
|
||
input_ids: [batch, seq_len]
|
||
pinyin_ids: [batch, 24]
|
||
attention_mask: [batch, seq_len]
|
||
|
||
Returns:
|
||
context_H, pinyin_P, context_mask, pinyin_mask
|
||
"""
|
||
# 准备输入
|
||
inputs = {
|
||
"input_ids": input_ids.numpy() if isinstance(input_ids, torch.Tensor) else input_ids,
|
||
"pinyin_ids": pinyin_ids.numpy() if isinstance(pinyin_ids, torch.Tensor) else pinyin_ids,
|
||
"attention_mask": attention_mask.numpy() if isinstance(attention_mask, torch.Tensor) else attention_mask,
|
||
}
|
||
|
||
# 运行推理
|
||
outputs = self.context_encoder_session.run(self.context_output_names, inputs)
|
||
|
||
# 解包输出
|
||
context_H, pinyin_P, context_mask, pinyin_mask = outputs
|
||
|
||
return (
|
||
torch.from_numpy(context_H),
|
||
torch.from_numpy(pinyin_P),
|
||
torch.from_numpy(context_mask),
|
||
torch.from_numpy(pinyin_mask),
|
||
)
|
||
|
||
def run_decoder(self, context_H, pinyin_P, history_slot_ids, context_mask, pinyin_mask):
|
||
"""
|
||
运行解码器
|
||
|
||
Args:
|
||
context_H: [batch, seq_len, 512]
|
||
pinyin_P: [batch, 24, 512]
|
||
history_slot_ids: [batch, 8]
|
||
context_mask: [batch, seq_len]
|
||
pinyin_mask: [batch, 24]
|
||
|
||
Returns:
|
||
logits: [batch, vocab_size]
|
||
"""
|
||
# 准备输入
|
||
inputs = {
|
||
"context_H": context_H.numpy() if isinstance(context_H, torch.Tensor) else context_H,
|
||
"pinyin_P": pinyin_P.numpy() if isinstance(pinyin_P, torch.Tensor) else pinyin_P,
|
||
"history_slot_ids": history_slot_ids.numpy() if isinstance(history_slot_ids, torch.Tensor) else history_slot_ids,
|
||
"context_mask": context_mask.numpy() if isinstance(context_mask, torch.Tensor) else context_mask,
|
||
"pinyin_mask": pinyin_mask.numpy() if isinstance(pinyin_mask, torch.Tensor) else pinyin_mask,
|
||
}
|
||
|
||
# 运行推理
|
||
outputs = self.decoder_session.run(self.decoder_output_names, inputs)
|
||
|
||
# 解包输出
|
||
logits = outputs[0]
|
||
return torch.from_numpy(logits)
|
||
|
||
def beam_search(self, context_H, pinyin_P, context_mask, pinyin_mask,
|
||
beam_size=5, max_length=10, vocab_size=10019):
|
||
"""
|
||
束搜索算法示例
|
||
|
||
Args:
|
||
context_H: 上下文编码
|
||
pinyin_P: 拼音编码
|
||
context_mask: 上下文掩码
|
||
pinyin_mask: 拼音掩码
|
||
beam_size: 束大小
|
||
max_length: 最大生成长度
|
||
vocab_size: 词汇表大小
|
||
|
||
Returns:
|
||
最佳序列列表
|
||
"""
|
||
# 初始束：空序列，分数为0
|
||
beams = [([], 0.0)] # (序列, 对数概率)
|
||
|
||
for step in range(max_length):
|
||
new_beams = []
|
||
|
||
for seq, score in beams:
|
||
# 构建history_slot_ids：已确认的字符ID
|
||
if len(seq) < 8:
|
||
history = seq + [0] * (8 - len(seq))
|
||
else:
|
||
history = seq[-8:] # 只保留最近8个
|
||
|
||
history_tensor = torch.tensor([history], dtype=torch.long)
|
||
|
||
# 运行解码器
|
||
logits = self.run_decoder(
|
||
context_H, pinyin_P, history_tensor,
|
||
context_mask, pinyin_mask
|
||
)
|
||
|
||
# 获取概率
|
||
probs = F.softmax(logits[0], dim=-1)
|
||
|
||
# 获取top-k候选
|
||
top_probs, top_indices = torch.topk(probs, beam_size)
|
||
|
||
# 扩展束
|
||
for prob, idx in zip(top_probs, top_indices):
|
||
new_seq = seq + [idx.item()]
|
||
new_score = score + torch.log(prob).item()
|
||
new_beams.append((new_seq, new_score))
|
||
|
||
# 剪枝：保留beam_size个最佳候选
|
||
new_beams.sort(key=lambda x: x[1], reverse=True)
|
||
beams = new_beams[:beam_size]
|
||
|
||
# 检查是否所有序列都已结束（以结束符0结尾）
|
||
all_ended = all(seq[-1] == 0 for seq, _ in beams if seq)
|
||
if all_ended:
|
||
break
|
||
|
||
return beams
|
||
|
||
def predict_single(self, input_ids, pinyin_ids, attention_mask, history_slot_ids):
|
||
"""
|
||
单步预测
|
||
|
||
Args:
|
||
input_ids: 输入token IDs
|
||
pinyin_ids: 拼音IDs
|
||
attention_mask: 注意力掩码
|
||
history_slot_ids: 历史槽位IDs
|
||
|
||
Returns:
|
||
预测logits
|
||
"""
|
||
# 1. 运行上下文编码器
|
||
context_H, pinyin_P, context_mask, pinyin_mask = self.run_context_encoder(
|
||
input_ids, pinyin_ids, attention_mask
|
||
)
|
||
|
||
# 2. 运行解码器
|
||
logits = self.run_decoder(
|
||
context_H, pinyin_P, history_slot_ids,
|
||
context_mask, pinyin_mask
|
||
)
|
||
|
||
return logits
|
||
|
||
|
||
def main():
|
||
"""示例主函数"""
|
||
print("ONNX模型推理示例")
|
||
print("=" * 60)
|
||
|
||
# 初始化推理器
|
||
context_encoder_path = "context_encoder.onnx"
|
||
decoder_path = "decoder.onnx"
|
||
|
||
if not os.path.exists(context_encoder_path) or not os.path.exists(decoder_path):
|
||
print("错误: 找不到ONNX模型文件")
|
||
print("请先运行export_onnx.py导出模型")
|
||
return
|
||
|
||
inference = ONNXInference(context_encoder_path, decoder_path)
|
||
|
||
print("✅ ONNX推理器初始化完成")
|
||
print("请参考此示例实现完整的输入法推理流程")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|