基于评测模板对 Qwen3-4B 模型进行性能测试

import json
from transformers import AutoTokenizer

TOKENIZER_PATH = "./tokenizer"
tokenizer = AutoTokenizer.from_pretrained(TOKENIZER_PATH, trust_remote_code=True)

# ── 正式评测集配置 ──────────────────────────────────────────
TARGETS = {
    "eval_1k":   {"min": 900,   "max": 1100,  "count": 200, "results": []},
    "eval_3_5k": {"min": 3200,  "max": 3800,  "count": 150, "results": []},
}

# ── 预热评测集配置 ──────────────────────────────────────────
WARMUP_COUNT = {
    "eval_1k":   20,
    "eval_3_5k": 20,
}

# ── 读取原始数据 ────────────────────────────────────────────
print("开始读取 Belle_open_source_0.5M.json ...")

items = []
with open('评测数据集/Belle_open_source_0.5M.json', 'r', encoding='utf-8', errors='ignore') as f:
    for i, line in enumerate(f):
        line = line.strip()
        if not line:
            continue
        try:
            item = json.loads(line)
        except json.JSONDecodeError:
            continue

        instruction = item.get('instruction', '').strip()
        inp = item.get('input', '').strip()

        # 压测只需要 prompt，不要 output
        user_text = (instruction + '\\n' + inp).strip() if inp else instruction
        if not user_text:
            continue

        token_count = len(tokenizer.encode(user_text))
        items.append((user_text, token_count))

        if i % 50000 == 0:
            print(f"  已读取 {i} 条...")

print(f"共读取 {len(items)} 条，开始按档位拼接筛选...")

# ── 生成正式评测集 ──────────────────────────────────────────
# 预热从短到长，正式压测也建议从短到长
for key in ["eval_1k", "eval_3_5k", "eval_16k"]:
    cfg = TARGETS[key]
    i = 0
    while i < len(items) and len(cfg["results"]) < cfg["count"]:
        combined_text = items[i][0]
        combined_tokens = items[i][1]
        j = i + 1
        while combined_tokens < cfg["min"] and j < len(items):
            combined_text += "\\n" + items[j][0]
            combined_tokens += items[j][1]
            j += 1
        if cfg["min"] <= combined_tokens <= cfg["max"]:
            cfg["results"].append({
                "messages": [{"role": "user", "content": combined_text.strip()}]
            })
        i = j
    print(f"  {key}：{len(cfg['results'])} 条")

# ── 保存正式评测集 ──────────────────────────────────────────
for key, cfg in TARGETS.items():
    path = f'评测数据集/{key}.jsonl'
    with open(path, 'w', encoding='utf-8') as f:
        for item in cfg["results"]:
            f.write(json.dumps(item, ensure_ascii=False) + '\\n')
    print(f"已保存正式评测集：{path}（{len(cfg['results'])} 条）")

# ── 生成并保存预热评测集 ────────────────────────────────────
print("\\n开始生成预热评测集...")

for key, n in WARMUP_COUNT.items():
    src_path    = f'评测数据集/{key}.jsonl'
    warmup_path = f'评测数据集/{key}_warmup.jsonl'

    samples = []
    with open(src_path, 'r', encoding='utf-8') as f:
        for i, line in enumerate(f):
            if i >= n:
                break
            samples.append(json.loads(line))

    with open(warmup_path, 'w', encoding='utf-8') as f:
        for item in samples:
            f.write(json.dumps(item, ensure_ascii=False) + '\\n')

    print(f"已保存预热评测集：{warmup_path}（{len(samples)} 条）")

print("全部完成")