12345678910111213141516171819202122232425262728293031323334 |
- import faiss
- import os
- from sentence_transformers import SentenceTransformer
- from dotenv import load_dotenv
- from typing import List, Dict, AsyncGenerator
- import numpy as np
- # 加载环境变量
- load_dotenv()
- # DeepSeek API配置
- EMBEDDING_MODEL = os.getenv("EMBEDDING_MODEL")
- class FaissHelper:
- def __init__(self, dimension: int):
- self.model = SentenceTransformer(EMBEDDING_MODEL)
- def generate_embeddings(self,texts: List[str]) -> np.ndarray:
- return self.model.encode(texts)
- def create_index(self,dimension: int):
- index = faiss.IndexFlatL2(dimension)
- return index
-
- def get_index_size(self,index):
- return index.ntotal
- def index_add(self,index, embeddings):
- index.add(embeddings)
- def load_faiss_index(self,index_path):
- index = faiss.read_index(index_path)
- return index
- def search_faiss_index(self, index, query, top_k):
- query_vector = query.reshape(1, -1)
- _, indices = index.search(query_vector, top_k)
- return indices[0]
|