gyt/opendan
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add sentence-transformer local text embedding supports

Browse Source
This commit is contained in:
liyaxing
2023-09-25 16:22:15 +08:00
committed by tsukasa
parent 1de94a4d06
commit f07976366f
3 changed files with 7 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/aios_kernel/__init__.py
+1
View File
@@ -25,5 +25,6 @@ from .local_stability_node import Local_Stability_ComputeNode
from .stability_node import Stability_ComputeNode
from .local_st_compute_node import LocalSentenceTransformer_Text_ComputeNode,LocalSentenceTransformer_Image_ComputeNode
AIOS_Version = "0.5.1, build 2023-9-27"
src/aios_kernel/local_st_compute_node.py
+4 -4
View File
@@ -58,7 +58,7 @@ class LocalSentenceTransformer_Text_ComputeNode(Queue_ComputeNode):
logger.debug(
f"LocalSentenceTransformer_Text_ComputeNode task input: {input}"
)
sentence_embeddings = self.model.encode(input).tolist()
sentence_embeddings = self.model.encode(input, show_progress_bar=False).tolist()
# logger.debug(f"LocalSentenceTransformer_Text_ComputeNode task sentence_embeddings: {sentence_embeddings}")
return {
"state": ComputeTaskState.DONE,
@@ -87,7 +87,7 @@ class LocalSentenceTransformer_Text_ComputeNode(Queue_ComputeNode):
pass
def is_support(self, task: ComputeTask) -> bool:
return task.task_type == ComputeTaskType.TEXT_EMBEDDING
return task.task_type == ComputeTaskType.TEXT_EMBEDDING and task.params["model_name"] == "all-MiniLM-L6-v2"
def is_local(self) -> bool:
return True
@@ -187,7 +187,7 @@ class LocalSentenceTransformer_Image_ComputeNode(Queue_ComputeNode):
logger.debug(
f"LocalSentenceTransformer_Text_ComputeNode task text input: {input}"
)
sentence_embeddings = self.multi_model.encode(input)
sentence_embeddings = self.multi_model.encode(input, show_progress_bar=False).tolist()
# logger.debug(f"LocalSentenceTransformer_Text_ComputeNode task sentence_embeddings: {sentence_embeddings}")
return {
"state": ComputeTaskState.DONE,
@@ -238,7 +238,7 @@ class LocalSentenceTransformer_Image_ComputeNode(Queue_ComputeNode):
def is_support(self, task: ComputeTask) -> bool:
return (
task.task_type == ComputeTaskType.TEXT_EMBEDDING
(task.task_type == ComputeTaskType.TEXT_EMBEDDING and task.params["model_name"] == "clip-ViT-B-32")
or task.task_type == ComputeTaskType.IMAGE_EMBEDDING
)