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

add asr and tts with settings

Browse Source
This commit is contained in:
0Xiao0
2024-08-22 15:26:27 +08:00
parent abb5195e55
commit c247d611a0
28 changed files with 626 additions and 82 deletions
Download Patch File Download Diff File Expand all files Collapse all files

18
src/blackbox/chroma_query.py
View File

@ -20,9 +20,9 @@ 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="/home/administrator/Workspace/Models/BAAI/bge-large-zh-v1.5", device = "cuda")
self.embedding_model_2 = embedding_functions.SentenceTransformerEmbeddingFunction(model_name="/home/administrator/Workspace/Models/BAAI/bge-large-en-v1.5", device = "cuda")
self.client_1 = chromadb.HttpClient(host='172.16.5.8', port=7000)
self.embedding_model_1 = embedding_functions.SentenceTransformerEmbeddingFunction(model_name="/home/gpu/Workspace/Models/BAAI/bge-large-zh-v1.5", device = "cuda:1")
self.embedding_model_2 = embedding_functions.SentenceTransformerEmbeddingFunction(model_name="/home/gpu/Workspace/Models/BAAI/bge-small-en-v1.5", device = "cuda:1")
self.client_1 = chromadb.HttpClient(host='10.6.81.119', port=7000)
# self.client_2 = chromadb.HttpClient(host='10.6.82.192', port=8000)
def __call__(self, *args, **kwargs):
@ -51,10 +51,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 = "/home/administrator/Workspace/Models/BAAI/bge-large-zh-v1.5"
chroma_embedding_model = "/home/gpu/Workspace/Models/BAAI/bge-large-zh-v1.5"
if chroma_host is None or chroma_host.isspace() or chroma_host == "":
chroma_host = "172.16.5.8"
chroma_host = "10.6.81.119"
if chroma_port is None or chroma_port.isspace() or chroma_port == "":
chroma_port = "7000"
@ -66,17 +66,17 @@ class ChromaQuery(Blackbox):
chroma_n_results = 3
# load client and embedding model from init
if re.search(r"172.16.5.8", chroma_host) and re.search(r"7000", chroma_port):
if re.search(r"10.6.81.119", chroma_host) and re.search(r"7000", chroma_port):
client = self.client_1
else:
client = chromadb.HttpClient(host=chroma_host, port=chroma_port)
if re.search(r"/home/administrator/Workspace/Models/BAAI/bge-large-zh-v1.5", chroma_embedding_model):
if re.search(r"/home/gpu/Workspace/Models/BAAI/bge-large-zh-v1.5", chroma_embedding_model):
embedding_model = self.embedding_model_1
elif re.search(r"/home/administrator/Workspace/Models/BAAI/bge-large-en-v1.5", chroma_embedding_model):
elif re.search(r"/home/gpu/Workspace/Models/BAAI/bge-small-en-v1.5", chroma_embedding_model):
embedding_model = self.embedding_model_2
else:
embedding_model = embedding_functions.SentenceTransformerEmbeddingFunction(model_name=chroma_embedding_model, device = "cuda")
embedding_model = embedding_functions.SentenceTransformerEmbeddingFunction(model_name=chroma_embedding_model, device = "cuda:1")
# load collection
collection = client.get_collection(chroma_collection_id, embedding_function=embedding_model)