amz_review_analyse/向量化.py

"""
从 voc_structured.sqlite 最新 job 展开 audience / pain_point / aspect_opinion
（与聚类.py 一致，不向量化单独的 aspect、opinion），使用本地 MLX 写入 voc_embeddings.sqlite。

溯源：source_row 与结构化时一致（CSV 第 1 条数据行=1）；对应 merged_reviews_cleaned.csv
物理行号 = source_row + 1（第 1 行为表头），content 取自该数据行。

用法（项目根目录，推荐 310py 虚拟环境 Python 3.10+）::

  ./310py/bin/python 向量化.py
  ./310py/bin/python 向量化.py --batch-size 16
  ./310py/bin/python 向量化.py --job-id 4 --csv merged_reviews_cleaned.csv

环境变量：VOC_EMBED_MODEL_PATH、VOC_EMBED_BATCH_SIZE（默认 16）、VOC_EMBED_MAX_TEXT_CHARS（默认 512）。
本地 MLX 推理串行执行，--workers 仅保留兼容、固定为 1。
"""
from __future__ import annotations

import argparse
import json
import logging
import os
import sqlite3
import struct
import sys
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Dict, List, Sequence, Tuple

from csv import DictReader

from local_embedding import (
    DEFAULT_BATCH_SIZE as LOCAL_DEFAULT_BATCH,
    embed_texts,
    embedding_dimensions,
)

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
    stream=sys.stderr,
)
logger = logging.getLogger("voc_embed")

PROJECT_ROOT = Path(__file__).resolve().parent
STRUCTURED_DB = PROJECT_ROOT / "voc_structured.sqlite"
EMBED_DB = PROJECT_ROOT / "voc_embeddings.sqlite"
DEFAULT_CSV = PROJECT_ROOT / "merged_reviews_cleaned.csv"
EMBEDDING_MODEL = os.environ.get(
    "VOC_EMBED_MODEL_PATH",
    str(PROJECT_ROOT / "Qwen3-Embedding-4B-mxfp8"),
).strip() or "Qwen3-Embedding-4B-mxfp8"
EMBED_BATCH_SIZE = LOCAL_DEFAULT_BATCH
EMBED_DEFAULT_WORKERS = 1  # 本地 MLX 模型不可多线程并行推理


def _resolve_embed_workers(explicit: int | None = None) -> int:
    if explicit is not None and explicit > 1:
        logger.warning("本地 embedding 仅支持串行，--workers 已忽略（使用 1）")
    return 1
ENTITY_TYPES = (
    "audience",
    "pain_point",
    "aspect_opinion",
)


@dataclass
class EmbedTask:
    job_id: int
    extraction_id: int
    source_row: int
    entity_type: str
    entity_index: int
    embed_text: str
    audience: str
    aspect: str | None
    opinion: str | None
    category: str | None
    sentiment: str | None
    content: str


def _latest_job_id(conn: sqlite3.Connection) -> int:
    row = conn.execute(
        "SELECT id FROM analysis_jobs ORDER BY id DESC LIMIT 1"
    ).fetchone()
    if not row:
        raise RuntimeError("analysis_jobs 为空，请先运行结构化")
    return int(row[0])


def _load_content_by_source_row(csv_path: Path) -> Dict[int, str]:
    """source_row（1 起）-> content；与结构化_server._load_csv_reviews 行号一致。"""
    with csv_path.open(encoding="utf-8-sig", newline="") as f:
        reader = DictReader(f)
        if not reader.fieldnames:
            raise ValueError(f"CSV 无表头: {csv_path}")
        headers = [h.strip() for h in reader.fieldnames]
        lower_map = {h.lower(): h for h in headers}
        col = None
        for name in ("content", "review", "评论"):
            if name.lower() in lower_map:
                col = lower_map[name.lower()]
                break
        if col is None:
            if len(headers) == 1:
                col = headers[0]
            else:
                raise ValueError(f"CSV 缺少 content 列: {headers}")
        out: Dict[int, str] = {}
        for idx, row in enumerate(reader, start=1):
            cell = row.get(col)
            t = str(cell or "").strip()
            if t:
                out[idx] = t
    return out


def _pack_embedding(vec: Sequence[float]) -> bytes:
    return struct.pack(f"{len(vec)}f", *vec)


