From 15645bbf56b719dc1a241167d2241463e3748498 Mon Sep 17 00:00:00 2001 From: baihaochen <1291308544@qq.com> Date: Sat, 16 May 2026 11:27:49 +0800 Subject: [PATCH] [feat] add GQA vs MLA evaluation script (PPL + Q&A) --- scripts/eval_compare.py | 181 ++++++++++++++++++++++++++++++++++++++++ 1 file changed, 181 insertions(+) create mode 100644 scripts/eval_compare.py diff --git a/scripts/eval_compare.py b/scripts/eval_compare.py new file mode 100644 index 0000000..c69379e --- /dev/null +++ b/scripts/eval_compare.py @@ -0,0 +1,181 @@ +""" +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()