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:25:30 +08:00
parent 01f2317bcc
commit 1ff999085f
1 changed files with 8 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
scripts/rag_build_query.py
View File

@ -21,12 +21,18 @@ class CustomEmbedding(BaseEmbedding):
from sentence_transformers import SentenceTransformer from sentence_transformers import SentenceTransformer
self.model = SentenceTransformer(model_name) self.model = SentenceTransformer(model_name)
def embed(self, text: str): def _get_text_embedding(self, text: str) -> list[float]:
return self.model.encode(text).tolist() return self.model.encode(text).tolist()
def embed_batch(self, texts: list): def _get_query_embedding(self, query: str) -> list[float]:
return self.model.encode(query).tolist()
def _get_text_embeddings(self, texts: list[str]) -> list[list[float]]:
return self.model.encode(texts).tolist() return self.model.encode(texts).tolist()
def _get_query_embeddings(self, queries: list[str]) -> list[list[float]]:
return self.model.encode(queries).tolist()
def build_user_index(user_id: str): def build_user_index(user_id: str):
doc_dir = os.path.join(USER_DOC_PATH, user_id) doc_dir = os.path.join(USER_DOC_PATH, user_id)
if not os.path.exists(doc_dir): if not os.path.exists(doc_dir):