#!/usr/bin/env python3 """Remove translated duplicate works from filtered JSONL and web/data assets.""" from __future__ import annotations import argparse import json import mmap from collections import defaultdict from datetime import datetime, timezone from pathlib import Path from typing import Any, Iterable def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser( description="Prune translated duplicates when the original Japanese work is present." ) parser.add_argument( "--filtered-input", default="asmr_works.filtered.jsonl", help="Filtered JSONL file to prune.", ) parser.add_argument("--data-dir", default="web/data", help="Directory containing web data assets.") parser.add_argument("--dry-run", action="store_true", help="Show removals without writing files.") parser.add_argument( "--diff-limit", type=int, default=50, help="Maximum removal rows to print. Use 0 for all rows.", ) return parser.parse_args() def clean_text(value: Any) -> str: return value.strip() if isinstance(value, str) else "" def source_id(work: dict[str, Any]) -> str: return clean_text(work.get("source_id")) or clean_text(work.get("sourceId")) def identity_keys(work: dict[str, Any]) -> list[str]: keys: list[str] = [] work_source_id = source_id(work) if work_source_id: keys.append(f"source:{work_source_id.lower()}") work_id = work.get("id") if work_id is not None: keys.append(f"id:{work_id}") return keys def primary_key(work: dict[str, Any]) -> str: keys = identity_keys(work) if keys: return keys[0] return f"title:{clean_text(work.get('title')).lower()}" def display_id(work: dict[str, Any]) -> str: return source_id(work) or str(work.get("id") or "unknown") def display_title(work: dict[str, Any]) -> str: title = clean_text(work.get("title")) or "title unknown" return title if len(title) <= 90 else title[:87] + "..." def translation_info(work: dict[str, Any]) -> dict[str, Any]: value = work.get("translation_info") return value if isinstance(value, dict) else {} def work_language(work: dict[str, Any]) -> str: info = translation_info(work) lang = clean_text(info.get("lang")) if lang: return lang if info.get("is_original") is True: return "JPN" editions = work.get("language_editions") work_source_id = source_id(work) if isinstance(editions, list) and work_source_id: for item in editions: if isinstance(item, dict) and clean_text(item.get("workno")) == work_source_id: return clean_text(item.get("lang")) or clean_text(item.get("label")) return "" def original_workno(work: dict[str, Any]) -> str: info = translation_info(work) for value in (info.get("original_workno"), work.get("original_workno")): text = clean_text(value) if text: return text editions = work.get("language_editions") if isinstance(editions, list): for item in editions: if not isinstance(item, dict): continue if item.get("lang") == "JPN": workno = clean_text(item.get("workno")) if workno: return workno other_editions = work.get("other_language_editions_in_db") if isinstance(other_editions, list): for item in other_editions: if not isinstance(item, dict): continue if item.get("is_original") is True: workno = clean_text(item.get("source_id")) if workno: return workno if info.get("is_original") is True: return source_id(work) return "" def translation_group_key(work: dict[str, Any]) -> str: original = original_workno(work) if original: return f"work:{original}" work_source_id = source_id(work) if work_source_id: return f"work:{work_source_id}" return f"id:{work.get('id')}" def read_jsonl(path: Path) -> list[dict[str, Any]]: works: list[dict[str, Any]] = [] with path.open(encoding="utf-8") as file: for line_no, line in enumerate(file, 1): if not line.strip(): continue try: value = json.loads(line) except json.JSONDecodeError as exc: raise SystemExit(f"{path}:{line_no}: invalid JSON: {exc}") from exc if not isinstance(value, dict): raise SystemExit(f"{path}:{line_no}: expected JSON object") value["__line_no"] = line_no works.append(value) return works def write_jsonl_atomic(path: Path, works: Iterable[dict[str, Any]]) -> None: temp_path = path.with_name(path.name + ".tmp") try: with temp_path.open("w", encoding="utf-8") as file: for work in works: clean_work = {key: value for key, value in work.items() if key != "__line_no"} file.write(json.dumps(clean_work, ensure_ascii=False, separators=(",", ":"))) file.write("\n") temp_path.replace(path) except Exception: temp_path.unlink(missing_ok=True) raise def read_json(path: Path) -> Any: with path.open(encoding="utf-8") as file: return json.load(file) def write_json_atomic(path: Path, value: Any, *, indent: int | None = None) -> None: temp_path = path.with_name(path.name + ".tmp") try: with temp_path.open("w", encoding="utf-8") as file: json.dump(value, file, ensure_ascii=False, indent=indent, separators=None if indent else (",", ":")) if indent: file.write("\n") temp_path.replace(path) except Exception: temp_path.unlink(missing_ok=True) raise def find_translated_duplicates(works: list[dict[str, Any]]) -> tuple[set[str], list[dict[str, Any]], int]: groups: dict[str, list[dict[str, Any]]] = defaultdict(list) for work in works: groups[translation_group_key(work)].append(work) remove_keys: set[str] = set() removals: list[dict[str, Any]] = [] duplicate_groups = 0 for key, group in groups.items(): originals = [work for work in group if work_language(work) == "JPN" and source_id(work)] translated = [work for work in group if work_language(work) not in ("", "JPN")] if not originals or not translated: continue duplicate_groups += 1 keep = originals[0] for work in translated: for identity_key in identity_keys(work): remove_keys.add(identity_key) removals.append( { "remove": work, "keep": keep, "group": key, "lang": work_language(work) or "unknown", } ) return remove_keys, removals, duplicate_groups def should_remove(work: dict[str, Any], remove_keys: set[str]) -> bool: return any(key in remove_keys for key in identity_keys(work)) def normalized_web_works(works: list[dict[str, Any]], remove_keys: set[str]) -> list[dict[str, Any]]: kept: list[dict[str, Any]] = [] for work in works: if should_remove(work, remove_keys): continue updated = dict(work) updated["embeddingIndex"] = len(kept) kept.append(updated) return kept def rebuild_embeddings( *, data_dir: Path, old_works: list[dict[str, Any]], kept_old_indexes: list[int], dimensions: int, ) -> None: embeddings_path = data_dir / "embeddings.f32" vector_size = dimensions * 4 expected_size = len(old_works) * vector_size actual_size = embeddings_path.stat().st_size if actual_size != expected_size: raise RuntimeError(f"embedding size mismatch: expected {expected_size} bytes, got {actual_size} bytes") temp_path = data_dir / "embeddings.f32.tmp" try: with embeddings_path.open("rb") as old_file, temp_path.open("wb") as new_file: with mmap.mmap(old_file.fileno(), 0, access=mmap.ACCESS_READ) as old_map: for old_index in kept_old_indexes: offset = old_index * vector_size new_file.write(old_map[offset : offset + vector_size]) temp_path.replace(embeddings_path) except Exception: temp_path.unlink(missing_ok=True) raise def print_removals(removals: list[dict[str, Any]], limit: int) -> None: rows = removals if limit == 0 else removals[:limit] for item in rows: remove = item["remove"] keep = item["keep"] print( f" - {display_id(remove)} -> keep {display_id(keep)} " f"/ {item['lang']} / {display_title(remove)}" ) remaining = 0 if limit == 0 else max(0, len(removals) - limit) if remaining: print(f" ... {remaining} more") def validate_args(args: argparse.Namespace) -> None: if args.diff_limit < 0: raise SystemExit("--diff-limit must be greater than or equal to 0") def main() -> int: args = parse_args() validate_args(args) filtered_path = Path(args.filtered_input) data_dir = Path(args.data_dir) manifest_path = data_dir / "manifest.json" works_path = data_dir / "works.json" filtered_works = read_jsonl(filtered_path) remove_keys, removals, duplicate_groups = find_translated_duplicates(filtered_works) pruned_filtered = [work for work in filtered_works if not should_remove(work, remove_keys)] manifest = read_json(manifest_path) web_works = read_json(works_path) if not isinstance(web_works, list): raise SystemExit(f"{works_path} must contain a JSON array") dimensions = int(manifest["dimensions"]) kept_old_indexes = [index for index, work in enumerate(web_works) if not should_remove(work, remove_keys)] pruned_web_works = normalized_web_works(web_works, remove_keys) print(f"Translation duplicate groups: {duplicate_groups}") print(f"Works to remove: {len(removals)}") print(f"Filtered works: {len(filtered_works)} -> {len(pruned_filtered)}") print(f"web/data works: {len(web_works)} -> {len(pruned_web_works)}") print(f"Embeddings to keep: {len(kept_old_indexes)} / {len(web_works)}") print_removals(removals, args.diff_limit) if args.dry_run: print("Dry run: no files were written.") return 0 write_jsonl_atomic(filtered_path, pruned_filtered) rebuild_embeddings( data_dir=data_dir, old_works=web_works, kept_old_indexes=kept_old_indexes, dimensions=dimensions, ) new_manifest = dict(manifest) new_manifest.update( { "count": len(pruned_web_works), "embeddingFile": "embeddings.f32", "worksFile": "works.json", "generatedAt": datetime.now(timezone.utc).isoformat(), } ) write_json_atomic(works_path, pruned_web_works) write_json_atomic(manifest_path, new_manifest, indent=2) print(f"Updated {filtered_path}") print(f"Updated {data_dir}") return 0 if __name__ == "__main__": raise SystemExit(main())