diff --git a/src/aios_kernel/compute_task.py b/src/aios_kernel/compute_task.py index 05125ee..12d320c 100644 --- a/src/aios_kernel/compute_task.py +++ b/src/aios_kernel/compute_task.py @@ -25,6 +25,7 @@ class ComputeTaskType(Enum): VOICE_2_TEXT = "voice_2_text" TEXT_2_VOICE = "text_2_voice" TEXT_EMBEDDING ="text_embedding" + IMAGE_EMBEDDING ="image_embedding" class ComputeTask: diff --git a/src/aios_kernel/local_st_compute_node.py b/src/aios_kernel/local_st_compute_node.py index e73b9f9..87fe37d 100644 --- a/src/aios_kernel/local_st_compute_node.py +++ b/src/aios_kernel/local_st_compute_node.py @@ -5,6 +5,7 @@ from pydantic import BaseModel from .compute_task import ComputeTask, ComputeTaskState, ComputeTaskType from .queue_compute_node import Queue_ComputeNode +from knowledge import ObjectID logger = logging.getLogger(__name__) @@ -13,19 +14,20 @@ This is a custom implementation, it should be redesigned. """ -class LocalSentenceTransformer_ComputeNode(Queue_ComputeNode): +class LocalSentenceTransformer_Text_ComputeNode(Queue_ComputeNode): + # For valid pretrained models, see https://www.sbert.net/docs/pretrained_models.html def __init__(self, model_name: str = "all-MiniLM-L6-v2"): super().__init__() - self.node_id = "local_sentence_transformer_node" + self.node_id = "local_sentence_transformer_text_embedding_node" self.model_name = model_name self.model = None def initial(self) -> bool: logger.info( - f"LocalSentenceTransformer_ComputeNode init, model_name: {self.model_name}" + f"LocalSentenceTransformer_Text_ComputeNode init, model_name: {self.model_name}" ) - + assert self.model_name is not None assert self.model is None try: @@ -35,9 +37,9 @@ class LocalSentenceTransformer_ComputeNode(Queue_ComputeNode): except Exception as err: logger.error(f"load model {self.model} failed: {err}") return False - + return True - + async def execute_task( self, task: ComputeTask ) -> { @@ -51,14 +53,14 @@ class LocalSentenceTransformer_ComputeNode(Queue_ComputeNode): }, }: try: - # logger.debug(f"LocalSentenceTransformer_ComputeNode task: {task}") + # logger.debug(f"LocalSentenceTransformer_Text_ComputeNode task: {task}") if task.task_type == ComputeTaskType.TEXT_EMBEDDING: input = task.params["input"] logger.debug( - f"LocalSentenceTransformer_ComputeNode task input: {input}" + f"LocalSentenceTransformer_Text_ComputeNode task input: {input}" ) sentence_embeddings = self.model.encode(input) - # logger.debug(f"LocalSentenceTransformer_ComputeNode task sentence_embeddings: {sentence_embeddings}") + # logger.debug(f"LocalSentenceTransformer_Text_ComputeNode task sentence_embeddings: {sentence_embeddings}") return { "state": ComputeTaskState.DONE, "content": sentence_embeddings, @@ -80,9 +82,7 @@ class LocalSentenceTransformer_ComputeNode(Queue_ComputeNode): } def display(self) -> str: - return ( - f"LocalSentenceTransformer_ComputeNode: {self.node_id}, {self.model_name}" - ) + return f"LocalSentenceTransformer_Text_ComputeNode: {self.node_id}, {self.model_name}" def get_capacity(self): pass @@ -94,3 +94,152 @@ class LocalSentenceTransformer_ComputeNode(Queue_ComputeNode): def is_local(self) -> bool: return True + +from typing import Union +from PIL import Image +import io + +class LocalSentenceTransformer_Image_ComputeNode(Queue_ComputeNode): + # For valid pretrained models, see https://www.sbert.net/docs/pretrained_models.html + def __init__(self, model_name: str = "clip-ViT-B-32", multi_model_name: str = "clip-ViT-B-32-multilingual-v1"): + super().__init__() + + self.node_id = "local_sentence_transformer_image_embedding_node" + self.model_name = model_name + self.multi_model_name = multi_model_name + self.model = None + self.multi_model = None + + def initial(self) -> bool: + logger.info( + f"LocalSentenceTransformer_Image_ComputeNode init, model_name: {self.model_name} {self.multi_model_name}" + ) + + assert self.model_name is not None + assert self.multi_model_name is not None + assert self.model is None + assert self.multi_model is None + + try: + from sentence_transformers import SentenceTransformer + + self.model = SentenceTransformer(self.model_name) + self.multi_model = SentenceTransformer(self.multi_model_name) + except Exception as err: + logger.error(f"load model {self.model} failed: {err}") + return False + + return True + + def _load_image(self, source: Union[ObjectID, bytes] ) -> Optional[Image]: + image_data = None + if isinstance(source, ObjectID): + from knowledge import KnowledgeStore, ImageObject + + buf = KnowledgeStore().get_object_store().get_object(source) + if buf is None: + logger.error(f"load image object but not found! {source}") + return None + + try: + image_obj= ImageObject.decode(buf) + except Exception as err: + logger.error(f"decode ImageObject from buffer failed: {source}, {err}") + return None + + file_size = image_obj.get_file_size() + print(f"got image object: {source.to_base58()}, size: {file_size}") + + image_data = KnowledgeStore().get_chunk_reader().read_chunk_list_to_single_bytes(image_obj.get_chunk_list()) + + elif isinstance(source, bytes): + image_data = source + else: + logger.error(f"unsupport image source type: {type(source)}, {source}") + return None + + try: + + img = Image.open(io.BytesIO(image_data)) + + return img + except Exception as err: + logger.error(f"load image from buffer failed: {source}, {err}") + return None + + 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_Text_ComputeNode task: {task}") + if task.task_type == ComputeTaskType.TEXT_EMBEDDING: + input = task.params["input"] + logger.debug( + f"LocalSentenceTransformer_Text_ComputeNode task text input: {input}" + ) + sentence_embeddings = self.multi_model.encode(input) + # logger.debug(f"LocalSentenceTransformer_Text_ComputeNode task sentence_embeddings: {sentence_embeddings}") + return { + "state": ComputeTaskState.DONE, + "content": sentence_embeddings, + "message": None, + } + elif task.task_type == ComputeTaskType.IMAGE_EMBEDDING: + input = task.params["input"] + logger.debug( + f"LocalSentenceTransformer_Image_ComputeNode task image input: {input}" + ) + + img = self._load_image(input) + if img is None: + return { + "state": ComputeTaskState.ERROR, + "error": {"code": -1, "message": "load image failed"}, + } + + sentence_embeddings = self.model.encode(img, convert_to_tensor=True) + + # logger.debug(f"LocalSentenceTransformer_Text_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_Image_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 + or task.task_type == ComputeTaskType.IMAGE_EMBEDDING + ) + + def is_local(self) -> bool: + return True diff --git a/src/knowledge/object/blob.py b/src/knowledge/object/blob.py index b005e62..08fb0bf 100644 --- a/src/knowledge/object/blob.py +++ b/src/knowledge/object/blob.py @@ -1,6 +1,8 @@ import os import shutil from .object import ObjectID +import logging +logger = logging.getLogger(__name__) class FileBlobStorage: @@ -38,16 +40,23 @@ class FileBlobStorage: def put(self, object_id: ObjectID, contents: bytes): full_path = self.get_full_path(object_id) + if os.path.exists(full_path): + logger.warning(f"will replace object: {object_id}") + self.write_sync(full_path, contents) def get(self, object_id: ObjectID) -> bytes: full_path = self.get_full_path(object_id) + if not os.path.exists(full_path): + return None + with open(full_path, "rb") as f: return f.read() def delete(self, object_id: ObjectID): full_path = self.get_full_path(object_id) - os.remove(full_path) + if os.path.exists(full_path): + os.remove(full_path) def exists(self, object_id: ObjectID) -> bool: full_path = self.get_full_path(object_id) diff --git a/test/test_embedding.py b/test/test_embedding.py index 7d9f649..e9454f7 100644 --- a/test/test_embedding.py +++ b/test/test_embedding.py @@ -20,6 +20,7 @@ root.addHandler(handler) def test_st(): + image_model = SentenceTransformer('clip-ViT-B-32-multilingual-v1') model = SentenceTransformer("all-MiniLM-L6-v2") # Our sentences we like to encode diff --git a/test/test_object.py b/test/test_object.py index 53dc6b0..94d99ca 100644 --- a/test/test_object.py +++ b/test/test_object.py @@ -74,6 +74,15 @@ class TestVectorSTorage(unittest.TestCase): image_data = KnowledgeStore().get_chunk_reader().read_chunk_list_to_single_bytes(image_obj.get_chunk_list()) self.assertEqual(file_size, len(image_data)) + from PIL import Image + import io + image = Image.open(io.BytesIO(image_data)) + image.show() + + from sentence_transformers import SentenceTransformer + #model = SentenceTransformer('clip-ViT-B-32-multilingual-v1') + model = SentenceTransformer('clip-ViT-B-32') + model.encode(image, convert_to_tensor=True) def test_relation(self): obj1 = ObjectID.hash_data("12345".encode("utf-8"))