liyaxing
125 lines
4.1 KiB
Python
125 lines
4.1 KiB
Python
import sys
|
|
import os
|
|
import logging
|
|
|
|
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)
|
|
|
|
|
|
from knowledge import (
|
|
ObjectID,
|
|
HashValue,
|
|
EmailObjectBuilder,
|
|
ObjectRelationStore,
|
|
KnowledgeStore,
|
|
EmailObject,
|
|
ImageObject,
|
|
)
|
|
from aios_kernel import LocalSentenceTransformer_Image_ComputeNode, ComputeTask
|
|
import asyncio
|
|
import unittest
|
|
|
|
|
|
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())
|
|
|
|
data2 = HashValue.from_base58(data.to_base58())
|
|
self.assertEqual(data.to_base36(), data2.to_base36())
|
|
|
|
data2 = HashValue.from_base36(data.to_base36())
|
|
self.assertEqual(data.to_base58(), data2.to_base58())
|
|
|
|
email_folder = "F:\\system\\Downloads\\8081ffdb80925f5bff9c6ab9c4756c7d"
|
|
email_object = EmailObjectBuilder({}, email_folder).build()
|
|
|
|
id = email_object.calculate_id()
|
|
print(f"got email object: {id.to_base58()}")
|
|
|
|
# test encode & decode
|
|
ret = email_object.encode()
|
|
obj = EmailObject.decode(ret)
|
|
id2 = obj.calculate_id()
|
|
print(f"got email object: {id2.to_base58()}")
|
|
self.assertEqual(id.to_base58(), id2.to_base58())
|
|
|
|
ret2 = obj.encode()
|
|
self.assertEqual(ret, ret2)
|
|
|
|
images = email_object.get_rich_text().get_images()
|
|
image_keys = list(images.keys())
|
|
print("got image list: ", image_keys)
|
|
|
|
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()
|
|
print(f"got image object: {image_id.to_base58()}, size: {file_size}")
|
|
|
|
|
|
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"))
|
|
obj2 = ObjectID.hash_data("67890".encode("utf-8"))
|
|
obj3 = ObjectID.hash_data("abcde".encode("utf-8"))
|
|
obj4 = ObjectID.hash_data("fghij".encode("utf-8"))
|
|
print(obj1.to_base58(), obj2.to_base58(), obj3.to_base58())
|
|
relation_store = KnowledgeStore().get_relation_store()
|
|
relation_store.add_relation(obj1, obj2)
|
|
relation_store.add_relation(obj1, obj2)
|
|
relation_store.add_relation(obj2, obj3)
|
|
|
|
relation_store.add_relation(obj1, obj3)
|
|
relation_store.add_relation(obj1, obj4)
|
|
|
|
objs = relation_store.get_related_objects(obj2)
|
|
self.assertEqual(len(objs), 1)
|
|
self.assertEqual(objs[0], obj3)
|
|
|
|
objs = relation_store.get_related_root_objects(obj1)
|
|
self.assertEqual(len(objs), 2)
|
|
self.assertEqual(obj3 in objs, True)
|
|
self.assertEqual(obj4 in objs, True)
|
|
# self.assertCountEqual(objs, [obj3, obj4])
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|