feat(ai-labeling): AI 打标实验工作区初始化 + 模型选型完成

2026-06-11 18:34:55 +08:00
parent 9a86c936ee
commit 202e02c254
22 changed files with 2762 additions and 0 deletions
@@ -0,0 +1,148 @@
+"""
+run_labeling.py - 单期或批量 AI 打标脚本
+用法:
+  单期: python run_labeling.py --ep 4 --model m3
+  批量: python run_labeling.py --all --model m3
+"""
+
+import sys
+sys.stdout.reconfigure(encoding='utf-8')
+sys.stderr.reconfigure(encoding='utf-8')
+
+import os
+import re
+import json
+import argparse
+from pathlib import Path
+from datetime import datetime
+from openai import OpenAI
+from dotenv import load_dotenv
+
+load_dotenv()
+
+BASE_DIR = Path(__file__).parent.parent
+TRANSCRIPTS_DIR = BASE_DIR / "benchmark-set" / "transcripts"
+PROMPTS_DIR = BASE_DIR / "prompts"
+EXPERIMENTS_DIR = BASE_DIR / "experiments"
+GROUND_TRUTH = BASE_DIR / "benchmark-set" / "ground-truth.json"
+
+MODEL_CONFIG = {
+    "m3": {
+        "base_url": "https://api.minimaxi.com/v1",
+        "model_name": "MiniMax-M3",
+        "api_key_env": "MINIMAX_API_KEY",
+    },
+    "deepseek-v4-pro": {
+        "base_url": "https://api.deepseek.com",
+        "model_name": "deepseek-v4-pro",
+        "api_key_env": "DEEPSEEK_API_KEY",
+    },
+    "mimo-v2.5-pro": {
+        "base_url": "https://api.xiaomimimo.com/v1",
+        "model_name": "mimo-v2.5-pro",
+        "api_key_env": "MIMO_API_KEY",
+    },
+}
+
+ALL_EPISODES = [3, 4, 5, 7, 10, 11, 12, 13, 14, 15]
+
+
+def load_prompt(field):
+    if field == "narrative":
+        return (PROMPTS_DIR / "prompt2_narrative.md").read_text(encoding="utf-8")
+    raise ValueError(f"Unknown field: {field}")
+
+
+def load_transcript(ep):
+    pattern = f"ep{ep:02d}_*.md"
+    files = list(TRANSCRIPTS_DIR.glob(pattern))
+    if not files:
+        raise FileNotFoundError(f"No transcript found for ep{ep:02d} in {TRANSCRIPTS_DIR}")
+    return files[0].read_text(encoding="utf-8"), files[0].name
+
+
+def load_ground_truth(ep):
+    data = json.loads(GROUND_TRUTH.read_text(encoding="utf-8"))
+    for episode in data["episodes"]:
+        if episode["ep"] == ep:
+            return episode
+    return None
+
+
+def parse_prompt(template, transcript):
+    parts = template.split("## USER")
+    system_prompt = parts[0].replace("# Prompt 2:叙事结构判别", "").strip()
+    user_prompt = parts[1].strip().replace("{transcript}", transcript)
+    return system_prompt, user_prompt
+
+
+def extract_json_from_response(raw: str) -> dict:
+    """从模型响应中提取 JSON，兼容推理模型的<think>...输出。"""
+    # 先去掉<think>...标签及其内容
+    text = re.sub(r'<think>.*?', '', raw, flags=re.DOTALL)
+    text = text.strip()
+    # 去掉markdown代码块
+    text = re.sub(r'^```(?:json)?\s*', '', text)
+    text = re.sub(r'\s*```$', '', text)
+    text = text.strip()
+    # 从第一个 { 开始，到最后一个 } 结束
+    first_brace = text.find('{')
+    last_brace = text.rfind('}')
+    if first_brace != -1 and last_brace != -1 and last_brace >= first_brace:
+        json_str = text[first_brace:last_brace + 1]
+        return json.loads(json_str)
+    # 兜底：直接尝试解析
+    return json.loads(text)
+
+
+def call_model(model_key, system_prompt, user_prompt):
+    config = MODEL_CONFIG[model_key]
+    client = OpenAI(
+        api_key=os.environ[config["api_key_env"]],
+        base_url=config["base_url"],
+    )
+    response = client.chat.completions.create(
+        model=config["model_name"],
+        messages=[
+            {"role": "system", "content": system_prompt},
+            {"role": "user", "content": user_prompt},
+        ],
+        temperature=0.0,
+    )
+    raw = response.choices[0].message.content
+    return extract_json_from_response(raw)
+
+
+def run_labeling(ep, model_key):
+    transcript, fname = load_transcript(ep)
+    template = load_prompt("narrative")
+    system_prompt, user_prompt = parse_prompt(template, transcript)
+    result = call_model(model_key, system_prompt, user_prompt)
+    gt = load_ground_truth(ep)
+    ts = datetime.now().strftime("%Y%m%d_%H%M%S")
+    out = EXPERIMENTS_DIR / f"{ts}_{model_key}_ep{ep:02d}.json"
+    out.write_text(
+        json.dumps({"episode": ep, "filename": fname, "result": result, "ground_truth": gt}, ensure_ascii=False, indent=2),
+        encoding="utf-8",
+    )
+    print(f"完成 ep{ep:02d} -> {out.name}")
+    return result
+
+
+def main():
+    parser = argparse.ArgumentParser(description="AI 打标脚本")
+    parser.add_argument("--ep", type=int, help="单期编号")
+    parser.add_argument("--all", action="store_true", help="跑全部")
+    parser.add_argument("--model", default="m3", help="模型键名")
+    args = parser.parse_args()
+    if args.all:
+        for ep in ALL_EPISODES:
+            run_labeling(ep, args.model)
+    elif args.ep:
+        run_labeling(args.ep, args.model)
+    else:
+        parser.print_help()
+
+
+if __name__ == "__main__":
+    main()