""" GQA vs MLA 模型对比评估脚本 - Perplexity（在 SFT 数据子集上） - 指定题目人工对比 """ import torch, json, argparse, random, sys, os sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))) from model.model_minimind import MiniMindConfig, MiniMindForCausalLM from model.model_minimind_mla import MiniMindMLAConfig, MiniMindMLAForCausalLM from transformers import AutoTokenizer TEST_QUESTIONS = [ "请用Python写一个计算斐波那契数列的函数", "解释一下机器学习和深度学习的区别", "推荐三本经典的科幻小说，并说明理由", "如果我想在北京开一家小餐馆，需要注意哪些方面", "请解释什么是光合作用", "比较一下燃油车和电动车的优缺点", "写一首关于秋天的五言绝句", "如何给孩子解释为什么天空是蓝色的", "请用简单的语言说明什么是API", "介绍一道你最喜欢的菜的做法", "什么是区块链技术，它可以应用在哪些领域", "如何有效地管理时间", "请翻译：Artificial intelligence is transforming every industry", "三个开关分别控制三个灯泡，你只能进房间一次，如何确定对应关系", "请用200字以内介绍中国的四大发明", ] def load_model(weight_path, config, device): """加载模型权重""" if isinstance(config, MiniMindMLAConfig): model = MiniMindMLAForCausalLM(config) else: model = MiniMindForCausalLM(config) model = model.to(device) if weight_path and os.path.exists(weight_path): state_dict = torch.load(weight_path, map_location=device, weights_only=True) model.load_state_dict(state_dict, strict=False) print(f" 加载权重: {weight_path}") model.eval() return model def compute_ppl(model, tokenizer, data_path, device, max_samples=200, max_length=512): """在数据集子集上计算 Perplexity""" print(f"\n 数据: {data_path} (最多 {max_samples} 条, max_len={max_length})") try: with open(data_path) as f: lines = f.readlines() except FileNotFoundError: print(f" [WARN] 数据文件不存在，跳过 PPL") return None random.shuffle(lines) total_loss = 0.0 total_tokens = 0 count = 0 for line in lines[:max_samples]: item = json.loads(line) # SFT 格式：有 conversations 字段 if "conversations" in item: text = tokenizer.apply_chat_template(item["conversations"], tokenize=False, add_generation_prompt=False) else: text = item.get("text", "") if not text.strip(): continue tokens = tokenizer(text, truncation=True, max_length=max_length, return_tensors="pt") input_ids = tokens["input_ids"].to(device) with torch.no_grad(): output = model(input_ids, labels=input_ids) loss = output.loss if loss is not None: total_loss += loss.item() * input_ids.size(1) total_tokens += input_ids.size(1) count += 1 if total_tokens == 0: return None avg_loss = total_loss / total_tokens ppl = torch.exp(torch.tensor(avg_loss)).item() print(f" 处理 {count} 条, 总 token {total_tokens}, loss={avg_loss:.4f}, ppl={ppl:.2f}") return ppl @torch.inference_mode() def generate_answer(model, tokenizer, prompt, device, max_new_tokens=512): """生成回答""" messages = [{"role": "user", "content": prompt}] text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) input_ids = tokenizer(text, return_tensors="pt")["input_ids"].to(device) output = model.generate( inputs=input_ids, max_new_tokens=max_new_tokens, use_cache=True, ) response = tokenizer.decode(output[0][input_ids.size(1):], skip_special_tokens=True) return response.strip() def main(): parser = argparse.ArgumentParser() parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu") parser.add_argument("--gqa_weight", type=str, default="../out/full_sft_768.pth") parser.add_argument("--mla_weight", type=str, default="../out/full_sft_768_mla.pth") parser.add_argument("--tokenizer_path", type=str, default="../model") parser.add_argument("--sft_data", type=str, default="../dataset/sft_t2t_mini.jsonl") parser.add_argument("--ppl_samples", type=int, default=200, help="PPL 采样条数") parser.add_argument("--max_new_tokens", type=int, default=256) parser.add_argument("--skip_ppl", action="store_true", help="跳过 PPL 计算") parser.add_argument("--skip_qna", action="store_true", help="跳过问答对比") parser.add_argument("--only_ppl", action="store_true", help="只算 PPL") args = parser.parse_args() device = args.device tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_path) print("=" * 80) print(" GQA vs MLA 模型对比评估") print("=" * 80) # ── 1. 加载模型 ── print("\n[1] 加载 GQA 模型...") gqa_config = MiniMindConfig() gqa_model = load_model(args.gqa_weight, gqa_config, device) print("\n[2] 加载 MLA 模型...") mla_config = MiniMindMLAConfig() mla_model = load_model(args.mla_weight, mla_config, device) # ── 2. PPL ── if not args.skip_ppl: print("\n" + "=" * 80) print(" Perplexity 对比") print("=" * 80) print("\n[GQA]") gqa_ppl = compute_ppl(gqa_model, tokenizer, args.sft_data, device, max_samples=args.ppl_samples) print("\n[MLA]") mla_ppl = compute_ppl(mla_model, tokenizer, args.sft_data, device, max_samples=args.ppl_samples) if gqa_ppl and mla_ppl: better = "GQA" if gqa_ppl < mla_ppl else "MLA" print(f"\n 结果: GQA PPL={gqa_ppl:.2f} MLA PPL={mla_ppl:.2f} → {better} 更低({abs(gqa_ppl - mla_ppl):.2f})") if args.only_ppl: return # ── 3. 问答对比 ── if not args.skip_qna: print("\n" + "=" * 80) print(" 问答对比") print("=" * 80) for i, question in enumerate(TEST_QUESTIONS, 1): print(f"\n{'─' * 80}") print(f" [{i}/{len(TEST_QUESTIONS)}] {question}") print(f"{'─' * 80}") print(f"\n [GQA]:") gqa_answer = generate_answer(gqa_model, tokenizer, question, device, max_new_tokens=args.max_new_tokens) for line in gqa_answer.split('\n'): print(f" {line}") print(f"\n [MLA]:") mla_answer = generate_answer(mla_model, tokenizer, question, device, max_new_tokens=args.max_new_tokens) for line in mla_answer.split('\n'): print(f" {line}") print("\n" + "=" * 80) print(" 评估完成") print("=" * 80) if __name__ == "__main__": main()