gyt/opendan
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add chunk store and chunk list helper

Browse Source
This commit is contained in:
liyaxing
2023-09-08 15:52:05 +08:00
parent 1e5b9afd0f
commit e8e3812d28
6 changed files with 164 additions and 28 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/knowledge/data/__init__.py
+2
View File
@@ -1,2 +1,4 @@
from .chunk import ChunkID, PositionType, PositionFileRange
from .tracker import ChunkTracker
from .chunk_store import ChunkStore
from .gen import ChunkListGenerator
src/knowledge/data/chunk.py
+1 -1
View File
@@ -8,7 +8,7 @@ class PositionType(Enum):
Device = 2
File = 3
FileRange = 4
ChunkManager = 5
ChunkStore = 5
class PositionFileRange:
src/knowledge/data/chunk_store.py
+26
View File
@@ -0,0 +1,26 @@
import os
import logging
from ..object import FileBlobStorage
from .chunk import ChunkID
class ChunkStore:
def __init__(self, root_dir: str):
logging.info(f"will init chunk store, root_dir={root_dir}")
if not os.path.exists(root_dir):
os.makedirs(root_dir)
self.root = root_dir
self.blob = FileBlobStorage(root_dir)
def put_chunk(self, chunk_id: ChunkID, contents: bytes):
self.blob.put(chunk_id, contents)
def get_chunk(self, chunk_id: ChunkID) -> bytes:
return self.blob.get(chunk_id)
def delete_chunk(self, chunk_id: ChunkID):
self.blob.delete(chunk_id)
src/knowledge/data/gen.py
+94
View File
@@ -0,0 +1,94 @@
import os
import hashlib
import re
from typing import Tuple, List
from .chunk_store import ChunkStore
from .chunk import ChunkID, PositionFileRange, PositionType
from ..object import HashValue
from .tracker import ChunkTracker
class ChunkListGenerator:
def __init__(self, chunk_store: ChunkStore, chunk_tracker: ChunkTracker):
self.chunk_store = chunk_store
self.chunk_tracker = chunk_tracker
def create_chunk_list_from_file(
self, file_path: str, chunk_size: int, restore: bool
) -> Tuple[List[ChunkID], HashValue]:
assert (
chunk_size % (1024 * 1024) == 0
), "chunk size should be an integral multiple of 1MB"
chunk_list = []
hash_obj = hashlib.sha256()
with open(file_path, "rb") as file:
while True:
chunk = file.read(chunk_size)
if not chunk:
break
chunk_id = ChunkID.hash_data(chunk)
chunk_list.append(chunk_id)
hash_obj.update(chunk)
if restore:
self.chunk_tracker.add_position(
chunk_id, file_path, PositionType.ChunkStore
)
self.chunk_store.put_chunk(chunk_id, chunk)
else:
file_range = PositionFileRange(
file_path, file.tell() - chunk_size, chunk_size
)
self.chunk_tracker.add_position(
chunk_id, file_range, PositionType.FileRange
)
file_hash = HashValue(hash_obj.digest())
print(f"calc file hash: {file_path}, {file_hash}")
return chunk_list, file_hash
def create_chunk_list_from_text(
self, text: str, chunk_max_words: int, separator_chars: str = ".,"
) -> Tuple[List[ChunkID], HashValue]:
text_list = self._split_text_list(text, chunk_max_words, separator_chars)
chunk_list = []
hash_obj = hashlib.sha256()
for text in text_list:
chunk_bytes = text.encode("utf-8")
hash_obj.update(chunk_bytes)
chunk_id = ChunkID.hash_data(chunk_bytes)
chunk_list.append(chunk_id)
self.chunk_tracker.add_position(
chunk_id, "", PositionType.ChunkStore
)
self.chunk_store.put_chunk(chunk_id, chunk_bytes)
hash = HashValue(hash_obj.digest())
return chunk_list, hash
@staticmethod
def _split_text_list(
text: str, chunk_max_words: int, separator_chars: str = ".,"
) -> List[str]:
sentences = re.split(f"[{separator_chars}]", text)
chunk_list = []
chunk = []
word_count = 0
for sentence in sentences:
words = sentence.split()
for word in words:
if word_count < chunk_max_words:
chunk.append(word)
word_count += 1
else:
chunk_list.append(" ".join(chunk))
chunk = [word]
word_count = 1
if chunk:
chunk_list.append(" ".join(chunk))
return chunk_list
src/knowledge/data/tracker.py
+17 -23
View File
@@ -2,31 +2,17 @@ import sqlite3
import time
import logging
import os
from .chunk import ChunkID, PositionType, PositionFileRange
from .chunk import ChunkID, PositionType, PositionFileRange
class ChunkTracker:
_instance = None
def __init__(self, root_dir: str):
if not os.path.exists(root_dir):
os.makedirs(root_dir)
file = os.path.join(root_dir, "chunk_tracker.db")
logging.info(f"will init chunk tracker, db={file}")
def __new__(cls):
if cls._instance is None:
cls._instance = super().__new__(cls)
directory = os.path.join(
os.path.dirname(__file__), "../../../rootfs/data/chunk/"
)
if not os.path.exists(directory):
os.mkdir(directory)
file = os.path.join(directory, "chunk_tracker.db")
logging.info(f"will init chunk tracker, db={file}")
cls._instance.__singleton_init__(file)
return cls._instance
def __singleton_init__(self, db_path: str):
logging.info(f"will init chunk tracker, db={db_path}")
self.conn = sqlite3.connect(db_path)
self.conn = sqlite3.connect(file)
self.cursor = self.conn.cursor()
self.cursor.execute(
"""
@@ -43,7 +29,9 @@ class ChunkTracker:
)
self.conn.commit()
def add_position(self, chunk_id: ChunkID, position: str, position_type: PositionType):
def add_position(
self, chunk_id: ChunkID, position: str, position_type: PositionType
):
logging.debug(f"add chunk position: {chunk_id}, {position}, {position_type}")
insert_time = update_time = int(time.time())
@@ -52,7 +40,13 @@ class ChunkTracker:
INSERT OR REPLACE INTO chunks (id, pos, pos_type, insert_time, update_time)
VALUES (?, ?, ?, ?, ?)
""",
(str(chunk_id), position, position_type.value, insert_time, update_time,),
(
str(chunk_id),
position,
position_type.value,
insert_time,
update_time,
),
)
self.conn.commit()
test/test_chunk.py
+22 -2
View File
@@ -7,16 +7,27 @@ print(dir_path)
sys.path.append("{}/../src/".format(dir_path))
print(sys.path)
from knowledge import ChunkTracker, ChunkID, HashValue, PositionType
import logging
import sys
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 ChunkTracker, ChunkID, HashValue, PositionType, KnowledgeStore, ChunkListGenerator
import asyncio
import unittest
class TestChunk(unittest.TestCase):
def test_chunk_tracker(self):
tracker = ChunkTracker()
tracker = KnowledgeStore().get_chunk_tracker()
hash = HashValue.hash_data("1234567890".encode("utf-8"));
cid = ChunkID.new_chunk_id(hash)
@@ -30,7 +41,16 @@ class TestChunk(unittest.TestCase):
ret = tracker.get_position(cid)
self.assertEqual(ret, None)
def test_chunk(self):
gen = ChunkListGenerator(KnowledgeStore().get_chunk_store(), KnowledgeStore().get_chunk_tracker())
gen.create_chunk_list_from_file("H:/test", 1024 * 1024, True)
# Read the file
text_file = "H:/test.txt"
with open(text_file, 'r', encoding = "utf-8") as file:
text = file.read()
gen.create_chunk_list_from_text(text, 1024)
if __name__ == "__main__":
unittest.main()