112 lines
3.4 KiB
Python
112 lines
3.4 KiB
Python
"""
|
|
FAISS向量匹配器
|
|
基于向量相似度从候选集中检索最匹配的菜品名
|
|
"""
|
|
import json
|
|
import os
|
|
from typing import List, Tuple
|
|
|
|
import faiss
|
|
import numpy as np
|
|
|
|
|
|
class DishNameMatcher:
|
|
"""
|
|
菜品名向量匹配器
|
|
加载预构建的FAISS索引，执行相似度检索
|
|
"""
|
|
|
|
def __init__(self, index_dir: str):
|
|
"""
|
|
参数：
|
|
index_dir: 索引目录，应包含：
|
|
- dish_names.json: 菜品名列表
|
|
- dish_embeddings.npy: 归一化后的向量
|
|
- faiss_index.bin: FAISS索引文件
|
|
"""
|
|
self.index_dir = index_dir
|
|
|
|
# 加载菜品名列表
|
|
names_path = os.path.join(index_dir, "dish_names.json")
|
|
with open(names_path, "r", encoding="utf-8") as f:
|
|
self.dish_names: List[str] = json.load(f)
|
|
|
|
# 加载向量（可选，用于调试）
|
|
embeddings_path = os.path.join(index_dir, "dish_embeddings.npy")
|
|
self.embeddings = np.load(embeddings_path)
|
|
|
|
# 加载FAISS索引
|
|
index_path = os.path.join(index_dir, "faiss_index.bin")
|
|
self.index = faiss.read_index(index_path)
|
|
|
|
print(
|
|
f"[DishNameMatcher] Loaded index_dir={index_dir} "
|
|
f"num_dishes={len(self.dish_names)} "
|
|
f"embeddings_shape={self.embeddings.shape} "
|
|
f"index_ntotal={self.index.ntotal}"
|
|
)
|
|
|
|
def match(
|
|
self,
|
|
query_embedding: np.ndarray,
|
|
top_k: int = 3
|
|
) -> List[Tuple[str, float]]:
|
|
"""
|
|
检索最相似的菜品名
|
|
|
|
参数：
|
|
query_embedding: 查询向量 (1D或2D数组)
|
|
top_k: 返回的候选数量
|
|
|
|
返回：
|
|
[(dish_name, similarity_score), ...] 按相似度降序排列
|
|
"""
|
|
# 确保是2D数组 (1, embedding_dim)
|
|
if query_embedding.ndim == 1:
|
|
query_embedding = query_embedding.reshape(1, -1)
|
|
|
|
# 归一化（因为索引是IndexFlatIP，需要归一化才能表示余弦相似度）
|
|
norm = np.linalg.norm(query_embedding, axis=1, keepdims=True)
|
|
if norm[0, 0] > 0:
|
|
query_embedding = query_embedding / norm
|
|
|
|
# 执行检索
|
|
scores, indices = self.index.search(
|
|
query_embedding.astype(np.float32),
|
|
min(top_k, len(self.dish_names))
|
|
)
|
|
|
|
# 构建结果
|
|
results = []
|
|
for score, idx in zip(scores[0], indices[0]):
|
|
if idx == -1: # FAISS在结果不足时会返回-1
|
|
break
|
|
dish_name = self.dish_names[idx]
|
|
results.append((dish_name, float(score)))
|
|
|
|
print(
|
|
f"[DishNameMatcher] Query shape={query_embedding.shape} "
|
|
f"top_k={top_k} results={len(results)}"
|
|
)
|
|
for rank, (name, score) in enumerate(results, 1):
|
|
print(f" [{rank}] {name} (score={score:.4f})")
|
|
|
|
return results
|
|
|
|
def match_text(
|
|
self,
|
|
text: str,
|
|
embedder,
|
|
top_k: int = 3
|
|
) -> List[Tuple[str, float]]:
|
|
"""
|
|
便捷方法：直接从文本检索（自动调用embedder）
|
|
|
|
参数：
|
|
text: 查询文本（菜品名）
|
|
embedder: OllamaEmbedder实例
|
|
top_k: 返回的候选数量
|
|
"""
|
|
query_emb = embedder.encode_single(text)
|
|
return self.match(query_emb, top_k)
|