初始化
This commit is contained in:
@@ -0,0 +1,35 @@
|
||||
"""向量生成模块 (M5)。
|
||||
|
||||
公共 API:
|
||||
- load_embedding_config / make_embedding_client
|
||||
- compose_text / embed_article / embed_articles
|
||||
- embed_all_events（批量管道）
|
||||
- EmbeddingResult / EmbeddingConfig / EmbeddingError
|
||||
"""
|
||||
|
||||
from embedding.client import (
|
||||
EmbeddingConfig,
|
||||
embed_batch,
|
||||
load_embedding_config,
|
||||
make_embedding_client,
|
||||
)
|
||||
from embedding.embedder import compose_text, embed_article, embed_articles
|
||||
from embedding.models import EmbeddingError, EmbeddingResult
|
||||
from embedding.pipeline import embed_all_events
|
||||
|
||||
__all__ = [
|
||||
# 客户端
|
||||
"EmbeddingConfig",
|
||||
"embed_batch",
|
||||
"load_embedding_config",
|
||||
"make_embedding_client",
|
||||
# 嵌入
|
||||
"compose_text",
|
||||
"embed_article",
|
||||
"embed_articles",
|
||||
# 批量管道
|
||||
"embed_all_events",
|
||||
# 模型
|
||||
"EmbeddingError",
|
||||
"EmbeddingResult",
|
||||
]
|
||||
@@ -0,0 +1,198 @@
|
||||
"""DashScope Embedding 客户端。
|
||||
|
||||
通过 OpenAI 兼容接口调用阿里百炼 text-embedding-v3：
|
||||
base_url: https://dashscope.aliyuncs.com/compatible-mode/v1
|
||||
model: text-embedding-v3（1024 维）
|
||||
|
||||
配置来源:
|
||||
- .env → DASHSCOPE_API_KEY / QWEN_BASE_URL（OpenAI 兼容端点）
|
||||
- configs/system.yaml → embedding 段（model / dimension / batch_size / timeout）
|
||||
"""
|
||||
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
|
||||
import yaml
|
||||
from openai import OpenAI
|
||||
|
||||
from embedding.models import EmbeddingError
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# 默认值
|
||||
DASHSCOPE_DEFAULT_BASE = "https://dashscope.aliyuncs.com/compatible-mode/v1"
|
||||
DASHSCOPE_DEFAULT_MODEL = "text-embedding-v3"
|
||||
DASHSCOPE_DEFAULT_DIM = 1024
|
||||
DASHSCOPE_BATCH_LIMIT = 10 # 百炼实测单批上限
|
||||
|
||||
# 重试
|
||||
DEFAULT_MAX_ATTEMPTS = 3
|
||||
RETRY_BASE_WAIT_SEC = 1.0
|
||||
RETRY_MAX_WAIT_SEC = 8.0
|
||||
|
||||
|
||||
def _load_embedding_config() -> dict:
|
||||
"""从 system.yaml 加载 embedding 段配置。"""
|
||||
config_path = Path("configs/system.yaml")
|
||||
if config_path.exists():
|
||||
try:
|
||||
with open(config_path, encoding="utf-8") as f:
|
||||
raw = yaml.safe_load(f)
|
||||
return raw.get("embedding", {})
|
||||
except Exception:
|
||||
logger.warning("加载 embedding 配置失败")
|
||||
return {}
|
||||
|
||||
|
||||
@dataclass
|
||||
class EmbeddingConfig:
|
||||
"""Embedding 调用配置。"""
|
||||
|
||||
provider: str = "dashscope"
|
||||
model: str = DASHSCOPE_DEFAULT_MODEL
|
||||
api_key: str = ""
|
||||
base_url: str = DASHSCOPE_DEFAULT_BASE
|
||||
dimension: int = DASHSCOPE_DEFAULT_DIM
|
||||
batch_size: int = DASHSCOPE_BATCH_LIMIT
|
||||
timeout_sec: float = 30.0
|
||||
max_attempts: int = DEFAULT_MAX_ATTEMPTS
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
if not self.api_key:
|
||||
raise EmbeddingError("DASHSCOPE_API_KEY 未配置，请检查 .env")
|
||||
|
||||
|
||||
def load_embedding_config(
|
||||
*,
|
||||
model: str | None = None,
|
||||
) -> EmbeddingConfig:
|
||||
"""根据配置构造 EmbeddingConfig。
|
||||
|
||||
优先级: system.yaml > .env 默认值 > 硬编码默认值
|
||||
"""
|
||||
sys_cfg = _load_embedding_config()
|
||||
|
||||
# API key: 从环境变量读取
|
||||
api_key = os.environ.get("DASHSCOPE_API_KEY", "")
|
||||
if not api_key:
|
||||
api_key = os.environ.get("QWEN_API_KEY", "")
|
||||
|
||||
# Base URL: 优先用 QWEN_BASE_URL（OpenAI 兼容），
|
||||
# DASHSCOPE_BASE_URL 通常是旧版非兼容端点，不作为默认
|
||||
base_url = (
|
||||
os.environ.get("QWEN_BASE_URL")
|
||||
or os.environ.get("DASHSCOPE_EMBEDDING_BASE_URL")
|
||||
or DASHSCOPE_DEFAULT_BASE
|
||||
)
|
||||
|
||||
m = model or sys_cfg.get("dashscope_model", DASHSCOPE_DEFAULT_MODEL)
|
||||
dimension = int(sys_cfg.get("dimension", DASHSCOPE_DEFAULT_DIM))
|
||||
batch_size = min(int(sys_cfg.get("batch_size", DASHSCOPE_BATCH_LIMIT)), DASHSCOPE_BATCH_LIMIT)
|
||||
timeout = float(sys_cfg.get("timeout_sec", 30.0))
|
||||
|
||||
if not api_key:
|
||||
raise EmbeddingError("DASHSCOPE_API_KEY 未配置，请检查 .env")
|
||||
|
||||
return EmbeddingConfig(
|
||||
provider="dashscope",
|
||||
model=m,
|
||||
api_key=api_key,
|
||||
base_url=base_url,
|
||||
dimension=dimension,
|
||||
batch_size=batch_size,
|
||||
timeout_sec=timeout,
|
||||
)
|
||||
|
||||
|
||||
def make_embedding_client(config: EmbeddingConfig) -> OpenAI:
|
||||
"""构造同步 OpenAI 客户端（指向 DashScope 兼容端点）。"""
|
||||
logger.info(
|
||||
"初始化 Embedding 客户端: provider=%s model=%s base_url=%s dim=%d",
|
||||
config.provider, config.model, config.base_url, config.dimension,
|
||||
)
|
||||
return OpenAI(
|
||||
api_key=config.api_key,
|
||||
base_url=config.base_url,
|
||||
timeout=config.timeout_sec,
|
||||
)
|
||||
|
||||
|
||||
def _chunked(items: list[str], size: int) -> list[list[str]]:
|
||||
"""把列表按 size 分块。"""
|
||||
return [items[i : i + size] for i in range(0, len(items), size)]
|
||||
|
||||
|
||||
def embed_batch(
|
||||
client: OpenAI,
|
||||
config: EmbeddingConfig,
|
||||
texts: list[str],
|
||||
) -> list[list[float]]:
|
||||
"""批量嵌入，自动分块+重试。
|
||||
|
||||
Args:
|
||||
client: OpenAI 客户端
|
||||
config: Embedding 配置
|
||||
texts: 待嵌入文本列表
|
||||
|
||||
Returns:
|
||||
与 texts 等长的向量列表，每个为 1024 维 float 列表
|
||||
"""
|
||||
if not texts:
|
||||
return []
|
||||
|
||||
all_results: list[list[float]] = []
|
||||
chunks = _chunked(texts, config.batch_size)
|
||||
|
||||
for chunk_idx, chunk in enumerate(chunks):
|
||||
result = _call_with_retry(client, config, chunk, chunk_idx, len(chunks))
|
||||
all_results.extend(result)
|
||||
|
||||
return all_results
|
||||
|
||||
|
||||
def _call_with_retry(
|
||||
client: OpenAI,
|
||||
config: EmbeddingConfig,
|
||||
batch: list[str],
|
||||
chunk_idx: int,
|
||||
total_chunks: int,
|
||||
) -> list[list[float]]:
|
||||
"""单批嵌入调用，带指数退避重试。"""
|
||||
last_err: Exception | None = None
|
||||
|
||||
for attempt in range(1, config.max_attempts + 1):
|
||||
try:
|
||||
resp = client.embeddings.create(model=config.model, input=batch)
|
||||
vectors = [d.embedding for d in resp.data]
|
||||
|
||||
# 维度校验
|
||||
if vectors and len(vectors[0]) != config.dimension:
|
||||
logger.warning(
|
||||
"实际维度 %d 与预期 %d 不一致",
|
||||
len(vectors[0]), config.dimension,
|
||||
)
|
||||
|
||||
logger.debug(
|
||||
"Embedding chunk %d/%d 完成（%d 条，attempt %d）",
|
||||
chunk_idx + 1, total_chunks, len(batch), attempt,
|
||||
)
|
||||
return vectors
|
||||
|
||||
except Exception as e:
|
||||
last_err = e
|
||||
logger.warning(
|
||||
"DashScope embed 失败 chunk %d/%d 尝试 %d/%d: %s: %s",
|
||||
chunk_idx + 1, total_chunks, attempt, config.max_attempts,
|
||||
type(e).__name__, e,
|
||||
)
|
||||
if attempt < config.max_attempts:
|
||||
wait = min(RETRY_BASE_WAIT_SEC * (2 ** (attempt - 1)), RETRY_MAX_WAIT_SEC)
|
||||
time.sleep(wait)
|
||||
|
||||
raise EmbeddingError(
|
||||
f"DashScope embed 放弃（{config.max_attempts} 次）: {last_err}",
|
||||
attempts=config.max_attempts,
|
||||
)
|
||||
@@ -0,0 +1,141 @@
|
||||
"""文本组装与向量生成。
|
||||
|
||||
核心逻辑:
|
||||
1. compose_text: 从 EnTranslatedArticle 拼接中文文本供嵌入
|
||||
2. embed_article: 单篇文章嵌入
|
||||
3. embed_articles: 批量嵌入（自动分块）
|
||||
"""
|
||||
|
||||
import logging
|
||||
|
||||
from openai import OpenAI
|
||||
|
||||
from embedding.client import EmbeddingConfig, embed_batch
|
||||
from embedding.models import EmbeddingResult
|
||||
from llm.models import EnTranslatedArticle
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# 嵌入文本最大字符数（DashScope text-embedding-v3 支持 8192 token，保守取 4000 字符）
|
||||
MAX_EMBED_CHARS = 4000
|
||||
|
||||
|
||||
def compose_text(article: EnTranslatedArticle, max_chars: int = MAX_EMBED_CHARS) -> str:
|
||||
"""把 EnTranslatedArticle 组装成单段嵌入文本。
|
||||
|
||||
拼接策略:
|
||||
- 中文标题（最高信号）
|
||||
- 事件摘要（语义浓缩）
|
||||
- 中文正文（截断）
|
||||
|
||||
Args:
|
||||
article: M4 输出的双语文章
|
||||
max_chars: 整段最大字符数
|
||||
|
||||
Returns:
|
||||
拼接后的中文嵌入文本
|
||||
"""
|
||||
parts: list[str] = []
|
||||
|
||||
# 1. 标题（中文）
|
||||
if article.title_zh:
|
||||
parts.append(f"标题: {article.title_zh}")
|
||||
|
||||
# 2. 事件标签（语义浓缩）
|
||||
for ev in article.events:
|
||||
event_parts = [
|
||||
f"[{ev.sentiment.value}]",
|
||||
f"{ev.event_type}",
|
||||
f"重要度{ev.importance}",
|
||||
]
|
||||
if ev.stock_codes:
|
||||
event_parts.append("代码:" + ",".join(ev.stock_codes[:5]))
|
||||
if ev.summary_zh:
|
||||
event_parts.append(ev.summary_zh)
|
||||
parts.append("事件: " + " ".join(event_parts))
|
||||
|
||||
# 3. 中文正文
|
||||
if article.content_zh:
|
||||
body = article.content_zh
|
||||
parts.append(f"正文: {body}")
|
||||
|
||||
text = "\n".join(parts)
|
||||
|
||||
# 截断保护
|
||||
if len(text) > max_chars:
|
||||
logger.debug("嵌入文本超长 %d → %d", len(text), max_chars)
|
||||
text = text[:max_chars]
|
||||
|
||||
return text
|
||||
|
||||
|
||||
def embed_article(
|
||||
client: OpenAI,
|
||||
config: EmbeddingConfig,
|
||||
article: EnTranslatedArticle,
|
||||
) -> EmbeddingResult:
|
||||
"""单篇文章嵌入。
|
||||
|
||||
Args:
|
||||
client: OpenAI 客户端
|
||||
config: Embedding 配置
|
||||
article: M4 输出的双语文章
|
||||
|
||||
Returns:
|
||||
EmbeddingResult 含向量 + 元信息
|
||||
"""
|
||||
text = compose_text(article)
|
||||
|
||||
vectors = embed_batch(client, config, [text])
|
||||
if not vectors:
|
||||
raise ValueError(f"嵌入返回空结果: {article.url_hash}")
|
||||
|
||||
return EmbeddingResult(
|
||||
url_hash=article.url_hash,
|
||||
source_id=article.source_id,
|
||||
vector=vectors[0],
|
||||
dimension=len(vectors[0]),
|
||||
embedded_text=text,
|
||||
provider=config.provider,
|
||||
model=config.model,
|
||||
)
|
||||
|
||||
|
||||
def embed_articles(
|
||||
client: OpenAI,
|
||||
config: EmbeddingConfig,
|
||||
articles: list[EnTranslatedArticle],
|
||||
) -> list[EmbeddingResult]:
|
||||
"""批量嵌入多篇文章（自动分块+重试）。
|
||||
|
||||
Args:
|
||||
client: OpenAI 客户端
|
||||
config: Embedding 配置
|
||||
articles: M4 输出的双语文章列表
|
||||
|
||||
Returns:
|
||||
EmbeddingResult 列表
|
||||
"""
|
||||
if not articles:
|
||||
return []
|
||||
|
||||
# 组装所有嵌入文本
|
||||
texts = [compose_text(a) for a in articles]
|
||||
|
||||
# 批量嵌入
|
||||
vectors = embed_batch(client, config, texts)
|
||||
|
||||
# 组装结果
|
||||
results: list[EmbeddingResult] = []
|
||||
for article, vec in zip(articles, vectors):
|
||||
results.append(EmbeddingResult(
|
||||
url_hash=article.url_hash,
|
||||
source_id=article.source_id,
|
||||
vector=vec,
|
||||
dimension=len(vec),
|
||||
embedded_text=compose_text(article),
|
||||
provider=config.provider,
|
||||
model=config.model,
|
||||
))
|
||||
|
||||
return results
|
||||
@@ -0,0 +1,34 @@
|
||||
"""Embedding 向量生成数据模型 (M5)。"""
|
||||
|
||||
from datetime import datetime
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
|
||||
class EmbeddingResult(BaseModel):
|
||||
"""单篇文章的嵌入向量结果，M5 最终落盘格式。"""
|
||||
|
||||
# 来源标识
|
||||
url_hash: str
|
||||
source_id: str
|
||||
|
||||
# 嵌入向量（1024 维 float 列表）
|
||||
vector: list[float] = Field(..., description="1024 维浮点向量")
|
||||
dimension: int = 1024
|
||||
|
||||
# 嵌入文本（用于检索时调试/可视化）
|
||||
embedded_text: str = Field(default="", description="拼接后送入 embedder 的文本")
|
||||
|
||||
# 调用元信息
|
||||
provider: str = "dashscope"
|
||||
model: str = ""
|
||||
embedded_at: datetime = Field(default_factory=datetime.now)
|
||||
|
||||
|
||||
class EmbeddingError(Exception):
|
||||
"""Embedding 调用失败。"""
|
||||
|
||||
def __init__(self, reason: str, *, attempts: int = 0) -> None:
|
||||
super().__init__(reason)
|
||||
self.reason = reason
|
||||
self.attempts = attempts
|
||||
@@ -0,0 +1,179 @@
|
||||
"""批量向量生成管道。
|
||||
|
||||
输入: data/events/{YYYYMMDD}/{url_hash}.json（M4 翻译+事件输出）
|
||||
输出: data/embeddings/{YYYYMMDD}/{url_hash}.json
|
||||
"""
|
||||
|
||||
import json
|
||||
import logging
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
|
||||
from crawler.utils import get_news_day
|
||||
from embedding.client import (
|
||||
EmbeddingConfig,
|
||||
load_embedding_config,
|
||||
make_embedding_client,
|
||||
)
|
||||
from embedding.embedder import embed_articles
|
||||
from embedding.models import EmbeddingError
|
||||
from llm.models import EnTranslatedArticle
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _load_event_articles(date_str: str) -> list[EnTranslatedArticle]:
|
||||
"""加载指定日期的翻译+事件文章。
|
||||
|
||||
Args:
|
||||
date_str: 日期 YYYYMMDD
|
||||
|
||||
Returns:
|
||||
EnTranslatedArticle 列表
|
||||
"""
|
||||
base_dir = Path(f"data/events/{date_str}")
|
||||
if not base_dir.exists():
|
||||
return []
|
||||
|
||||
articles: list[EnTranslatedArticle] = []
|
||||
for json_file in sorted(base_dir.glob("*.json")):
|
||||
if json_file.name == "index.json":
|
||||
continue
|
||||
try:
|
||||
data = json.loads(json_file.read_text(encoding="utf-8"))
|
||||
articles.append(EnTranslatedArticle(**data))
|
||||
except (json.JSONDecodeError, Exception) as e:
|
||||
logger.warning("解析事件文章失败 %s: %s", json_file, e)
|
||||
|
||||
return articles
|
||||
|
||||
|
||||
def embed_all_events(
|
||||
date_str: str | None = None,
|
||||
*,
|
||||
model: str | None = None,
|
||||
) -> dict:
|
||||
"""对所有 M4 输出的文章执行向量化。
|
||||
|
||||
Args:
|
||||
date_str: 日期 YYYYMMDD，默认当前新闻日
|
||||
model: Embedding 模型名，默认从 system.yaml 读取
|
||||
|
||||
Returns:
|
||||
统计 dict
|
||||
"""
|
||||
if date_str is None:
|
||||
date_str = get_news_day()
|
||||
|
||||
logger.info("══════ 开始向量生成，日期: %s ══════", date_str)
|
||||
|
||||
# 加载文章
|
||||
articles = _load_event_articles(date_str)
|
||||
if not articles:
|
||||
logger.warning("事件目录无文章: data/events/%s/", date_str)
|
||||
return {"date": date_str, "total": 0, "success": 0, "failed": 0, "elapsed_sec": 0}
|
||||
|
||||
# 初始化 Embedding 客户端
|
||||
config = load_embedding_config(model=model)
|
||||
client = make_embedding_client(config)
|
||||
|
||||
# 输出目录
|
||||
out_dir = Path(f"data/embeddings/{date_str}")
|
||||
out_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# 增量：跳过已向量化的文章
|
||||
new_articles = []
|
||||
skipped = 0
|
||||
for a in articles:
|
||||
if (out_dir / f"{a.url_hash}.json").exists():
|
||||
skipped += 1
|
||||
else:
|
||||
new_articles.append(a)
|
||||
if skipped > 0:
|
||||
logger.info("增量跳过 %d 篇已向量化，剩余 %d 篇待处理", skipped, len(new_articles))
|
||||
articles = new_articles
|
||||
|
||||
start_time = datetime.now()
|
||||
success = 0
|
||||
failed = 0
|
||||
|
||||
# 批量嵌入（按 batch_size 分块，每批输出进度）
|
||||
batch_size = config.batch_size
|
||||
total = len(articles)
|
||||
logger.info("开始向量化 %d 篇文章（batch_size=%d, model=%s）",
|
||||
total, batch_size, config.model)
|
||||
|
||||
for batch_start in range(0, total, batch_size):
|
||||
batch_end = min(batch_start + batch_size, total)
|
||||
batch = articles[batch_start:batch_end]
|
||||
|
||||
try:
|
||||
results = embed_articles(client, config, batch)
|
||||
for result in results:
|
||||
out_file = out_dir / f"{result.url_hash}.json"
|
||||
out_file.write_text(
|
||||
result.model_dump_json(indent=2, ensure_ascii=False),
|
||||
encoding="utf-8",
|
||||
)
|
||||
success += 1
|
||||
|
||||
logger.info(" [%d/%d] ✅ %d 篇 → %d 维向量",
|
||||
batch_end, total, len(results), config.dimension)
|
||||
|
||||
except EmbeddingError as e:
|
||||
failed += len(batch)
|
||||
logger.error("批量嵌入失败 [%d-%d]: %s", batch_start, batch_end, e.reason)
|
||||
except Exception as e:
|
||||
failed += len(batch)
|
||||
logger.exception("批量嵌入异常 [%d-%d]: %s", batch_start, batch_end, e)
|
||||
|
||||
elapsed = (datetime.now() - start_time).total_seconds()
|
||||
|
||||
# 写入索引
|
||||
_write_embedding_index(date_str, success, failed, elapsed, config)
|
||||
|
||||
logger.info(
|
||||
"══════ 向量生成完成: 成功 %d / 失败 %d / 总计 %d，耗时 %.1f 秒 ══════",
|
||||
success, failed, len(articles), elapsed,
|
||||
)
|
||||
|
||||
return {
|
||||
"date": date_str,
|
||||
"total": len(articles),
|
||||
"success": success,
|
||||
"failed": failed,
|
||||
"elapsed_sec": elapsed,
|
||||
"provider": config.provider,
|
||||
"model": config.model,
|
||||
"dimension": config.dimension,
|
||||
}
|
||||
|
||||
|
||||
def _write_embedding_index(
|
||||
date_str: str,
|
||||
success: int,
|
||||
failed: int,
|
||||
elapsed_sec: float,
|
||||
config: EmbeddingConfig,
|
||||
) -> None:
|
||||
"""写入向量索引文件。"""
|
||||
out_dir = Path(f"data/embeddings/{date_str}")
|
||||
out_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
index_data = {
|
||||
"date": date_str,
|
||||
"success": success,
|
||||
"failed": failed,
|
||||
"elapsed_sec": round(elapsed_sec, 1),
|
||||
"provider": config.provider,
|
||||
"model": config.model,
|
||||
"dimension": config.dimension,
|
||||
"generated_at": datetime.now().isoformat(),
|
||||
}
|
||||
|
||||
index_path = out_dir / "index.json"
|
||||
index_path.write_text(
|
||||
json.dumps(index_data, indent=2, ensure_ascii=False),
|
||||
encoding="utf-8",
|
||||
)
|
||||
logger.info("向量索引已写入: %s", index_path)
|
||||
Reference in New Issue
Block a user