2441c9119c
包含七步编排入口、结构化/向量化/聚类/词频/报告模块与 prompts 配置；忽略原始 CSV 与本地密钥。 Co-authored-by: Cursor <cursoragent@cursor.com>
534 lines
17 KiB
Python
534 lines
17 KiB
Python
"""
|
||
从 voc_structured.sqlite 最新 job 展开 audience / pain_points / aspect / opinion /
|
||
aspect_opinion，调用 DashScope text-embedding-v4（256 维）写入 voc_embeddings.sqlite。
|
||
|
||
溯源：source_row 与结构化时一致（CSV 第 1 条数据行=1）；对应 merged_reviews_cleaned.csv
|
||
物理行号 = source_row + 1（第 1 行为表头），content 取自该数据行。
|
||
|
||
用法（项目根目录）::
|
||
|
||
python3 向量化.py
|
||
python3 向量化.py --workers 8
|
||
python3 向量化.py --job-id 4 --csv merged_reviews_cleaned.csv
|
||
|
||
说明：DashScope text-embedding-v4 单次请求最多 10 条文本；脚本按批调用 API，
|
||
默认多线程并行多批（--workers），并非逐条请求。
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
import json
|
||
import logging
|
||
import os
|
||
import sqlite3
|
||
import struct
|
||
import sys
|
||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||
from dataclasses import dataclass
|
||
from pathlib import Path
|
||
from typing import Any, Dict, List, Sequence, Tuple
|
||
|
||
from csv import DictReader
|
||
from openai import OpenAI
|
||
|
||
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"
|
||
DASHSCOPE_BASE_URL = "https://dashscope.aliyuncs.com/compatible-mode/v1"
|
||
EMBEDDING_MODEL = "text-embedding-v4"
|
||
EMBEDDING_DIMENSIONS = 256
|
||
EMBED_BATCH_SIZE = 10 # DashScope text-embedding-v4 单请求 input 数组上限 10
|
||
EMBED_DEFAULT_WORKERS = 6 # 并行批次数（每批最多 10 条）
|
||
ENTITY_TYPES = (
|
||
"audience",
|
||
"pain_point",
|
||
"aspect",
|
||
"opinion",
|
||
"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 _resolve_api_key() -> str:
|
||
v = os.environ.get("DASHSCOPE_API_KEY", "").strip()
|
||
if v:
|
||
return v
|
||
fp = os.environ.get("DASHSCOPE_API_KEY_FILE", "").strip()
|
||
if fp:
|
||
p = Path(fp).expanduser()
|
||
if p.is_file():
|
||
return p.read_text(encoding="utf-8").strip().strip('"').strip("'")
|
||
local = PROJECT_ROOT / ".dashscope_key"
|
||
if local.is_file():
|
||
return local.read_text(encoding="utf-8").strip().strip('"').strip("'")
|
||
return ""
|
||
|
||
|
||
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 not opinion:
|
||
continue
|
||
if aspect:
|
||
tasks.append(
|
||
EmbedTask(
|
||
job_id=job_id,
|
||
extraction_id=ext_id,
|
||
source_row=source_row,
|
||
entity_type="aspect",
|
||
entity_index=i,
|
||
embed_text=aspect,
|
||
audience=aud_norm,
|
||
aspect=aspect,
|
||
opinion=opinion or None,
|
||
category=category,
|
||
sentiment=sentiment,
|
||
content=content,
|
||
)
|
||
)
|
||
if opinion:
|
||
tasks.append(
|
||
EmbedTask(
|
||
job_id=job_id,
|
||
extraction_id=ext_id,
|
||
source_row=source_row,
|
||
entity_type="opinion",
|
||
entity_index=i,
|
||
embed_text=opinion,
|
||
audience=aud_norm,
|
||
aspect=aspect or None,
|
||
opinion=opinion,
|
||
category=category,
|
||
sentiment=sentiment,
|
||
content=content,
|
||
)
|
||
)
|
||
if aspect and opinion:
|
||
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_one_api_batch(api_key: str, texts: List[str]) -> List[bytes]:
|
||
"""单次 API 调用（最多 batch_size 条文本）。"""
|
||
client = OpenAI(api_key=api_key, base_url=DASHSCOPE_BASE_URL)
|
||
resp = client.embeddings.create(
|
||
model=EMBEDDING_MODEL,
|
||
input=texts,
|
||
dimensions=EMBEDDING_DIMENSIONS,
|
||
)
|
||
if len(resp.data) != len(texts):
|
||
raise RuntimeError(
|
||
f"embedding 返回条数 {len(resp.data)} != 请求 {len(texts)}"
|
||
)
|
||
ordered = sorted(resp.data, key=lambda d: d.index)
|
||
blobs: List[bytes] = []
|
||
for item in ordered:
|
||
vec = item.embedding
|
||
if len(vec) != EMBEDDING_DIMENSIONS:
|
||
raise RuntimeError(f"维度 {len(vec)} != 期望 {EMBEDDING_DIMENSIONS}")
|
||
blobs.append(_pack_embedding(vec))
|
||
return blobs
|
||
|
||
|
||
def _embed_batches(
|
||
api_key: str,
|
||
tasks: List[EmbedTask],
|
||
*,
|
||
batch_size: int = EMBED_BATCH_SIZE,
|
||
workers: int = EMBED_DEFAULT_WORKERS,
|
||
) -> List[bytes]:
|
||
"""
|
||
将任务切成每批最多 batch_size 条，并行请求 DashScope。
|
||
返回与 tasks 顺序一致的 embedding BLOB 列表。
|
||
"""
|
||
if batch_size < 1 or batch_size > 10:
|
||
raise ValueError("batch_size 须在 1–10 之间（DashScope 单请求上限 10）")
|
||
workers = max(1, workers)
|
||
|
||
chunks: List[List[EmbedTask]] = [
|
||
tasks[i : i + batch_size] for i in range(0, len(tasks), batch_size)
|
||
]
|
||
n_chunks = len(chunks)
|
||
if n_chunks == 0:
|
||
return []
|
||
|
||
logger.info(
|
||
"共 %s 条文本，%s 批（每批≤%s 条），并行 workers=%s",
|
||
len(tasks),
|
||
n_chunks,
|
||
batch_size,
|
||
min(workers, n_chunks),
|
||
)
|
||
|
||
# 单线程：逻辑简单，便于限流环境
|
||
if workers == 1:
|
||
all_blobs: List[bytes] = []
|
||
for i, chunk in enumerate(chunks, start=1):
|
||
blobs = _embed_one_api_batch(api_key, [t.embed_text for t in chunk])
|
||
all_blobs.extend(blobs)
|
||
if i == n_chunks or i % 20 == 0:
|
||
logger.info("Embedding 进度 %s/%s 批", i, n_chunks)
|
||
return all_blobs
|
||
|
||
results: List[List[bytes] | None] = [None] * n_chunks
|
||
done = 0
|
||
with ThreadPoolExecutor(max_workers=min(workers, n_chunks)) as pool:
|
||
future_map = {
|
||
pool.submit(
|
||
_embed_one_api_batch,
|
||
api_key,
|
||
[t.embed_text for t in chunk],
|
||
): idx
|
||
for idx, chunk in enumerate(chunks)
|
||
}
|
||
for fut in as_completed(future_map):
|
||
idx = future_map[fut]
|
||
results[idx] = fut.result()
|
||
done += 1
|
||
if done == n_chunks or done % 20 == 0:
|
||
logger.info("Embedding 进度 %s/%s 批", done, n_chunks)
|
||
|
||
return [blob for batch in results for blob in batch] # type: ignore[union-attr]
|
||
|
||
|
||
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 = EMBED_DEFAULT_WORKERS,
|
||
reset_db: bool = True,
|
||
) -> Dict[str, Any]:
|
||
api_key = _resolve_api_key()
|
||
if not api_key:
|
||
raise RuntimeError(
|
||
"缺少 DASHSCOPE_API_KEY 或项目根 .dashscope_key"
|
||
)
|
||
|
||
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_batches(
|
||
api_key,
|
||
tasks,
|
||
batch_size=batch_size,
|
||
workers=workers,
|
||
)
|
||
|
||
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,
|
||
EMBEDDING_DIMENSIONS,
|
||
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": EMBEDDING_DIMENSIONS,
|
||
"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="每批 API 请求条数，最大 10（DashScope 限制）",
|
||
)
|
||
parser.add_argument(
|
||
"--workers",
|
||
type=int,
|
||
default=EMBED_DEFAULT_WORKERS,
|
||
help="并行批次数；设为 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=args.workers,
|
||
)
|
||
print(json.dumps(summary, ensure_ascii=False, indent=2))
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|