ivanwu/jarvis-models
1
0
Fork 0
mirror of https://github.com/BoardWare-Genius/jarvis-models.git synced 2025-12-13 16:53:24 +00:00
Code Issues Packages Projects Releases Wiki Activity

update blackbox chroma and chat

Browse Source
This commit is contained in:
ACBBZ
2024-05-28 08:23:10 +00:00
parent 347eea2f9c
commit ca34962376
4 changed files with 113 additions and 48 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
src/blackbox/chroma_query.py
View File

@ -51,7 +51,7 @@ class ChromaQuery(Blackbox):
return JSONResponse(content={"error": "question is required"}, status_code=status.HTTP_400_BAD_REQUEST)
if chroma_embedding_model is None or chroma_embedding_model.isspace() or chroma_embedding_model == "":
chroma_embedding_model = "bge-large-zh-v1.5"
chroma_embedding_model = "/model/Weight/BAAI/bge-large-zh-v1.5"
if chroma_host is None or chroma_host.isspace() or chroma_host == "":
chroma_host = "10.6.82.192"
@ -71,10 +71,9 @@ class ChromaQuery(Blackbox):
else:
client = chromadb.HttpClient(host=chroma_host, port=chroma_port)
if re.search(r"bge-large-zh-v1.5", chroma_embedding_model):
if re.search(r"/model/Weight/BAAI/bge-large-zh-v1.5", chroma_embedding_model):
embedding_model = self.embedding_model_1
else:
chroma_embedding_model = "/model/Weight/BAAI/" + chroma_embedding_model
embedding_model = embedding_functions.SentenceTransformerEmbeddingFunction(model_name=chroma_embedding_model, device = "cuda")
# load collection