Add sentence-transformer local text embedding supports

src/aios_kernel/__init__.py

@@ -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

@@ -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
         )
 

test/test_embedding.py

@@ -67,11 +67,13 @@ def test_st():
     ]
 
     # Compute embeddings
+
     #embeddings = model.encode(sentences, convert_to_tensor=True)
     embeddings = model.encode(sentences)
     print("embeddings as follows: ")
     print(embeddings)
 
+
     # Compute cosine-similarities for each sentence with each other sentence
     cosine_scores = util.cos_sim(embeddings, embeddings)