148 lines
6.2 KiB
Python
148 lines
6.2 KiB
Python
"""
|
||
Qualitative eval of a MiniMind pretrain (or post-pretrain) checkpoint by sampling
|
||
text continuations from a fixed prompt set spanning multiple domains.
|
||
|
||
For pretrain (base) models — does NOT use chat template. Use eval_llm.py for SFT models.
|
||
|
||
Usage:
|
||
# Default: 64M pretrain ckpt
|
||
python eval_completion.py
|
||
|
||
# 1B ckpt
|
||
python eval_completion.py --hidden_size 2048 --num_hidden_layers 20 --weight pretrain
|
||
|
||
# Custom prompt
|
||
python eval_completion.py --prompt "今天天气不错，我打算"
|
||
|
||
# Save outputs to markdown for design doc
|
||
python eval_completion.py --output_md eval_results.md
|
||
|
||
What to look for:
|
||
✓ Grammatically coherent Chinese, no repetition loops
|
||
✓ Stays roughly on topic of the prefix
|
||
✓ Variety across runs (different temperature should give different outputs)
|
||
✗ Repetitive: "的的的的" / 同一句重复 → 训练或解码有问题
|
||
✗ 全部生成都一模一样 → 模型坍缩 / temperature 太低
|
||
✗ 乱码 / 非中文片段（如果是纯中文 prefix）→ tokenizer 问题
|
||
"""
|
||
import os
|
||
import sys
|
||
|
||
__package__ = "scripts"
|
||
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
|
||
|
||
import argparse
|
||
import warnings
|
||
import time
|
||
import torch
|
||
from transformers import AutoTokenizer, TextStreamer
|
||
from model.model_minimind import MiniMindConfig, MiniMindForCausalLM
|
||
|
||
warnings.filterwarnings('ignore')
|
||
|
||
|
||
DEFAULT_PROMPTS = [
|
||
("叙事", "在很久很久以前，有一个小村庄，"),
|
||
("日常", "今天天气不错，我打算"),
|
||
("技术", "Python 是一种编程语言，它的特点是"),
|
||
("事实", "中华人民共和国成立于"),
|
||
("代码", "def fibonacci(n):\n if n < 2:"),
|
||
("科学", "牛顿第一定律说的是"),
|
||
("文学", "李白是唐代著名的"),
|
||
("数学", "二的十次方等于"),
|
||
("常识", "水的沸点在标准大气压下是"),
|
||
("地理", "中国的首都是"),
|
||
]
|
||
|
||
|
||
def load_model(args):
|
||
cfg = MiniMindConfig(
|
||
hidden_size=args.hidden_size,
|
||
num_hidden_layers=args.num_hidden_layers,
|
||
use_moe=bool(args.use_moe),
|
||
)
|
||
model = MiniMindForCausalLM(cfg)
|
||
moe_suffix = '_moe' if args.use_moe else ''
|
||
ckp = f'{args.save_dir}/{args.weight}_{args.hidden_size}{moe_suffix}.pth'
|
||
state_dict = torch.load(ckp, map_location='cpu', weights_only=True)
|
||
model.load_state_dict(state_dict, strict=False)
|
||
model = model.half().eval().to(args.device)
|
||
return model, cfg
|
||
|
||
|
||
@torch.no_grad()
|
||
def main():
|
||
parser = argparse.ArgumentParser(description="MiniMind Pretrain Completion Eval")
|
||
parser.add_argument('--hidden_size', type=int, default=768)
|
||
parser.add_argument('--num_hidden_layers', type=int, default=8)
|
||
parser.add_argument('--use_moe', type=int, default=0, choices=[0, 1])
|
||
parser.add_argument('--save_dir', type=str, default='../out')
|
||
parser.add_argument('--weight', type=str, default='pretrain')
|
||
parser.add_argument('--tokenizer_path', type=str, default='../model')
|
||
parser.add_argument('--prompt', type=str, default=None,
|
||
help='单条自定义 prompt；不指定则跑 DEFAULT_PROMPTS 全集')
|
||
parser.add_argument('--max_new_tokens', type=int, default=80)
|
||
parser.add_argument('--temperature', type=float, default=0.7)
|
||
parser.add_argument('--top_p', type=float, default=0.9)
|
||
parser.add_argument('--top_k', type=int, default=50)
|
||
parser.add_argument('--repetition_penalty', type=float, default=1.1)
|
||
parser.add_argument('--num_samples_per_prompt', type=int, default=1,
|
||
help='每个 prompt 采样几次（验证多样性）')
|
||
parser.add_argument('--output_md', type=str, default=None,
|
||
help='输出 markdown 报告路径（可选）')
|
||
parser.add_argument('--device', type=str, default='cuda' if torch.cuda.is_available() else 'cpu')
|
||
parser.add_argument('--seed', type=int, default=42)
|
||
args = parser.parse_args()
|
||
|
||
torch.manual_seed(args.seed)
|
||
|
||
print(f'[load] ckpt={args.save_dir}/{args.weight}_{args.hidden_size}.pth')
|
||
model, cfg = load_model(args)
|
||
tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_path)
|
||
|
||
prompts = [('custom', args.prompt)] if args.prompt else DEFAULT_PROMPTS
|
||
|
||
md_lines = []
|
||
md_lines.append(f'# Completion Eval: `{args.weight}_{args.hidden_size}.pth`')
|
||
md_lines.append(f'temperature={args.temperature}, top_p={args.top_p}, top_k={args.top_k}, '
|
||
f'rep_penalty={args.repetition_penalty}, max_new={args.max_new_tokens}\n')
|
||
|
||
for label, prompt in prompts:
|
||
print(f'\n{"=" * 60}')
|
||
print(f'[{label}] Prompt: {prompt!r}')
|
||
md_lines.append(f'## [{label}] `{prompt}`\n')
|
||
for sample_i in range(args.num_samples_per_prompt):
|
||
print(f'\n--- Sample {sample_i + 1} ---')
|
||
# 必须加 BOS：训练时所有样本以 [BOS] 开头，eval 时不加 BOS 模型会产生 degenerate
|
||
# 输出（"的的的"/"是是是" 重复），因为它从没见过"无 BOS 开头"的序列
|
||
body = tokenizer(prompt, add_special_tokens=False).input_ids
|
||
ids = torch.tensor([[tokenizer.bos_token_id] + body], device=args.device)
|
||
t0 = time.time()
|
||
output_ids = model.generate(
|
||
ids,
|
||
max_new_tokens=args.max_new_tokens,
|
||
temperature=args.temperature,
|
||
top_p=args.top_p,
|
||
top_k=args.top_k,
|
||
repetition_penalty=args.repetition_penalty,
|
||
do_sample=True,
|
||
eos_token_id=tokenizer.eos_token_id,
|
||
)
|
||
elapsed = time.time() - t0
|
||
text = tokenizer.decode(output_ids[0].tolist(), skip_special_tokens=True)
|
||
continuation = text[len(prompt):]
|
||
print(continuation)
|
||
new_tokens = output_ids.shape[1] - ids.shape[1]
|
||
print(f' [{new_tokens} new tokens in {elapsed:.2f}s = {new_tokens/elapsed:.1f} tok/s]')
|
||
md_lines.append(f'**Sample {sample_i + 1}**: {continuation.strip()}\n')
|
||
md_lines.append('')
|
||
|
||
if args.output_md:
|
||
with open(args.output_md, 'w') as f:
|
||
f.write('\n'.join(md_lines))
|
||
print(f'\n[saved] {args.output_md}')
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|