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

style: add path to yaml

Browse Source
This commit is contained in:
0Xiao0
2025-04-03 18:10:32 +08:00
parent 1ee838327e
commit eb9ec3c0bf
8 changed files with 98 additions and 49 deletions
Download Patch File Download Diff File Expand all files Collapse all files

33
src/blackbox/chroma_query.py
View File

@ -12,22 +12,29 @@ from ..log.logging_time import logging_time
import re
from sentence_transformers import CrossEncoder
from pathlib import Path
from ..configuration import Configuration
from ..configuration import PathConf
logger = logging.getLogger
DEFAULT_COLLECTION_ID = "123"
from injector import singleton
@singleton
class ChromaQuery(Blackbox):
def __init__(self, *args, **kwargs) -> None:
# config = read_yaml(args[0])
# load chromadb and embedding model
self.embedding_model_1 = embedding_functions.SentenceTransformerEmbeddingFunction(model_name="/Workspace/Models/BAAI/bge-large-zh-v1.5", device = "cuda:0")
self.embedding_model_2 = embedding_functions.SentenceTransformerEmbeddingFunction(model_name="/Workspace/Models/BAAI/bge-small-en-v1.5", device = "cuda:0")
self.embedding_model_3 = embedding_functions.SentenceTransformerEmbeddingFunction(model_name="/Workspace/Models/BAAI/bge-m3", device = "cuda:0")
self.client_1 = chromadb.HttpClient(host='10.6.44.141', port=7000)
path = PathConf(Configuration())
self.model_path = Path(path.chroma_rerank_embedding_model)
self.embedding_model_1 = embedding_functions.SentenceTransformerEmbeddingFunction(model_name=str(self.model_path / "bge-large-zh-v1.5"), device = "cuda:0")
self.embedding_model_2 = embedding_functions.SentenceTransformerEmbeddingFunction(model_name=str(self.model_path / "bge-small-en-v1.5"), device = "cuda:0")
self.embedding_model_3 = embedding_functions.SentenceTransformerEmbeddingFunction(model_name=str(self.model_path / "bge-m3"), device = "cuda:0")
self.client_1 = chromadb.HttpClient(host='localhost', port=7000)
# self.client_2 = chromadb.HttpClient(host='10.6.82.192', port=8000)
self.reranker_model_1 = CrossEncoder("/Workspace/Models/BAAI/bge-reranker-v2-m3", max_length=512, device = "cuda")
self.reranker_model_1 = CrossEncoder(str(self.model_path / "bge-reranker-v2-m3"), max_length=512, device = "cuda")
def __call__(self, *args, **kwargs):
return self.processing(*args, **kwargs)
@ -57,10 +64,10 @@ 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 = "/Workspace/Models/BAAI/bge-large-zh-v1.5"
chroma_embedding_model = str(self.model_path / "bge-large-zh-v1.5")
if chroma_host is None or chroma_host.isspace() or chroma_host == "":
chroma_host = "10.6.44.141"
chroma_host = "localhost"
if chroma_port is None or chroma_port.isspace() or chroma_port == "":
chroma_port = "7000"
@ -72,7 +79,7 @@ class ChromaQuery(Blackbox):
chroma_n_results = 10
# load client and embedding model from init
if re.search(r"10.6.44.141", chroma_host) and re.search(r"7000", chroma_port):
if re.search(r"localhost", chroma_host) and re.search(r"7000", chroma_port):
client = self.client_1
else:
try:
@ -80,11 +87,11 @@ class ChromaQuery(Blackbox):
except:
return JSONResponse(content={"error": "chroma client not found"}, status_code=status.HTTP_400_BAD_REQUEST)
if re.search(r"/Workspace/Models/BAAI/bge-large-zh-v1.5", chroma_embedding_model):
if re.search(str(self.model_path / "bge-large-zh-v1.5"), chroma_embedding_model):
embedding_model = self.embedding_model_1
elif re.search(r"/Workspace/Models/BAAI/bge-small-en-v1.5", chroma_embedding_model):
elif re.search(str(self.model_path / "bge-small-en-v1.5"), chroma_embedding_model):
embedding_model = self.embedding_model_2
elif re.search(r"/Workspace/Models/BAAI/bge-m3", chroma_embedding_model):
elif re.search(str(self.model_path / "bge-m3"), chroma_embedding_model):
embedding_model = self.embedding_model_3
else:
try:
@ -123,7 +130,7 @@ class ChromaQuery(Blackbox):
final_result = str(results["documents"])
if chroma_reranker_model:
if re.search(r"/Workspace/Models/BAAI/bge-reranker-v2-m3", chroma_reranker_model):
if re.search(str(self.model_path / "bge-reranker-v2-m3"), chroma_reranker_model):
reranker_model = self.reranker_model_1
else:
try: