#!/usr/bin/env python3 """ 评估预训练模型的推理效果 用于测试不同实验中训练出来的模型在eval_data.json上的表现 """ import os import json import argparse import torch import torch.nn.functional as F from transformers import AutoTokenizer from model.LMConfig import LMConfig def load_model(model_path, model_type, device, config_params=None): """ 加载模型和tokenizer Args: model_path: 模型权重文件路径 model_type: 模型类型 (model/model_original/model_no_feed) device: 运行设备 config_params: 模型配置参数字典 Returns: model: 加载好的模型 tokenizer: tokenizer实例 """ # 初始化配置 if config_params: lm_config = LMConfig(**config_params) else: lm_config = LMConfig() # 打印配置信息 print(f"模型配置:") print(f" dim: {lm_config.dim}") print(f" n_layers: {lm_config.n_layers}") print(f" n_heads: {lm_config.n_heads}") print(f" vocab_size: {lm_config.vocab_size}") print(f" max_seq_len: {lm_config.max_seq_len}") if hasattr(lm_config, 'knowledge_num'): print(f" knowledge_num: {lm_config.knowledge_num}") print(f" knowledge_length: {lm_config.knowledge_length}") print(f" knowledge_dim: {lm_config.knowledge_dim}") print() # 加载tokenizer tokenizer = AutoTokenizer.from_pretrained('./model/minimind_tokenizer') # 根据模型类型导入对应的模型类 if model_type == "model": from model.model import MiniMindLM elif model_type == "model_original": from model.model_original import MiniMindLM elif model_type == "model_no_feed": from model.model_no_feed import MiniMindLM else: raise ValueError(f"不支持的模型类型: {model_type}") # 初始化模型 model = MiniMindLM(lm_config) # 加载权重 if os.path.exists(model_path): print(f"正在从 {model_path} 加载模型权重...") # 加载权重文件 state_dict = torch.load(model_path, map_location=device) # 获取模型的参数名称 model_keys = set(model.state_dict().keys()) checkpoint_keys = set(state_dict.keys()) # 统计权重匹配情况 matched_keys = model_keys & checkpoint_keys missing_keys = model_keys - checkpoint_keys unexpected_keys = checkpoint_keys - model_keys print(f"\n权重加载详情:") print(f" 模型总参数数量: {len(model_keys)}") print(f" 权重文件参数数量: {len(checkpoint_keys)}") print(f" 成功匹配参数: {len(matched_keys)}") print(f" 缺失参数: {len(missing_keys)}") print(f" 多余参数: {len(unexpected_keys)}") # 详细列出缺失和多余的参数 if missing_keys: print(f"\n❌ 缺失的参数 ({len(missing_keys)}):") for key in sorted(missing_keys): print(f" - {key}") if unexpected_keys: print(f"\n⚠️ 权重文件中多余的参数 ({len(unexpected_keys)}):") for key in sorted(unexpected_keys): print(f" + {key}") # 加载权重（允许部分匹配） try: incompatible_keys = model.load_state_dict(state_dict, strict=False) # 检查加载结果 if len(incompatible_keys.missing_keys) == 0 and len(incompatible_keys.unexpected_keys) == 0: print(f"\n✅ 权重加载完全成功！") elif len(incompatible_keys.missing_keys) == 0: print(f"\n✅ 权重加载成功（忽略多余参数）") else: print(f"\n⚠️ 权重加载部分成功，存在缺失参数") print(f" 这可能影响模型性能，请检查模型配置参数是否正确") # 计算加载成功率 success_rate = len(matched_keys) / len(model_keys) * 100 print(f" 参数加载成功率: {success_rate:.1f}%") if success_rate < 90: print(f" ❌ 警告：加载成功率过低，模型可能无法正常工作！") elif success_rate < 100: print(f" ⚠️ 警告：存在缺失参数，可能影响模型性能") except Exception as e: raise RuntimeError(f"权重加载失败: {e}") # 验证关键层的形状 print("🔍 验证关键层形状:") key_layers = [ 'tok_embeddings.weight', 'output.weight', 'norm.weight', ] # 添加每一层的验证 for i in range(lm_config.n_layers): key_layers.extend([ f'layers.{i}.attention_norm.weight', f'layers.{i}.ffn_norm.weight', f'layers.{i}.self_attention.wq.weight', f'layers.{i}.self_attention.wk.weight', f'layers.{i}.self_attention.wv.weight', f'layers.{i}.self_attention.wo.weight', ]) # FFN层的验证（model_original有FFN，其他模型可能没有） if f'layers.{i}.feed_forward.w1.weight' in model_keys: key_layers.extend([ f'layers.{i}.feed_forward.w1.weight', f'layers.{i}.feed_forward.w2.weight', f'layers.{i}.feed_forward.w3.weight', ]) # 验证KnowledgeDataset相关层（仅model和model_no_feed） if model_type in ['model', 'model_no_feed']: key_layers.extend([ 'knowledge_dataset.to_queries.0.weight', 'knowledge_dataset.keys', 'knowledge_dataset.knowledge_dataset', ]) # 添加CrossAttention层 for i in range(lm_config.n_layers): key_layers.extend([ f'layers.{i}.cross_attention.to_q.weight', f'layers.{i}.cross_attention.to_k.weight', f'layers.{i}.cross_attention.to_v.weight', f'layers.{i}.cross_attention.to_out.weight', ]) # 检查关键层 verified_layers = 0 total_key_layers = 0 for layer_name in key_layers: if layer_name in model_keys: # 只检查模型中实际存在的层 total_key_layers += 1 if layer_name in matched_keys: verified_layers += 1 expected_shape = model.state_dict()[layer_name].shape actual_shape = state_dict[layer_name].shape if layer_name in state_dict else "缺失" if layer_name in state_dict and expected_shape == actual_shape: print(f" ✅ {layer_name}: {actual_shape}") else: print(f" ❌ {layer_name}: 期望 {expected_shape}, 实际 {actual_shape}") else: print(f" ❌ {layer_name}: 缺失") print(f"\n关键层验证结果: {verified_layers}/{total_key_layers} 层验证成功") if verified_layers == total_key_layers: print("✅ 所有关键层验证通过！") elif verified_layers / total_key_layers >= 0.9: print("⚠️ 大部分关键层验证通过，模型应该可以正常工作") else: print("❌ 关键层验证失败过多，模型可能无法正常工作！") print() else: raise FileNotFoundError(f"模型文件不存在: {model_path}") model.to(device) model.eval() return model, tokenizer def load_eval_data(data_path, num_samples=20): """ 加载评估数据集 Args: data_path: 数据文件路径 num_samples: 要评估的样本数量 Returns: samples: 数据样本列表 """ data = [] with open(data_path, 'r', encoding='utf-8') as f: for line_num, line in enumerate(f): line = line.strip() if line: # 跳过空行 try: sample = json.loads(line) data.append(sample) if len(data) >= num_samples: break except json.JSONDecodeError as e: print(f"警告：第{line_num+1}行JSON解析失败: {e}") continue # 只取前num_samples条数据 samples = data[:num_samples] print(f"加载了 {len(samples)} 条评估数据") return samples def evaluate_sample(model, tokenizer, text, input_length=100, predict_length=100, device='cuda'): """ 评估单个样本 Args: model: 模型实例 tokenizer: tokenizer实例 text: 输入文本 input_length: 输入token数量 predict_length: 预测token数量 device: 运行设备 Returns: input_text: 输入文本 predicted_text: 预测文本 ground_truth_text: 真实文本 loss: 预测损失（如果可计算） """ # 对文本进行分词 tokens = tokenizer.encode(text, add_special_tokens=False) # 确保有足够的token if len(tokens) < input_length + predict_length: print(f"警告：文本长度不足，只有 {len(tokens)} 个token") return None, None, None, None # 分割输入和目标 input_tokens = tokens[:input_length] target_tokens = tokens[input_length:input_length + predict_length] # 转换为张量 input_ids = torch.tensor([input_tokens], dtype=torch.long).to(device) # 生成预测 with torch.no_grad(): # 使用generate方法生成，调整参数改善生成质量 generated = model.generate( input_ids, max_new_tokens=predict_length, temperature=1.0, top_p=0.95, eos_token_id=tokenizer.eos_token_id, pad_token_id=tokenizer.pad_token_id ) # 提取生成的token（去掉输入部分） # generated包含完整序列，需要从input_length位置开始提取新生成的部分 full_generated_tokens = generated[0].tolist() if len(full_generated_tokens) > input_length: predicted_tokens = full_generated_tokens[input_length:] else: # 如果生成序列长度不够，说明没有新生成内容 predicted_tokens = [] # 检查是否因EOS token提前结束生成 eos_found = False eos_position = -1 actual_predicted_length = len(predicted_tokens) if predicted_tokens and tokenizer.eos_token_id is not None: try: eos_position = predicted_tokens.index(tokenizer.eos_token_id) eos_found = True # 只保留EOS token之前的内容 predicted_tokens = predicted_tokens[:eos_position] actual_predicted_length = len(predicted_tokens) except ValueError: # 没有找到EOS token pass # 计算loss（使用forward方法） # 准备用于loss计算的输入 loss_input_ids = torch.tensor([tokens[:input_length + predict_length]], dtype=torch.long).to(device) outputs = model(loss_input_ids) # 移除logits_to_keep参数 # 计算loss logits = outputs.logits loss = None if logits is not None: # 重塑logits和目标 - 修复：使用正确的位置切片 shift_logits = logits[0, input_length:input_length + predict_length, :].contiguous() shift_labels = torch.tensor(target_tokens, dtype=torch.long).to(device) # 计算交叉熵损失 loss = F.cross_entropy(shift_logits, shift_labels, reduction='mean') loss = loss.item() # 解码文本 input_text = tokenizer.decode(input_tokens, skip_special_tokens=True) # 只解码实际生成的token，限制在predict_length内 actual_predicted_tokens = predicted_tokens[:predict_length] if predicted_tokens else [] predicted_text = tokenizer.decode(actual_predicted_tokens, skip_special_tokens=True) if actual_predicted_tokens else "[未生成内容]" ground_truth_text = tokenizer.decode(target_tokens, skip_special_tokens=True) # 返回额外的生成统计信息 generation_stats = { 'requested_length': predict_length, 'actual_length': actual_predicted_length, 'eos_found': eos_found, 'eos_position': eos_position if eos_found else None, 'truncated_by_eos': eos_found and eos_position < predict_length } return input_text, predicted_text, ground_truth_text, loss, generation_stats def main(): parser = argparse.ArgumentParser(description='评估预训练模型') parser.add_argument('--model_path', type=str, default='out/experiment_1_4_0/pretrain_512.pth', help='模型权重文件路径') parser.add_argument('--model_type', type=str, default='model', choices=['model', 'model_original', 'model_no_feed'], help='模型类型') parser.add_argument('--data_path', type=str, default='dataset/stable/eval_data.json', help='评估数据集路径') parser.add_argument('--num_samples', type=int, default=20, help='评估样本数量') parser.add_argument('--input_length', type=int, default=100, help='输入token长度') parser.add_argument('--predict_length', type=int, default=100, help='预测token长度') parser.add_argument('--device', type=str, default='cuda' if torch.cuda.is_available() else 'cpu', help='运行设备') # 模型架构参数 parser.add_argument('--dim', type=int, default=512, help='模型维度') parser.add_argument('--n_layers', type=int, default=8, help='Transformer层数') parser.add_argument('--n_heads', type=int, default=32, help='注意力头数') parser.add_argument('--n_kv_heads', type=int, default=8, help='KV注意力头数') parser.add_argument('--vocab_size', type=int, default=6400, help='词汇表大小') parser.add_argument('--max_seq_len', type=int, default=512, help='最大序列长度') parser.add_argument('--dropout', type=float, default=0.0, help='Dropout率') parser.add_argument('--norm_eps', type=float, default=1e-5, help='层归一化epsilon') parser.add_argument('--rope_theta', type=float, default=1e6, help='RoPE theta参数') # KnowledgeDataset相关参数（仅model和model_no_feed使用） parser.add_argument('--knowledge_num', type=int, default=1048576, help='知识条目数量') parser.add_argument('--knowledge_length', type=int, default=32, help='单条知识长度') parser.add_argument('--knowledge_dim', type=int, default=128, help='知识维度') # MOE相关参数 parser.add_argument('--use_moe', action='store_true', help='是否使用MOE') parser.add_argument('--num_experts_per_tok', type=int, default=2, help='每个token激活的专家数') parser.add_argument('--n_routed_experts', type=int, default=4, help='路由专家数量') args = parser.parse_args() print(f"评估配置:") print(f" 模型路径: {args.model_path}") print(f" 模型类型: {args.model_type}") print(f" 数据路径: {args.data_path}") print(f" 样本数量: {args.num_samples}") print(f" 输入长度: {args.input_length} tokens") print(f" 预测长度: {args.predict_length} tokens") print(f" 运行设备: {args.device}") print() # 构建配置参数字典 config_params = { 'dim': args.dim, 'n_layers': args.n_layers, 'n_heads': args.n_heads, 'n_kv_heads': args.n_kv_heads, 'vocab_size': args.vocab_size, 'max_seq_len': args.max_seq_len, 'dropout': args.dropout, 'norm_eps': args.norm_eps, 'rope_theta': args.rope_theta, 'use_moe': args.use_moe, 'num_experts_per_tok': args.num_experts_per_tok, 'n_routed_experts': args.n_routed_experts, } # 只有model和model_no_feed需要KnowledgeDataset参数 if args.model_type in ['model', 'model_no_feed']: config_params.update({ 'knowledge_num': args.knowledge_num, 'knowledge_length': args.knowledge_length, 'knowledge_dim': args.knowledge_dim, }) # 加载模型 model, tokenizer = load_model(args.model_path, args.model_type, args.device, config_params) # 加载数据 samples = load_eval_data(args.data_path, args.num_samples) # 评估每个样本 total_loss = 0 valid_samples = 0 total_requested_tokens = 0 total_actual_tokens = 0 samples_with_eos = 0 samples_truncated_by_eos = 0 for i, sample in enumerate(samples): print(f"\n{'='*60}") print(f"样本 {i+1}/{len(samples)}") print(f"{'='*60}") text = sample['text'] # 评估样本 input_text, predicted_text, ground_truth_text, loss, generation_stats = evaluate_sample( model, tokenizer, text, args.input_length, args.predict_length, args.device ) if input_text is None: print("跳过该样本（文本长度不足）") continue # 打印结果 print(f"\n输入 ({args.input_length} tokens):") print(f" {input_text}") print(f"\n预测输出 (请求{generation_stats['requested_length']}个token, 实际生成{generation_stats['actual_length']}个):") print(f" {predicted_text}") print(f"\n真实值 ({args.predict_length} tokens):") print(f" {ground_truth_text}") # 打印生成统计信息 print(f"\n生成统计:") print(f" 请求生成: {generation_stats['requested_length']} tokens") print(f" 实际生成: {generation_stats['actual_length']} tokens") if generation_stats['eos_found']: print(f" ✅ 发现EOS token在位置 {generation_stats['eos_position']}") if generation_stats['truncated_by_eos']: print(f" ⚠️ 因EOS token提前结束生成") else: print(f" ✅ EOS token出现在预期位置") else: print(f" ❌ 未发现EOS token (可能达到最大长度限制)") if loss is not None: print(f"\nLoss: {loss:.4f}") total_loss += loss valid_samples += 1 # 更新生成统计 total_requested_tokens += generation_stats['requested_length'] total_actual_tokens += generation_stats['actual_length'] if generation_stats['eos_found']: samples_with_eos += 1 if generation_stats['truncated_by_eos']: samples_truncated_by_eos += 1 # 打印总体统计 if valid_samples > 0: print(f"\n{'='*60}") print(f"总体统计:") print(f" 有效样本数: {valid_samples}") print(f" 平均Loss: {total_loss / valid_samples:.4f}") print() print(f"生成统计:") print(f" 请求生成总tokens: {total_requested_tokens}") print(f" 实际生成总tokens: {total_actual_tokens}") print(f" 生成完成率: {total_actual_tokens / total_requested_tokens * 100:.1f}%" if total_requested_tokens > 0 else " 生成完成率: N/A") print(f" 发现EOS的样本: {samples_with_eos}/{len(samples)} ({samples_with_eos/len(samples)*100:.1f}%)" if len(samples) > 0 else " 发现EOS的样本: N/A") print(f" 被EOS截断的样本: {samples_truncated_by_eos}/{len(samples)} ({samples_truncated_by_eos/len(samples)*100:.1f}%)" if len(samples) > 0 else " 被EOS截断的样本: N/A") print(f" 平均每样本生成长度: {total_actual_tokens/len(samples):.1f} tokens" if len(samples) > 0 else " 平均每样本生成长度: N/A") print(f"{'='*60}") if __name__ == "__main__": main()