"""Build annotation units from a DMHY prefix DAG. The DAG producer keeps repeated suffix structure shared across many raw prefixes. This tool turns those shared nodes into compact, traceable units for LLM or human review without calling any remote service. """ from __future__ import annotations import argparse import json from dataclasses import dataclass from pathlib import Path from typing import Any, Iterable from tools.annotate_dmhy_prefix_graph import ( heuristic_patch, string_list, ) DEFAULT_DAG = Path("datasets/AnimeName/dmhy_prefix_dag.json") DEFAULT_OUTPUT = Path("datasets/AnimeName/dmhy_prefix_dag.annotation_units.jsonl") @dataclass(frozen=True) class Args: dag: Path output: Path min_reachable_terminals: int min_incoming_count: int limit: int | None example_count: int include_terminal_sink_units: bool def parse_args() -> Args: parser = argparse.ArgumentParser( description="Emit DAG-aware DMHY prefix annotation units as JSONL" ) parser.add_argument("--dag", type=Path, default=DEFAULT_DAG, help="Input dmhy_prefix_dag.json") parser.add_argument( "--output", type=Path, default=DEFAULT_OUTPUT, help="Output annotation unit JSONL", ) parser.add_argument( "--min-reachable-terminals", type=int, default=2, help="Select non-root nodes reaching at least this many terminals", ) parser.add_argument( "--min-incoming-count", type=int, default=2, help="Select nodes with at least this many incoming DAG edges", ) parser.add_argument("--limit", type=int, default=None, help="Maximum units to write") parser.add_argument( "--example-count", type=int, default=5, help="Maximum examples retained per example field", ) parser.add_argument( "--include-terminal-sink-units", action="store_true", help="Include terminal nodes with no outgoing edges as annotation units", ) ns = parser.parse_args() return Args( dag=ns.dag, output=ns.output, min_reachable_terminals=max(1, ns.min_reachable_terminals), min_incoming_count=max(1, ns.min_incoming_count), limit=ns.limit, example_count=max(1, ns.example_count), include_terminal_sink_units=ns.include_terminal_sink_units, ) def load_dag(path: Path) -> dict[str, Any]: if not path.exists(): raise SystemExit(f"DAG not found: {path}") try: dag = json.loads(path.read_text(encoding="utf-8")) except json.JSONDecodeError as exc: raise SystemExit(f"invalid DAG JSON in {path}: {exc}") from exc if not isinstance(dag, dict): raise SystemExit(f"invalid DAG schema in {path}: root must be an object") if not isinstance(dag.get("nodes"), list): raise SystemExit(f"invalid DAG schema in {path}: missing nodes list") if not isinstance(dag.get("terminals"), list): raise SystemExit(f"invalid DAG schema in {path}: missing terminals list") return dag def node_id(node: dict[str, Any], fallback: int) -> int: value = node.get("id", fallback) try: return int(value) except (TypeError, ValueError): raise SystemExit(f"invalid node id: {value!r}") from None def terminal_id(terminal: dict[str, Any], fallback: int) -> str: value = terminal.get("terminal_id", terminal.get("id", fallback)) return str(value) def int_field(row: dict[str, Any], key: str, default: int = 0) -> int: try: return int(row.get(key, default) or default) except (TypeError, ValueError): return default def build_indexes(dag: dict[str, Any]) -> tuple[dict[int, dict[str, Any]], dict[int, list[dict[str, Any]]]]: nodes: dict[int, dict[str, Any]] = {} for fallback, node in enumerate(dag["nodes"]): if not isinstance(node, dict): continue nodes[node_id(node, fallback)] = node terminals_by_node: dict[int, list[dict[str, Any]]] = {} for fallback, terminal in enumerate(dag["terminals"]): if not isinstance(terminal, dict): continue try: terminal_node_id = int(terminal.get("node_id")) except (TypeError, ValueError): continue terminal = dict(terminal) terminal["_terminal_id"] = terminal_id(terminal, fallback) terminal["_terminal_index"] = fallback terminals_by_node.setdefault(terminal_node_id, []).append(terminal) return nodes, terminals_by_node class ReachableTerminalIndex: def __init__( self, nodes: dict[int, dict[str, Any]], terminals_by_node: dict[int, list[dict[str, Any]]], ) -> None: self.nodes = nodes self.terminals_by_node = terminals_by_node self.memo: dict[int, list[dict[str, Any]]] = {} self.visiting: set[int] = set() def get(self, start_node_id: int) -> list[dict[str, Any]]: if start_node_id in self.memo: return self.memo[start_node_id] if start_node_id in self.visiting: raise SystemExit(f"cycle detected while traversing DAG at node {start_node_id}") self.visiting.add(start_node_id) found = list(self.terminals_by_node.get(start_node_id, [])) node = self.nodes.get(start_node_id, {}) for edge in node.get("children") or []: if not isinstance(edge, dict): continue try: target = int(edge.get("target")) except (TypeError, ValueError): continue found.extend(self.get(target)) self.visiting.remove(start_node_id) deduped = dedupe_terminals(found) self.memo[start_node_id] = deduped return deduped def dedupe_terminals(terminals: Iterable[dict[str, Any]]) -> list[dict[str, Any]]: seen: set[str] = set() result: list[dict[str, Any]] = [] for terminal in terminals: tid = str(terminal.get("_terminal_id") or terminal.get("terminal_id") or "") if not tid or tid in seen: continue seen.add(tid) result.append(terminal) return result def limited_unique(values: Iterable[str], limit: int) -> list[str]: seen: set[str] = set() result: list[str] = [] for value in values: if not value or not value.strip() or value in seen: continue seen.add(value) result.append(value) if len(result) >= limit: break return result def edge_labels(node: dict[str, Any], limit: int) -> list[str]: labels = [] for edge in node.get("children") or []: if isinstance(edge, dict) and edge.get("label") is not None: labels.append(str(edge["label"])) return limited_unique(labels, limit) def has_outgoing_edges(node: dict[str, Any]) -> bool: return any(isinstance(edge, dict) for edge in node.get("children") or []) def aggregate_examples(terminals: list[dict[str, Any]], key: str, limit: int) -> list[str]: values: list[str] = [] for terminal in terminals: if key == "prefix": value = terminal.get("prefix") if value is not None: values.append(str(value)) else: values.extend(string_list(terminal.get(key))) return limited_unique(values, limit) def aggregate_needs_review(terminals: list[dict[str, Any]]) -> bool: for index, terminal in enumerate(terminals): annotations = terminal.get("annotations") if isinstance(annotations, dict) and annotations.get("needs_llm_review") is not None: if bool(annotations["needs_llm_review"]): return True continue if heuristic_patch(terminal, index)["needs_llm_review"]: return True return False def annotation_template() -> dict[str, Any]: return { "episode_title_suffixes": [], "media_suffixes": [], "title_candidates": [], "notes": None, } def recommended_prompt(kind: str, terminal_count: int) -> str: if kind == "shared_suffix": return ( "Review the shared DAG suffix examples and mark episode-title text, media metadata, " f"and possible title candidates for {terminal_count} linked terminals." ) return "Review this terminal cluster and mark episode-title text, media metadata, and title candidates." def make_unit( node: dict[str, Any], node_id_value: int, terminals: list[dict[str, Any]], kind: str, example_count: int, ) -> dict[str, Any]: terminal_ids = [str(terminal["_terminal_id"]) for terminal in terminals] reachable_weight = int_field(node, "reachable_weight") if reachable_weight <= 0: reachable_weight = sum(int_field(terminal, "weight", int_field(terminal, "count", 1)) for terminal in terminals) return { "unit_id": f"dag-node-{node_id_value}", "node_id": node_id_value, "kind": kind, "incoming_count": int_field(node, "incoming_count"), "reachable_terminals": len(terminals), "reachable_weight": reachable_weight, "terminal_ids": terminal_ids, "prefix_examples": aggregate_examples(terminals, "prefix", example_count), "value_examples": aggregate_examples(terminals, "value_examples", example_count), "suffix_examples": aggregate_examples(terminals, "suffix_examples", example_count), "common_edge_labels": edge_labels(node, example_count), "needs_llm_review": aggregate_needs_review(terminals), "recommended_prompt": recommended_prompt(kind, len(terminals)), "annotations": annotation_template(), } def node_reachable_count(node: dict[str, Any]) -> int: return int_field(node, "reachable_terminals") def node_reachable_weight(node: dict[str, Any]) -> int: return int_field(node, "reachable_weight") def selected_units(dag: dict[str, Any], args: Args) -> list[dict[str, Any]]: nodes, terminals_by_node = build_indexes(dag) reachable_index = ReachableTerminalIndex(nodes, terminals_by_node) root = int(dag.get("root", 0) or 0) shared_candidates: list[tuple[tuple[int, int, int, int, int], int, dict[str, Any]]] = [] for current_id, node in nodes.items(): if not args.include_terminal_sink_units and not has_outgoing_edges(node): continue incoming_count = int_field(node, "incoming_count") reachable_count = node_reachable_count(node) by_shared_incoming = incoming_count >= args.min_incoming_count by_reachable = current_id != root and reachable_count >= args.min_reachable_terminals if by_shared_incoming or by_reachable: sort_key = ( 0, -node_reachable_weight(node), -reachable_count, -incoming_count, current_id, ) shared_candidates.append((sort_key, current_id, node)) shared_candidates.sort(key=lambda item: item[0]) candidates: list[tuple[tuple[int, int, int, int, int], dict[str, Any]]] = [] for sort_key, current_id, node in shared_candidates: if args.limit is not None and len(candidates) >= args.limit: break terminals = reachable_index.get(current_id) if not terminals: continue unit = make_unit(node, current_id, terminals, "shared_suffix", args.example_count) candidates.append((sort_key, unit)) if args.limit is not None and len(candidates) >= args.limit: return [unit for _sort_key, unit in candidates[: max(0, args.limit)]] covered_terminal_ids = { terminal_id for _sort_key, unit in candidates for terminal_id in unit["terminal_ids"] } for current_id, terminals in terminals_by_node.items(): node = nodes.get(current_id, {"id": current_id}) if not args.include_terminal_sink_units and not has_outgoing_edges(node): continue uncovered = [terminal for terminal in terminals if terminal["_terminal_id"] not in covered_terminal_ids] if not uncovered: continue unit = make_unit(node, current_id, uncovered, "terminal_cluster", args.example_count) sort_key = ( 1, -unit["reachable_weight"], -unit["reachable_terminals"], -unit["incoming_count"], current_id, ) candidates.append((sort_key, unit)) candidates.sort(key=lambda item: item[0]) units = [unit for _sort_key, unit in candidates] if args.limit is not None: units = units[: max(0, args.limit)] return units def write_jsonl(path: Path, rows: Iterable[dict[str, Any]]) -> int: path.parent.mkdir(parents=True, exist_ok=True) count = 0 with path.open("w", encoding="utf-8", newline="\n") as handle: for row in rows: handle.write(json.dumps(row, ensure_ascii=False, separators=(",", ":")) + "\n") count += 1 return count def main() -> None: args = parse_args() dag = load_dag(args.dag) units = selected_units(dag, args) count = write_jsonl(args.output, units) summary = { "dag": str(args.dag), "output": str(args.output), "annotation_units": count, "min_reachable_terminals": args.min_reachable_terminals, "min_incoming_count": args.min_incoming_count, "example_count": args.example_count, } print(json.dumps(summary, ensure_ascii=False, indent=2)) if __name__ == "__main__": main()