| |
| |
| """Prepare and assemble iter006 data from iter005 Chinese assets plus new English assets.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import shutil |
| from pathlib import Path |
| from typing import Any, Dict, Iterable, List, Set |
|
|
|
|
| DEFAULT_ITER005 = Path("/workspace/echoloc/Dataset/Novel/query_data/supervisor_iterations/iter_005") |
| DEFAULT_ITER006 = Path("/workspace/echoloc/Dataset/Novel/query_data/supervisor_iterations/iter_006") |
|
|
|
|
| def read_jsonl(path: Path) -> List[Dict[str, Any]]: |
| rows: List[Dict[str, Any]] = [] |
| if not path.exists(): |
| return rows |
| with path.open("r", encoding="utf-8") as handle: |
| for line in handle: |
| line = line.strip() |
| if line: |
| try: |
| rows.append(json.loads(line)) |
| except json.JSONDecodeError: |
| rows.append(json.loads(line.lstrip("\ufeff；;"))) |
| return rows |
|
|
|
|
| def write_jsonl(path: Path, rows: Iterable[Dict[str, Any]]) -> int: |
| path.parent.mkdir(parents=True, exist_ok=True) |
| n = 0 |
| with path.open("w", encoding="utf-8") as handle: |
| for row in rows: |
| handle.write(json.dumps(row, ensure_ascii=False) + "\n") |
| n += 1 |
| return n |
|
|
|
|
| def qid(row: Dict[str, Any]) -> str: |
| return str(row.get("qid") or row.get("id") or "") |
|
|
|
|
| def is_vstyle_en(row: Dict[str, Any]) -> bool: |
| cur = qid(row) |
| if cur.startswith("vstyle_en_"): |
| return True |
| language = str(row.get("language") or ((row.get("source") or {}).get("language")) or "").lower() |
| return language == "en" and cur.startswith("vstyle_") |
|
|
|
|
| def is_vstyle_zh(row: Dict[str, Any]) -> bool: |
| cur = qid(row) |
| if cur.startswith("vstyle_en_"): |
| return False |
| if cur.startswith("vstyle_zh_"): |
| return True |
| language = str(row.get("language") or ((row.get("source") or {}).get("language")) or "").lower() |
| return language == "zh" and cur.startswith("vstyle_") |
|
|
|
|
| def is_story_zh(row: Dict[str, Any]) -> bool: |
| cur = qid(row) |
| if cur.startswith("vstyle_") or cur.startswith("story_en_"): |
| return False |
| if "vstyle_en_" in cur: |
| return False |
| language = str(row.get("language") or ((row.get("source") or {}).get("language")) or "zh").lower() |
| return language == "zh" |
|
|
|
|
| def is_zh_reuse(row: Dict[str, Any]) -> bool: |
| return is_story_zh(row) or is_vstyle_zh(row) |
|
|
|
|
| def mark_language(row: Dict[str, Any], language: str) -> Dict[str, Any]: |
| row = dict(row) |
| row["language"] = language |
| source = row.get("source") |
| if isinstance(source, dict): |
| source = dict(source) |
| source["language"] = language |
| row["source"] = source |
| source_candidate = row.get("source_query_candidate") |
| if isinstance(source_candidate, dict): |
| source_candidate = dict(source_candidate) |
| source_candidate["language"] = language |
| row["source_query_candidate"] = source_candidate |
| return row |
|
|
|
|
| def by_qid(rows: Iterable[Dict[str, Any]]) -> Dict[str, Dict[str, Any]]: |
| out: Dict[str, Dict[str, Any]] = {} |
| for row in rows: |
| cur = qid(row) |
| if cur: |
| out[cur] = row |
| return out |
|
|
|
|
| def read_first_jsonl(paths: Iterable[Path]) -> List[Dict[str, Any]]: |
| for path in paths: |
| rows = read_jsonl(path) |
| if rows: |
| return rows |
| return [] |
|
|
|
|
| def copy_if_exists(src: Path, dst: Path) -> bool: |
| if not src.exists(): |
| return False |
| dst.parent.mkdir(parents=True, exist_ok=True) |
| shutil.copy2(src, dst) |
| return True |
|
|
|
|
| def extract_vstyle_en(iter005: Path, iter006: Path) -> Dict[str, Any]: |
| out_dir = iter006 / "english_rebuild" / "vstyle_en" |
| out_dir.mkdir(parents=True, exist_ok=True) |
|
|
| candidates = [mark_language(row, "en") for row in read_jsonl(iter005 / "filtered/query_candidates_keep.jsonl") if is_vstyle_en(row)] |
| controls = [mark_language(row, "en") for row in read_jsonl(iter005 / "filtered/query_tts_qwen3_input_structured.jsonl") if is_vstyle_en(row)] |
| eval_rows = [dict(row, language="en") for row in read_jsonl(iter005 / "eval_inputs/e2e_input_aligned.jsonl") if is_vstyle_en(row)] |
|
|
| counts = { |
| "vstyle_en_candidates": write_jsonl(out_dir / "query_candidates.jsonl", candidates), |
| "vstyle_en_candidates_keep": write_jsonl(out_dir / "filtered/query_candidates_keep.jsonl", candidates), |
| "vstyle_en_query_tts_controls": write_jsonl(out_dir / "filtered/query_tts_qwen3_input_structured.jsonl", controls), |
| "vstyle_en_existing_eval_rows": write_jsonl(out_dir / "existing_eval_audio_rows.jsonl", eval_rows), |
| } |
| (out_dir / "README_iter006_vstyle_en.txt").write_text( |
| "Use filtered/query_candidates_keep.jsonl as RESPONSE_IN for English response/thinker target regeneration. " |
| "existing_eval_audio_rows.jsonl preserves iter005 query_audio_path for vstyle_en.\n", |
| encoding="utf-8", |
| ) |
| return {"out_dir": str(out_dir), "counts": counts} |
|
|
|
|
| def copy_zh_reuse(iter005: Path, iter006: Path) -> Dict[str, Any]: |
| out_dir = iter006 / "zh_reuse_from_iter005" |
| out_dir.mkdir(parents=True, exist_ok=True) |
| specs = [ |
| ("query_candidates.jsonl", "query_candidates.jsonl"), |
| ("filtered/query_candidates_keep.jsonl", "filtered/query_candidates_keep.jsonl"), |
| ("filtered/query_tts_qwen3_input_structured.jsonl", "filtered/query_tts_qwen3_input_structured.jsonl"), |
| ("thinker_targets/response_controls.jsonl", "thinker_targets/response_controls.jsonl"), |
| ("thinker_targets/thinker_targets.jsonl", "thinker_targets/thinker_targets.jsonl"), |
| ("eval_inputs/e2e_input_aligned.jsonl", "eval_inputs/e2e_input_aligned.jsonl"), |
| ("eval_inputs/talker_input_all.jsonl", "eval_inputs/talker_input_all.jsonl"), |
| ("eval_inputs/response_tts_input_aligned.jsonl", "eval_inputs/response_tts_input_aligned.jsonl"), |
| ] |
| counts: Dict[str, Any] = {} |
| for rel_src, rel_dst in specs: |
| rows = [row for row in read_jsonl(iter005 / rel_src) if is_zh_reuse(row)] |
| counts[rel_dst] = write_jsonl(out_dir / rel_dst, rows) |
| for rel in ["sft/thinker_emochange/train.jsonl", "sft/thinker_emochange/val.jsonl", "sft/thinker_emochange/manifest.json"]: |
| counts[rel] = "not_copied_final_sft_rebuilt_from_merged_eval" |
| return {"out_dir": str(out_dir), "counts": counts} |
|
|
|
|
| def merge_unique(parts: Iterable[Iterable[Dict[str, Any]]]) -> List[Dict[str, Any]]: |
| out: Dict[str, Dict[str, Any]] = {} |
| for rows in parts: |
| for row in rows: |
| cur = qid(row) |
| if cur: |
| out[cur] = row |
| return [out[key] for key in sorted(out)] |
|
|
|
|
| def merge_eval_with_existing_audio(new_eval_rows: List[Dict[str, Any]], existing_audio_rows: List[Dict[str, Any]]) -> List[Dict[str, Any]]: |
| audio_by_qid = by_qid(existing_audio_rows) |
| merged = [] |
| for row in new_eval_rows: |
| cur = qid(row) |
| audio = audio_by_qid.get(cur) |
| if audio: |
| row = dict(row) |
| row["query_audio_path"] = audio.get("query_audio_path", row.get("query_audio_path", "")) |
| row["query_audio_control_path"] = audio.get("query_audio_control_path", row.get("query_audio_control_path", "")) |
| merged.append(row) |
| return merged |
|
|
|
|
| def merge_final(iter005: Path, iter006: Path, story_en_root: Path, vstyle_en_root: Path) -> Dict[str, Any]: |
| zh_root = iter006 / "zh_reuse_from_iter005" |
| final_eval = iter006 / "eval_inputs" |
| final_targets = iter006 / "thinker_targets" |
| final_filtered = iter006 / "filtered" |
|
|
| if not zh_root.exists(): |
| copy_zh_reuse(iter005, iter006) |
|
|
| zh_candidates = read_jsonl(zh_root / "query_candidates.jsonl") |
| zh_keep = read_jsonl(zh_root / "filtered/query_candidates_keep.jsonl") |
| zh_controls = read_jsonl(zh_root / "filtered/query_tts_qwen3_input_structured.jsonl") |
| zh_responses = read_jsonl(zh_root / "thinker_targets/response_controls.jsonl") |
| zh_targets = read_jsonl(zh_root / "thinker_targets/thinker_targets.jsonl") |
| zh_e2e = read_jsonl(zh_root / "eval_inputs/e2e_input_aligned.jsonl") |
| zh_talker = read_jsonl(zh_root / "eval_inputs/talker_input_all.jsonl") |
| zh_response_aligned = read_jsonl(zh_root / "eval_inputs/response_tts_input_aligned.jsonl") |
|
|
| story_candidates = read_jsonl(story_en_root / "query_candidates.jsonl") |
| story_keep = read_jsonl(story_en_root / "filtered/query_candidates_keep.jsonl") |
| story_controls = read_jsonl(story_en_root / "filtered/query_tts_qwen3_input_structured.jsonl") |
| story_responses = read_jsonl(story_en_root / "thinker_targets/response_controls.jsonl") |
| story_targets = read_jsonl(story_en_root / "thinker_targets/thinker_targets.jsonl") |
| story_e2e = read_jsonl(story_en_root / "eval_inputs/e2e_input_aligned.jsonl") |
| story_talker = read_jsonl(story_en_root / "eval_inputs/talker_input_all.jsonl") |
| story_response_aligned = read_jsonl(story_en_root / "eval_inputs/response_tts_input_aligned.jsonl") |
|
|
| vstyle_candidates = read_first_jsonl([vstyle_en_root / "filtered/query_candidates_keep.jsonl", vstyle_en_root / "query_candidates_keep.jsonl"]) |
| vstyle_controls = read_first_jsonl([vstyle_en_root / "filtered/query_tts_qwen3_input_structured.jsonl", vstyle_en_root / "query_tts_qwen3_input_structured.jsonl"]) |
| vstyle_responses = read_jsonl(vstyle_en_root / "thinker_targets/response_controls.jsonl") |
| vstyle_targets = read_jsonl(vstyle_en_root / "thinker_targets/thinker_targets.jsonl") |
| vstyle_export = read_jsonl(vstyle_en_root / "eval_inputs/e2e_input_aligned.jsonl") |
| vstyle_existing_audio = read_jsonl(vstyle_en_root / "existing_eval_audio_rows.jsonl") |
| if not vstyle_export and vstyle_responses and vstyle_targets and vstyle_existing_audio: |
| response_by_qid = by_qid(vstyle_responses) |
| audio_by_qid = by_qid(vstyle_existing_audio) |
| vstyle_export = [] |
| for target in vstyle_targets: |
| cur = qid(target) |
| response = response_by_qid.get(cur, {}) |
| audio = audio_by_qid.get(cur, {}) |
| style = str((target.get("combined") or {}).get("instruct") or "").strip() |
| text = str((target.get("combined") or {}).get("txt") or "").strip() |
| visible_query = response.get("visible_query") or (response.get("source_query_candidate") or {}).get("visible_query") or target.get("source_query") or {} |
| vstyle_export.append( |
| { |
| "id": cur, |
| "qid": cur, |
| "query_type": target.get("query_type") or response.get("query_type"), |
| "language": "en", |
| "query_text": visible_query.get("text", audio.get("query_text", "")) if isinstance(visible_query, dict) else "", |
| "query_audio_path": audio.get("query_audio_path", ""), |
| "query_audio_control_path": audio.get("query_audio_control_path", ""), |
| "oracle_thinker_style": style, |
| "oracle_thinker_text": text, |
| "oracle_response_text": response.get("audio_content", text.replace(" <|EMO_CHANGE|> ", " ")), |
| "response_control": response.get("final_generated_control") or response.get("response_generated_control"), |
| "source_event": response.get("source_event"), |
| } |
| ) |
| vstyle_e2e = merge_eval_with_existing_audio(vstyle_export, vstyle_existing_audio) |
| vstyle_talker = read_jsonl(vstyle_en_root / "eval_inputs/talker_input_all.jsonl") |
| vstyle_response_aligned = read_jsonl(vstyle_en_root / "eval_inputs/response_tts_input_aligned.jsonl") |
| if not vstyle_talker and vstyle_targets: |
| vstyle_talker = [] |
| for target in vstyle_targets: |
| cur = qid(target) |
| combined = target.get("combined") or {} |
| vstyle_talker.append( |
| { |
| "id": cur, |
| "qid": cur, |
| "language": "en", |
| "ability": "constructed_thinker_to_talker", |
| "query_type": target.get("query_type"), |
| "thinker_style": str(combined.get("instruct") or "").strip(), |
| "thinker_text": str(combined.get("txt") or "").strip(), |
| "need_emochange": target.get("need_emochange"), |
| "source_query": target.get("source_query"), |
| "aligned_with_query_audio": cur in {qid(row) for row in vstyle_e2e}, |
| } |
| ) |
| if not vstyle_response_aligned and vstyle_responses and vstyle_e2e: |
| e2e_qids = {qid(row) for row in vstyle_e2e} |
| vstyle_response_aligned = [row for row in vstyle_responses if qid(row) in e2e_qids] |
|
|
| counts = { |
| "query_candidates": write_jsonl(iter006 / "query_candidates.jsonl", merge_unique([zh_candidates, story_candidates, vstyle_candidates])), |
| "query_candidates_keep": write_jsonl(final_filtered / "query_candidates_keep.jsonl", merge_unique([zh_keep, story_keep, vstyle_candidates])), |
| "query_tts_controls": write_jsonl(final_filtered / "query_tts_qwen3_input_structured.jsonl", merge_unique([zh_controls, story_controls, vstyle_controls])), |
| "response_controls": write_jsonl(final_targets / "response_controls.jsonl", merge_unique([zh_responses, story_responses, vstyle_responses])), |
| "thinker_targets": write_jsonl(final_targets / "thinker_targets.jsonl", merge_unique([zh_targets, story_targets, vstyle_targets])), |
| "e2e_input_aligned": write_jsonl(final_eval / "e2e_input_aligned.jsonl", merge_unique([zh_e2e, story_e2e, vstyle_e2e])), |
| "talker_input_all": write_jsonl(final_eval / "talker_input_all.jsonl", merge_unique([zh_talker, story_talker, vstyle_talker])), |
| "response_tts_input_aligned": write_jsonl(final_eval / "response_tts_input_aligned.jsonl", merge_unique([zh_response_aligned, story_response_aligned, vstyle_response_aligned])), |
| } |
| (iter006 / "iter006_merge_summary.json").write_text(json.dumps(counts, ensure_ascii=False, indent=2), encoding="utf-8") |
| return {"iter006": str(iter006), "counts": counts} |
|
|
|
|
| def counts(iter006: Path) -> Dict[str, Any]: |
| paths = [ |
| "query_candidates.jsonl", |
| "filtered/query_candidates_keep.jsonl", |
| "filtered/query_tts_qwen3_input_structured.jsonl", |
| "thinker_targets/response_controls.jsonl", |
| "thinker_targets/thinker_targets.jsonl", |
| "eval_inputs/e2e_input_aligned.jsonl", |
| "eval_inputs/talker_input_all.jsonl", |
| "eval_inputs/response_tts_input_aligned.jsonl", |
| "sft/thinker_emochange/train.jsonl", |
| "sft/thinker_emochange/val.jsonl", |
| ] |
| return {rel: len(read_jsonl(iter006 / rel)) for rel in paths} |
|
|
|
|
| def main() -> int: |
| parser = argparse.ArgumentParser() |
| parser.add_argument("mode", choices=["copy-zh", "extract-vstyle-en", "merge-final", "counts"]) |
| parser.add_argument("--iter005", type=Path, default=DEFAULT_ITER005) |
| parser.add_argument("--iter006", type=Path, default=DEFAULT_ITER006) |
| parser.add_argument("--story_en_root", type=Path, default=DEFAULT_ITER006 / "english_rebuild/story_en") |
| parser.add_argument("--vstyle_en_root", type=Path, default=DEFAULT_ITER006 / "english_rebuild/vstyle_en") |
| args = parser.parse_args() |
|
|
| if args.mode == "copy-zh": |
| result = copy_zh_reuse(args.iter005, args.iter006) |
| elif args.mode == "extract-vstyle-en": |
| result = extract_vstyle_en(args.iter005, args.iter006) |
| elif args.mode == "merge-final": |
| result = merge_final(args.iter005, args.iter006, args.story_en_root, args.vstyle_en_root) |
| else: |
| result = {"iter006": str(args.iter006), "counts": counts(args.iter006)} |
| print(json.dumps(result, ensure_ascii=False, indent=2)) |
| return 0 |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|