def _init_embed_schema(conn: sqlite3.Connection) -> None:
    conn.executescript(
        """
        CREATE TABLE IF NOT EXISTS embedding_items (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            job_id INTEGER NOT NULL,
            extraction_id INTEGER NOT NULL,
            source_row INTEGER NOT NULL,
            entity_type TEXT NOT NULL,
            entity_index INTEGER NOT NULL DEFAULT 0,
            embed_text TEXT NOT NULL,
            audience TEXT,
            aspect TEXT,
            opinion TEXT,
            category TEXT,
            sentiment TEXT,
            content TEXT NOT NULL,
            dimensions INTEGER NOT NULL,
            embedding BLOB NOT NULL
        );
        CREATE INDEX IF NOT EXISTS idx_embed_job_type
            ON embedding_items(job_id, entity_type);
        CREATE INDEX IF NOT EXISTS idx_embed_source_row
            ON embedding_items(job_id, source_row);
        CREATE INDEX IF NOT EXISTS idx_embed_extraction
            ON embedding_items(extraction_id, entity_type, entity_index);
        """
    )


def _should_skip_extraction(audience: str, pain_points: List[Any], feedback: List[Any]) -> bool:
    aud = (audience or "").strip().lower()
    return aud == "unknown" and not pain_points and not feedback


def _build_tasks(
    job_id: int,
    rows: List[sqlite3.Row],
    content_map: Dict[int, str],
) -> List[EmbedTask]:
    tasks: List[EmbedTask] = []
    for row in rows:
        ext_id = int(row["id"])
        source_row = int(row["source_row"])
        try:
            data = json.loads(row["extraction_json"])
        except json.JSONDecodeError as e:
            logger.warning("跳过无效 JSON extraction_id=%s: %s", ext_id, e)
            continue

        audience = str(data.get("audience", "unknown")).strip() or "unknown"
        pain_points = data.get("pain_points") or []
        feedback = data.get("product_feedback") or []
        if not isinstance(pain_points, list):
            pain_points = []
        if not isinstance(feedback, list):
            feedback = []

        if _should_skip_extraction(audience, pain_points, feedback):
            continue

        content = content_map.get(source_row, "")
        if not content:
            logger.warning(
                "source_row=%s 在 CSV 中无 content，仍写入向量但 content 为空",
                source_row,
            )

        aud_norm = audience.strip()
        if aud_norm.lower() != "unknown":
            tasks.append(
                EmbedTask(
                    job_id=job_id,
                    extraction_id=ext_id,
                    source_row=source_row,
                    entity_type="audience",
                    entity_index=0,
                    embed_text=aud_norm,
                    audience=aud_norm,
                    aspect=None,
                    opinion=None,
                    category=None,
                    sentiment=None,
                    content=content,
                )
            )

        for i, p in enumerate(pain_points):
            text = str(p).strip()
            if not text:
                continue
            tasks.append(
                EmbedTask(
                    job_id=job_id,
                    extraction_id=ext_id,
                    source_row=source_row,
                    entity_type="pain_point",
                    entity_index=i,
                    embed_text=text,
                    audience=aud_norm,
                    aspect=None,
                    opinion=None,
                    category=None,
                    sentiment=None,
                    content=content,
                )
            )

        for i, item in enumerate(feedback):
            if not isinstance(item, dict):
                continue
            aspect = str(item.get("aspect", "")).strip()
            opinion = str(item.get("opinion", "")).strip()
            category = str(item.get("category", "")).strip() or None
            sentiment = str(item.get("sentiment", "")).strip() or None
            if not (aspect and opinion):
                continue
            merged = f"{aspect}, {opinion}"
            tasks.append(
                EmbedTask(
                    job_id=job_id,
                    extraction_id=ext_id,
                    source_row=source_row,
                    entity_type="aspect_opinion",
                    entity_index=i,
                    embed_text=merged,
                    audience=aud_norm,
                    aspect=aspect,
                    opinion=opinion,
                    category=category,
                    sentiment=sentiment,
                    content=content,
                )
            )
    return tasks


def _embed_batches(
    tasks: List[EmbedTask],
    *,
    batch_size: int = EMBED_BATCH_SIZE,
) -> Tuple[List[bytes], int]:
    """本地 MLX 串行分批 embedding，返回 BLOB 列表与向量维度。"""
    if batch_size < 1:
        raise ValueError("batch_size 须 ≥ 1")
    if not tasks:
        return [], 0

    texts = [t.embed_text for t in tasks]
    n_chunks = (len(texts) + batch_size - 1) // batch_size
    logger.info(
        "共 %s 条文本，%s 批（每批≤%s 条），本地 MLX 串行",
        len(texts),
        n_chunks,
        batch_size,
    )

    vecs = embed_texts(texts, batch_size=batch_size)
    dim = len(vecs[0]) if vecs else embedding_dimensions()
    blobs = [_pack_embedding(v) for v in vecs]
    for i, v in enumerate(vecs):
        if len(v) != dim:
            raise RuntimeError(f"第 {i} 条维度 {len(v)} != {dim}")
    return blobs, dim


