探索搜索领域中AI搜索的无限可能

    from transformers import BertModel, BertTokenizer
    import torch
    
    # 加载预训练模型和分词器
    model_name = 'bert-base-uncased'
    tokenizer = BertTokenizer.from_pretrained(model_name)
    model = BertModel.from_pretrained(model_name)
    
    def get_bert_embedding(text):
    # 分词和转换为模型输入格式
    inputs = tokenizer(text, return_tensors='pt', truncation=True, padding=True, max_length=512)
    
    # 获取模型输出
    with torch.no_grad():
        outputs = model(**inputs)
    
    # 使用[CLS]标记的隐藏状态作为整个文本的表示
    embedding = outputs.last_hidden_state[:, 0, :].squeeze()
    
    return embedding.numpy()
    
    # 示例使用
    query = "best Italian restaurant near me"
    embedding = get_bert_embedding(query)
    print(f"生成的向量维度: {embedding.shape}")
    
    
    python
    
    
![](https://ad.itadn.com/c/weblog/blog-img/images/2025-08-19/tNylDrzOeaWkqvmxREb81IQAdPXL.png)

    import faiss
    import numpy as np
    
    # 假设我们有一组文档向量
    num_docs = 10000
    dim = 768  # BERT-base的向量维度
    doc_vectors = np.random.rand(num_docs, dim).astype('float32')
    
    # 构建FAISS索引
    index = faiss.IndexFlatIP(dim)  # 内积作为相似度度量
    index.add(doc_vectors)
    
    # 搜索示例
    query_vector = np.random.rand(1, dim).astype('float32')
    k = 5  # 返回前5个最相似结果
    distances, indices = index.search(query_vector, k)
    
    print("最相似的文档索引:", indices)
    print("相似度分数:", distances)
    
    
    python
    
    
![](https://ad.itadn.com/c/weblog/blog-img/images/2025-08-19/2fjlXDhVZmIBQSJw7Uyou6T8dHr3.png)

    from sklearn.ensemble import RandomForestClassifier
    from sklearn.model_selection import train_test_split
    import numpy as np
    
    # 假设我们有以下特征(实际中会更复杂)
    # 第一阶段：BM25分数
    # 第二阶段：语义相似度、点击率预测、个性化分数等
    
    # 生成模拟数据
    num_samples = 1000
    X = np.random.rand(num_samples, 4)  # 4个特征
    y = np.random.randint(2, size=num_samples)  # 二元标签(相关或不相关)
    
    # 分割训练测试集
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2)
    
    # 训练排序模型
    model = RandomForestClassifier(n_estimators=100)
    model.fit(X_train, y_train)
    
    # 评估模型
    print("测试集准确率:", model.score(X_test, y_test))
    
    # 使用模型进行排序
    query_features = np.random.rand(5, 4)  # 5个候选文档的特征
    scores = model.predict_proba(query_features)[:, 1]  # 获取正类概率
    ranked_indices = np.argsort(-scores)  # 按分数降序排列
    
    print("排序后的文档索引:", ranked_indices)
    print("对应的相关性分数:", scores[ranked_indices])
    
    
    python
    
    
![](https://ad.itadn.com/c/weblog/blog-img/images/2025-08-19/0s9qf2TB6IDEAyaJO3HwlidLnKjt.png)

    conda create -n ai_search python=3.8
    conda activate ai_search
    pip install torch transformers faiss-cpu scikit-learn pandas
    
    
    bash

    import numpy as np
    import faiss
    from transformers import BertModel, BertTokenizer
    import torch
    from typing import List, Dict
    import json
    
    class AISearchSystem:
    def __init__(self, model_name: str = 'bert-base-uncased'):
        # 初始化BERT模型和分词器
        self.tokenizer = BertTokenizer.from_pretrained(model_name)
        self.model = BertModel.from_pretrained(model_name)
        self.model.eval()  # 设置为评估模式
        
        # 初始化FAISS索引
        self.dim = 768  # BERT-base的维度
        self.index = faiss.IndexFlatIP(self.dim)
        self.documents = []  # 存储原始文档
        
        # 缓存文档向量(实际应用中应使用持久化存储)
        self.doc_vectors = np.zeros((0, self.dim), dtype='float32')
    
    def add_document(self, text: str, metadata: Dict = None):
        """添加文档到搜索系统"""
        # 生成文档向量
        doc_embedding = self._get_embedding(text)
        
        # 添加到FAISS索引
        if len(self.doc_vectors) == 0:
            self.doc_vectors = doc_embedding.reshape(1, -1)
        else:
            self.doc_vectors = np.vstack([self.doc_vectors, doc_embedding])
        
        # 重置FAISS索引
        self.index.reset()
        self.index.add(self.doc_vectors)
        
        # 存储文档和元数据
        doc_id = len(self.documents)
        self.documents.append({
            'id': doc_id,
            'text': text,
            'metadata': metadata or {}
        })
        
        return doc_id
    
    def search(self, query: str, k: int = 5) -> List[Dict]:
        """执行搜索"""
        # 生成查询向量
        query_embedding = self._get_embedding(query)
        
        # 搜索相似文档
        distances, indices = self.index.search(query_embedding, k)
        
        # 准备结果
        results = []
        for i, idx in enumerate(indices[0]):
            if idx >= 0:  # FAISS可能返回-1
                doc = self.documents[idx]
                results.append({
                    'document': doc,
                    'score': float(distances[0][i]),
                    'rank': i + 1
                })
        
        return results
    
    def _get_embedding(self, text: str) -> np.ndarray:
        """使用BERT生成文本嵌入"""
        inputs = self.tokenizer(
            text, 
            return_tensors='pt', 
            truncation=True, 
            padding=True, 
            max_length=512
        )
        
        with torch.no_grad():
            outputs = self.model(**inputs)
        
        # 使用[CLS]标记的隐藏状态作为文本表示
        embedding = outputs.last_hidden_state[:, 0, :].squeeze()
        
        # 归一化向量(使点积等于余弦相似度)
        embedding = embedding / torch.norm(embedding, p=2)
        
        return embedding.numpy().astype('float32')
    
    # 示例使用
    if __name__ == "__main__":
    # 初始化搜索系统
    search_system = AISearchSystem()
    
    # 添加一些文档
    docs = [
        "The quick brown fox jumps over the lazy dog",
        "Artificial intelligence is transforming search technology",
        "BERT is a powerful language model developed by Google",
        "Python is a popular programming language for AI and data science",
        "Neural networks can learn complex patterns from data"
    ]
    
    for doc in docs:
        search_system.add_document(doc)
    
    # 执行搜索
    query = "models for language processing"
    results = search_system.search(query, k=3)
    
    # 打印结果
    print(f"查询: '{query}'")
    print("搜索结果:")
    for result in results:
        print(f"[Rank {result['rank']}, Score: {result['score']:.4f}] {result['document']['text']}")
    
    
    python
    
    
![](https://ad.itadn.com/c/weblog/blog-img/images/2025-08-19/aUjX9qm4eNG1hzJITPnYf2grByov.png)

 * 使用[CLS]标记的隐藏状态作为整个文本的表示
 * 对向量进行L2归一化，使点积等于余弦相似度

 * IndexFlatIP使用内积作为相似度度量
 * 支持动态添加文档

 * 返回结果包括文档内容、相似度分数和排名