audit_hk_songs_l2.py 12.3 KB
#!/usr/bin/env python3
"""L2 歌词去重库内审计脚本

分两阶段运行:
  Phase 1 (--phase download): 从目标库并行下载歌词 → 本地 JSONL 缓存(断点续传)
  Phase 2 (--phase compare):  读缓存 → 倒排索引召回 → DuplicateChecker 判定 → 输出重复对 CSV

用法:
  python audit_hk_songs_l2.py --phase download --workers 32
  python audit_hk_songs_l2.py --phase compare
  python audit_hk_songs_l2.py --phase all --workers 32
"""

from __future__ import annotations

import argparse
import csv
import json
import os
import sys
import time
from collections import defaultdict
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path

import pymysql
from dotenv import load_dotenv
from tqdm import tqdm

sys.path.insert(0, str(Path(__file__).resolve().parent))
from import_hk_songs import (
    L2CandidateIndex,
    _decode_lyric_content,
    _is_lyrics_effectively_empty,
    download_file,
)
from lyric_dedup import DuplicateChecker, DuplicateDecision, LyricRecord

load_dotenv()

TARGET_DB_CONFIG = {
    "host": os.getenv("TARGET_DB_HOST"),
    "port": int(os.getenv("TARGET_DB_PORT", 3306)),
    "user": os.getenv("TARGET_DB_USER"),
    "password": os.getenv("TARGET_DB_PASSWORD"),
    "database": os.getenv("TARGET_DB_NAME"),
    "charset": "utf8mb4",
    "cursorclass": pymysql.cursors.DictCursor,
}
TARGET_TABLE = os.getenv("TARGET_TABLE_NAME", "hk_songs")

OUTPUT_DIR = Path(__file__).resolve().parent / "output"
CACHE_DIR = OUTPUT_DIR / "cache"
REPORT_DIR = OUTPUT_DIR / "reports"
CACHE_FILE = CACHE_DIR / f"{TARGET_TABLE}_l2_audit_lyrics.jsonl"

# ── 阶段一:下载 ─────────────────────────────────────────────────────────────

