Add sentence-transformer local text embedding supports

src/aios_kernel/__init__.py

@@ -23,6 +23,7 @@ from .text_to_speech_function import TextToSpeechFunction
 from .workspace_env import WorkspaceEnvironment
 from .local_stability_node import Local_Stability_ComputeNode
 from .stability_node import Stability_ComputeNode
+from .local_st_compute_node import LocalSentenceTransformer_ComputeNode
 
 AIOS_Version = "0.5.1, build 2023-9-26"
 

src/aios_kernel/local_st_compute_node.py

@@ -0,0 +1,87 @@
+import logging
+import requests
+from typing import Optional, List
+from pydantic import BaseModel
+
+from .compute_task import ComputeTask, ComputeTaskState, ComputeTaskType
+from .queue_compute_node import Queue_ComputeNode
+
+logger = logging.getLogger(__name__)
+
+"""
+This is a custom implementation, it should be redesigned.
+"""
+
+
+class LocalSentenceTransformer_ComputeNode(Queue_ComputeNode):
+    def __init__(self, model_name: str = "all-MiniLM-L6-v2"):
+        super().__init__()
+
+        logger.info(
+            f"LocalSentenceTransformer_ComputeNode init, model_name: {model_name}"
+        )
+        self.model_name = model_name
+
+        try:
+            from sentence_transformers import SentenceTransformer
+
+            self.model = SentenceTransformer(self.model)
+        except Exception as err:
+            logger.error(f"load model {self.model} failed: {err}")
+
+    async def execute_task(
+        self, task: ComputeTask
+    ) -> {
+        "task_type": str,
+        "content": str,
+        "message": str,
+        "state": ComputeTaskState,
+        "error": {
+            "code": int,
+            "message": str,
+        },
+    }:
+        try:
+            # logger.debug(f"LocalSentenceTransformer_ComputeNode task: {task}")
+            if task.task_type == ComputeTaskType.TEXT_EMBEDDING:
+                input = task.params["input"]
+                logger.debug(
+                    f"LocalSentenceTransformer_ComputeNode task input: {input}"
+                )
+                sentence_embeddings = self.model.encode(input)
+                # logger.debug(f"LocalSentenceTransformer_ComputeNode task sentence_embeddings: {sentence_embeddings}")
+                return {
+                    "state": ComputeTaskState.DONE,
+                    "content": sentence_embeddings,
+                    "message": None,
+                }
+            else:
+                return {
+                    "state": ComputeTaskState.ERROR,
+                    "error": {"code": -1, "message": "unsupport embedding task type"},
+                }
+        except Exception as err:
+            import traceback
+
+            logger.error(f"{traceback.format_exc()}, error: {err}")
+
+            return {
+                "state": ComputeTaskState.ERROR,
+                "error": {"code": -1, "message": "unknown exception: " + str(err)},
+            }
+
+    def display(self) -> str:
+        return (
+            f"LocalSentenceTransformer_ComputeNode: {self.node_id}, {self.model_name}"
+        )
+
+    def get_capacity(self):
+        pass
+
+    def is_support(self, task: ComputeTask) -> bool:
+        return task.task_type == ComputeTaskType.TEXT_EMBEDDING and (
+            not task.params["model_name"] or task.params["model_name"] == "llama"
+        )
+
+    def is_local(self) -> bool:
+        return True