data_construction/generate_thinker_targets.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""Generate thinker supervision targets from query candidates."""

from __future__ import annotations

import argparse
import json
import re
from concurrent.futures import FIRST_COMPLETED, ThreadPoolExecutor, wait
from pathlib import Path
from typing import Any, Dict, Iterable, List

from common_llm import DEFAULT_BASE_URL, chat_completion, extract_json_object


DEFAULT_INPUT = "/workspace/echoloc/Dataset/Novel/query_data/v2_2000/query_candidates.jsonl"
DEFAULT_OUTPUT = "/workspace/echoloc/Dataset/Novel/query_data/v2_2000/thinker_targets.jsonl"
DEFAULT_CONTROL_DIR = "/workspace/echoloc/Dataset/Novel/query_data/v2_2000/thinker_control2instruct"
EMO_TOKEN = "<|EMO_CHANGE|>"


def sanitize_id(qid: str) -> str:
    return re.sub(r"[^A-Za-z0-9_.-]+", "_", qid)[:180]


def iter_items(path: Path) -> Iterable[Dict[str, Any]]:
    with path.open("r", encoding="utf-8") as f:
        for line in f:
            if line.strip():
                yield json.loads(line)


def load_existing_qids(path: Path) -> set:
    qids = set()
    if not path.exists():
        return qids
    with path.open("r", encoding="utf-8") as f:
        for line in f:
            try:
                qid = json.loads(line).get("qid")
            except json.JSONDecodeError:
                continue
            if qid:
                qids.add(qid)
    return qids


def normalize_internal_instruct(text: Any) -> str:
    text = str(text or "").strip()
    text = re.sub(r"^\s*请(用|使用|以)?", "", text).strip()
    text = re.sub(r"^\s*以", "", text).strip()
    text = re.sub(r"，?进行(表达|朗读|演绎)\s*$", "", text).strip()
    return text


def build_messages(item: Dict[str, Any]) -> List[Dict[str, str]]:
    system = f"""
你是 Omni thinker 监督数据生成 Agent。输入是一条用户 query candidate；你要生成 thinker 应输出给 talker 的控制文本：
combined.instruct + combined.txt，以及 segments。

重要约束：
1. 你只能根据 visible_query.text 和未来 query 音频可能真实承载的声学信息来推断，不要把 hidden_context/source_event 里的细节当成用户一定说了出来。
2. hidden_context、source_event、target_contract 只能作为弱参考，帮助理解数据构造意图；如果 visible_query 没表达，就不要幻觉具体人物、关系、事件细节。
3. dialogue 类型：生成高情商 AI 助手回复。回复要接住用户情绪、清晰、有边界，不要角色续写，不要替用户编造事实。
4. instruction 类型：解析用户 TTS 指令，输出要朗读/表演的文本和对应声学控制；不要把“请用...说”整句当作要播报的内容。
5. 如果情绪/语气应发生变化，在 combined.txt 中用 {EMO_TOKEN} 标出自然分段位置。segments[].txt 不包含该 token。
6. 是否加入 {EMO_TOKEN} 由你根据 query 决定；不要为了形式强行加入。
7. combined.instruct 必须是中文自然语言声学控制，但应像 thinker 内部控制描述，不要像用户指令。格式接近：“年轻女性嗓音；平稳温和开场，逐渐转为坚定，最终以明亮上扬收束”。
8. instruct 字段不要使用祈使句，不要以“请、请用、请使用、请以”开头。
9. segments 建议 1-4 段。每段包含 instruct 和 txt。
10. 输出 thinker_reasoning，但它只用于分析，不一定进训练；要简洁说明你如何从 query 可见信息推断。

只输出 JSON，不要 Markdown。

