| |
| |
| """Create jsonl shards for unfinished e2e Omni inference.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| from pathlib import Path |
| from typing import Any, Dict, Iterable, List |
|
|
|
|
| BASE_DIR = Path("/workspace/echoloc/Dataset/Novel/query_data/v2_2000") |
| DEFAULT_INPUT = BASE_DIR / "eval_inputs/e2e_input_aligned.jsonl" |
| DEFAULT_OUT_DIR = BASE_DIR / "eval_inputs/e2e_remaining_shards" |
| DEFAULT_WAV_DIR = BASE_DIR / "model_outputs/qwen3omni_e2e_from_query_audio/wavs" |
|
|
|
|
| def iter_jsonl(path: Path) -> Iterable[Dict[str, Any]]: |
| with path.open("r", encoding="utf-8") as f: |
| for line in f: |
| line = line.strip() |
| if line: |
| yield json.loads(line) |
|
|
|
|
| def write_jsonl(path: Path, rows: List[Dict[str, Any]]) -> None: |
| with path.open("w", encoding="utf-8") as f: |
| for row in rows: |
| f.write(json.dumps(row, ensure_ascii=False) + "\n") |
|
|
|
|
| def main() -> None: |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--input_jsonl", type=Path, default=DEFAULT_INPUT) |
| ap.add_argument("--out_dir", type=Path, default=DEFAULT_OUT_DIR) |
| ap.add_argument("--wav_dir", type=Path, default=DEFAULT_WAV_DIR) |
| ap.add_argument("--num_shards", type=int, default=4) |
| args = ap.parse_args() |
|
|
| args.out_dir.mkdir(parents=True, exist_ok=True) |
| rows = list(iter_jsonl(args.input_jsonl)) |
| remaining = [ |
| row for row in rows |
| if not (args.wav_dir / f"{row['qid']}.wav").exists() |
| ] |
|
|
| shards: List[List[Dict[str, Any]]] = [[] for _ in range(args.num_shards)] |
| for i, row in enumerate(remaining): |
| shards[i % args.num_shards].append(row) |
|
|
| for i, shard in enumerate(shards): |
| write_jsonl(args.out_dir / f"e2e_remaining_shard_{i:02d}.jsonl", shard) |
|
|
| summary = { |
| "input_total": len(rows), |
| "done_wavs": len(rows) - len(remaining), |
| "remaining": len(remaining), |
| "num_shards": args.num_shards, |
| "shard_sizes": [len(x) for x in shards], |
| "out_dir": str(args.out_dir), |
| } |
| with (args.out_dir / "summary.json").open("w", encoding="utf-8") as f: |
| json.dump(summary, f, ensure_ascii=False, indent=2) |
| print(json.dumps(summary, ensure_ascii=False, indent=2)) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
