code/02-prompting/prompt_test.py

"""批量测试提示词：每个版本把每条用例跑 3 次，统计通过率，找出时对时错的用例。

用法：python prompt_test.py prompts/classify_v1.txt prompts/classify_v2.txt
结果会另存一份到 results.jsonl，方便以后对比。
"""
import json
import os
import sys
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path

from openai import OpenAI

client = OpenAI(
    api_key=os.environ["LLM_API_KEY"],
    base_url=os.environ.get("LLM_BASE_URL", "https://api.deepseek.com"),
)
MODEL = os.environ.get("LLM_MODEL", "deepseek-flash")
RUNS = 3
HERE = Path(__file__).parent

cases = [json.loads(line) for line in (HERE / "prompts/cases.jsonl").read_text().splitlines() if line.strip()]


def classify(system, text):
    response = client.chat.completions.create(
        model=MODEL,
        messages=[{"role": "system", "content": system}, {"role": "user", "content": f"留言：{text}\n类别："}],
        extra_body={"thinking": {"type": "disabled"}},
    )
    return response.choices[0].message.content.strip()


records = []
for prompt_path in sys.argv[1:]:
    system = (HERE / prompt_path).read_text()
    jobs = [case for case in cases for _ in range(RUNS)]
    with ThreadPoolExecutor(10) as pool:
        outputs = list(pool.map(lambda c: classify(system, c["text"]), jobs))

    passed = sum(out == case["label"] for out, case in zip(outputs, jobs))
    print(f"{prompt_path}：{passed}/{len(jobs)} 通过（{passed / len(jobs):.0%}）")
    for i, case in enumerate(cases):
        answers = outputs[i * RUNS:(i + 1) * RUNS]
        right = sum(a == case["label"] for a in answers)
        records.append({"prompt": prompt_path, "text": case["text"], "label": case["label"], "outputs": answers})
        if right == 0:
            print(f"    全错  {case['text']}  标注={case['label']}  模型={answers}")
        elif right < RUNS:
            print(f"    不稳  {case['text']}  标注={case['label']}  模型={answers}")

with open(HERE / "results.jsonl", "w") as f:
    for r in records:
        f.write(json.dumps(r, ensure_ascii=False) + "\n")