def load_cached_ids(cache_path: Path) -> set[str]:
    """读取已缓存的 record_id 集合,用于断点续传。"""
    ids: set[str] = set()
    if not cache_path.exists():
        return ids
    with cache_path.open("r", encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            try:
                item = json.loads(line)
                if item.get("id"):
                    ids.add(str(item["id"]))
            except json.JSONDecodeError:
                continue
    return ids


def fetch_db_rows(table: str) -> list[dict]:
    """从目标库拉取所有有歌词 URL 的记录元数据。"""
    conn = pymysql.connect(**TARGET_DB_CONFIG)
    try:
        with conn.cursor() as cur:
            cur.execute(
                f"SELECT id, name, lyricist, composer, singer, lyrics_url "
                f"FROM `{table}` "
                f"WHERE deleted = '0' AND lyrics_url IS NOT NULL AND lyrics_url LIKE 'http%'"
            )
            return list(cur.fetchall())
    finally:
        conn.close()


def _download_one(row: dict) -> dict | None:
    """下载单条歌词,返回缓存条目或 None(失败时)。"""
    record_id = str(row["id"])
    url = str(row["lyrics_url"])
    content, _ = download_file(url, timeout=15, max_retries=3)
    if not content:
        return None
    lyrics = _decode_lyric_content(content)
    if _is_lyrics_effectively_empty(lyrics):
        # 歌词实质为空:写占位符,让 compare 阶段跳过,但记录已处理避免重下
        return {
            "id": record_id,
            "url": url,
            "lyrics": "",
            "name": row.get("name"),
            "singer": row.get("singer"),
            "lyricist": row.get("lyricist"),
            "composer": row.get("composer"),
        }
    return {
        "id": record_id,
        "url": url,
        "lyrics": lyrics,
        "name": row.get("name"),
        "singer": row.get("singer"),
        "lyricist": row.get("lyricist"),
        "composer": row.get("composer"),
    }


def phase_download(workers: int) -> None:
    CACHE_DIR.mkdir(parents=True, exist_ok=True)

    print("从目标库加载记录列表...")
    rows = fetch_db_rows(TARGET_TABLE)
    print(f"  共 {len(rows)} 条有歌词 URL 的记录")

    cached_ids = load_cached_ids(CACHE_FILE)
    pending = [r for r in rows if str(r["id"]) not in cached_ids]
    print(f"  已缓存: {len(cached_ids)} | 待下载: {len(pending)}")

    if not pending:
        print("所有歌词已缓存,跳过下载阶段。")
        return

    failed = 0
    with CACHE_FILE.open("a", encoding="utf-8") as cache_f:
        with ThreadPoolExecutor(max_workers=workers) as executor:
            futures = {executor.submit(_download_one, row): row for row in pending}
            with tqdm(total=len(pending), desc="下载歌词", unit="条") as pbar:
                for future in as_completed(futures):
                    row = futures[future]
                    try:
                        item = future.result()
                    except Exception as e:
                        item = None
                        tqdm.write(f"[错误] id={row['id']} 下载异常: {e}")
                    if item is None:
                        failed += 1
                        tqdm.write(f"[失败] id={row['id']} url={row.get('lyrics_url', '')[:80]}")
                    else:
                        cache_f.write(json.dumps(item, ensure_ascii=False) + "\n")
                    pbar.update(1)
                    pbar.set_postfix(failed=failed)

    total_cached = len(load_cached_ids(CACHE_FILE))
    print(f"\n下载完成: 成功缓存 {total_cached} 条,失败 {failed} 条")
    print(f"缓存文件: {CACHE_FILE}")


# ── 阶段二:比对 ─────────────────────────────────────────────────────────────

def load_cache(cache_path: Path) -> list[dict]:
    """读取全部缓存条目(过滤空歌词)。"""
    items = []
    if not cache_path.exists():
        return items
    with cache_path.open("r", encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            try:
                item = json.loads(line)
            except json.JSONDecodeError:
                continue
            # 空歌词占位符跳过
            if item.get("lyrics") and str(item["lyrics"]).strip():
                items.append(item)
    return items


def phase_compare(
    output_path: Path,
    threshold: float,
    top_k: int,
    decisions: set[str],
) -> None:
    REPORT_DIR.mkdir(parents=True, exist_ok=True)

    print(f"读取本地歌词缓存: {CACHE_FILE}")
    items = load_cache(CACHE_FILE)
    if not items:
        print("缓存为空或不存在,请先运行 --phase download")
        return
    print(f"  有效歌词条目: {len(items)}")

    checker = DuplicateChecker(duplicate_jaccard_threshold=threshold)
    index = L2CandidateIndex(checker, mode="topk", top_k=top_k)

    print("构建 L2 倒排索引...")
    for item in tqdm(items, desc="建索引", unit="条"):
        index.add(LyricRecord(
            record_id=str(item["id"]),
            lyrics=item["lyrics"],
            title=item.get("name"),
            artist=item.get("singer"),
            lyricist=item.get("lyricist"),
            composer=item.get("composer"),
        ))
    print(f"  索引构建完成: {len(index)} 条")

    # 已输出的 pair 集合(用 frozenset 去重 A→B / B→A)
    seen_pairs: set[frozenset] = set()
    dup_count = 0
    review_count = 0

    fieldnames = [
        "id_a", "name_a", "lyricist_a", "composer_a",
        "id_b", "name_b", "lyricist_b", "composer_b",
        "decision", "confidence", "jaccard", "primary_jaccard",
        "line_coverage", "primary_line_coverage", "reason",
    ]

    # 构建 id → item 映射,方便写 CSV 时取元数据
    id_to_item = {str(item["id"]): item for item in items}

    with output_path.open("w", encoding="utf-8-sig", newline="") as f:
        writer = csv.DictWriter(f, fieldnames=fieldnames)
        writer.writeheader()

        for item in tqdm(items, desc="歌词比对", unit="条"):
            qid = str(item["id"])
            record = LyricRecord(
                record_id=qid,
                lyrics=item["lyrics"],
                title=item.get("name"),
                artist=item.get("singer"),
                lyricist=item.get("lyricist"),
                composer=item.get("composer"),
            )

            recalled = index.recall(record)
            # 过滤自身
            recalled = [r for r in recalled if r.record_id != qid]
            if not recalled:
                continue

            result = checker.check_record_against_candidates(record, recalled, max_candidates=10)

            for cm in result.candidates:
                if cm.decision.value not in decisions:
                    continue
                pair = frozenset({qid, cm.record_id})
                if pair in seen_pairs:
                    continue
                seen_pairs.add(pair)

                b = id_to_item.get(cm.record_id, {})
                row = {
                    "id_a": qid,
                    "name_a": item.get("name", ""),
                    "lyricist_a": item.get("lyricist", ""),
                    "composer_a": item.get("composer", ""),
                    "id_b": cm.record_id,
                    "name_b": b.get("name", ""),
                    "lyricist_b": b.get("lyricist", ""),
                    "composer_b": b.get("composer", ""),
                    "decision": cm.decision.value,
                    "confidence": f"{cm.confidence:.4f}",
                    "jaccard": f"{cm.jaccard:.4f}",
                    "primary_jaccard": f"{cm.primary_jaccard:.4f}",
                    "line_coverage": f"{cm.line_coverage:.4f}",
                    "primary_line_coverage": f"{cm.primary_line_coverage:.4f}",
                    "reason": cm.reason,
                }
                writer.writerow(row)
                f.flush()

                if cm.decision == DuplicateDecision.DUPLICATE:
                    dup_count += 1
                else:
                    review_count += 1

    print(f"\n比对完成:")
    print(f"  duplicate 对数: {dup_count}")
    print(f"  review    对数: {review_count}")
    print(f"  输出文件: {output_path}")


# ── 入口 ──────────────────────────────────────────────────────────────────────

def main() -> None:
    parser = argparse.ArgumentParser(description="hk_songs 库内 L2 歌词去重审计")
    parser.add_argument(
        "--phase",
        choices=["download", "compare", "all"],
        default="all",
        help="运行阶段:download=仅下载缓存,compare=仅比对,all=两阶段连续执行(默认 all)",
    )
    parser.add_argument("--workers", type=int, default=24, help="下载并发线程数(默认 24)")
    parser.add_argument(
        "--threshold", type=float, default=0.78,
        help="L2 Jaccard 判重阈值(默认 0.78,与导入脚本一致)",
    )
    parser.add_argument(
        "--top-k", type=int, default=200,
        help="每条记录最多召回候选数(默认 200)",
    )
    parser.add_argument(
        "--decisions",
        default="duplicate,review",
        help="输出哪些判定结果,逗号分隔,可选 duplicate/review/new(默认 duplicate,review)",
    )
    parser.add_argument(
        "--output", type=Path, default=None,
        help="CSV 输出路径(默认 output/reports/l2_audit_<timestamp>.csv)",
    )
    args = parser.parse_args()

    decisions = {d.strip() for d in args.decisions.split(",") if d.strip()}
    if args.output is None:
        from datetime import datetime
        ts = datetime.now().strftime("%Y%m%d_%H%M%S")
        args.output = REPORT_DIR / f"l2_audit_{ts}.csv"

    if args.phase in ("download", "all"):
        print("=" * 60)
        print("阶段一:下载歌词缓存")
        print("=" * 60)
        phase_download(workers=args.workers)

    if args.phase in ("compare", "all"):
        print()
        print("=" * 60)
        print("阶段二:L2 歌词比对")
        print("=" * 60)
        phase_compare(
            output_path=args.output,
            threshold=args.threshold,
            top_k=args.top_k,
            decisions=decisions,
        )


if __name__ == "__main__":
    main()