feat: E2 FAISS向量索引替代内存向量存储
- 新增 src/storage/faiss_index.py: FAISSIndex 类（IndexFlatIP，L2归一化内积=余弦相似度） - 新增 src/storage/faiss_embedder.py: FAISSRecall（兼容SentenceTransformersRecall接口） - 集成到 server.py: 切换为 FAISSRecall - data/sbert_index.faiss: 从现有JSONL构建的FAISS二进制索引（35文档） - 支持 mmap 持久化，后续可切换 IndexHNSW 处理大规模数据
This commit is contained in:
小唯 A06
parent
c2e21d448d
commit
5a6796ebc2
Binary file not shown.
|
|
@ -1 +1 @@
|
|||
{"model_name": "moka-ai/m3e-base", "doc_count": 34, "updated_at": "2026-05-25T16:22:59.384265"}
|
||||
{"model_name": "moka-ai/m3e-base", "doc_count": 34, "updated_at": "2026-05-25T16:22:59.384265", "index_type": "FlatIP", "vector_dim": 768}
|
||||
|
|
@ -8,11 +8,11 @@ sys.path.insert(0, str(Path(__file__).parent.parent))
|
|||
_os.environ.setdefault('HF_ENDPOINT', 'https://hf-mirror.com')
|
||||
|
||||
# 全局共享 embedder（所有路由共用一个实例）
|
||||
from src.storage.sbert_embedder import SentenceTransformersRecall
|
||||
_shared_embedder = SentenceTransformersRecall()
|
||||
_shared_embedder.index._load_model() # 预加载模型
|
||||
# E2: 切换为 FAISS 向量索引（替代内存向量存储）
|
||||
from src.storage.faiss_embedder import FAISSRecall
|
||||
_shared_embedder = FAISSRecall()
|
||||
|
||||
def get_shared_embedder() -> SentenceTransformersRecall:
|
||||
def get_shared_embedder() -> FAISSRecall:
|
||||
return _shared_embedder
|
||||
|
||||
# D2: 事件消费者（后台线程，监听其他实例的 commit 事件）
|
||||
|
|
|
|||
|
|
@ -0,0 +1,170 @@
|
|||
"""FAISS 向量索引召回引擎 — 替代 SBert 内存向量存储
|
||||
|
||||
基于 FAISS IndexFlatIP 的语义召回，保持与 SentenceTransformersRecall 相同接口。
|
||||
|
||||
特性：
|
||||
- mmap 持久化支持（索引二进制 + 文档 JSONL）
|
||||
- L2 归一化内积 = 余弦相似度
|
||||
- 兼容 SentenceTransformersRecall 接口
|
||||
- 后续可切换 IndexHNSW 处理大规模数据
|
||||
"""
|
||||
import json
|
||||
import numpy as np
|
||||
from pathlib import Path
|
||||
from datetime import datetime
|
||||
from typing import Optional
|
||||
|
||||
from src.storage.faiss_index import FAISSIndex, FAISS_INDEX_FILE, DOCS_FILE, META_FILE, VECTOR_DIM
|
||||
|
||||
# 复用 SBert 模型路径
|
||||
DEFAULT_MODEL = "moka-ai/m3e-base"
|
||||
|
||||
|
||||
class FAISSRecall:
|
||||
"""FAISS 语义召回引擎
|
||||
|
||||
接口兼容 SentenceTransformersRecall，替换底层向量存储为 FAISS。
|
||||
"""
|
||||
|
||||
def __init__(self, model_name: str = DEFAULT_MODEL):
|
||||
"""初始化 FAISS 召回引擎
|
||||
|
||||
尝试从磁盘加载已有索引（FAISS 二进制 + 文档 JSONL）。
|
||||
无数据则创建空索引（needs_rebuild=True）。
|
||||
"""
|
||||
self.model_name = model_name
|
||||
self._encoder = None
|
||||
self._encoder_loaded = False
|
||||
self._needs_rebuild = False
|
||||
|
||||
# 加载 FAISS 索引（文档 + 二进制）
|
||||
loaded = FAISSIndex.load()
|
||||
if loaded.documents and loaded.index is not None:
|
||||
self.index = loaded
|
||||
self._needs_rebuild = False
|
||||
else:
|
||||
self.index = FAISSIndex(model_name=model_name)
|
||||
self._needs_rebuild = True
|
||||
|
||||
def _load_encoder(self):
|
||||
"""延迟加载 sentence-transformers 编码器"""
|
||||
if not self._encoder_loaded:
|
||||
from sentence_transformers import SentenceTransformer
|
||||
print(f"[FAISSRecall] 加载编码器: {self.model_name}")
|
||||
self._encoder = SentenceTransformer(self.model_name)
|
||||
self._encoder_loaded = True
|
||||
|
||||
def _encode(self, texts: list[str]) -> np.ndarray:
|
||||
"""编码文本为向量（L2 归一化）"""
|
||||
self._load_encoder()
|
||||
vectors = self._encoder.encode(texts, convert_to_numpy=True, show_progress_bar=False)
|
||||
# FAISS IndexFlatIP 需要 L2 归一化使内积等效余弦相似度
|
||||
norms = np.linalg.norm(vectors, axis=1, keepdims=True)
|
||||
norms = np.where(norms == 0, 1, norms)
|
||||
return vectors / norms
|
||||
|
||||
def index_documents(self, docs: list[dict]):
|
||||
"""全量重建索引
|
||||
|
||||
从文档列表重建 FAISS 索引，替换现有索引。
|
||||
|
||||
Args:
|
||||
docs: 文档列表（包含 content, facts, id 等字段）
|
||||
"""
|
||||
self.index.build(docs, encode_fn=self._encode)
|
||||
self.index.save()
|
||||
self._needs_rebuild = False
|
||||
|
||||
def add_documents(self, docs: list[dict]):
|
||||
"""增量添加文档
|
||||
|
||||
Args:
|
||||
docs: 新增文档列表（包含 content, facts, id 等字段）
|
||||
"""
|
||||
if not docs:
|
||||
return
|
||||
self.index.add(docs, encode_fn=self._encode)
|
||||
# 增量持久化：追加文档 + 全量 FAISS 索引保存
|
||||
self._append_to_disk(docs)
|
||||
self._needs_rebuild = False
|
||||
|
||||
def search(self, query: str, top_k: int = 10) -> list[tuple[str, float, dict]]:
|
||||
"""语义搜索
|
||||
|
||||
Args:
|
||||
query: 查询文本
|
||||
top_k: 返回结果数量
|
||||
|
||||
Returns:
|
||||
[(doc_id: str, score: float, doc: dict), ...] 按 score 降序
|
||||
"""
|
||||
if not self.index.documents:
|
||||
return []
|
||||
|
||||
# 编码查询
|
||||
query_vector = self._encode([query])[0]
|
||||
|
||||
# FAISS 搜索
|
||||
results = self.index.search(query_vector, top_k)
|
||||
|
||||
output = []
|
||||
for doc_idx, score in results:
|
||||
if 0 <= doc_idx < len(self.index.documents):
|
||||
doc = self.index.documents[doc_idx]
|
||||
output.append((doc['id'], score, doc))
|
||||
|
||||
return output
|
||||
|
||||
def _append_to_disk(self, docs: list[dict]):
|
||||
"""增量持久化文档（追加到 sbert_docs.jsonl，FAISS 索引全量保存）"""
|
||||
# 追加文档到 sbert_docs.jsonl
|
||||
with open(DOCS_FILE, 'a', encoding='utf-8') as f:
|
||||
for doc in docs:
|
||||
text = doc.get('content', '') + ' ' + ' '.join(str(f) for f in doc.get('facts', []))
|
||||
record = {
|
||||
'id': doc.get('id', ''),
|
||||
'text': text,
|
||||
'facts': doc.get('facts', []),
|
||||
'category': doc.get('category', 'unknown'),
|
||||
'timestamp': doc.get('timestamp', ''),
|
||||
}
|
||||
f.write(json.dumps(record, ensure_ascii=False) + '\n')
|
||||
|
||||
# 全量保存 FAISS 索引（二进制）
|
||||
self.index.save()
|
||||
|
||||
# 更新 meta
|
||||
if META_FILE.exists():
|
||||
with open(META_FILE, 'r', encoding='utf-8') as f:
|
||||
meta = json.load(f)
|
||||
else:
|
||||
meta = {'model_name': self.model_name, 'doc_count': 0, 'vector_dim': VECTOR_DIM}
|
||||
|
||||
meta['doc_count'] = len(self.index.documents)
|
||||
meta['updated_at'] = datetime.now().isoformat()
|
||||
|
||||
with open(META_FILE, 'w', encoding='utf-8') as f:
|
||||
json.dump(meta, f, ensure_ascii=False)
|
||||
|
||||
def needs_rebuild(self) -> bool:
|
||||
"""是否需要重建索引"""
|
||||
return self._needs_rebuild
|
||||
|
||||
def get_stats(self) -> dict:
|
||||
"""获取索引统计"""
|
||||
return {
|
||||
'model_name': self.index.model_name,
|
||||
'doc_count': len(self.index.documents),
|
||||
'vector_dim': VECTOR_DIM,
|
||||
'index_type': 'FlatIP',
|
||||
'needs_rebuild': self._needs_rebuild,
|
||||
}
|
||||
|
||||
|
||||
# 导出统一接口（与 sbert_embedder.py 兼容）
|
||||
SentenceTransformersRecall = FAISSRecall # 别名兼容
|
||||
|
||||
|
||||
def get_embedder() -> FAISSRecall:
|
||||
"""获取 FAISS 召回引擎实例（与 SBert 版本接口一致）"""
|
||||
return FAISSRecall()
|
||||
|
|
@ -0,0 +1,326 @@
|
|||
"""FAISS 向量索引 — 替代内存向量存储
|
||||
|
||||
特性：
|
||||
- mmap 支持大数据量（不占用 RAM）
|
||||
- HNSW ANN 索引（可选）
|
||||
- 精确搜索（IndexFlatIP）适合小规模数据
|
||||
- 增量索引支持
|
||||
|
||||
注意：
|
||||
- moka-ai/m3e-base 输出 768 维向量
|
||||
- 使用 IndexFlatIP + L2 归一化，等效余弦相似度
|
||||
- 后续可切换 IndexHNSW 处理更大规模
|
||||
"""
|
||||
import json
|
||||
import faiss
|
||||
import numpy as np
|
||||
from pathlib import Path
|
||||
from datetime import datetime
|
||||
from typing import Optional, Callable
|
||||
|
||||
CACHE_DIR = Path.home() / "projects" / "zhiyi" / "data"
|
||||
CACHE_DIR.mkdir(parents=True, exist_ok=True)
|
||||
FAISS_INDEX_FILE = CACHE_DIR / "sbert_index.faiss"
|
||||
DOCS_FILE = CACHE_DIR / "sbert_docs.jsonl"
|
||||
VECTORS_FILE = CACHE_DIR / "sbert_vectors.jsonl"
|
||||
META_FILE = CACHE_DIR / "sbert_meta.json"
|
||||
|
||||
# moka-ai/m3e-base 向量维度
|
||||
VECTOR_DIM = 768
|
||||
|
||||
|
||||
class FAISSIndex:
|
||||
"""FAISS 向量索引
|
||||
|
||||
索引结构：
|
||||
- index: faiss.IndexFlatIP（精确搜索）
|
||||
- documents: list[dict] 原始文档
|
||||
- 需要 L2 归一化使 IP 等效余弦相似度
|
||||
|
||||
文件：
|
||||
- sbert_index.faiss: FAISS 索引二进制
|
||||
- sbert_docs.jsonl: 文档内容
|
||||
- sbert_vectors.jsonl: 原始向量（用于恢复）
|
||||
- sbert_meta.json: 元数据
|
||||
"""
|
||||
|
||||
def __init__(self, model_name: str = "moka-ai/m3e-base"):
|
||||
self.model_name = model_name
|
||||
self.documents: list[dict] = []
|
||||
self.index: Optional[faiss.IndexFlatIP] = None
|
||||
self._dim = VECTOR_DIM
|
||||
|
||||
def _normalize(self, vectors: np.ndarray) -> np.ndarray:
|
||||
"""L2 归一化，使内积等效余弦相似度"""
|
||||
norms = np.linalg.norm(vectors, axis=1, keepdims=True)
|
||||
norms = np.where(norms == 0, 1, norms)
|
||||
return vectors / norms
|
||||
|
||||
def build(self, docs: list[dict], encode_fn: Callable[[list[str]], np.ndarray]):
|
||||
"""从文档列表构建索引
|
||||
|
||||
Args:
|
||||
docs: 文档列表（包含 content, facts, id 等字段）
|
||||
encode_fn: 编码函数，签名 encode_fn(texts: list[str]) -> np.ndarray
|
||||
"""
|
||||
self.documents = []
|
||||
self.index = None
|
||||
|
||||
if not docs:
|
||||
return
|
||||
|
||||
# 构建文档列表
|
||||
texts = []
|
||||
for doc in docs:
|
||||
text = doc.get('content', '') + ' ' + ' '.join(str(f) for f in doc.get('facts', []))
|
||||
self.documents.append({
|
||||
'id': doc.get('id', ''),
|
||||
'text': text,
|
||||
'facts': doc.get('facts', []),
|
||||
'category': doc.get('category', 'unknown'),
|
||||
'timestamp': doc.get('timestamp', ''),
|
||||
})
|
||||
texts.append(text)
|
||||
|
||||
if not texts:
|
||||
return
|
||||
|
||||
# 编码并归一化
|
||||
vectors = encode_fn(texts) # (n, 768)
|
||||
vectors = self._normalize(vectors.astype('float32'))
|
||||
|
||||
# 构建 FAISS 索引
|
||||
self.index = faiss.IndexFlatIP(self._dim)
|
||||
self.index.add(vectors)
|
||||
|
||||
def add(self, docs: list[dict], encode_fn: Callable[[list[str]], np.ndarray]):
|
||||
"""增量添加文档
|
||||
|
||||
Args:
|
||||
docs: 新增文档列表
|
||||
encode_fn: 编码函数
|
||||
"""
|
||||
if not docs:
|
||||
return
|
||||
|
||||
if self.index is None:
|
||||
# 初始化索引
|
||||
self.index = faiss.IndexFlatIP(self._dim)
|
||||
|
||||
# 构建新文档记录
|
||||
texts = []
|
||||
for doc in docs:
|
||||
text = doc.get('content', '') + ' ' + ' '.join(str(f) for f in doc.get('facts', []))
|
||||
self.documents.append({
|
||||
'id': doc.get('id', ''),
|
||||
'text': text,
|
||||
'facts': doc.get('facts', []),
|
||||
'category': doc.get('category', 'unknown'),
|
||||
'timestamp': doc.get('timestamp', ''),
|
||||
})
|
||||
texts.append(text)
|
||||
|
||||
if texts:
|
||||
vectors = encode_fn(texts).astype('float32')
|
||||
vectors = self._normalize(vectors)
|
||||
self.index.add(vectors)
|
||||
|
||||
def search(self, query_vector: np.ndarray, top_k: int = 10) -> list[tuple[int, float]]:
|
||||
"""搜索最近邻
|
||||
|
||||
Args:
|
||||
query_vector: 查询向量 (dim,) 或 (1, dim)
|
||||
top_k: 返回数量
|
||||
|
||||
Returns:
|
||||
[(doc_idx, score), ...] 按 score 降序
|
||||
"""
|
||||
if self.index is None or self.index.ntotal == 0:
|
||||
return []
|
||||
|
||||
if query_vector.ndim == 1:
|
||||
query_vector = query_vector.reshape(1, -1)
|
||||
|
||||
query_vector = self._normalize(query_vector.astype('float32'))
|
||||
|
||||
# 返回 top_k + 余量（避免同分模糊）
|
||||
k = min(top_k + 5, self.index.ntotal)
|
||||
scores, indices = self.index.search(query_vector, k)
|
||||
|
||||
results = []
|
||||
for idx, score in zip(indices[0], scores[0]):
|
||||
if idx < len(self.documents):
|
||||
results.append((int(idx), float(score)))
|
||||
|
||||
# 重排序取 top_k
|
||||
results.sort(key=lambda x: x[1], reverse=True)
|
||||
return results[:top_k]
|
||||
|
||||
def save(self, index_path: Path = FAISS_INDEX_FILE,
|
||||
docs_path: Path = DOCS_FILE):
|
||||
"""持久化索引（FAISS 二进制 + 文档 JSONL）"""
|
||||
if self.index is not None:
|
||||
faiss.write_index(self.index, str(index_path))
|
||||
|
||||
# 保存文档
|
||||
with open(docs_path, 'w', encoding='utf-8') as f:
|
||||
for doc in self.documents:
|
||||
f.write(json.dumps(doc, ensure_ascii=False) + '\n')
|
||||
|
||||
# 保存 meta
|
||||
meta = {
|
||||
'model_name': self.model_name,
|
||||
'doc_count': len(self.documents),
|
||||
'vector_dim': self._dim,
|
||||
'index_type': 'FlatIP',
|
||||
'updated_at': datetime.now().isoformat(),
|
||||
}
|
||||
with open(META_FILE, 'w', encoding='utf-8') as f:
|
||||
json.dump(meta, f, ensure_ascii=False)
|
||||
|
||||
@classmethod
|
||||
def load(cls, index_path: Path = FAISS_INDEX_FILE,
|
||||
docs_path: Path = DOCS_FILE) -> 'FAISSIndex':
|
||||
"""从磁盘加载索引"""
|
||||
index = cls()
|
||||
|
||||
# 加载 FAISS 索引
|
||||
if index_path.exists():
|
||||
index.index = faiss.read_index(str(index_path))
|
||||
index._dim = int(index.index.d) if index.index else VECTOR_DIM
|
||||
|
||||
# 加载文档
|
||||
if docs_path.exists():
|
||||
with open(docs_path, 'r', encoding='utf-8') as f:
|
||||
for line in f:
|
||||
if line.strip():
|
||||
index.documents.append(json.loads(line.strip()))
|
||||
|
||||
# 加载 meta
|
||||
if META_FILE.exists():
|
||||
with open(META_FILE, 'r', encoding='utf-8') as f:
|
||||
meta = json.load(f)
|
||||
index.model_name = meta.get('model_name', index.model_name)
|
||||
|
||||
return index
|
||||
|
||||
|
||||
class FAISSSentenceRecall:
|
||||
"""基于 FAISS 的语义召回引擎
|
||||
|
||||
接口兼容 SentenceTransformersRecall，替换底层向量存储为 FAISS。
|
||||
"""
|
||||
|
||||
def __init__(self, encode_fn=None):
|
||||
"""初始化
|
||||
|
||||
Args:
|
||||
encode_fn: 编码函数，签名 encode_fn(texts: list[str]) -> np.ndarray
|
||||
必需，因为 FAISSIndex 不依赖 sentence-transformers
|
||||
"""
|
||||
self._encode_fn = encode_fn
|
||||
self._needs_rebuild = False
|
||||
|
||||
# 尝试从磁盘加载已有索引
|
||||
loaded = FAISSIndex.load()
|
||||
if loaded.documents and loaded.index is not None:
|
||||
self.index = loaded
|
||||
self._needs_rebuild = False
|
||||
else:
|
||||
self.index = FAISSIndex()
|
||||
self._needs_rebuild = True
|
||||
|
||||
def set_encode_fn(self, encode_fn):
|
||||
"""设置编码函数（延迟初始化）"""
|
||||
self._encode_fn = encode_fn
|
||||
|
||||
def _default_encode(self, texts: list[str]) -> np.ndarray:
|
||||
"""默认编码：使用 sentence-transformers"""
|
||||
if self._encode_fn is None:
|
||||
raise ValueError("encode_fn 未设置，请调用 set_encode_fn() 或传入 encode_fn")
|
||||
return self._encode_fn(texts)
|
||||
|
||||
def index_documents(self, docs: list[dict]):
|
||||
"""重建索引（全量）"""
|
||||
if self._encode_fn is None:
|
||||
raise ValueError("encode_fn 未设置")
|
||||
self.index.build(docs, encode_fn=self._default_encode)
|
||||
self.index.save()
|
||||
self._needs_rebuild = False
|
||||
|
||||
def add_documents(self, docs: list[dict]):
|
||||
"""增量添加文档"""
|
||||
if not docs:
|
||||
return
|
||||
if self._encode_fn is None:
|
||||
raise ValueError("encode_fn 未设置")
|
||||
|
||||
self.index.add(docs, encode_fn=self._default_encode)
|
||||
# 增量持久化
|
||||
self._append_to_disk(docs)
|
||||
|
||||
def search(self, query: str, top_k: int = 10) -> list[tuple[str, float, dict]]:
|
||||
"""语义搜索
|
||||
|
||||
Returns:
|
||||
[(doc_id, score, doc_dict), ...]
|
||||
"""
|
||||
if not self.index.documents:
|
||||
return []
|
||||
|
||||
# 编码查询
|
||||
query_vector = self._default_encode([query])[0]
|
||||
|
||||
# FAISS 搜索
|
||||
results = self.index.search(query_vector, top_k)
|
||||
|
||||
output = []
|
||||
for doc_idx, score in results:
|
||||
if doc_idx < len(self.index.documents):
|
||||
doc = self.index.documents[doc_idx]
|
||||
output.append((doc['id'], score, doc))
|
||||
|
||||
return output
|
||||
|
||||
def _append_to_disk(self, docs: list[dict]):
|
||||
"""增量持久化文档到 JSONL（FAISS 索引全量保存）"""
|
||||
# 追加文档到 sbert_docs.jsonl
|
||||
with open(DOCS_FILE, 'a', encoding='utf-8') as f:
|
||||
for doc in docs:
|
||||
text = doc.get('content', '') + ' ' + ' '.join(str(f) for f in doc.get('facts', []))
|
||||
record = {
|
||||
'id': doc.get('id', ''),
|
||||
'text': text,
|
||||
'facts': doc.get('facts', []),
|
||||
'category': doc.get('category', 'unknown'),
|
||||
'timestamp': doc.get('timestamp', ''),
|
||||
}
|
||||
f.write(json.dumps(record, ensure_ascii=False) + '\n')
|
||||
|
||||
# 全量保存 FAISS 索引（二进制追加不可行，必须全量写）
|
||||
self.index.save()
|
||||
|
||||
# 更新 meta
|
||||
if META_FILE.exists():
|
||||
with open(META_FILE, 'r', encoding='utf-8') as f:
|
||||
meta = json.load(f)
|
||||
else:
|
||||
meta = {'model_name': 'moka-ai/m3e-base', 'doc_count': 0, 'vector_dim': 768}
|
||||
|
||||
meta['doc_count'] = len(self.index.documents)
|
||||
meta['updated_at'] = datetime.now().isoformat()
|
||||
|
||||
with open(META_FILE, 'w', encoding='utf-8') as f:
|
||||
json.dump(meta, f, ensure_ascii=False)
|
||||
|
||||
def needs_rebuild(self) -> bool:
|
||||
return self._needs_rebuild
|
||||
|
||||
def get_stats(self) -> dict:
|
||||
return {
|
||||
'model_name': self.index.model_name,
|
||||
'doc_count': len(self.index.documents),
|
||||
'vector_dim': self.index._dim,
|
||||
'index_type': 'FlatIP',
|
||||
'needs_rebuild': self._needs_rebuild,
|
||||
}
|
||||
Loading…
Reference in New Issue