From 87f82ec9e2dcf06dd666fba41e88e388fb1f6064 Mon Sep 17 00:00:00 2001
From: liyaxing <liyaxing@buckyos.com>
Date: Wed, 27 Sep 2023 17:24:38 +0800
Subject: [PATCH] Add image embedding task and test case

---
 src/aios_kernel/compute_task.py | 15 ++++++++++++++-
 src/requirements.txt            |  1 +
 test/test_embedding.py          |  5 ++++-
 test/test_object.py             | 15 +++++++++++++--
 4 files changed, 32 insertions(+), 4 deletions(-)

diff --git a/src/aios_kernel/compute_task.py b/src/aios_kernel/compute_task.py
index 12d320c..6569fcb 100644
--- a/src/aios_kernel/compute_task.py
+++ b/src/aios_kernel/compute_task.py
@@ -2,6 +2,8 @@
 from enum import Enum
 import uuid
 import time
+from typing import Union
+from knowledge import ObjectID
 
 class ComputeTaskResultCode(Enum):
     OK = 0
@@ -61,7 +63,7 @@ class ComputeTask:
         if inner_functions is not None:
             self.params["inner_functions"] = inner_functions
 
-    def set_text_embedding_params(self, input, model_name=None, callchain_id = None):
+    def set_text_embedding_params(self, input: str, model_name=None, callchain_id = None):
         self.task_type = ComputeTaskType.TEXT_EMBEDDING
         self.create_time = time.time()
         self.task_id = uuid.uuid4().hex
@@ -71,6 +73,17 @@ class ComputeTask:
         else:
             self.params["model_name"] = "text-embedding-ada-002"
         self.params["input"] = input
+        
+    def set_image_embedding_params(self, input = Union[ObjectID, bytes], model_name=None, callchain_id = None):
+        self.task_type = ComputeTaskType.IMAGE_EMBEDDING
+        self.create_time = time.time()
+        self.task_id = uuid.uuid4().hex
+        self.callchain_id = callchain_id
+        if model_name is not None:
+            self.params["model_name"] = model_name
+        else:
+            self.params["model_name"] = None
+        self.params["input"] = input
     
     def set_text_2_image_params(self, prompt: str, model_name, callchain_id=None):
         self.task_type = ComputeTaskType.TEXT_2_IMAGE
diff --git a/src/requirements.txt b/src/requirements.txt
index ec9ef4f..c6128ee 100644
--- a/src/requirements.txt
+++ b/src/requirements.txt
@@ -137,3 +137,4 @@ python-telegram-bot
 pydub
 stability_sdk
 sentence-transformers==2.2.2
+tiktoken
diff --git a/test/test_embedding.py b/test/test_embedding.py
index e9454f7..5b14dda 100644
--- a/test/test_embedding.py
+++ b/test/test_embedding.py
@@ -67,7 +67,10 @@ def test_st():
     ]
 
     # Compute embeddings
-    embeddings = model.encode(sentences, convert_to_tensor=True)
+    #embeddings = model.encode(sentences, convert_to_tensor=True)
+    embeddings = model.encode(sentences)
+    print("embeddings as follows: ")
+    print(embeddings)
 
     # Compute cosine-similarities for each sentence with each other sentence
     cosine_scores = util.cos_sim(embeddings, embeddings)
diff --git a/test/test_object.py b/test/test_object.py
index 94d99ca..e464aa0 100644
--- a/test/test_object.py
+++ b/test/test_object.py
@@ -26,12 +26,13 @@ from knowledge import (
     EmailObject,
     ImageObject,
 )
+from aios_kernel import LocalSentenceTransformer_Image_ComputeNode, ComputeTask
 import asyncio
 import unittest
 
 
-class TestVectorSTorage(unittest.TestCase):
-    def test_object(self):
+class TestVectorSTorage(unittest.IsolatedAsyncioTestCase):
+    async def test_object(self):
         data = HashValue.hash_data("1233".encode("utf-8"))
         print(data.to_base58())
         print(data.to_base36())
@@ -65,6 +66,15 @@ class TestVectorSTorage(unittest.TestCase):
         image_id = images[image_keys[1]]
         print(f"got image object: {image_keys[1]} {image_id.to_base58()}")
         
+        node = LocalSentenceTransformer_Image_ComputeNode();
+        ret = node.initial()
+        self.assertEqual(ret, True)
+        
+        task = ComputeTask()
+        task.set_image_embedding_params(image_id)
+        ret = await node.execute_task(task)
+        print(ret)
+        '''
         buf = KnowledgeStore().get_object_store().get_object(image_id)
         image_obj= ImageObject.decode(buf)
         file_size = image_obj.get_file_size()
@@ -83,6 +93,7 @@ class TestVectorSTorage(unittest.TestCase):
         #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"))