#!/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())