🧬 Embedding模型完整使用指南

# ========== 完整语义搜索引擎 ==========
from sentence_transformers import SentenceTransformer
import faiss
import numpy as np
import json
from typing import List, Dict

class SemanticSearchEngine:
    def __init__(self, model_name='BAAI/bge-large-zh-v1.5'):
        """初始化搜索引擎"""
        self.model = SentenceTransformer(model_name)
        self.index = None
        self.documents = []
        self.dimension = 1024
        
    def build_index(self, documents: List[Dict]):
        """构建FAISS索引"""
        print("🔨 开始构建索引...")
        
        # 提取文本
        texts = [doc['text'] for doc in documents]
        self.documents = documents
        
        # 向量化
        embeddings = self.model.encode(
            texts,
            batch_size=64,
            show_progress_bar=True,
            normalize_embeddings=True  # L2规范化
        )
        
        # 创建FAISS索引（使用HNSW）
        self.index = faiss.IndexHNSWFlat(self.dimension, 32)
        self.index.hnsw.efConstruction = 40
        self.index.hnsw.efSearch = 16
        
        # 添加向量
        self.index.add(embeddings.astype('float32'))
        
        print(f"✅ 索引构建完成！共 {len(documents)} 条记录")
        
    def search(self, query: str, top_k: int = 10, filters: Dict = None):
        """语义搜索"""
        # 查询向量化
        query_vec = self.model.encode(
            query,
            normalize_embeddings=True
        ).astype('float32').reshape(1, -1)
        
        # 检索
        distances, indices = self.index.search(query_vec, top_k * 2)
        
        # 格式化结果
        results = []
        for dist, idx in zip(distances[0], indices[0]):
            if idx == -1:  # 无效索引
                continue
                
            doc = self.documents[idx].copy()
            doc['score'] = float(1 / (1 + dist))  # 转换为相似度分数
            
            # 应用过滤器
            if filters:
                if not self._apply_filters(doc, filters):
                    continue
            
            results.append(doc)
            
            if len(results) >= top_k:
                break
        
        return results
    
    def _apply_filters(self, doc: Dict, filters: Dict) -> bool:
        """应用过滤条件"""
        for key, value in filters.items():
            if key not in doc or doc[key] != value:
                return False
        return True
    
    def save_index(self, filepath: str):
        """保存索引"""
        faiss.write_index(self.index, f"{filepath}.index")
        with open(f"{filepath}.docs", 'w', encoding='utf-8') as f:
            json.dump(self.documents, f, ensure_ascii=False, indent=2)
        print(f"✅ 索引已保存到 {filepath}")
    
    def load_index(self, filepath: str):
        """加载索引"""
        self.index = faiss.read_index(f"{filepath}.index")
        with open(f"{filepath}.docs", 'r', encoding='utf-8') as f:
            self.documents = json.load(f)
        print(f"✅ 索引已加载，共 {len(self.documents)} 条记录")

# ========== 使用示例 ==========

# 准备数据
documents = [
    {
        "id": 1,
        "text": "机器学习是人工智能的核心技术",
        "category": "AI",
        "date": "2025-01-01"
    },
    {
        "id": 2,
        "text": "深度学习使用神经网络进行训练",
        "category": "AI",
        "date": "2025-01-02"
    },
    {
        "id": 3,
        "text": "Python是数据科学最流行的编程语言",
        "category": "编程",
        "date": "2025-01-03"
    },
    # ... 更多文档
]

# 初始化搜索引擎
engine = SemanticSearchEngine()

# 构建索引
engine.build_index(documents)

# 保存索引（可选）
engine.save_index("./search_index")

# ========== 搜索测试 ==========

# 1. 基础搜索
results = engine.search("人工智能技术", top_k=5)
print("\n🔍 搜索: 人工智能技术")
for result in results:
    print(f"  [{result['score']:.3f}] {result['text']}")

# 2. 带过滤器的搜索
results = engine.search(
    "编程语言",
    top_k=5,
    filters={"category": "编程"}
)
print("\n🔍 搜索: 编程语言 (仅限编程类)")
for result in results:
    print(f"  [{result['score']:.3f}] {result['text']}")

# 3. 模糊搜索（同义词）
queries = ["AI技术", "人工智能", "机器智能"]
for query in queries:
    results = engine.search(query, top_k=3)
    print(f"\n🔍 搜索: {query}")
    print(f"  最佳匹配: {results[0]['text']}")

print("\n✅ 搜索引擎测试完成！")