def run_embed(
    *,
    job_id: int | None = None,
    csv_path: Path = DEFAULT_CSV,
    structured_db: Path = STRUCTURED_DB,
    embed_db: Path = EMBED_DB,
    batch_size: int = EMBED_BATCH_SIZE,
    workers: int | None = None,
    reset_db: bool = True,
) -> Dict[str, Any]:
    embed_workers = _resolve_embed_workers(workers)

    csv_path = csv_path.expanduser().resolve()
    if not csv_path.is_file():
        raise FileNotFoundError(csv_path)
    if not structured_db.is_file():
        raise FileNotFoundError(structured_db)

    content_map = _load_content_by_source_row(csv_path)
    sconn = sqlite3.connect(structured_db)
    sconn.row_factory = sqlite3.Row
    try:
        jid = job_id if job_id is not None else _latest_job_id(sconn)
        ext_rows = sconn.execute(
            """
            SELECT id, source_row, extraction_json
            FROM comment_extractions
            WHERE job_id = ?
            ORDER BY source_row
            """,
            (jid,),
        ).fetchall()
        if not ext_rows:
            raise RuntimeError(f"job_id={jid} 无 comment_extractions 记录")
    finally:
        sconn.close()

    tasks = _build_tasks(jid, ext_rows, content_map)
    if not tasks:
        raise RuntimeError("没有可向量化的条目（检查过滤规则与结构化结果）")

    by_type: Dict[str, int] = {}
    for t in tasks:
        by_type[t.entity_type] = by_type.get(t.entity_type, 0) + 1
    logger.info(
        "job_id=%s：%s 条评论结构化记录 -> %s 条向量任务 %s",
        jid,
        len(ext_rows),
        len(tasks),
        by_type,
    )

    blobs, embed_dim = _embed_batches(tasks, batch_size=batch_size)

    econn = sqlite3.connect(embed_db)
    try:
        _init_embed_schema(econn)
        if reset_db:
            econn.execute("DELETE FROM embedding_items")
        else:
            econn.execute("DELETE FROM embedding_items WHERE job_id = ?", (jid,))
        econn.commit()
        insert_sql = """
            INSERT INTO embedding_items (
                job_id, extraction_id, source_row, entity_type, entity_index,
                embed_text, audience, aspect, opinion, category, sentiment,
                content, dimensions, embedding
            ) VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?,?)
        """
        rows = [
            (
                task.job_id,
                task.extraction_id,
                task.source_row,
                task.entity_type,
                task.entity_index,
                task.embed_text,
                task.audience,
                task.aspect,
                task.opinion,
                task.category,
                task.sentiment,
                task.content,
                embed_dim,
                blob,
            )
            for task, blob in zip(tasks, blobs)
        ]
        econn.executemany(insert_sql, rows)
        econn.commit()
        total = econn.execute("SELECT COUNT(*) FROM embedding_items").fetchone()[0]
    finally:
        econn.close()

    summary = {
        "job_id": jid,
        "structured_db": str(structured_db),
        "embed_db": str(embed_db),
        "csv": str(csv_path),
        "model": EMBEDDING_MODEL,
        "dimensions": embed_dim,
        "embed_workers": embed_workers,
        "extractions": len(ext_rows),
        "vectors": total,
        "by_entity_type": by_type,
    }
    logger.info("完成：%s", summary)
    return summary


def main() -> None:
    parser = argparse.ArgumentParser(description="VOC 结构化结果向量化")
    parser.add_argument("--job-id", type=int, default=None, help="默认最新 job")
    parser.add_argument("--csv", type=Path, default=DEFAULT_CSV)
    parser.add_argument("--structured-db", type=Path, default=STRUCTURED_DB)
    parser.add_argument("--embed-db", type=Path, default=EMBED_DB)
    parser.add_argument(
        "--batch-size",
        type=int,
        default=EMBED_BATCH_SIZE,
        help=f"每批本地推理条数（默认 {EMBED_BATCH_SIZE}；可用 VOC_EMBED_BATCH_SIZE）",
    )
    parser.add_argument(
        "--workers",
        type=int,
        default=None,
        help="保留兼容；本地 MLX 固定串行 workers=1",
    )
    args = parser.parse_args()
    summary = run_embed(
        job_id=args.job_id,
        csv_path=args.csv,
        structured_db=args.structured_db,
        embed_db=args.embed_db,
        batch_size=args.batch_size,
        workers=_resolve_embed_workers(args.workers),
    )
    print(json.dumps(summary, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()