| """Smoke tests for annotated DMHY graph dataset helpers.""" |
|
|
| from __future__ import annotations |
|
|
| import tempfile |
| import json |
| import subprocess |
| import sys |
| import unittest |
| from pathlib import Path |
|
|
| from tools.annotate_dmhy_prefix_graph import normalize_generated_tokens |
| from tools.annotate_dmhy_prefix_tree_agent import ( |
| Args as PrefixTreeAgentArgs, |
| annotate_terminal_with_llm, |
| annotate_selected, |
| enrich_patch, |
| request_body, |
| validate_strict_annotation, |
| ) |
| from tools.convert_annotated_dmhy_dataset import ( |
| iter_validated_jsonl, |
| validate_record, |
| ) |
| from tools.convert_to_char_dataset import convert_record |
|
|
|
|
| class AnnotatedDmhyWorkflowTests(unittest.TestCase): |
| def prefix_tree_agent_args(self) -> PrefixTreeAgentArgs: |
| return PrefixTreeAgentArgs( |
| graph=Path("graph.json"), |
| source_list=Path("source.jsonl"), |
| output=Path("out.jsonl"), |
| patch_output=Path("patch.jsonl"), |
| units_output=Path("units.jsonl"), |
| limit=None, |
| min_weight=None, |
| only_needs_review=False, |
| llm=True, |
| base_url="http://example.test/v1", |
| api_key="test-key", |
| model="gpt-5.4-mini", |
| max_requests=None, |
| http_timeout=1, |
| preserve_i_labels=False, |
| examples_only=False, |
| workers=8, |
| retries=2, |
| resume=False, |
| reasoning_effort="medium", |
| ) |
|
|
| def prefix_selected(self, count: int = 3) -> list[tuple[int, dict, dict]]: |
| selected: list[tuple[int, dict, dict]] = [] |
| for index in range(count): |
| terminal_id = f"t{index}" |
| terminal = { |
| "terminal_id": terminal_id, |
| "prefix": f"Show {index} - ", |
| "suffix_examples": [f"0{index} - Title [1080P]"], |
| "value_examples": [f"Show {index} - 0{index} - Title [1080P].mkv"], |
| "weight": index + 1, |
| } |
| patch = enrich_patch( |
| { |
| "terminal_id": terminal_id, |
| "needs_llm_review": True, |
| "episode_title_suffixes": [], |
| "media_suffixes": ["[1080P]"], |
| "title_candidates": [], |
| "llm_label": None, |
| "notes": "heuristic", |
| "source": "heuristic-v1", |
| } |
| ) |
| selected.append((index, terminal, patch)) |
| return selected |
|
|
| def test_prefix_tree_agent_max_requests_does_not_fallback_unprocessed_terminals(self) -> None: |
| args = self.prefix_tree_agent_args() |
| args = PrefixTreeAgentArgs(**{**args.__dict__, "max_requests": 1, "workers": 1}) |
| selected = self.prefix_selected(3) |
|
|
| def fake_annotate(_terminal, patch, _args): |
| merged = dict(patch) |
| merged.update({"source": "responses-v1", "status": "ok", "fallback": False, "errors": []}) |
| return merged, False |
|
|
| import tools.annotate_dmhy_prefix_tree_agent as agent |
|
|
| original_annotate = agent.annotate_terminal_with_llm |
| try: |
| agent.annotate_terminal_with_llm = fake_annotate |
| completed, resume_skipped, llm_requests, llm_fallbacks, pending_unprocessed = annotate_selected(selected, args) |
| finally: |
| agent.annotate_terminal_with_llm = original_annotate |
|
|
| self.assertEqual([patch["terminal_id"] for _ord, _idx, _term, patch in completed], ["t0"]) |
| self.assertEqual(resume_skipped, 0) |
| self.assertEqual(llm_requests, 1) |
| self.assertEqual(llm_fallbacks, 0) |
| self.assertEqual(pending_unprocessed, 2) |
| self.assertFalse(any(patch["fallback"] for _ord, _idx, _term, patch in completed)) |
|
|
| def test_prefix_tree_agent_resume_continues_next_batch(self) -> None: |
| with tempfile.TemporaryDirectory() as tmpdir: |
| patch_path = Path(tmpdir) / "patches.jsonl" |
| patch_path.write_text( |
| json.dumps( |
| { |
| "terminal_id": "t0", |
| "unit_id": "t0", |
| "terminal_ids": ["t0"], |
| "episode_title_suffixes": ["old"], |
| "media_suffixes": ["[1080P]"], |
| "title_candidates": ["old"], |
| "llm_label": "old", |
| "notes": "previous", |
| "source": "responses-v1", |
| "status": "ok", |
| "fallback": False, |
| "errors": [], |
| }, |
| ensure_ascii=False, |
| separators=(",", ":"), |
| ) |
| + "\n", |
| encoding="utf-8", |
| ) |
| args = self.prefix_tree_agent_args() |
| args = PrefixTreeAgentArgs( |
| **{**args.__dict__, "patch_output": patch_path, "resume": True, "max_requests": 1, "workers": 1} |
| ) |
| selected = self.prefix_selected(3) |
|
|
| def fake_annotate(_terminal, patch, _args): |
| merged = dict(patch) |
| merged.update( |
| { |
| "episode_title_suffixes": [patch["terminal_id"]], |
| "source": "responses-v1", |
| "status": "ok", |
| "fallback": False, |
| "errors": [], |
| } |
| ) |
| return merged, False |
|
|
| import tools.annotate_dmhy_prefix_tree_agent as agent |
|
|
| original_annotate = agent.annotate_terminal_with_llm |
| try: |
| agent.annotate_terminal_with_llm = fake_annotate |
| completed, resume_skipped, llm_requests, _llm_fallbacks, pending_unprocessed = annotate_selected( |
| selected, args |
| ) |
| finally: |
| agent.annotate_terminal_with_llm = original_annotate |
|
|
| patches = [patch for _ord, _idx, _term, patch in completed] |
| self.assertEqual([patch["terminal_id"] for patch in patches], ["t0", "t1"]) |
| self.assertEqual(patches[0]["notes"], "previous") |
| self.assertEqual(patches[1]["episode_title_suffixes"], ["t1"]) |
| self.assertEqual(resume_skipped, 1) |
| self.assertEqual(llm_requests, 1) |
| self.assertEqual(pending_unprocessed, 1) |
|
|
| def test_prefix_tree_agent_completed_order_is_deterministic(self) -> None: |
| args = self.prefix_tree_agent_args() |
| args = PrefixTreeAgentArgs(**{**args.__dict__, "max_requests": 3, "workers": 3}) |
| selected = self.prefix_selected(3) |
|
|
| def fake_annotate(_terminal, patch, _args): |
| merged = dict(patch) |
| merged.update({"source": "responses-v1", "status": "ok", "fallback": False, "errors": []}) |
| return merged, False |
|
|
| import tools.annotate_dmhy_prefix_tree_agent as agent |
|
|
| original_annotate = agent.annotate_terminal_with_llm |
| try: |
| agent.annotate_terminal_with_llm = fake_annotate |
| completed, _resume_skipped, _llm_requests, _llm_fallbacks, _pending_unprocessed = annotate_selected( |
| selected, args |
| ) |
| finally: |
| agent.annotate_terminal_with_llm = original_annotate |
|
|
| self.assertEqual([patch["terminal_id"] for _ord, _idx, _term, patch in completed], ["t0", "t1", "t2"]) |
|
|
| def test_prefix_tree_agent_writes_llm_failures_without_completing_patches(self) -> None: |
| with tempfile.TemporaryDirectory() as tmpdir: |
| failure_path = Path(tmpdir) / "failures.jsonl" |
| args = self.prefix_tree_agent_args() |
| args = PrefixTreeAgentArgs( |
| **{**args.__dict__, "max_requests": 1, "workers": 1, "failure_output": failure_path} |
| ) |
| selected = self.prefix_selected(2) |
|
|
| def fake_annotate(_terminal, patch, _args): |
| failed = dict(patch) |
| failed.update( |
| { |
| "status": "fallback", |
| "fallback": True, |
| "errors": ["Responses API HTTP 429: busy"], |
| } |
| ) |
| return failed, True |
|
|
| import tools.annotate_dmhy_prefix_tree_agent as agent |
|
|
| original_annotate = agent.annotate_terminal_with_llm |
| try: |
| agent.annotate_terminal_with_llm = fake_annotate |
| completed, _resume_skipped, llm_requests, llm_fallbacks, pending_unprocessed = annotate_selected( |
| selected, args |
| ) |
| finally: |
| agent.annotate_terminal_with_llm = original_annotate |
|
|
| self.assertEqual(completed, []) |
| self.assertEqual(llm_requests, 1) |
| self.assertEqual(llm_fallbacks, 1) |
| self.assertEqual(pending_unprocessed, 2) |
| failures = [ |
| json.loads(line) |
| for line in failure_path.read_text(encoding="utf-8").splitlines() |
| if line.strip() |
| ] |
| self.assertEqual(len(failures), 1) |
| self.assertEqual(failures[0]["terminal_id"], "t0") |
| self.assertEqual(failures[0]["status"], "pending") |
| self.assertEqual(failures[0]["source"], "responses-v1") |
| self.assertIn("Responses API HTTP 429", failures[0]["errors"][0]) |
|
|
| def test_generated_tokens_split_punctuation_and_use_b_only_labels(self) -> None: |
| tokens, labels = normalize_generated_tokens( |
| ["[ANi]", " ", "Title-Name", "07"], |
| ["B-GROUP", "O", "I-TITLE", "B-EPISODE"], |
| ) |
| self.assertEqual(tokens, ["[", "ANi", "]", " ", "Title", "-", "Name", "07"]) |
| self.assertEqual( |
| labels, |
| ["O", "B-GROUP", "O", "O", "B-TITLE", "O", "B-TITLE", "B-EPISODE"], |
| ) |
| self.assertTrue(all(label == "O" or label.startswith("B-") for label in labels)) |
|
|
| def test_preserve_i_labels_keeps_i_on_non_separator_pieces(self) -> None: |
| tokens, labels = normalize_generated_tokens( |
| ["Title-Name"], |
| ["I-TITLE"], |
| preserve_i_labels=True, |
| ) |
| self.assertEqual(tokens, ["Title", "-", "Name"]) |
| self.assertEqual(labels, ["I-TITLE", "O", "I-TITLE"]) |
|
|
| def test_prefix_tree_agent_strict_schema_payload(self) -> None: |
| args = self.prefix_tree_agent_args() |
| terminal = { |
| "terminal_id": "t0", |
| "prefix": "Show - 01", |
| "suffix_examples": [" [1080P][WEB-DL]"], |
| "value_examples": ["Show - 01 [1080P][WEB-DL].mkv"], |
| } |
| patch = { |
| "terminal_id": "t0", |
| "episode_title_suffixes": [], |
| "media_suffixes": ["[1080P]", "[WEB-DL]"], |
| "title_candidates": [], |
| "notes": "heuristic", |
| } |
| body = request_body(args=args, terminal=terminal, patch=patch, previous_error=None, mode="reasoning_json_schema") |
| self.assertEqual(body["model"], "gpt-5.4-mini") |
| self.assertEqual(body["reasoning"], {"effort": "medium"}) |
| schema_format = body["text"]["format"] |
| self.assertEqual(schema_format["type"], "json_schema") |
| self.assertTrue(schema_format["strict"]) |
| self.assertIn("terminal_id", schema_format["schema"]["required"]) |
| self.assertIn("episode_title_suffixes", schema_format["schema"]["required"]) |
|
|
| def test_prefix_tree_agent_reasoning_effort_is_configurable(self) -> None: |
| args = self.prefix_tree_agent_args() |
| args = PrefixTreeAgentArgs(**{**args.__dict__, "reasoning_effort": "xhigh"}) |
| body = request_body( |
| args=args, |
| terminal={"terminal_id": "t0"}, |
| patch={"terminal_id": "t0", "notes": ""}, |
| previous_error=None, |
| mode="reasoning_json_schema", |
| ) |
| self.assertEqual(body["reasoning"], {"effort": "xhigh"}) |
|
|
| def test_prefix_tree_agent_terminal_id_mismatch_validation(self) -> None: |
| with self.assertRaisesRegex(ValueError, "terminal_id mismatch"): |
| validate_strict_annotation( |
| { |
| "terminal_id": "wrong", |
| "episode_title_suffixes": [], |
| "media_suffixes": [], |
| "title_candidates": [], |
| "llm_label": "uncertain", |
| "notes": "bad id", |
| }, |
| expected_terminal_id="expected", |
| ) |
|
|
| def test_prefix_tree_agent_retries_after_terminal_id_mismatch(self) -> None: |
| args = self.prefix_tree_agent_args() |
| terminal = { |
| "terminal_id": "t0", |
| "prefix": "Show - 01", |
| "suffix_examples": [" - Finale [1080P]"], |
| "value_examples": ["Show - 01 - Finale [1080P].mkv"], |
| } |
| patch = { |
| "terminal_id": "t0", |
| "needs_llm_review": True, |
| "episode_title_suffixes": [], |
| "media_suffixes": ["[1080P]"], |
| "title_candidates": [], |
| "llm_label": None, |
| "notes": "heuristic", |
| "source": "heuristic-v1", |
| } |
| calls: list[dict] = [] |
|
|
| def fake_post(body, _args): |
| calls.append(body) |
| terminal_id = "wrong" if len(calls) == 1 else "t0" |
| return { |
| "output_text": json.dumps( |
| { |
| "terminal_id": terminal_id, |
| "episode_title_suffixes": ["Finale"], |
| "media_suffixes": ["[1080P]"], |
| "title_candidates": ["Finale"], |
| "llm_label": "mixed", |
| "notes": "validated", |
| } |
| ) |
| } |
|
|
| import tools.annotate_dmhy_prefix_tree_agent as agent |
|
|
| original_post = agent.post_responses |
| try: |
| agent.post_responses = fake_post |
| merged, used_fallback = annotate_terminal_with_llm(terminal, patch, args) |
| finally: |
| agent.post_responses = original_post |
|
|
| self.assertFalse(used_fallback) |
| self.assertEqual(len(calls), 2) |
| self.assertIn("Previous attempt failed validation", calls[1]["input"][0]["content"]) |
| self.assertEqual(merged["terminal_id"], "t0") |
| self.assertEqual(merged["source"], "responses-v1") |
| self.assertEqual(merged["episode_title_suffixes"], ["Finale"]) |
|
|
| def test_prefix_tree_agent_cli_writes_validator_compatible_units(self) -> None: |
| with tempfile.TemporaryDirectory() as tmpdir: |
| tmp = Path(tmpdir) |
| graph_path = tmp / "graph.json" |
| dataset_path = tmp / "dmhy_weak.prefix_sample.jsonl" |
| patch_path = tmp / "patches.jsonl" |
| units_path = tmp / "units.jsonl" |
| graph_path.write_text( |
| json.dumps( |
| { |
| "terminals": [ |
| { |
| "terminal_id": "t0", |
| "node_id": 42, |
| "prefix": "[ANi] Prefix Show - ", |
| "digit_skeleton": "[ANi] Prefix Show - ##", |
| "weight": 7, |
| "value_examples": [ |
| "[ANi] Prefix Show - 01 [1080P][WEB-DL].mkv" |
| ], |
| "suffix_examples": ["01 [1080P][WEB-DL]"], |
| } |
| ] |
| }, |
| ensure_ascii=False, |
| ), |
| encoding="utf-8", |
| ) |
|
|
| annotate = subprocess.run( |
| [ |
| sys.executable, |
| "-m", |
| "tools.annotate_dmhy_prefix_tree_agent", |
| "--graph", |
| str(graph_path), |
| "--output", |
| str(dataset_path), |
| "--patch-output", |
| str(patch_path), |
| "--units-output", |
| str(units_path), |
| "--examples-only", |
| "--workers", |
| "1", |
| ], |
| check=False, |
| capture_output=True, |
| text=True, |
| ) |
| self.assertEqual(annotate.returncode, 0, annotate.stderr) |
|
|
| units = [ |
| json.loads(line) |
| for line in units_path.read_text(encoding="utf-8").splitlines() |
| if line.strip() |
| ] |
| self.assertEqual(len(units), 1) |
| self.assertEqual( |
| set(units[0]), |
| { |
| "unit_id", |
| "source_kind", |
| "source_id", |
| "terminal_ids", |
| "weight", |
| "context", |
| "examples", |
| "expected_output", |
| }, |
| ) |
| self.assertEqual(units[0]["unit_id"], "t0") |
| self.assertEqual(units[0]["source_kind"], "prefix_tree") |
| self.assertEqual(units[0]["source_id"], 42) |
| self.assertEqual(units[0]["terminal_ids"], ["t0"]) |
| self.assertEqual(units[0]["weight"], 7) |
| self.assertEqual(units[0]["context"]["prefixes"], ["[ANi] Prefix Show - "]) |
| self.assertEqual(units[0]["context"]["digit_skeletons"], ["[ANi] Prefix Show - ##"]) |
| self.assertEqual(units[0]["context"]["edge_labels"], []) |
| self.assertEqual( |
| units[0]["examples"]["values"], |
| ["[ANi] Prefix Show - 01 [1080P][WEB-DL].mkv"], |
| ) |
| self.assertEqual(units[0]["examples"]["suffixes"], ["01 [1080P][WEB-DL]"]) |
| self.assertEqual(units[0]["expected_output"], {"schema_version": "dmhy-annotation-v1"}) |
|
|
| patches = [ |
| json.loads(line) |
| for line in patch_path.read_text(encoding="utf-8").splitlines() |
| if line.strip() |
| ] |
| self.assertEqual(patches[0]["unit_id"], "t0") |
| self.assertEqual(patches[0]["terminal_ids"], ["t0"]) |
| self.assertEqual(patches[0]["status"], "ok") |
| self.assertEqual(patches[0]["errors"], []) |
| self.assertFalse(patches[0]["fallback"]) |
|
|
| validate_units = subprocess.run( |
| [ |
| sys.executable, |
| "-m", |
| "tools.validate_dmhy_annotation_pipeline", |
| "validate-units", |
| str(units_path), |
| ], |
| check=False, |
| capture_output=True, |
| text=True, |
| ) |
| self.assertEqual(validate_units.returncode, 0, validate_units.stdout + validate_units.stderr) |
| validate_patches = subprocess.run( |
| [ |
| sys.executable, |
| "-m", |
| "tools.validate_dmhy_annotation_pipeline", |
| "validate-patches", |
| str(patch_path), |
| "--units", |
| str(units_path), |
| ], |
| check=False, |
| capture_output=True, |
| text=True, |
| ) |
| self.assertEqual(validate_patches.returncode, 0, validate_patches.stdout + validate_patches.stderr) |
|
|
| def test_validation_rejects_embedded_punctuation(self) -> None: |
| record = { |
| "filename": "Title-Name 07", |
| "tokens": ["Title-Name", "07"], |
| "labels": ["B-TITLE", "B-EPISODE"], |
| } |
| with self.assertRaisesRegex(ValueError, "contains punctuation"): |
| validate_record(record, Path("sample.jsonl"), 1) |
|
|
| def test_validation_rejects_embedded_symbol_separator(self) -> None: |
| record = { |
| "filename": "Title 1920×1080 07", |
| "tokens": ["Title", "1920×1080", "07"], |
| "labels": ["B-TITLE", "B-RESOLUTION", "B-EPISODE"], |
| } |
| with self.assertRaisesRegex(ValueError, "contains punctuation"): |
| validate_record(record, Path("sample.jsonl"), 1) |
|
|
| def test_b_only_input_converts_to_char_i_labels(self) -> None: |
| record = { |
| "filename": "Title-Name 07", |
| "tokens": ["Title", "-", "Name", " ", "07"], |
| "labels": ["B-TITLE", "O", "B-TITLE", "O", "B-EPISODE"], |
| } |
| validate_record(record, Path("sample.jsonl"), 1) |
| converted = convert_record(record) |
| self.assertIn("I-TITLE", converted["labels"]) |
| self.assertEqual(converted["tokens"][:5], ["T", "i", "t", "l", "e"]) |
|
|
| def test_iter_validated_jsonl_accepts_generated_shape(self) -> None: |
| with tempfile.TemporaryDirectory() as tmpdir: |
| path = Path(tmpdir) / "records.jsonl" |
| path.write_text( |
| '{"filename":"A 01","tokens":["A"," ","01"],"labels":["B-TITLE","O","B-EPISODE"]}\n', |
| encoding="utf-8", |
| ) |
| rows = list(iter_validated_jsonl(path)) |
| self.assertEqual(len(rows), 1) |
| self.assertEqual(rows[0]["filename"], "A 01") |
|
|
| def test_cli_smoke_annotate_then_convert_with_temp_files(self) -> None: |
| with tempfile.TemporaryDirectory() as tmpdir: |
| tmp = Path(tmpdir) |
| graph_path = tmp / "graph.json" |
| dataset_path = tmp / "dmhy_weak.generated.jsonl" |
| char_path = tmp / "dmhy_weak.generated_char.jsonl" |
| vocab_path = tmp / "vocab.generated.char.json" |
| manifest_path = tmp / "manifest.json" |
| graph_path.write_text( |
| json.dumps( |
| { |
| "terminals": [ |
| { |
| "terminal_id": "t0", |
| "weight": 1, |
| "value_examples": [ |
| "[ANi] Test Show - 01 [1080P][WEB-DL].mkv" |
| ], |
| "suffix_examples": [" [1080P][WEB-DL]"], |
| } |
| ] |
| }, |
| ensure_ascii=False, |
| ), |
| encoding="utf-8", |
| ) |
|
|
| annotate = subprocess.run( |
| [ |
| sys.executable, |
| "-m", |
| "tools.annotate_dmhy_prefix_graph", |
| "--graph", |
| str(graph_path), |
| "--output", |
| str(dataset_path), |
| "--patch-output", |
| "", |
| "--examples-only", |
| ], |
| check=False, |
| capture_output=True, |
| text=True, |
| ) |
| self.assertEqual(annotate.returncode, 0, annotate.stderr) |
|
|
| rows = [ |
| json.loads(line) |
| for line in dataset_path.read_text(encoding="utf-8").splitlines() |
| if line.strip() |
| ] |
| self.assertEqual(len(rows), 1) |
| self.assertIn("annotations", rows[0]) |
| self.assertEqual(rows[0]["tokens"][0], "[") |
| self.assertEqual(rows[0]["labels"][0], "O") |
|
|
| convert = subprocess.run( |
| [ |
| sys.executable, |
| "-m", |
| "tools.convert_annotated_dmhy_dataset", |
| "--input", |
| str(dataset_path), |
| "--output", |
| str(char_path), |
| "--vocab-output", |
| str(vocab_path), |
| "--manifest-output", |
| str(manifest_path), |
| "--progress", |
| "0", |
| ], |
| check=False, |
| capture_output=True, |
| text=True, |
| ) |
| self.assertEqual(convert.returncode, 0, convert.stderr) |
| self.assertTrue(char_path.exists()) |
| self.assertTrue(vocab_path.exists()) |
| self.assertTrue(manifest_path.exists()) |
|
|
| def test_cli_source_list_mode_expands_beyond_value_examples(self) -> None: |
| with tempfile.TemporaryDirectory() as tmpdir: |
| tmp = Path(tmpdir) |
| graph_path = tmp / "graph.json" |
| source_path = tmp / "dmhy_list.jsonl" |
| dataset_path = tmp / "dmhy_weak.generated.jsonl" |
| graph_path.write_text( |
| json.dumps( |
| { |
| "terminals": [ |
| { |
| "terminal_id": "t0", |
| "prefix": "[ANi] Full Show - ", |
| "weight": 10, |
| "value_examples": [ |
| "[ANi] Full Show - 01 [1080P][WEB-DL].mkv" |
| ], |
| "suffix_examples": ["01 [1080P][WEB-DL]"], |
| }, |
| { |
| "terminal_id": "t1", |
| "prefix": "[ANi] Other Show - ", |
| "weight": 10, |
| "value_examples": [ |
| "[ANi] Other Show - 01 [1080P][WEB-DL].mkv" |
| ], |
| "suffix_examples": ["01 [1080P][WEB-DL]"], |
| }, |
| ] |
| }, |
| ensure_ascii=False, |
| ), |
| encoding="utf-8", |
| ) |
| source_path.write_text( |
| "\n".join( |
| json.dumps({"value": value}, ensure_ascii=False) |
| for value in [ |
| "[ANi] Full Show - 01 [1080P][WEB-DL].mkv", |
| "[ANi] Full Show - 02 [1080P][WEB-DL].mkv", |
| "[ANi] Full Show - 03 [1080P][WEB-DL].mkv", |
| "[ANi] Other Show - 01 [1080P][WEB-DL].mkv", |
| ] |
| ) |
| + "\n", |
| encoding="utf-8", |
| ) |
|
|
| annotate = subprocess.run( |
| [ |
| sys.executable, |
| "-m", |
| "tools.annotate_dmhy_prefix_graph", |
| "--graph", |
| str(graph_path), |
| "--source-list", |
| str(source_path), |
| "--output", |
| str(dataset_path), |
| "--patch-output", |
| "", |
| "--limit", |
| "1", |
| ], |
| check=False, |
| capture_output=True, |
| text=True, |
| ) |
| self.assertEqual(annotate.returncode, 0, annotate.stderr) |
| rows = [ |
| json.loads(line) |
| for line in dataset_path.read_text(encoding="utf-8").splitlines() |
| if line.strip() |
| ] |
| self.assertEqual(len(rows), 3) |
| self.assertEqual([row["filename"] for row in rows], [ |
| "[ANi] Full Show - 01 [1080P][WEB-DL].mkv", |
| "[ANi] Full Show - 02 [1080P][WEB-DL].mkv", |
| "[ANi] Full Show - 03 [1080P][WEB-DL].mkv", |
| ]) |
| self.assertTrue(all(row["terminal_id"] == "t0" for row in rows)) |
|
|
| def test_cli_examples_only_uses_terminal_value_examples(self) -> None: |
| with tempfile.TemporaryDirectory() as tmpdir: |
| tmp = Path(tmpdir) |
| graph_path = tmp / "graph.json" |
| source_path = tmp / "dmhy_list.jsonl" |
| dataset_path = tmp / "dmhy_weak.generated.jsonl" |
| graph_path.write_text( |
| json.dumps( |
| { |
| "terminals": [ |
| { |
| "terminal_id": "t0", |
| "prefix": "[ANi] Example Show - ", |
| "weight": 10, |
| "value_examples": [ |
| "[ANi] Example Show - 01 [1080P][WEB-DL].mkv" |
| ], |
| "suffix_examples": ["01 [1080P][WEB-DL]"], |
| } |
| ] |
| }, |
| ensure_ascii=False, |
| ), |
| encoding="utf-8", |
| ) |
| source_path.write_text( |
| "\n".join( |
| json.dumps({"value": value}, ensure_ascii=False) |
| for value in [ |
| "[ANi] Example Show - 01 [1080P][WEB-DL].mkv", |
| "[ANi] Example Show - 02 [1080P][WEB-DL].mkv", |
| ] |
| ) |
| + "\n", |
| encoding="utf-8", |
| ) |
|
|
| annotate = subprocess.run( |
| [ |
| sys.executable, |
| "-m", |
| "tools.annotate_dmhy_prefix_graph", |
| "--graph", |
| str(graph_path), |
| "--source-list", |
| str(source_path), |
| "--output", |
| str(dataset_path), |
| "--patch-output", |
| "", |
| "--examples-only", |
| ], |
| check=False, |
| capture_output=True, |
| text=True, |
| ) |
| self.assertEqual(annotate.returncode, 0, annotate.stderr) |
| rows = [ |
| json.loads(line) |
| for line in dataset_path.read_text(encoding="utf-8").splitlines() |
| if line.strip() |
| ] |
| self.assertEqual(len(rows), 1) |
| self.assertEqual(rows[0]["filename"], "[ANi] Example Show - 01 [1080P][WEB-DL].mkv") |
|
|
| def test_cli_dag_annotation_units_include_shared_node_terminals(self) -> None: |
| with tempfile.TemporaryDirectory() as tmpdir: |
| tmp = Path(tmpdir) |
| dag_path = tmp / "dmhy_prefix_dag.json" |
| output_path = tmp / "dmhy_prefix_dag.annotation_units.jsonl" |
| dag_path.write_text( |
| json.dumps( |
| { |
| "meta": {"version": "prefix-dag-v1"}, |
| "root": 0, |
| "nodes": [ |
| { |
| "id": 0, |
| "terminal": False, |
| "children": [ |
| {"label": "A", "target": 1}, |
| {"label": "B", "target": 2}, |
| ], |
| "incoming_count": 0, |
| "reachable_terminals": 2, |
| "reachable_weight": 20, |
| }, |
| { |
| "id": 1, |
| "terminal": False, |
| "children": [{"label": " shared", "target": 3}], |
| "incoming_count": 1, |
| "reachable_terminals": 1, |
| "reachable_weight": 10, |
| }, |
| { |
| "id": 2, |
| "terminal": False, |
| "children": [{"label": " shared", "target": 3}], |
| "incoming_count": 1, |
| "reachable_terminals": 1, |
| "reachable_weight": 10, |
| }, |
| { |
| "id": 3, |
| "terminal": False, |
| "children": [ |
| {"label": " 01", "target": 4}, |
| {"label": " 02", "target": 5}, |
| ], |
| "incoming_count": 2, |
| "reachable_terminals": 2, |
| "reachable_weight": 20, |
| }, |
| { |
| "id": 4, |
| "terminal": True, |
| "children": [], |
| "incoming_count": 1, |
| "reachable_terminals": 1, |
| "reachable_weight": 10, |
| }, |
| { |
| "id": 5, |
| "terminal": True, |
| "children": [], |
| "incoming_count": 1, |
| "reachable_terminals": 1, |
| "reachable_weight": 10, |
| }, |
| ], |
| "terminals": [ |
| { |
| "terminal_id": "t0", |
| "node_id": 4, |
| "prefix": "Show A shared 01", |
| "digit_skeleton": "Show A shared <NUM>", |
| "count": 10, |
| "weight": 10, |
| "suffix_examples": [" [1080P][WEB-DL]"], |
| "value_examples": ["Show A shared 01 [1080P][WEB-DL].mkv"], |
| "annotations": {}, |
| }, |
| { |
| "terminal_id": "t1", |
| "node_id": 5, |
| "prefix": "Show B shared 02", |
| "digit_skeleton": "Show B shared <NUM>", |
| "count": 10, |
| "weight": 10, |
| "suffix_examples": [" [1080P][WEB-DL]"], |
| "value_examples": ["Show B shared 02 [1080P][WEB-DL].mkv"], |
| "annotations": {}, |
| }, |
| ], |
| }, |
| ensure_ascii=False, |
| ), |
| encoding="utf-8", |
| ) |
|
|
| annotate = subprocess.run( |
| [ |
| sys.executable, |
| "-m", |
| "tools.annotate_dmhy_prefix_dag", |
| "--dag", |
| str(dag_path), |
| "--output", |
| str(output_path), |
| "--min-reachable-terminals", |
| "2", |
| "--min-incoming-count", |
| "2", |
| "--limit", |
| "1", |
| ], |
| check=False, |
| capture_output=True, |
| text=True, |
| ) |
| self.assertEqual(annotate.returncode, 0, annotate.stderr) |
| rows = [ |
| json.loads(line) |
| for line in output_path.read_text(encoding="utf-8").splitlines() |
| if line.strip() |
| ] |
| self.assertEqual(len(rows), 1) |
| self.assertEqual(rows[0]["unit_id"], "dag-node-3") |
| self.assertEqual(rows[0]["kind"], "shared_suffix") |
| self.assertEqual(rows[0]["terminal_ids"], ["t0", "t1"]) |
| self.assertEqual( |
| rows[0]["prefix_examples"], |
| ["Show A shared 01", "Show B shared 02"], |
| ) |
| self.assertEqual(rows[0]["common_edge_labels"], [" 01", " 02"]) |
| self.assertIn("annotations", rows[0]) |
|
|
| def test_cli_dag_annotation_units_skip_terminal_sink_by_default(self) -> None: |
| with tempfile.TemporaryDirectory() as tmpdir: |
| tmp = Path(tmpdir) |
| dag_path = tmp / "dmhy_prefix_dag.json" |
| output_path = tmp / "dmhy_prefix_dag.annotation_units.jsonl" |
| dag_path.write_text( |
| json.dumps( |
| { |
| "meta": {"version": "prefix-dag-v1"}, |
| "root": 0, |
| "nodes": [ |
| { |
| "id": 0, |
| "terminal": False, |
| "children": [{"label": "shared", "target": 2}], |
| "incoming_count": 0, |
| "reachable_terminals": 1, |
| "reachable_weight": 12908, |
| }, |
| { |
| "id": 2, |
| "terminal": True, |
| "children": [], |
| "incoming_count": 12908, |
| "reachable_terminals": 1, |
| "reachable_weight": 12908, |
| }, |
| ], |
| "terminals": [ |
| { |
| "terminal_id": "t-sink", |
| "node_id": 2, |
| "prefix": "Shared Sink", |
| "digit_skeleton": "Shared Sink", |
| "count": 12908, |
| "weight": 12908, |
| "suffix_examples": [], |
| "value_examples": ["Shared Sink.mkv"], |
| "annotations": {}, |
| } |
| ], |
| }, |
| ensure_ascii=False, |
| ), |
| encoding="utf-8", |
| ) |
|
|
| annotate = subprocess.run( |
| [ |
| sys.executable, |
| "-m", |
| "tools.annotate_dmhy_prefix_dag", |
| "--dag", |
| str(dag_path), |
| "--output", |
| str(output_path), |
| "--min-reachable-terminals", |
| "1", |
| "--min-incoming-count", |
| "2", |
| ], |
| check=False, |
| capture_output=True, |
| text=True, |
| ) |
| self.assertEqual(annotate.returncode, 0, annotate.stderr) |
| rows = [ |
| json.loads(line) |
| for line in output_path.read_text(encoding="utf-8").splitlines() |
| if line.strip() |
| ] |
| self.assertEqual(rows, []) |
|
|
| def test_cli_dag_annotation_units_include_terminal_sink_when_requested(self) -> None: |
| with tempfile.TemporaryDirectory() as tmpdir: |
| tmp = Path(tmpdir) |
| dag_path = tmp / "dmhy_prefix_dag.json" |
| output_path = tmp / "dmhy_prefix_dag.annotation_units.jsonl" |
| dag_path.write_text( |
| json.dumps( |
| { |
| "meta": {"version": "prefix-dag-v1"}, |
| "root": 0, |
| "nodes": [ |
| { |
| "id": 0, |
| "terminal": False, |
| "children": [{"label": "shared", "target": 2}], |
| "incoming_count": 0, |
| "reachable_terminals": 1, |
| "reachable_weight": 12908, |
| }, |
| { |
| "id": 2, |
| "terminal": True, |
| "children": [], |
| "incoming_count": 12908, |
| "reachable_terminals": 1, |
| "reachable_weight": 12908, |
| }, |
| ], |
| "terminals": [ |
| { |
| "terminal_id": "t-sink", |
| "node_id": 2, |
| "prefix": "Shared Sink", |
| "digit_skeleton": "Shared Sink", |
| "count": 12908, |
| "weight": 12908, |
| "suffix_examples": [], |
| "value_examples": ["Shared Sink.mkv"], |
| "annotations": {}, |
| } |
| ], |
| }, |
| ensure_ascii=False, |
| ), |
| encoding="utf-8", |
| ) |
|
|
| annotate = subprocess.run( |
| [ |
| sys.executable, |
| "-m", |
| "tools.annotate_dmhy_prefix_dag", |
| "--dag", |
| str(dag_path), |
| "--output", |
| str(output_path), |
| "--min-reachable-terminals", |
| "1", |
| "--min-incoming-count", |
| "2", |
| "--include-terminal-sink-units", |
| ], |
| check=False, |
| capture_output=True, |
| text=True, |
| ) |
| self.assertEqual(annotate.returncode, 0, annotate.stderr) |
| rows = [ |
| json.loads(line) |
| for line in output_path.read_text(encoding="utf-8").splitlines() |
| if line.strip() |
| ] |
| self.assertEqual(len(rows), 1) |
| self.assertEqual(rows[0]["unit_id"], "dag-node-2") |
| self.assertEqual(rows[0]["kind"], "shared_suffix") |
| self.assertEqual(rows[0]["terminal_ids"], ["t-sink"]) |
| self.assertEqual(rows[0]["common_edge_labels"], []) |
|
|
|
|
| if __name__ == "__main__": |
| unittest.main() |
|
|