输出格式：
{{
  "qid": "...",
  "query_type": "dialogue|instruction",
  "thinker_reasoning": "简要说明：可见文本/可听声学线索是什么，哪些信息不能幻觉，为什么这样回复/控制",
  "need_emochange": true,
  "segments": [
    {{"instruct": "句级声学控制", "txt": "该段要说的话"}}
  ],
  "combined": {{
    "instruct": "整体声学控制",
    "txt": "完整文本，必要时包含 <|EMO_CHANGE|>"
  }},
  "combined_no_speaker": {{
    "instruct": "不含具体音色人设的整体声学控制",
    "txt": "完整文本，必要时包含 <|EMO_CHANGE|>"
  }},
  "safety_notes": []
}}
"""
    payload = {
        "qid": item.get("qid"),
        "query_type": item.get("query_type"),
        "visible_query": item.get("visible_query"),
        "query_voice": item.get("query_voice"),
        "source_event": item.get("source_event"),
        "hidden_context": item.get("hidden_context"),
        "target_contract": item.get("target_contract"),
    }
    return [
        {"role": "system", "content": system.strip()},
        {"role": "user", "content": json.dumps(payload, ensure_ascii=False, indent=2)},
    ]


def normalize_target(item: Dict[str, Any], parsed: Dict[str, Any], model: str) -> Dict[str, Any]:
    qid = item.get("qid")
    parsed["qid"] = qid
    parsed["query_type"] = item.get("query_type")
    parsed["source_query"] = item.get("visible_query")
    parsed["source_event"] = item.get("source_event")
    parsed["generation"] = {
        "model": model,
        "prompt_version": "thinker_target_v1_visible_query_grounded",
    }
    for segment in parsed.get("segments") or []:
        if isinstance(segment, dict):
            segment["instruct"] = normalize_internal_instruct(segment.get("instruct", ""))
    for key in ["combined", "combined_no_speaker"]:
        value = parsed.get(key)
        if isinstance(value, dict):
            value["instruct"] = normalize_internal_instruct(value.get("instruct", ""))
    combined = parsed.get("combined") or {}
    txt = combined.get("txt", "")
    parsed["need_emochange"] = bool(parsed.get("need_emochange")) or (EMO_TOKEN in txt)
    return parsed


def generate_one(item: Dict[str, Any], args: argparse.Namespace) -> Dict[str, Any]:
    try:
        raw = chat_completion(
            build_messages(item),
            model=args.model,
            base_url=args.base_url,
            api_key_env=args.api_key_env,
            temperature=args.temperature,
            enable_thinking=args.enable_thinking,
            timeout=args.http_timeout,
        )
        parsed = extract_json_object(raw)
        return normalize_target(item, parsed, args.model)
    except Exception as exc:
        return {
            "qid": item.get("qid"),
            "query_type": item.get("query_type"),
            "status": "failed",
            "error": repr(exc),
        }


def append_jsonl(path: Path, rows: List[Dict[str, Any]]) -> None:
    with path.open("a", encoding="utf-8") as f:
        for row in rows:
            f.write(json.dumps(row, ensure_ascii=False) + "\n")


def export_control(target: Dict[str, Any], control_dir: Path) -> None:
    if target.get("status") == "failed":
        return
    qid = target.get("qid")
    if not qid:
        return
    out = control_dir / sanitize_id(qid)
    out.mkdir(parents=True, exist_ok=True)
    control = {
        "segments": target.get("segments", []),
        "combined": target.get("combined", {}),
        "combined_no_speaker": target.get("combined_no_speaker", {}),
        "meta": {
            "qid": qid,
            "query_type": target.get("query_type"),
            "thinker_reasoning": target.get("thinker_reasoning", ""),
            "source_query": target.get("source_query"),
        },
    }
    (out / "control2instruct.json").write_text(json.dumps(control, ensure_ascii=False, indent=2), encoding="utf-8")


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--input", default=DEFAULT_INPUT)
    parser.add_argument("--output", default=DEFAULT_OUTPUT)
    parser.add_argument("--control_dir", default=DEFAULT_CONTROL_DIR)
    parser.add_argument("--model", default="deepseek-v4-pro")
    parser.add_argument("--base_url", default=DEFAULT_BASE_URL)
    parser.add_argument("--api_key_env", default="DASHSCOPE_API_KEY")
    parser.add_argument("--temperature", type=float, default=0.4)
    parser.add_argument("--http_timeout", type=int, default=180)
    parser.add_argument("--num_workers", type=int, default=4)
    parser.add_argument("--limit", type=int, default=0)
    parser.add_argument("--resume", action="store_true")
    parser.add_argument("--enable_thinking", action="store_true")
    parser.add_argument("--export_controls", action="store_true")
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    input_path = Path(args.input)
    output_path = Path(args.output)
    output_path.parent.mkdir(parents=True, exist_ok=True)
    control_dir = Path(args.control_dir)
    if args.export_controls:
        control_dir.mkdir(parents=True, exist_ok=True)

    done = load_existing_qids(output_path) if args.resume else set()
    items = [x for x in iter_items(input_path) if x.get("qid") not in done]
    if args.limit > 0:
        items = items[: args.limit]
    print(f"[INFO] target_items={len(items)} resume_done={len(done)} output={output_path}", flush=True)

    with ThreadPoolExecutor(max_workers=max(1, args.num_workers)) as executor:
        pending = {}
        it = iter(items)

        def submit_next() -> bool:
            try:
                item = next(it)
            except StopIteration:
                return False
            pending[executor.submit(generate_one, item, args)] = item.get("qid")
            return True

        while len(pending) < args.num_workers and submit_next():
            pass
        done_count = 0
        while pending:
            done_futures, _ = wait(pending, return_when=FIRST_COMPLETED)
            rows = []
            for fut in done_futures:
                pending.pop(fut, None)
                row = fut.result()
                rows.append(row)
                if args.export_controls:
                    export_control(row, control_dir)
                done_count += 1
                submit_next()
            append_jsonl(output_path, rows)
            print(f"[PROGRESS] generated={done_count}/{len(items)}", flush=True)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())