""" 推理脚本 可视化扩散翻译过程 """ import os import argparse import torch import torch.nn.functional as F from typing import Optional, Tuple, List from config import Config from tokenizer import Tokenizer from embedding import DualLanguageEmbedding, DualOutputProjection from model import create_model from diffusion import get_diffusion from switcher import create_switcher class Translator: """翻译器""" def __init__(self, config: Config, checkpoint_path: Optional[str] = None): self.config = config self.device = torch.device("cpu") # 加载分词器 cache_dir = os.path.join(config.project_dir, config.data.cache_dir) self.zh_tokenizer = Tokenizer.load(os.path.join(cache_dir, "tokenizer_zh.json")) self.en_tokenizer = Tokenizer.load(os.path.join(cache_dir, "tokenizer_en.json")) # 初始化模型组件 self.embedding = DualLanguageEmbedding( vocab_size_zh=self.zh_tokenizer.vocab_size_actual, vocab_size_en=self.en_tokenizer.vocab_size_actual, d_model=config.model.d_model, max_len=config.model.max_len, dropout=0.0, # 推理时不使用dropout ) self.output_proj = DualOutputProjection( d_model=config.model.d_model, vocab_size_zh=self.zh_tokenizer.vocab_size_actual, vocab_size_en=self.en_tokenizer.vocab_size_actual, ) self.model = create_model(config) self.switcher = create_switcher(config) self.diffusion, self.ddim_sampler = get_diffusion(config) # 加载权重 if checkpoint_path: self._load_checkpoint(checkpoint_path) def _load_checkpoint(self, path: str): """加载检查点""" state = torch.load(path, map_location=self.device, weights_only=False) self.embedding.load_state_dict(state['embedding']) self.output_proj.load_state_dict(state['output_proj']) self.model.load_state_dict(state['model']) self.switcher.load_state_dict(state['switcher']) print(f"已加载检查点: {path}") def _encode(self, text: str, lang: str) -> torch.Tensor: """编码文本""" if lang == "zh": ids = self.zh_tokenizer.encode(text, add_sos=True, add_eos=True) return torch.tensor(ids, dtype=torch.long).unsqueeze(0) else: ids = self.en_tokenizer.encode(text, add_sos=True, add_eos=True) return torch.tensor(ids, dtype=torch.long).unsqueeze(0) def _decode(self, ids: torch.Tensor, lang: str) -> str: """解码为文本""" ids = ids[0].tolist() if lang == "zh": return self.zh_tokenizer.decode(ids, skip_special=True) else: return self.en_tokenizer.decode(ids, skip_special=True) def _embed_to_tokens(self, x: torch.Tensor, lang: str) -> torch.Tensor: """从嵌入空间解码到token""" logits = self.output_proj(x, lang) ids = logits.argmax(dim=-1) return ids @torch.no_grad() def translate( self, text: str, source_lang: str, verbose: bool = True, ddim: bool = True, ) -> str: """翻译文本 Args: text: 输入文本 source_lang: 源语言 "zh" 或 "en" verbose: 是否打印扩散过程 ddim: 是否使用DDIM加速 Returns: 翻译结果 """ self.model.eval() self.embedding.eval() self.output_proj.eval() self.switcher.eval() target_lang = "en" if source_lang == "zh" else "zh" if verbose: print(f"\n翻译模式: {source_lang.upper()} → {target_lang.upper()}") print(f"输入: {text}") print(f"\n扩散过程:") # 编码源语言 source_ids = self._encode(text, source_lang) source_len = torch.tensor([source_ids.size(1)]) # 嵌入源语言 source_emb = self.embedding(source_ids, source_lang, source_len) # 完整前向扩散到纯噪声 if verbose: print(f" 前向扩散: {source_lang} → 噪声空间") batch_size = source_emb.size(0) t_full = torch.full((batch_size,), self.config.diffusion.timesteps - 1, dtype=torch.long) noise = torch.randn_like(source_emb) x_t, _ = self.diffusion.q_sample(source_emb, t_full, noise) # DDIM反向扩散 if ddim: result = self._ddim_reverse( x_t, source_lang, target_lang, verbose ) else: result = self._ddpm_reverse( x_t, source_lang, target_lang, verbose ) return result def _ddim_reverse( self, x_t: torch.Tensor, source_lang: str, target_lang: str, verbose: bool, ) -> str: """DDIM反向扩散""" ddim_steps = self.config.diffusion.ddim_steps timesteps = self.ddim_sampler.ddim_timesteps total_steps = len(timesteps) switch_point = total_steps // 2 # 在中间切换语言 for i, t in enumerate(timesteps[:-1]): t_prev = timesteps[i + 1] # 根据进度决定用哪种语言去噪和显示 # 前半段：源语言，后半段：目标语言 if i < switch_point: current_lang = source_lang else: current_lang = target_lang # 预测噪声 t_tensor = torch.full((x_t.size(0),), t, dtype=torch.long) predicted_noise = self.model(x_t, t_tensor, lang=current_lang) if verbose: # 显示当前语言的解码结果 current_ids = self._embed_to_tokens(x_t, current_lang) current_text = self._decode(current_ids, current_lang) if len(current_text) > 50: current_text = current_text[:50] + "..." print(f" Step {t:4d} → {current_text}") # DDIM步骤 x_t = self.ddim_sampler.ddim_step(x_t, t, t_prev, predicted_noise, eta=0.0) # 最终解码 final_ids = self._embed_to_tokens(x_t, target_lang) result = self._decode(final_ids, target_lang) if verbose: print(f"\n输出: {result}") return result def _ddpm_reverse( self, x_t: torch.Tensor, source_lang: str, target_lang: str, verbose: bool, ) -> str: """DDPM反向扩散（标准方法，较慢）""" total_steps = self.config.diffusion.timesteps switch_point = total_steps // 2 # 在中间切换语言 for t in range(total_steps - 1, -1, -1): # 根据时间步决定用哪种语言 if t > switch_point: current_lang = source_lang else: current_lang = target_lang t_tensor = torch.full((x_t.size(0),), t, dtype=torch.long) # 预测噪声 predicted_noise = self.model(x_t, t_tensor, lang=current_lang) if verbose: current_ids = self._embed_to_tokens(x_t, current_lang) current_text = self._decode(current_ids, current_lang) if len(current_text) > 50: current_text = current_text[:50] + "..." print(f" Step {t:4d} → {current_text}") # DDPM步骤 x_t = self.diffusion.p_sample(x_t, t_tensor, predicted_noise) # 解码 final_ids = self._embed_to_tokens(x_t, target_lang) result = self._decode(final_ids, target_lang) if verbose: print(f"\n输出: {result}") return result def interactive(self): """交互模式""" print("\n" + "=" * 50) print("Diffutslator 交互翻译模式") print("=" * 50) print("输入 'zh: 文本' 翻译中文到英文") print("输入 'en: text' 翻译英文到中文") print("输入 'quit' 或 'exit' 退出") print("=" * 50 + "\n") while True: try: user_input = input(">>> ").strip() if user_input.lower() in ['quit', 'exit', 'q']: print("再见!") break if not user_input: continue # 解析输入 if user_input.lower().startswith('zh:'): text = user_input[3:].strip() source_lang = "zh" elif user_input.lower().startswith('en:'): text = user_input[3:].strip() source_lang = "en" else: # 自动检测（简单判断） if any('\u4e00' <= c <= '\u9fff' for c in user_input): text = user_input source_lang = "zh" else: text = user_input source_lang = "en" # 翻译 result = self.translate(text, source_lang, verbose=True) except KeyboardInterrupt: print("\n再见!") break except Exception as e: print(f"错误: {e}") def main(): parser = argparse.ArgumentParser(description="Diffutslator 推理脚本") parser.add_argument("--checkpoint", type=str, default=None, help="检查点路径") parser.add_argument("--text", type=str, default=None, help="要翻译的文本") parser.add_argument("--zh", action="store_true", help="输入是中文") parser.add_argument("--en", action="store_true", help="输入是英文") parser.add_argument("--interactive", "-i", action="store_true", help="交互模式") parser.add_argument("--quiet", "-q", action="store_true", help="安静模式，不打印过程") parser.add_argument("--ddim-steps", type=int, default=50, help="DDIM步数") args = parser.parse_args() # 配置 config = Config() config.diffusion.ddim_steps = args.ddim_steps # 找检查点 checkpoint_path = args.checkpoint if checkpoint_path is None: checkpoint_dir = os.path.join(config.project_dir, config.training.checkpoint_dir) best_path = os.path.join(checkpoint_dir, "best.pt") if os.path.exists(best_path): checkpoint_path = best_path else: # 找最新的检查点 checkpoints = [f for f in os.listdir(checkpoint_dir) if f.endswith('.pt')] if checkpoints: checkpoint_path = os.path.join(checkpoint_dir, checkpoints[-1]) if checkpoint_path is None: print("错误: 未找到检查点，请先训练模型") return # 创建翻译器 translator = Translator(config, checkpoint_path) # 模式 if args.interactive: translator.interactive() elif args.text: if args.zh: source_lang = "zh" elif args.en: source_lang = "en" else: # 自动检测 if any('\u4e00' <= c <= '\u9fff' for c in args.text): source_lang = "zh" else: source_lang = "en" result = translator.translate(args.text, source_lang, verbose=not args.quiet) if args.quiet: print(result) else: # 默认交互模式 translator.interactive() if __name__ == "__main__": main()