hailin
/
faiss_rag_enterprise
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

.

This commit is contained in:
hailin 2025-05-09 20:36:18 +08:00
parent 7872cb5197
commit afdbdd0fe2
1 changed files with 3 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
scripts/rag_build_query.py
View File

@ -15,6 +15,7 @@ import faiss
from typing import List from typing import List
import asyncio import asyncio
from sentence_transformers import SentenceTransformer from sentence_transformers import SentenceTransformer
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
USER_INDEX_PATH = "index_data" USER_INDEX_PATH = "index_data"
USER_DOC_PATH = "docs" USER_DOC_PATH = "docs"
@ -54,7 +55,7 @@ def build_user_index(user_id: str):
raise FileNotFoundError(f"文档目录不存在: {doc_dir}") raise FileNotFoundError(f"文档目录不存在: {doc_dir}")
documents = SimpleDirectoryReader(doc_dir).load_data() documents = SimpleDirectoryReader(doc_dir).load_data()
embed_model = CustomEmbedding(model_name=settings.MODEL_NAME) embed_model = HuggingFaceEmbedding(model_name=settings.MODEL_NAME)
service_context = ServiceContext.from_defaults(embed_model=embed_model) service_context = ServiceContext.from_defaults(embed_model=embed_model)
@ -69,7 +70,7 @@ def build_user_index(user_id: str):
print(f"[BUILD] 为用户 {user_id} 构建并保存了索引 → {index_path}") print(f"[BUILD] 为用户 {user_id} 构建并保存了索引 → {index_path}")
def query_user_rag(user_id: str, question: str, top_k: int = 4) -> str: def query_user_rag(user_id: str, question: str, top_k: int = 4) -> str:
embed_model = CustomEmbedding(model_name=settings.MODEL_NAME) embed_model = HuggingFaceEmbedding(model_name=settings.MODEL_NAME)
service_context = ServiceContext.from_defaults(embed_model=embed_model) service_context = ServiceContext.from_defaults(embed_model=embed_model)
all_nodes = [] all_nodes = []