"""Run a small DAG-aware LLM annotation pass for DMHY filename suffixes. This runner is intentionally separate from the prefix-tree annotation scripts. It selects shared DAG suffix units, asks an OpenAI-compatible Responses API for strict JSON annotations, falls back to deterministic heuristics on failure, and emits both terminal patch rows and dmhy_weak-compatible sample records. """ from __future__ import annotations import argparse import json import os import sys import time import urllib.error import urllib.request from concurrent.futures import ThreadPoolExecutor, as_completed from dataclasses import dataclass from datetime import datetime, timezone from pathlib import Path from typing import Any, Iterable from anifilebert.tokenizer import AnimeTokenizer from tools.annotate_dmhy_prefix_graph import ( heuristic_patch, normalize_generated_tokens, string_list, unique_keep_order, write_jsonl, ) from tools.dmhy_dataset import weak_label_filename DEFAULT_DAG = Path("datasets/AnimeName/dmhy_prefix_dag.json") DEFAULT_OUTPUT = Path("datasets/AnimeName/dmhy_weak.dag_sample.jsonl") DEFAULT_PATCH_OUTPUT = Path("datasets/AnimeName/dmhy_prefix_dag.annotations.dag_sample.jsonl") DEFAULT_UNITS_OUTPUT = Path("datasets/AnimeName/dmhy_prefix_dag.annotation_units.dag_sample.jsonl") DEFAULT_BASE_URL = "http://10.137.32.209/v1" DEFAULT_MODEL = "gpt-5.4-mini" LLM_SOURCE = "responses-dag-v1" HEURISTIC_SOURCE = "heuristic-dag-v1" COVERAGE_SHARED_ONLY = "shared-only" COVERAGE_ALL_TERMINALS = "all-terminals" @dataclass(frozen=True) class Args: dag: Path output: Path patch_output: Path units_output: Path | None limit: int max_requests: int units_per_request: int max_context_units: int workers: int coverage_mode: str min_incoming_count: int min_reachable_terminals: int max_reachable_terminals: int | None example_count: int max_records: int records_per_terminal: int terminal_sink_threshold: int terminal_cluster_size: int llm: bool base_url: str api_key: str | None model: str http_timeout: int retries: int preserve_i_labels: bool resume: bool def parse_args() -> Args: parser = argparse.ArgumentParser( description="Annotate a small sample of shared DMHY prefix-DAG suffix units" ) parser.add_argument("--dag", type=Path, default=DEFAULT_DAG) parser.add_argument("--output", type=Path, default=DEFAULT_OUTPUT) parser.add_argument("--patch-output", type=Path, default=DEFAULT_PATCH_OUTPUT) parser.add_argument( "--units-output", default=str(DEFAULT_UNITS_OUTPUT), help="Optional selected annotation units JSONL; use empty string to disable", ) parser.add_argument("--limit", type=int, default=20, help="Maximum DAG units to annotate") parser.add_argument("--max-requests", type=int, default=20, help="Maximum LLM requests") parser.add_argument("--units-per-request", type=int, default=10, help="DAG units per Responses API request") parser.add_argument( "--max-context-units", type=int, default=20, help="Maximum unit contexts included in one Responses API request", ) parser.add_argument("--workers", type=int, default=8, help="Concurrent LLM workers") parser.add_argument( "--coverage-mode", choices=[COVERAGE_SHARED_ONLY, COVERAGE_ALL_TERMINALS], default=COVERAGE_SHARED_ONLY, help="Select only shared DAG units, or cover every terminal with terminal_cluster fallback units", ) parser.add_argument("--min-incoming-count", type=int, default=2) parser.add_argument("--min-reachable-terminals", type=int, default=2) parser.add_argument( "--max-reachable-terminals", type=int, default=500, help="Skip very large shared units for sample runs; use 0 to disable", ) parser.add_argument("--example-count", type=int, default=8) parser.add_argument("--max-records", type=int, default=200, help="Maximum dataset records to emit") parser.add_argument("--records-per-terminal", type=int, default=1) parser.add_argument( "--terminal-sink-threshold", type=int, default=1000, help="Do not collect terminals from no-outgoing sink nodes above this size", ) parser.add_argument( "--terminal-cluster-size", type=int, default=1, help="Number of uncovered terminals per all-terminals fallback cluster", ) parser.add_argument("--llm", action="store_true", help="Use the Responses API; otherwise heuristics only") parser.add_argument( "--base-url", default=os.environ.get("ANIFILEBERT_LLM_BASE_URL", DEFAULT_BASE_URL), help="OpenAI-compatible API base URL", ) parser.add_argument( "--api-key", default=os.environ.get("ANIFILEBERT_LLM_API_KEY"), help="API key; defaults to ANIFILEBERT_LLM_API_KEY", ) parser.add_argument("--model", default=DEFAULT_MODEL) parser.add_argument("--http-timeout", type=int, default=120) parser.add_argument("--retries", type=int, default=3) parser.add_argument("--preserve-i-labels", action="store_true") parser.add_argument("--resume", action="store_true", help="Skip ok/fallback unit_ids already in patch output") ns = parser.parse_args() units_arg = str(ns.units_output).strip() return Args( dag=ns.dag, output=ns.output, patch_output=ns.patch_output, units_output=Path(units_arg) if units_arg else None, limit=max(0, ns.limit), max_requests=max(0, ns.max_requests), units_per_request=max(1, ns.units_per_request), max_context_units=max(1, ns.max_context_units), workers=max(1, ns.workers), coverage_mode=ns.coverage_mode, min_incoming_count=max(1, ns.min_incoming_count), min_reachable_terminals=max(1, ns.min_reachable_terminals), max_reachable_terminals=None if ns.max_reachable_terminals <= 0 else ns.max_reachable_terminals, example_count=max(1, ns.example_count), max_records=max(0, ns.max_records), records_per_terminal=max(1, ns.records_per_terminal), terminal_sink_threshold=max(1, ns.terminal_sink_threshold), terminal_cluster_size=max(1, ns.terminal_cluster_size), llm=ns.llm, base_url=ns.base_url, api_key=ns.api_key, model=ns.model, http_timeout=max(1, ns.http_timeout), retries=max(1, ns.retries), preserve_i_labels=ns.preserve_i_labels, resume=ns.resume, ) 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) or 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 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 node_id(node: dict[str, Any], fallback: int) -> int: try: return int(node.get("id", fallback)) except (TypeError, ValueError): raise SystemExit(f"invalid node id: {node.get('id')!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 has_outgoing_edges(node: dict[str, Any]) -> bool: return any(isinstance(edge, dict) for edge in node.get("children") or []) def edge_labels(node: dict[str, Any], limit: int) -> list[str]: labels: list[str] = [] for edge in node.get("children") or []: if isinstance(edge, dict) and edge.get("label") is not None: labels.append(str(edge["label"])) return unique_keep_order(labels)[:limit] 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 index, node in enumerate(dag["nodes"]): if isinstance(node, dict): nodes[node_id(node, index)] = node terminals_by_node: dict[int, list[dict[str, Any]]] = {} for index, terminal in enumerate(dag["terminals"]): if not isinstance(terminal, dict): continue try: terminal_node_id = int(terminal.get("node_id")) except (TypeError, ValueError): continue row = dict(terminal) row["_terminal_id"] = terminal_id(row, index) row["_terminal_index"] = index terminals_by_node.setdefault(terminal_node_id, []).append(row) return nodes, terminals_by_node class ReachableTerminals: def __init__( self, nodes: dict[int, dict[str, Any]], terminals_by_node: dict[int, list[dict[str, Any]]], terminal_sink_threshold: int, ) -> None: self.nodes = nodes self.terminals_by_node = terminals_by_node self.terminal_sink_threshold = terminal_sink_threshold self.memo: dict[int, list[dict[str, Any]]] = {} self.visiting: set[int] = set() def get(self, start: int) -> list[dict[str, Any]]: if start in self.memo: return self.memo[start] if start in self.visiting: raise SystemExit(f"cycle detected while traversing DAG at node {start}") self.visiting.add(start) node = self.nodes.get(start, {}) local_terminals = self.terminals_by_node.get(start, []) is_large_terminal_sink = not has_outgoing_edges(node) and len(local_terminals) >= self.terminal_sink_threshold found = [] if is_large_terminal_sink else list(local_terminals) for edge in node.get("children") or []: if not isinstance(edge, dict): continue try: found.extend(self.get(int(edge.get("target")))) except (TypeError, ValueError): continue self.visiting.remove(start) deduped = dedupe_terminals(found) self.memo[start] = 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 aggregate_examples(terminals: Iterable[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 unique_keep_order(values)[:limit] def aggregate_terminal_field(terminals: Iterable[dict[str, Any]], key: str, limit: int) -> list[str]: values: list[str] = [] for terminal in terminals: value = terminal.get(key) if value is not None: values.append(str(value)) return unique_keep_order(values)[:limit] def heuristic_unit_patch(unit: dict[str, Any]) -> dict[str, Any]: episode_titles: list[str] = [] media_suffixes: list[str] = [] title_candidates: list[str] = [] needs_review = False for terminal in unit["_terminals"]: patch = heuristic_patch(terminal, int(terminal.get("_terminal_index", 0))) episode_titles.extend(patch["episode_title_suffixes"]) media_suffixes.extend(patch["media_suffixes"]) title_candidates.extend(patch["title_candidates"]) needs_review = needs_review or bool(patch["needs_llm_review"]) return { "unit_id": unit["unit_id"], "terminal_ids": unit["terminal_ids"], "episode_title_suffixes": unique_keep_order(episode_titles), "media_suffixes": unique_keep_order(media_suffixes), "title_candidates": unique_keep_order(title_candidates), "llm_label": None, "notes": f"heuristic fallback; terminals={len(unit['terminal_ids'])}; needs_review={needs_review}", } def make_unit( node: dict[str, Any], node_id_value: int, terminals: list[dict[str, Any]], 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": "shared_suffix", "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), "digit_skeleton_examples": aggregate_terminal_field(terminals, "digit_skeleton", 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), "_terminals": terminals, } def make_terminal_cluster_unit( terminals: list[dict[str, Any]], cluster_index: int, example_count: int, ) -> dict[str, Any]: terminal_ids = [str(terminal["_terminal_id"]) for terminal in terminals] reachable_weight = sum(int_field(terminal, "weight", int_field(terminal, "count", 1)) for terminal in terminals) return { "unit_id": f"terminal-cluster-{cluster_index}", "node_id": None, "kind": "terminal_cluster", "incoming_count": 0, "reachable_terminals": len(terminals), "reachable_weight": reachable_weight, "terminal_ids": terminal_ids, "prefix_examples": aggregate_examples(terminals, "prefix", example_count), "digit_skeleton_examples": aggregate_terminal_field(terminals, "digit_skeleton", example_count), "value_examples": aggregate_examples(terminals, "value_examples", example_count), "suffix_examples": aggregate_examples(terminals, "suffix_examples", example_count), "common_edge_labels": [], "_terminals": terminals, } def all_terminals(dag: dict[str, Any]) -> list[dict[str, Any]]: terminals: list[dict[str, Any]] = [] for index, terminal in enumerate(dag["terminals"]): if not isinstance(terminal, dict): continue row = dict(terminal) row["_terminal_id"] = terminal_id(row, index) row["_terminal_index"] = index terminals.append(row) return dedupe_terminals(terminals) def candidate_shared_units( dag: dict[str, Any], args: Args, ) -> list[dict[str, Any]]: nodes, terminals_by_node = build_indexes(dag) reachable = ReachableTerminals(nodes, terminals_by_node, args.terminal_sink_threshold) root = int(dag.get("root", 0) or 0) candidates: list[tuple[tuple[int, int, int, int], int, dict[str, Any]]] = [] for current_id, node in nodes.items(): if current_id == root or not has_outgoing_edges(node): continue incoming_count = int_field(node, "incoming_count") reachable_count = int_field(node, "reachable_terminals") if incoming_count < args.min_incoming_count: continue if reachable_count < args.min_reachable_terminals: continue if args.max_reachable_terminals is not None and reachable_count > args.max_reachable_terminals: continue sort_key = (-incoming_count, -int_field(node, "reachable_weight"), -reachable_count, current_id) candidates.append((sort_key, current_id, node)) candidates.sort(key=lambda item: item[0]) candidates.sort(key=lambda item: item[0]) units: list[dict[str, Any]] = [] for _sort_key, current_id, node in candidates: terminals = reachable.get(current_id) if not terminals: continue units.append(make_unit(node, current_id, terminals, args.example_count)) return units def unique_terminal_coverage(units: Iterable[dict[str, Any]]) -> set[str]: covered: set[str] = set() for unit in units: covered.update(str(terminal_id_value) for terminal_id_value in unit.get("terminal_ids") or []) return covered def shared_only_units(dag: dict[str, Any], args: Args) -> list[dict[str, Any]]: limit = args.limit if args.limit > 0 else 0 return candidate_shared_units(dag, args)[:limit] def all_terminal_units(dag: dict[str, Any], args: Args) -> list[dict[str, Any]]: terminals = all_terminals(dag) terminals_by_id = {str(terminal["_terminal_id"]): terminal for terminal in terminals} selected: list[dict[str, Any]] = [] covered: set[str] = set() for unit in candidate_shared_units(dag, args): unit_terminal_ids = set(unit["terminal_ids"]) if not unit_terminal_ids: continue if unit_terminal_ids & covered: continue selected.append(unit) covered.update(unit_terminal_ids) uncovered = [terminal for terminal in terminals if str(terminal["_terminal_id"]) not in covered] for start in range(0, len(uncovered), args.terminal_cluster_size): cluster_terminals = uncovered[start : start + args.terminal_cluster_size] selected.append( make_terminal_cluster_unit( cluster_terminals, cluster_index=(start // args.terminal_cluster_size) + 1, example_count=args.example_count, ) ) covered.update(str(terminal["_terminal_id"]) for terminal in cluster_terminals) expected_ids = set(terminals_by_id) if covered != expected_ids: missing = sorted(expected_ids - covered)[:10] extra = sorted(covered - expected_ids)[:10] raise SystemExit( "all-terminals coverage mismatch: " f"expected={len(expected_ids)} covered={len(covered)} missing={missing} extra={extra}" ) return selected def select_units(dag: dict[str, Any], args: Args) -> list[dict[str, Any]]: if args.coverage_mode == COVERAGE_ALL_TERMINALS: return all_terminal_units(dag, args) return shared_only_units(dag, args) def annotation_schema() -> dict[str, Any]: return { "type": "object", "additionalProperties": False, "required": [ "unit_id", "terminal_ids", "episode_title_suffixes", "media_suffixes", "title_candidates", "llm_label", "notes", ], "properties": { "unit_id": {"type": "string"}, "terminal_ids": {"type": "array", "items": {"type": "string"}}, "episode_title_suffixes": {"type": "array", "items": {"type": "string"}}, "media_suffixes": {"type": "array", "items": {"type": "string"}}, "title_candidates": {"type": "array", "items": {"type": "string"}}, "llm_label": {"type": ["string", "null"]}, "notes": {"type": "string"}, }, } def batch_annotation_schema() -> dict[str, Any]: return { "type": "object", "additionalProperties": False, "required": ["annotations"], "properties": { "annotations": { "type": "array", "items": annotation_schema(), } }, } def responses_url(base_url: str) -> str: return base_url.rstrip("/") + "/responses" def unit_request_context(unit: dict[str, Any]) -> dict[str, Any]: return { "unit_id": unit["unit_id"], "node_id": unit["node_id"], "terminal_ids": unit["terminal_ids"], "incoming_count": unit["incoming_count"], "reachable_terminals": unit["reachable_terminals"], "reachable_weight": unit["reachable_weight"], "prefix_examples": unit["prefix_examples"], "value_examples": unit["value_examples"], "suffix_examples": unit["suffix_examples"], "common_edge_labels": unit["common_edge_labels"], "heuristic_patch": heuristic_unit_patch(unit), } def request_variants(units: list[dict[str, Any]], args: Args, previous_error: str | None = None) -> list[dict[str, Any]]: if not units: raise ValueError("request_variants requires at least one unit") instructions = ( "You annotate anime release filename DAG suffix units. Classify only suffix text that is " "a human episode title into episode_title_suffixes. Put resolution, codec, source, audio, " "subtitle, language, hash, release, and file-format fragments into media_suffixes. " "Use title_candidates for possible anime title strings found in prefix/value examples. Return " "strict JSON matching the supplied schema with one annotation per input unit. Keep unit_id and " "terminal_ids exactly unchanged for every unit." ) if previous_error: instructions += f" Previous attempt failed validation: {previous_error}. Correct only that issue." input_obj = { "units": [unit_request_context(unit) for unit in units[: args.max_context_units]], } schema = batch_annotation_schema() base = { "model": args.model, "instructions": instructions, "input": json.dumps(input_obj, ensure_ascii=False), } return [ { **base, "reasoning": {"effort": "medium"}, "text": { "format": { "type": "json_schema", "name": "dmhy_dag_batch_annotation", "strict": True, "schema": schema, } }, }, { **base, "text": { "format": { "type": "json_schema", "name": "dmhy_dag_batch_annotation", "strict": True, "schema": schema, } }, }, { **base, "tools": [ { "type": "function", "name": "submit_dmhy_dag_batch_annotation", "strict": True, "parameters": schema, "description": "Submit strict DMHY DAG suffix annotations for all input units.", } ], "tool_choice": {"type": "function", "name": "submit_dmhy_dag_batch_annotation"}, }, ] def http_post_json(payload: dict[str, Any], args: Args) -> dict[str, Any]: if not args.api_key: raise RuntimeError("--llm requires --api-key or ANIFILEBERT_LLM_API_KEY") request = urllib.request.Request( responses_url(args.base_url), data=json.dumps(payload, ensure_ascii=False).encode("utf-8"), headers={ "Authorization": f"Bearer {args.api_key}", "Content-Type": "application/json", }, method="POST", ) try: with urllib.request.urlopen(request, timeout=args.http_timeout) as response: raw = response.read().decode("utf-8") except urllib.error.HTTPError as exc: body = exc.read().decode("utf-8", errors="replace") raise RuntimeError(f"Responses API HTTP {exc.code}: {body[:500]}") from exc except (urllib.error.URLError, OSError) as exc: raise RuntimeError(f"Responses API request failed: {exc}") from exc try: data = json.loads(raw) except json.JSONDecodeError as exc: raise RuntimeError(f"Responses API returned invalid JSON envelope: {raw[:500]}") from exc if not isinstance(data, dict): raise RuntimeError("Responses API envelope must be a JSON object") return data def extract_annotation_payload(data: dict[str, Any]) -> dict[str, Any]: output_text = data.get("output_text") if isinstance(output_text, str) and output_text.strip(): return json.loads(strip_json_fence(output_text)) text_chunks: list[str] = [] for item in data.get("output") or []: if not isinstance(item, dict): continue if item.get("type") in {"function_call", "tool_call"}: arguments = item.get("arguments") if isinstance(arguments, str) and arguments.strip(): return json.loads(arguments) if isinstance(arguments, dict): return arguments for content in item.get("content") or []: if not isinstance(content, dict): continue text = content.get("text") if isinstance(text, str): text_chunks.append(text) if content.get("type") in {"function_call", "tool_call"}: arguments = content.get("arguments") if isinstance(arguments, str) and arguments.strip(): return json.loads(arguments) if isinstance(arguments, dict): return arguments if text_chunks: return json.loads(strip_json_fence("\n".join(text_chunks))) raise RuntimeError("Responses API did not include annotation text or tool arguments") def strip_json_fence(text: str) -> str: text = text.strip() if text.startswith("```"): lines = text.splitlines() if lines: lines = lines[1:] if lines and lines[-1].strip() == "```": lines = lines[:-1] return "\n".join(lines).strip() return text def validate_annotation(annotation: Any, unit: dict[str, Any]) -> dict[str, Any]: if not isinstance(annotation, dict): raise RuntimeError("annotation must be a JSON object") missing = [key for key in annotation_schema()["required"] if key not in annotation] if missing: raise RuntimeError(f"annotation missing required keys: {', '.join(missing)}") if annotation["unit_id"] != unit["unit_id"]: raise RuntimeError(f"unit_id mismatch: {annotation['unit_id']!r}") returned_ids = [str(item) for item in annotation.get("terminal_ids") or []] if returned_ids != unit["terminal_ids"]: raise RuntimeError("terminal_ids mismatch") cleaned = { "unit_id": unit["unit_id"], "terminal_ids": unit["terminal_ids"], "episode_title_suffixes": unique_keep_order(str(item) for item in annotation["episode_title_suffixes"]), "media_suffixes": unique_keep_order(str(item) for item in annotation["media_suffixes"]), "title_candidates": unique_keep_order(str(item) for item in annotation["title_candidates"]), "llm_label": annotation["llm_label"] if annotation["llm_label"] is None else str(annotation["llm_label"]), "notes": str(annotation["notes"]), } return cleaned def validate_batch_annotation(payload: Any, units: list[dict[str, Any]]) -> list[dict[str, Any]]: if not isinstance(payload, dict): raise RuntimeError("batch annotation must be a JSON object") annotations = payload.get("annotations") if not isinstance(annotations, list): raise RuntimeError("batch annotation missing annotations list") expected_ids = [unit["unit_id"] for unit in units] returned_ids = [ str(annotation.get("unit_id")) for annotation in annotations if isinstance(annotation, dict) and annotation.get("unit_id") is not None ] if sorted(returned_ids) != sorted(expected_ids) or len(returned_ids) != len(expected_ids): raise RuntimeError(f"unit_id set mismatch: expected={expected_ids!r} returned={returned_ids!r}") annotations_by_id: dict[str, Any] = {} for annotation in annotations: if not isinstance(annotation, dict): raise RuntimeError("annotations entries must be JSON objects") unit_id_value = str(annotation.get("unit_id")) if unit_id_value in annotations_by_id: raise RuntimeError(f"duplicate annotation unit_id: {unit_id_value}") annotations_by_id[unit_id_value] = annotation cleaned: list[dict[str, Any]] = [] for unit in units: cleaned.append(validate_annotation(annotations_by_id[unit["unit_id"]], unit)) return cleaned def annotate_batch(units: list[dict[str, Any]], args: Args, use_llm: bool) -> list[dict[str, Any]]: if not use_llm: patches: list[dict[str, Any]] = [] for unit in units: patch = heuristic_unit_patch(unit) patch.update({"source": HEURISTIC_SOURCE, "fallback": True, "attempts": 0, "errors": []}) patches.append(patch) return patches errors: list[str] = [] for attempt in range(args.retries): variants = request_variants(units, args, errors[-1] if errors else None) payload = variants[min(attempt, len(variants) - 1)] try: data = http_post_json(payload, args) annotations = validate_batch_annotation(extract_annotation_payload(data), units) for annotation in annotations: annotation.update({"source": LLM_SOURCE, "fallback": False, "attempts": attempt + 1, "errors": []}) return annotations except (RuntimeError, json.JSONDecodeError, TypeError, ValueError) as exc: errors.append(str(exc)) time.sleep(min(2.0, 0.25 * (attempt + 1))) patches = [] for unit in units: patch = heuristic_unit_patch(unit) patch.update( { "source": HEURISTIC_SOURCE, "fallback": True, "attempts": args.retries, "errors": errors[-3:], "notes": f"{patch['notes']}; llm_failed_after={args.retries}", } ) patches.append(patch) return patches def annotate_unit(unit: dict[str, Any], args: Args, use_llm: bool) -> dict[str, Any]: return annotate_batch([unit], args, use_llm)[0] def chunk_units(units: list[dict[str, Any]], size: int) -> list[list[dict[str, Any]]]: return [units[index : index + size] for index in range(0, len(units), size)] def load_resume_patches(path: Path) -> dict[str, dict[str, Any]]: if not path.exists(): return {} completed: dict[str, dict[str, Any]] = {} for line_number, line in enumerate(path.read_text(encoding="utf-8").splitlines(), 1): if not line.strip(): continue try: row = json.loads(line) except json.JSONDecodeError as exc: raise SystemExit(f"invalid resume patch JSON in {path}:{line_number}: {exc}") from exc if not isinstance(row, dict): continue unit_id_value = row.get("unit_id") if not unit_id_value: continue status = row.get("status") is_legacy_completed = status is None and ("source" in row or "fallback" in row) is_completed = status in {"ok", "fallback"} or is_legacy_completed if is_completed: completed[str(unit_id_value)] = row return completed def patch_from_resume_row(unit: dict[str, Any], row: dict[str, Any]) -> dict[str, Any]: return { "unit_id": unit["unit_id"], "terminal_ids": unit["terminal_ids"], "episode_title_suffixes": unique_keep_order(str(item) for item in row.get("episode_title_suffixes") or []), "media_suffixes": unique_keep_order(str(item) for item in row.get("media_suffixes") or []), "title_candidates": unique_keep_order(str(item) for item in row.get("title_candidates") or []), "llm_label": row.get("llm_label"), "notes": str(row.get("notes", "")), "source": str(row.get("source") or HEURISTIC_SOURCE), "fallback": bool(row.get("fallback", row.get("status") == "fallback")), "attempts": int_field(row, "attempts"), "errors": list(row.get("errors") or []), } def patch_row(unit: dict[str, Any], patch: dict[str, Any]) -> dict[str, Any]: fallback = bool(patch["fallback"]) return { "unit_id": unit["unit_id"], "node_id": unit["node_id"], "kind": unit["kind"], "incoming_count": unit["incoming_count"], "reachable_terminals": unit["reachable_terminals"], "reachable_weight": unit["reachable_weight"], "terminal_ids": patch["terminal_ids"], "episode_title_suffixes": patch["episode_title_suffixes"], "media_suffixes": patch["media_suffixes"], "title_candidates": patch["title_candidates"], "llm_label": patch["llm_label"], "notes": patch["notes"], "source": patch["source"], "fallback": fallback, "status": "fallback" if fallback else "ok", "attempts": patch["attempts"], "errors": patch["errors"], "annotated_at": datetime.now(timezone.utc).isoformat(), } def unit_source_id(unit: dict[str, Any]) -> str: node_id_value = unit.get("node_id") if node_id_value is None: return str(unit["unit_id"]) return str(node_id_value) def unit_public_row(unit: dict[str, Any]) -> dict[str, Any]: return { "unit_id": unit["unit_id"], "source_kind": "prefix_dag", "source_id": unit_source_id(unit), "terminal_ids": unit["terminal_ids"], "weight": int(unit["reachable_weight"]), "context": { "prefixes": unit["prefix_examples"], "digit_skeletons": unit["digit_skeleton_examples"], "edge_labels": unit["common_edge_labels"], "notes": ( f"kind={unit['kind']}; incoming_count={unit['incoming_count']}; " f"reachable_terminals={unit['reachable_terminals']}" ), }, "examples": { "values": unit["value_examples"], "suffixes": unit["suffix_examples"], }, "expected_output": { "schema_version": "dmhy-annotation-v1", }, } def dataset_records( units: list[dict[str, Any]], patches: list[dict[str, Any]], args: Args, ) -> list[dict[str, Any]]: tokenizer = AnimeTokenizer() records: list[dict[str, Any]] = [] seen: set[str] = set() for unit, patch in zip(units, patches): per_terminal_counts: dict[str, int] = {} for terminal in unit["_terminals"]: terminal_id_value = str(terminal["_terminal_id"]) for source_index, filename in enumerate(string_list(terminal.get("value_examples"))): if len(records) >= args.max_records: return records if per_terminal_counts.get(terminal_id_value, 0) >= args.records_per_terminal: break if filename in seen: continue seen.add(filename) sample = weak_label_filename(filename, tokenizer) if sample is None: continue tokens, labels = normalize_generated_tokens( sample["tokens"], sample["labels"], preserve_i_labels=args.preserve_i_labels, ) records.append( { "file_id": f"prefix-dag:{unit['unit_id']}:{terminal_id_value}:{source_index}", "filename": filename, "tokens": tokens, "labels": labels, "terminal_id": terminal_id_value, "terminal_index": terminal["_terminal_index"], "unit_id": unit["unit_id"], "source": patch["source"], "needs_llm_review": patch["fallback"], "episode_title_suffixes": patch["episode_title_suffixes"], "media_suffixes": patch["media_suffixes"], "title_candidates": patch["title_candidates"], "annotations": { "unit_id": unit["unit_id"], "terminal_id": terminal_id_value, "terminal_index": terminal["_terminal_index"], "source": patch["source"], "fallback": patch["fallback"], "episode_title_suffixes": patch["episode_title_suffixes"], "media_suffixes": patch["media_suffixes"], "title_candidates": patch["title_candidates"], "llm_label": patch["llm_label"], "notes": patch["notes"], }, } ) per_terminal_counts[terminal_id_value] = per_terminal_counts.get(terminal_id_value, 0) + 1 return records def annotate_units( units: list[dict[str, Any]], args: Args, ) -> tuple[list[tuple[dict[str, Any], dict[str, Any]]], int, int, int]: resume_rows = load_resume_patches(args.patch_output) if args.resume else {} results: list[dict[str, Any] | None] = [None] * len(units) pending: list[tuple[int, dict[str, Any]]] = [] skipped = 0 for index, unit in enumerate(units): row = resume_rows.get(unit["unit_id"]) if row is not None and [str(item) for item in row.get("terminal_ids") or []] == unit["terminal_ids"]: results[index] = patch_from_resume_row(unit, row) skipped += 1 else: pending.append((index, unit)) batch_size = max(1, min(args.units_per_request, args.max_context_units)) batches = [ ([index for index, _unit in batch], [unit for _index, unit in batch]) for batch in chunk_units(pending, batch_size) ] if args.llm: runnable_batches = batches[: args.max_requests] llm_requested = len(runnable_batches) else: runnable_batches = batches llm_requested = 0 with ThreadPoolExecutor(max_workers=args.workers) as executor: futures = { executor.submit(annotate_batch, batch_units, args, args.llm): indices for indices, batch_units in runnable_batches } for future in as_completed(futures): indices = futures[future] for index, patch in zip(indices, future.result()): results[index] = patch completed = [ (unit, patch) for unit, patch in zip(units, results) if patch is not None ] pending_unprocessed = len(units) - len(completed) return completed, skipped, llm_requested, pending_unprocessed def main() -> None: args = parse_args() if args.llm and not args.api_key: raise SystemExit("--llm requires --api-key or ANIFILEBERT_LLM_API_KEY") dag = load_dag(args.dag) units = select_units(dag, args) if not units: raise SystemExit("no DAG units selected; adjust --limit/--min-incoming-count") completed_pairs, resume_skipped, llm_requested, pending_unprocessed = annotate_units(units, args) completed_units = [unit for unit, _patch in completed_pairs] patches = [patch for _unit, patch in completed_pairs] patch_rows = [patch_row(unit, patch) for unit, patch in completed_pairs] records = dataset_records(completed_units, patches, args) patch_count = write_jsonl(args.patch_output, patch_rows) record_count = write_jsonl(args.output, records) unit_count = 0 if args.units_output is not None: unit_count = write_jsonl(args.units_output, (unit_public_row(unit) for unit in units)) fallback_count = sum(1 for patch in patches if patch["fallback"]) terminal_count = len(all_terminals(dag)) unique_coverage_count = len(unique_terminal_coverage(units)) summary = { "dag": str(args.dag), "coverage_mode": args.coverage_mode, "units_output": str(args.units_output) if args.units_output is not None else None, "patch_output": str(args.patch_output), "output": str(args.output), "selected_units": len(units), "dag_terminals": terminal_count, "unique_terminal_coverage": unique_coverage_count, "written_units": unit_count, "written_patches": patch_count, "written_records": record_count, "llm_requested": llm_requested, "pending_unprocessed": pending_unprocessed, "units_per_request": args.units_per_request, "max_context_units": args.max_context_units, "resume": args.resume, "resume_skipped": resume_skipped, "fallback_count": fallback_count, "workers": args.workers, "model": args.model if args.llm else None, } print(json.dumps(summary, ensure_ascii=False, indent=2)) if fallback_count: print(f"warning: {fallback_count} unit(s) used heuristic fallback", file=sys.stderr) if __name__ == "__main__": main()