82 lines
2.6 KiB
Python
82 lines
2.6 KiB
Python
|
#!/usr/bin/env python
|
|||
|
|
# -*- coding: utf-8 -*-
|
|||
|
|
"""
|
|||
|
|
批量为已有知识库条目生成 Embedding 向量（本地模型）
|
|||
|
|
运行方式: python scripts/migrate_embeddings.py
|
|||
|
|
|
|||
|
|
首次运行会自动下载模型（~95MB），之后走本地缓存
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
import sys
|
|||
|
|
import os
|
|||
|
|
import json
|
|||
|
|
import logging
|
|||
|
|
|
|||
|
|
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
|||
|
|
|
|||
|
|
from src.config.unified_config import get_config
|
|||
|
|
from src.core.database import db_manager
|
|||
|
|
from src.core.models import KnowledgeEntry
|
|||
|
|
from src.core.embedding_client import EmbeddingClient
|
|||
|
|
from src.core.vector_store import vector_store
|
|||
|
|
|
|||
|
|
logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
|
|||
|
|
logger = logging.getLogger(__name__)
|
|||
|
|
|
|||
|
|
|
|||
|
|
def migrate():
|
|||
|
|
config = get_config()
|
|||
|
|
if not config.embedding.enabled:
|
|||
|
|
logger.warning("Embedding 功能未启用，请在 .env 中设置 EMBEDDING_ENABLED=True")
|
|||
|
|
return
|
|||
|
|
|
|||
|
|
client = EmbeddingClient()
|
|||
|
|
|
|||
|
|
# 测试模型加载
|
|||
|
|
logger.info("正在加载本地 embedding 模型（首次运行需下载）...")
|
|||
|
|
if not client.test_connection():
|
|||
|
|
logger.error("Embedding 模型加载失败，请检查 sentence-transformers 是否安装")
|
|||
|
|
return
|
|||
|
|
logger.info("模型加载成功")
|
|||
|
|
|
|||
|
|
# 获取所有需要生成 embedding 的条目
|
|||
|
|
with db_manager.get_session() as session:
|
|||
|
|
entries = session.query(KnowledgeEntry).filter(
|
|||
|
|
KnowledgeEntry.is_active == True
|
|||
|
|
).all()
|
|||
|
|
|
|||
|
|
# 筛选出没有 embedding 的条目
|
|||
|
|
to_process = []
|
|||
|
|
for entry in entries:
|
|||
|
|
if not entry.vector_embedding or entry.vector_embedding.strip() == '':
|
|||
|
|
to_process.append(entry)
|
|||
|
|
|
|||
|
|
logger.info(f"共 {len(entries)} 条活跃知识，{len(to_process)} 条需要生成 embedding")
|
|||
|
|
|
|||
|
|
if not to_process:
|
|||
|
|
logger.info("所有条目已有 embedding，无需迁移")
|
|||
|
|
return
|
|||
|
|
|
|||
|
|
# 批量生成
|
|||
|
|
texts = [e.question + " " + e.answer for e in to_process]
|
|||
|
|
logger.info(f"开始批量生成 embedding...")
|
|||
|
|
vectors = client.embed_batch(texts)
|
|||
|
|
|
|||
|
|
success_count = 0
|
|||
|
|
for i, entry in enumerate(to_process):
|
|||
|
|
vec = vectors[i]
|
|||
|
|
if vec:
|
|||
|
|
entry.vector_embedding = json.dumps(vec)
|
|||
|
|
success_count += 1
|
|||
|
|
|
|||
|
|
session.commit()
|
|||
|
|
logger.info(f"Embedding 生成完成: 成功 {success_count}/{len(to_process)}")
|
|||
|
|
|
|||
|
|
# 重建向量索引
|
|||
|
|
vector_store.load_from_db()
|
|||
|
|
logger.info(f"向量索引重建完成: {vector_store.size} 条")
|
|||
|
|
|
|||
|
|
|
|||
|
|
if __name__ == "__main__":
|
|||
|
|
migrate()
|