diff --git a/src/aios_kernel/__init__.py b/src/aios_kernel/__init__.py index 36ef512..2be558b 100644 --- a/src/aios_kernel/__init__.py +++ b/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" diff --git a/src/aios_kernel/local_st_compute_node.py b/src/aios_kernel/local_st_compute_node.py new file mode 100644 index 0000000..135797f --- /dev/null +++ b/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 diff --git a/src/requirements.txt b/src/requirements.txt index 7bd02d6..3f19c10 100644 --- a/src/requirements.txt +++ b/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 \ No newline at end of file +stability_sdk +sentence-transformers==2.2 diff --git a/test/test_embedding.py b/test/test_embedding.py new file mode 100644 index 0000000..7d9f649 --- /dev/null +++ b/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()