#!/usr/bin/env python3 """ 束搜索算法演示 展示如何使用导出的两个ONNX模型进行束搜索推理， 模拟输入法场景：给定上下文、拼音和已确认字符，生成候选汉字序列。 """ import argparse import os import sys from pathlib import Path from typing import List, Tuple, Dict, Any import numpy as np import onnxruntime as ort import torch import torch.nn.functional as F # 添加src目录到路径 sys.path.insert(0, str(Path(__file__).parent)) class ONNXBeamSearch: """ 基于ONNX模型的束搜索解码器 """ def __init__( self, context_encoder_path: str, decoder_path: str, device: str = "cpu" ): """ 初始化 Args: context_encoder_path: 上下文编码器ONNX模型路径 decoder_path: 解码器ONNX模型路径 device: 推理设备（cpu或cuda） """ self.device = device # 配置ONNX Runtime提供程序 if device == "cuda": providers = ["CUDAExecutionProvider", "CPUExecutionProvider"] else: providers = ["CPUExecutionProvider"] # 创建ONNX Runtime会话 self.context_session = ort.InferenceSession( context_encoder_path, providers=providers ) self.decoder_session = ort.InferenceSession(decoder_path, providers=providers) # 获取输入输出信息 self.context_input_info = { input.name: input for input in self.context_session.get_inputs() } self.context_output_info = { output.name: output for output in self.context_session.get_outputs() } self.decoder_input_info = { input.name: input for input in self.decoder_session.get_inputs() } self.decoder_output_info = { output.name: output for output in self.decoder_session.get_outputs() } print("📦 ONNX模型加载完成") print(f" 上下文编码器输入: {list(self.context_input_info.keys())}") print(f" 上下文编码器输出: {list(self.context_output_info.keys())}") print(f" 解码器输入: {list(self.decoder_input_info.keys())}") print(f" 解码器输出: {list(self.decoder_output_info.keys())}") def prepare_inputs( self, text_before: str, text_after: str, pinyin: str, context_prompts: List[str] = None, tokenizer=None, query_engine=None, max_seq_len: int = 128, ) -> Dict[str, np.ndarray]: """ 准备模型输入 注意: 这是一个简化的版本，实际应用中需要实现完整的预处理逻辑 这里使用随机数据模拟，实际应使用tokenizer和查询引擎 Args: text_before: 光标前文本 text_after: 光标后文本 pinyin: 拼音输入 context_prompts: 上下文提示 tokenizer: 分词器（未实现） query_engine: 查询引擎（未实现） max_seq_len: 最大序列长度 Returns: 输入字典，包含input_ids, pinyin_ids, attention_mask """ # 在实际应用中，这里应该实现: # 1. 使用tokenizer将文本转换为input_ids # 2. 使用text_to_pinyin_ids将拼音转换为pinyin_ids # 3. 构建attention_mask # 简化为随机数据（用于演示） batch_size = 1 seq_len = min(max_seq_len, 64) # 简化长度 # 模拟输入 input_ids = np.random.randint(0, 1000, (batch_size, seq_len), dtype=np.int64) pinyin_ids = np.random.randint( 0, 30, (batch_size, 24), dtype=np.int64 ) # 固定24长度 attention_mask = np.ones((batch_size, seq_len), dtype=np.int64) # 在实际应用中，应该这样处理拼音: # from src.model.dataset import text_to_pinyin_ids # pinyin_ids_list = text_to_pinyin_ids(pinyin) # pinyin_ids_array = np.array(pinyin_ids_list, dtype=np.int64).reshape(1, -1) return { "input_ids": input_ids, "pinyin_ids": pinyin_ids, "attention_mask": attention_mask, } def run_context_encoder( self, inputs: Dict[str, np.ndarray] ) -> Tuple[np.ndarray, ...]: """ 运行上下文编码器 Args: inputs: 输入字典 Returns: 上下文编码器输出: (context_H, pinyin_P, context_mask, pinyin_mask) """ # 准备ONNX输入（确保顺序正确） onnx_inputs = {} for input_name in self.context_input_info.keys(): if input_name in inputs: onnx_inputs[input_name] = inputs[input_name] else: # 对于缺失的输入，使用默认值 shape = self.context_input_info[input_name].shape dtype = self.context_input_info[input_name].type # 简化处理：创建零数组 if "int" in dtype: onnx_inputs[input_name] = np.zeros(shape, dtype=np.int64) else: onnx_inputs[input_name] = np.zeros(shape, dtype=np.float32) # 运行推理 outputs = self.context_session.run(None, onnx_inputs) return tuple(outputs) 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: """ 运行解码器 Args: context_H: 上下文编码 [batch, seq_len, 512] pinyin_P: 拼音编码 [batch, 24, 512] history_slot_ids: 历史槽位IDs [batch, 8] context_mask: 上下文掩码 [batch, seq_len] pinyin_mask: 拼音掩码 [batch, 24] Returns: logits [batch, vocab_size] """ 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(None, inputs) return outputs[0] def beam_search( self, context_H: np.ndarray, pinyin_P: np.ndarray, context_mask: np.ndarray, pinyin_mask: np.ndarray, beam_size: int = 5, max_length: int = 8, vocab_size: int = 10019, temperature: float = 1.0, length_penalty: float = 0.6, ) -> List[Tuple[List[int], float]]: """ 束搜索算法 Args: context_H: 上下文编码 pinyin_P: 拼音编码 context_mask: 上下文掩码 pinyin_mask: 拼音掩码 beam_size: 束大小 max_length: 最大生成长度 vocab_size: 词汇表大小 temperature: 温度参数（控制随机性） length_penalty: 长度惩罚系数 Returns: 排序后的候选序列列表，每个元素为(序列, 分数) """ # 初始束：空序列，对数概率为0 beams = [([], 0.0)] # (序列, 累计对数概率) for step in range(max_length): new_beams = [] for seq, score in beams: # 构建history_slot_ids if len(seq) < 8: # 填充到8个槽位 history = seq + [0] * (8 - len(seq)) else: # 只保留最近8个字符 history = seq[-8:] history_array = np.array([history], dtype=np.int64) # 运行解码器 logits = self.run_decoder( context_H, pinyin_P, history_array, context_mask, pinyin_mask ) # 应用温度 logits = logits / temperature # 转换为概率 probs = F.softmax(torch.from_numpy(logits[0]), dim=-1).numpy() # 获取top-k候选（k = beam_size * 2 以增加多样性） top_k = min(beam_size * 2, vocab_size) top_indices = np.argsort(probs)[-top_k:][::-1] top_probs = probs[top_indices] # 扩展束 for idx, prob in zip(top_indices, top_probs): new_seq = seq + [int(idx)] # 更新分数：累计对数概率 + log(prob) new_score = score + np.log(prob + 1e-10) # 应用长度惩罚 length_penalized_score = new_score / ( (5 + len(new_seq)) ** length_penalty ) new_beams.append((new_seq, length_penalized_score, new_score)) # 剪枝：保留beam_size个最佳候选 new_beams.sort(key=lambda x: x[1], reverse=True) # 按惩罚后分数排序 beams = [(seq, orig_score) for seq, _, orig_score in new_beams[:beam_size]] # 检查是否所有序列都已结束（结束符ID为0） all_ended = all(seq[-1] == 0 for seq, _ in beams if len(seq) > 0) if all_ended: break # 返回原始分数（未应用长度惩罚） return beams def interactive_beam_search(self): """交互式束搜索演示""" print("\n" + "=" * 60) print("束搜索交互演示") print("=" * 60) print("\n📝 输入法场景模拟:") print(" 假设用户正在输入拼音 'nihao'，已确认第一个字 '你'") print(" 上下文: 光标前文本 '今天天气很好'，光标后文本 '我们去公园玩'") print(" 上下文提示: '张三,李四'（模型不掌握的专有名词）") print("-" * 60) # 模拟输入（实际应用中应从用户获取） text_before = "今天天气很好" text_after = "我们去公园玩" pinyin = "hao" # 继续输入"hao" slot_chars = ["你"] # 已确认的字符 context_prompts = ["张三", "李四"] print(f"\n📋 输入参数:") print(f" 光标前文本: '{text_before}'") print(f" 光标后文本: '{text_after}'") print(f" 拼音: '{pinyin}'") print(f" 槽位历史: {slot_chars}") print(f" 上下文提示: {context_prompts}") # 准备输入（简化版，使用随机数据） print(f"\n🔧 准备模型输入...") inputs = self.prepare_inputs( text_before=text_before, text_after=text_after, pinyin=pinyin, context_prompts=context_prompts, ) # 运行上下文编码器 print(f"🧠 运行上下文编码器...") context_outputs = self.run_context_encoder(inputs) context_H, pinyin_P, context_mask, pinyin_mask = context_outputs print(f"✅ 上下文编码完成") print(f" context_H形状: {context_H.shape}") print(f" pinyin_P形状: {pinyin_P.shape}") # 将槽位历史转换为ID（简化：使用随机ID） # 实际应用中应使用query_engine将汉字转换为ID slot_ids = [42] if slot_chars else [0] # 假设'你'的ID是42 if len(slot_ids) < 8: slot_ids = slot_ids + [0] * (8 - len(slot_ids)) # 运行束搜索 print(f"\n🔍 运行束搜索 (beam_size=3, max_length=4)...") beams = self.beam_search( context_H, pinyin_P, context_mask, pinyin_mask, beam_size=3, max_length=4, vocab_size=10019, ) # 显示结果 print(f"\n🏆 束搜索结果:") print("-" * 50) for i, (seq, score) in enumerate(beams): # 将ID序列转换为汉字（简化：显示ID） seq_str = " ".join([f"ID:{id}" if id != 0 else "END" for id in seq]) print(f"{i + 1}. 序列: [{seq_str}]") print(f" 对数概率: {score:.4f}") print(f" 概率: {np.exp(score):.6f}") print() print("📝 说明:") print(" - 'END' 表示结束符 (ID: 0)") print(" - 实际应用中应将ID转换为汉字") print(" - 最高分数的序列作为最终预测") return beams def main(): parser = argparse.ArgumentParser(description="束搜索算法演示") parser.add_argument( "--context-encoder", type=str, default="./exported_models/context_encoder.onnx", help="上下文编码器ONNX路径（默认: ./exported_models/context_encoder.onnx）", ) parser.add_argument( "--decoder", type=str, default="./exported_models/decoder.onnx", help="解码器ONNX路径（默认: ./exported_models/decoder.onnx）", ) parser.add_argument( "--device", type=str, default="cpu", choices=["cpu", "cuda"], help="推理设备（默认: cpu）", ) parser.add_argument("--beam-size", type=int, default=3, help="束大小（默认: 3）") parser.add_argument( "--max-length", type=int, default=4, help="最大生成长度（默认: 4）" ) parser.add_argument( "--interactive", action="store_true", default=True, help="交互模式（默认: True）", ) args = parser.parse_args() # 检查模型文件是否存在 if not os.path.exists(args.context_encoder): print(f"❌ 上下文编码器文件不存在: {args.context_encoder}") print(f" 请先运行 export_onnx.py 导出模型") return if not os.path.exists(args.decoder): print(f"❌ 解码器文件不存在: {args.decoder}") print(f" 请先运行 export_onnx.py 导出模型") return print("🚀 束搜索算法演示") print("=" * 60) print(f"上下文编码器: {args.context_encoder}") print(f"解码器: {args.decoder}") print(f"设备: {args.device}") print(f"束大小: {args.beam_size}") print(f"最大长度: {args.max_length}") # 初始化束搜索器 beam_searcher = ONNXBeamSearch( context_encoder_path=args.context_encoder, decoder_path=args.decoder, device=args.device, ) if args.interactive: # 运行交互演示 beam_searcher.interactive_beam_search() print("\n" + "=" * 60) print("🎉 演示完成") print("\n💡 下一步:") print(" 1. 实现完整的输入预处理（tokenizer和拼音转换）") print(" 2. 集成查询引擎以将ID转换为汉字") print(" 3. 根据实际场景调整束搜索参数") print(" 4. 性能优化：批量处理、缓存等") if __name__ == "__main__": main()