AniFileBERT / tools /annotate_dmhy_prefix_tree_agent.py
ModerRAS's picture
Add DMHY annotation workflow helpers
7509455
Raw
History Blame Contribute Delete
26.7 kB
"""Agentic strict LLM annotator for DMHY prefix-tree terminals.
This runner is intentionally separate from ``annotate_dmhy_prefix_graph.py`` so
prefix-tree and DAG experiments can evolve independently. It reads the existing
prefix graph plus the full DMHY source list, emits terminal annotation patches,
and writes dmhy_weak-compatible dataset JSONL records.
"""
from __future__ import annotations
import argparse
import concurrent.futures
import json
import os
import re
import sys
import urllib.error
import urllib.request
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Iterable
from anifilebert.tokenizer import AnimeTokenizer
from tools.annotate_dmhy_prefix_graph import (
DEFAULT_GRAPH,
DEFAULT_MODEL,
DEFAULT_OUTPUT,
DEFAULT_PATCH_OUTPUT,
DEFAULT_SOURCE_LIST,
LLM_SOURCE,
dataset_records,
heuristic_patch,
load_graph,
selected_terminals,
source_list_matches,
string_list,
)
DEFAULT_BASE_URL = "http://10.137.32.209/v1"
DEFAULT_UNITS_OUTPUT = Path("datasets/AnimeName/dmhy_prefix_graph.annotation_units.prefix_sample.jsonl")
NON_TITLE_NUMERIC_FRAGMENT_RE = re.compile(
r"(?ix)^\s*"
r"(?:"
r"[\[(【《]?\s*(?:EP?|\#)?\d{1,4}(?:\s*(?:v|ver|version|rev)\s*\d{1,3})?\s*[\])】》]?|"
r"(?:v|ver|version|rev)\s*\d{1,3}|"
r"vol(?:ume)?\.?\s*\d{1,3}|"
r"S\d{1,2}E\d{1,4}|"
r"\d{1,2}x\d{1,4}"
r")"
r"\s*$"
)
ANNOTATION_SCHEMA: dict[str, Any] = {
"type": "object",
"additionalProperties": False,
"required": [
"terminal_id",
"episode_title_suffixes",
"media_suffixes",
"title_candidates",
"llm_label",
"notes",
],
"properties": {
"terminal_id": {"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"],
"description": "Short classification such as accept, media_only, episode_title, mixed, or uncertain.",
},
"notes": {"type": "string"},
},
}
@dataclass(frozen=True)
class Args:
graph: Path
source_list: Path
output: Path
patch_output: Path | None
units_output: Path | None
limit: int | None
min_weight: int | None
only_needs_review: bool
llm: bool
base_url: str
api_key: str | None
model: str
max_requests: int | None
http_timeout: int
preserve_i_labels: bool
examples_only: bool
workers: int
retries: int
resume: bool
failure_output: Path | None = None
reasoning_effort: str = "medium"
def parse_args() -> Args:
parser = argparse.ArgumentParser(
description="Strict Responses API agent for annotating DMHY prefix-tree terminals"
)
parser.add_argument("--graph", type=Path, default=DEFAULT_GRAPH)
parser.add_argument("--source-list", type=Path, default=DEFAULT_SOURCE_LIST)
parser.add_argument("--output", type=Path, default=DEFAULT_OUTPUT)
parser.add_argument("--patch-output", default=str(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=None)
parser.add_argument("--min-weight", type=int, default=None)
parser.add_argument("--only-needs-review", action="store_true")
parser.add_argument("--llm", action="store_true")
parser.add_argument(
"--base-url",
default=os.environ.get("ANIFILEBERT_LLM_BASE_URL", DEFAULT_BASE_URL),
help="OpenAI-compatible Responses API base URL",
)
parser.add_argument(
"--api-key",
default=os.environ.get("ANIFILEBERT_LLM_API_KEY"),
help="API key; prefer ANIFILEBERT_LLM_API_KEY",
)
parser.add_argument("--model", default=DEFAULT_MODEL)
parser.add_argument("--max-requests", type=int, default=None)
parser.add_argument("--http-timeout", type=int, default=120)
parser.add_argument("--preserve-i-labels", action="store_true")
parser.add_argument("--examples-only", action="store_true")
parser.add_argument("--workers", type=int, default=8, help="Concurrent LLM terminal patch workers")
parser.add_argument("--retries", type=int, default=3, help="LLM validation retry attempts per terminal")
parser.add_argument("--resume", action="store_true", help="Skip ok LLM terminal patches already in patch output")
parser.add_argument(
"--reasoning-effort",
choices=("minimal", "low", "medium", "high", "xhigh"),
default="medium",
help="Responses API reasoning effort when the target model supports reasoning",
)
parser.add_argument(
"--failure-output",
default="",
help="Optional JSONL path for attempted terminals that failed LLM validation/request retries",
)
ns = parser.parse_args()
patch_arg = str(ns.patch_output).strip()
units_arg = str(ns.units_output).strip()
failure_arg = str(ns.failure_output).strip()
return Args(
graph=ns.graph,
source_list=ns.source_list,
output=ns.output,
patch_output=Path(patch_arg) if patch_arg else None,
units_output=Path(units_arg) if units_arg else None,
limit=ns.limit,
min_weight=ns.min_weight,
only_needs_review=ns.only_needs_review,
llm=ns.llm,
base_url=ns.base_url,
api_key=ns.api_key,
model=ns.model,
max_requests=max(0, ns.max_requests) if ns.max_requests is not None else None,
http_timeout=ns.http_timeout,
preserve_i_labels=ns.preserve_i_labels,
examples_only=ns.examples_only,
workers=max(1, ns.workers),
retries=max(1, ns.retries),
resume=ns.resume,
failure_output=Path(failure_arg) if failure_arg else None,
reasoning_effort=ns.reasoning_effort,
)
def responses_url(base_url: str) -> str:
return base_url.rstrip("/") + "/responses"
def redact_secrets(text: str) -> str:
return re.sub(r"sk-[A-Za-z0-9._-]+", "sk-REDACTED", text)
def unique_strings(value: Any, *, field: str) -> list[str]:
if not isinstance(value, list):
raise ValueError(f"{field} must be a list")
seen: set[str] = set()
result: list[str] = []
for item in value:
if not isinstance(item, str):
raise ValueError(f"{field} items must be strings")
cleaned = " ".join(item.split()).strip()
if cleaned and cleaned not in seen:
seen.add(cleaned)
result.append(cleaned)
return result
def validate_strict_annotation(annotation: Any, *, expected_terminal_id: str) -> dict[str, Any]:
if not isinstance(annotation, dict):
raise ValueError("annotation must be a JSON object")
terminal_id = annotation.get("terminal_id")
if not isinstance(terminal_id, str):
raise ValueError("terminal_id must be a string")
if terminal_id != expected_terminal_id:
raise ValueError(f"terminal_id mismatch: expected {expected_terminal_id}, got {terminal_id}")
llm_label = annotation.get("llm_label")
if llm_label is not None and not isinstance(llm_label, str):
raise ValueError("llm_label must be a string or null")
notes = annotation.get("notes")
if not isinstance(notes, str):
raise ValueError("notes must be a string")
validated = {
"terminal_id": terminal_id,
"episode_title_suffixes": unique_strings(annotation.get("episode_title_suffixes"), field="episode_title_suffixes"),
"media_suffixes": unique_strings(annotation.get("media_suffixes"), field="media_suffixes"),
"title_candidates": unique_strings(annotation.get("title_candidates"), field="title_candidates"),
"llm_label": llm_label,
"notes": " ".join(notes.split()).strip(),
}
errors = annotation_quality_errors(validated)
if errors:
raise ValueError("; ".join(errors))
return validated
def annotation_quality_errors(annotation: dict[str, Any]) -> list[str]:
errors: list[str] = []
for field in ("episode_title_suffixes", "title_candidates"):
values = annotation.get(field)
if not isinstance(values, list):
continue
bad_values = [
str(value)
for value in values
if isinstance(value, str) and NON_TITLE_NUMERIC_FRAGMENT_RE.fullmatch(value)
]
if bad_values:
errors.append(
f"{field} must not contain pure episode/version/volume fragments: {bad_values[:5]}"
)
return errors
def extract_response_annotation(data: dict[str, Any]) -> Any:
output_text = data.get("output_text")
if isinstance(output_text, str) and output_text.strip():
return json.loads(output_text)
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):
return json.loads(arguments)
if isinstance(arguments, dict):
return arguments
for content in item.get("content") or []:
if not isinstance(content, dict):
continue
if content.get("type") in {"output_json", "json"} and "json" in content:
return content["json"]
text = content.get("text")
if isinstance(text, str) and text.strip():
return json.loads(text)
raise ValueError("Responses API returned no annotation payload")
def format_error_hint(error: str | None) -> str:
if not error:
return ""
return (
"\nPrevious attempt failed validation. Correct the issue exactly and submit a new schema-valid "
f"annotation. Error: {error}"
)
def request_body(
*,
args: Args,
terminal: dict[str, Any],
patch: dict[str, Any],
previous_error: str | None,
mode: str,
) -> dict[str, Any]:
task = {
"terminal_id": patch["terminal_id"],
"prefix": terminal.get("prefix"),
"digit_skeleton": terminal.get("digit_skeleton"),
"suffix_examples": terminal.get("suffix_examples") or [],
"value_examples": terminal.get("value_examples") or [],
"heuristic_patch": {
"episode_title_suffixes": patch.get("episode_title_suffixes") or [],
"media_suffixes": patch.get("media_suffixes") or [],
"title_candidates": patch.get("title_candidates") or [],
"notes": patch.get("notes") or "",
},
}
instructions = (
"Annotate one anime filename prefix-tree terminal. Distinguish human episode-title suffix text "
"from media metadata such as resolution, source, codec, audio, subtitles, language, hashes, and "
"release tags. Do not put pure episode numbers, episode versions like 01v2, version tokens like "
"ver1, or volume markers like Vol.1 in episode_title_suffixes or title_candidates; those are "
"episode/version structure, not human title text. Keep the same terminal_id. Return only the "
"required strict schema data."
+ format_error_hint(previous_error)
)
body: dict[str, Any] = {
"model": args.model,
"input": [
{"role": "developer", "content": instructions},
{"role": "user", "content": json.dumps(task, ensure_ascii=False)},
],
}
if "reasoning" in mode:
body["reasoning"] = {"effort": args.reasoning_effort}
if "json_schema" in mode:
body["text"] = {
"format": {
"type": "json_schema",
"name": "dmhy_prefix_terminal_annotation",
"strict": True,
"schema": ANNOTATION_SCHEMA,
}
}
else:
body["tools"] = [
{
"type": "function",
"name": "submit_annotation",
"description": "Submit the validated DMHY terminal annotation.",
"strict": True,
"parameters": ANNOTATION_SCHEMA,
}
]
body["tool_choice"] = {"type": "function", "name": "submit_annotation"}
return body
def post_responses(body: 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(body, 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:
return json.loads(response.read().decode("utf-8"))
except urllib.error.HTTPError as exc:
detail = exc.read().decode("utf-8", errors="replace")
raise RuntimeError(f"Responses API HTTP {exc.code}: {redact_secrets(detail[:500])}") from exc
except urllib.error.URLError as exc:
raise RuntimeError(f"Responses API request failed: {redact_secrets(str(exc.reason))}") from exc
except json.JSONDecodeError as exc:
raise RuntimeError("Responses API returned invalid response JSON") from exc
def next_mode(mode: str, error: str) -> str:
lowered = error.lower()
if mode == "reasoning_json_schema" and "reasoning" in lowered:
return "json_schema"
if "json_schema" in mode and any(token in lowered for token in ("text.format", "json_schema", "schema")):
return "function_tool"
if mode == "reasoning_json_schema":
return "json_schema"
return mode
def merge_llm_annotation(patch: dict[str, Any], annotation: dict[str, Any]) -> dict[str, Any]:
merged = dict(patch)
merged.update(
{
"episode_title_suffixes": annotation["episode_title_suffixes"],
"media_suffixes": annotation["media_suffixes"],
"title_candidates": annotation["title_candidates"],
"llm_label": annotation["llm_label"],
"notes": annotation["notes"],
"source": LLM_SOURCE,
"status": "ok",
"fallback": False,
"errors": [],
}
)
return merged
def annotate_terminal_with_llm(
terminal: dict[str, Any],
patch: dict[str, Any],
args: Args,
) -> tuple[dict[str, Any], bool]:
mode = "reasoning_json_schema"
previous_error: str | None = None
errors: list[str] = []
for _attempt in range(args.retries):
try:
body = request_body(args=args, terminal=terminal, patch=patch, previous_error=previous_error, mode=mode)
data = post_responses(body, args)
raw_annotation = extract_response_annotation(data)
annotation = validate_strict_annotation(raw_annotation, expected_terminal_id=str(patch["terminal_id"]))
return merge_llm_annotation(patch, annotation), False
except (RuntimeError, ValueError, json.JSONDecodeError) as exc:
error = redact_secrets(str(exc))
errors.append(error)
mode = next_mode(mode, error)
previous_error = error
fallback = dict(patch)
fallback["notes"] = f"{fallback.get('notes') or ''}; llm_error={' | '.join(errors)}"
fallback["status"] = "fallback"
fallback["fallback"] = True
fallback["errors"] = errors[-3:]
return fallback, True
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 is_completed_resume_row(row: dict[str, Any]) -> bool:
status = row.get("status")
source = str(row.get("source") or "")
if status == "ok" and source.startswith("responses"):
return not annotation_quality_errors(row)
return status is None and source == LLM_SOURCE
def load_resume_patches(path: Path | None) -> dict[str, dict[str, Any]]:
if path is None or 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) or not is_completed_resume_row(row):
continue
keys: list[str] = []
for key in ("terminal_id", "unit_id"):
value = row.get(key)
if isinstance(value, (str, int)) and not isinstance(value, bool):
keys.append(str(value))
terminal_ids = row.get("terminal_ids")
if isinstance(terminal_ids, list):
keys.extend(str(item) for item in terminal_ids if isinstance(item, (str, int)) and not isinstance(item, bool))
for key in keys:
completed[key] = row
return completed
def patch_from_resume_row(patch: dict[str, Any], row: dict[str, Any]) -> dict[str, Any]:
merged = dict(patch)
for key in ("episode_title_suffixes", "media_suffixes", "title_candidates"):
merged[key] = unique_strings(row.get(key) or [], field=key)
if "llm_label" in row:
merged["llm_label"] = row.get("llm_label")
if "notes" in row:
merged["notes"] = str(row.get("notes") or "")
merged["source"] = str(row.get("source") or patch.get("source") or LLM_SOURCE)
merged["fallback"] = False
merged["status"] = "ok"
errors = row.get("errors")
merged["errors"] = errors if isinstance(errors, list) else []
return enrich_patch(merged)
def resume_row_for_patch(
resume_rows: dict[str, dict[str, Any]],
patch: dict[str, Any],
) -> dict[str, Any] | None:
for key in (patch.get("terminal_id"), patch.get("unit_id")):
if key is None:
continue
row = resume_rows.get(str(key))
if row is not None:
return row
return None
def annotate_selected(
selected: list[tuple[int, dict[str, Any], dict[str, Any]]],
args: Args,
) -> tuple[list[tuple[int, int, dict[str, Any], dict[str, Any]]], int, int, int, int]:
resume_rows = load_resume_patches(args.patch_output) if args.resume else {}
results: list[dict[str, Any] | None] = [None] * len(selected)
pending_llm: list[tuple[int, dict[str, Any], dict[str, Any]]] = []
resume_skipped = 0
for ordinal, (_index, _terminal, patch) in enumerate(selected):
row = resume_row_for_patch(resume_rows, patch)
if row is not None:
results[ordinal] = patch_from_resume_row(patch, row)
resume_skipped += 1
continue
if not args.llm:
results[ordinal] = patch
continue
pending_llm.append((ordinal, selected[ordinal][1], patch))
candidates = pending_llm
if args.max_requests is not None:
candidates = candidates[: args.max_requests]
fallback_count = 0
failure_rows: list[dict[str, Any]] = []
if candidates:
with concurrent.futures.ThreadPoolExecutor(max_workers=args.workers) as executor:
future_to_ordinal = {
executor.submit(annotate_terminal_with_llm, terminal, patch, args): ordinal
for ordinal, terminal, patch in candidates
}
for future in concurrent.futures.as_completed(future_to_ordinal):
ordinal = future_to_ordinal[future]
try:
patch, used_fallback = future.result()
except Exception as exc: # defensive; worker should already convert errors to fallback
patch = dict(selected[ordinal][2])
patch["errors"] = [redact_secrets(str(exc))]
used_fallback = True
if used_fallback:
fallback_count += 1
terminal_id = selected[ordinal][2]["terminal_id"]
failure_rows.append(llm_failure_row(selected[ordinal][1], patch, args))
print(f"warning: terminal {terminal_id}: LLM failed; leaving pending", file=sys.stderr)
continue
results[ordinal] = patch
if args.failure_output is not None:
write_jsonl(args.failure_output, failure_rows)
completed = [
(ordinal, selected[ordinal][0], selected[ordinal][1], patch)
for ordinal, patch in enumerate(results)
if patch is not None
]
pending_unprocessed = len(selected) - len(completed)
return completed, resume_skipped, len(candidates), fallback_count, pending_unprocessed
def llm_failure_row(terminal: dict[str, Any], patch: dict[str, Any], args: Args) -> dict[str, Any]:
terminal_id = str(patch.get("terminal_id") or terminal.get("terminal_id") or "")
errors = patch.get("errors")
return {
"terminal_id": terminal_id,
"unit_id": str(patch.get("unit_id") or terminal_id),
"model": args.model,
"status": "pending",
"source": LLM_SOURCE,
"errors": errors if isinstance(errors, list) else [],
"prefix": terminal.get("prefix") if isinstance(terminal.get("prefix"), str) else None,
"digit_skeleton": terminal.get("digit_skeleton") if isinstance(terminal.get("digit_skeleton"), str) else None,
"suffix_examples": string_list(terminal.get("suffix_examples")),
"value_examples": string_list(terminal.get("value_examples")),
}
def source_id(terminal: dict[str, Any], index: int, terminal_id: str) -> str | int:
for key in ("node_id", "id"):
value = terminal.get(key)
if isinstance(value, (str, int)) and not isinstance(value, bool):
return value
return terminal_id or index
def int_weight(terminal: dict[str, Any]) -> int:
try:
return int(terminal.get("weight") or terminal.get("count") or 0)
except (TypeError, ValueError):
return 0
def enrich_patch(patch: dict[str, Any]) -> dict[str, Any]:
terminal_id = str(patch["terminal_id"])
enriched = dict(patch)
enriched.setdefault("unit_id", terminal_id)
enriched.setdefault("terminal_ids", [terminal_id])
enriched.setdefault("status", "ok")
enriched.setdefault("fallback", False)
enriched.setdefault("errors", [])
return enriched
def unit_public_row(index: int, terminal: dict[str, Any], patch: dict[str, Any]) -> dict[str, Any]:
terminal_id = str(patch["terminal_id"])
return {
"unit_id": str(patch.get("unit_id") or terminal_id),
"source_kind": "prefix_tree",
"source_id": source_id(terminal, index, terminal_id),
"terminal_ids": [terminal_id],
"weight": int_weight(terminal),
"context": {
"prefixes": string_list([terminal.get("prefix")] if terminal.get("prefix") is not None else []),
"digit_skeletons": string_list(
[terminal.get("digit_skeleton")] if terminal.get("digit_skeleton") is not None else []
),
"edge_labels": [],
"notes": patch.get("notes") if isinstance(patch.get("notes"), str) else None,
},
"examples": {
"values": string_list(terminal.get("value_examples")),
"suffixes": string_list(terminal.get("suffix_examples")),
},
"expected_output": {"schema_version": "dmhy-annotation-v1"},
}
def main() -> None:
args = parse_args()
graph = load_graph(args.graph)
selected = selected_terminals(graph, args)
if not selected:
raise SystemExit("no terminals selected; adjust --limit/--min-weight/--only-needs-review")
selected = [(index, terminal, enrich_patch(patch)) for index, terminal, patch in selected]
completed, resume_skipped, llm_requests, llm_fallbacks, pending_unprocessed = annotate_selected(selected, args)
tokenizer = AnimeTokenizer()
source_matches = None if args.examples_only else source_list_matches(args.source_list, selected)
records: list[dict[str, Any]] = []
patches: list[dict[str, Any]] = []
completed_by_ordinal = {ordinal: patch for ordinal, _index, _terminal, patch in completed}
for ordinal, index, terminal, patch in completed:
patches.append(patch)
records.extend(
dataset_records(
terminal,
index,
patch,
tokenizer,
filenames=None if args.examples_only else source_matches.get(ordinal, []),
preserve_i_labels=args.preserve_i_labels,
)
)
record_count = write_jsonl(args.output, records)
patch_count = 0
if args.patch_output is not None:
patch_count = write_jsonl(args.patch_output, patches)
unit_count = 0
if args.units_output is not None:
unit_count = write_jsonl(
args.units_output,
(
unit_public_row(index, terminal, completed_by_ordinal.get(ordinal, patch))
for ordinal, (index, terminal, patch) in enumerate(selected)
),
)
summary = {
"graph": str(args.graph),
"source_list": None if args.examples_only else str(args.source_list),
"output": str(args.output),
"patch_output": str(args.patch_output) if args.patch_output is not None else None,
"units_output": str(args.units_output) if args.units_output is not None else None,
"selected_terminals": len(selected),
"examples_only": args.examples_only,
"dataset_records": record_count,
"patches": patch_count,
"units": unit_count,
"llm_enabled": args.llm,
"llm_requests": llm_requests,
"llm_fallbacks": llm_fallbacks,
"resume": args.resume,
"resume_skipped": resume_skipped,
"pending_unprocessed": pending_unprocessed,
"workers": args.workers,
}
print(json.dumps(summary, ensure_ascii=False, indent=2))
if __name__ == "__main__":
main()