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

test/test_embedding.py

@@ -0,0 +1,93 @@
+import sys
+import os
+import logging
+from sentence_transformers import SentenceTransformer, util
+
+
+dir_path = os.path.dirname(os.path.realpath(__file__))
+print(dir_path)
+
+sys.path.append("{}/../src/".format(dir_path))
+print(sys.path)
+
+root = logging.getLogger()
+root.setLevel(logging.DEBUG)
+handler = logging.StreamHandler(sys.stdout)
+handler.setLevel(logging.DEBUG)
+formatter = logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s")
+handler.setFormatter(formatter)
+root.addHandler(handler)
+
+
+def test_st():
+    model = SentenceTransformer("all-MiniLM-L6-v2")
+
+    # Our sentences we like to encode
+    sentences = [
+        "This framework generates embeddings for each input sentence",
+        "Sentences are passed as a list of string.",
+        "The quick brown fox jumps over the lazy dog.",
+    ]
+
+    # Sentences are encoded by calling model.encode()
+    sentence_embeddings = model.encode(sentences)
+
+    # Print the embeddings
+    for sentence, embedding in zip(sentences, sentence_embeddings):
+        print("Sentence:", sentence)
+        print("Embedding:", embedding)
+        print("")
+
+    # Single list of sentences
+
+    """
+    sentences = [
+        "The cat sits outside",
+        "A man is playing guitar",
+        "I love pasta",
+        "The new movie is awesome",
+        "The cat plays in the garden",
+        "A woman watches TV",
+        "The new movie is so great",
+        "Do you like pizza?",
+    ]
+    """
+    sentences = [
+        "猫坐在外面",
+        "狗坐在上面",
+        "狗坐在里面",
+        "一个男人在弹吉他",
+        "我爱意大利面",
+        "新电影太精彩了",
+        "猫在花园里玩耍",
+        "一个女人在看电视",
+        "新电影太棒了",
+        "你喜欢披萨吗？",
+    ]
+
+    # Compute embeddings
+    embeddings = model.encode(sentences, convert_to_tensor=True)
+
+    # Compute cosine-similarities for each sentence with each other sentence
+    cosine_scores = util.cos_sim(embeddings, embeddings)
+
+    # Find the pairs with the highest cosine similarity scores
+    pairs = []
+    for i in range(len(cosine_scores) - 1):
+        for j in range(i + 1, len(cosine_scores)):
+            pairs.append({"index": [i, j], "score": cosine_scores[i][j]})
+
+    # Sort scores in decreasing order
+    pairs = sorted(pairs, key=lambda x: x["score"], reverse=True)
+
+    for pair in pairs[0:10]:
+        i, j = pair["index"]
+        print(
+            "{} \t\t {} \t\t Score: {:.4f}".format(
+                sentences[i], sentences[j], pair["score"]
+            )
+        )
+
+
+if __name__ == "__main__":
+    test_st()