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

src/requirements.txt

@@ -80,10 +80,60 @@ toml>=0.10.0
 protobuf
 grpcio
 grpcio-status
+h11==0.14.0
+httpcore==0.17.3
+httptools==0.6.0
+httpx==0.24.1
+huggingface-hub==0.16.4
+humanfriendly==10.0
+idna==3.4
+imageio==2.31.3
+imageio-ffmpeg==0.4.8
+importlib-resources==6.0.1
+mail-parser==3.15.0
+monotonic==1.6
+moviepy==1.0.0
+mpmath==1.3.0
+multidict==6.0.4
+numpy==1.25.2
+onnxruntime==1.15.1
+openai==0.28.0
+overrides==7.4.0
+packaging==23.1
+pandas==2.1.0
+Pillow==10.0.0
+posthog==3.0.2
+proglog==0.1.10
+prompt-toolkit==3.0.39
+proto-plus==1.22.3
+protobuf
+pulsar-client==3.3.0
+pyasn1==0.5.0
+pyasn1-modules==0.3.0
+pydantic==1.10.12
+PyPika==0.48.9
+pyreadline3==3.4.1
+python-dateutil==2.8.2
+python-dotenv==1.0.0
+python-telegram-bot==20.5
+pytz==2023.3.post1
+PyYAML==6.0.1
+requests==2.31.0
+rsa==4.9
+simplejson==3.19.1
+six==1.16.0
+sniffio==1.3.0
+soupsieve==2.5
+starlette==0.27.0
+sympy==1.12
+telegram==0.0.1
+tokenizers==0.14.0
+toml==0.10.0
 pysocks
 chardet
 pydub
 aiosqlite
 python-telegram-bot
 pydub
-stability_sdk
+stability_sdk
+sentence-transformers==2.2