refactor read and chunk operators with no side effects

Author: ChenZiHong-Gavin

…phgen/models/storage/networkx_storage.py …models/storage/graph/networkx_storage.pygraphgen/models/storage/networkx_storage.py renamed to graphgen/models/storage/graph/networkx_storage.py graphgen/models/storage/networkx_storage.py renamed to graphgen/models/storage/graph/networkx_storage.py

@@ -1,7 +1,7 @@
 import os
 from dataclasses import dataclass
 
-from graphgen.bases.base_storage import BaseKVStorage, BaseListStorage
+from graphgen.bases.base_storage import BaseKVStorage
 from graphgen.utils import load_json, logger, write_json
 
 
@@ -54,41 +54,3 @@ def upsert(self, data: dict):
     def drop(self):
         if self._data:
             self._data.clear()
-
-
-@dataclass
-class JsonListStorage(BaseListStorage):
-    working_dir: str = None
-    namespace: str = None
-    _data: list = None
-
-    def __post_init__(self):
-        self._file_name = os.path.join(self.working_dir, f"{self.namespace}.json")
-        self._data = load_json(self._file_name) or []
-        logger.info("Load List %s with %d data", self.namespace, len(self._data))
-
-    @property
-    def data(self):
-        return self._data
-
-    def all_items(self) -> list:
-        return self._data
-
-    def index_done_callback(self):
-        write_json(self._data, self._file_name)
-
-    def get_by_index(self, index: int):
-        if index < 0 or index >= len(self._data):
-            return None
-        return self._data[index]
-
-    def append(self, data):
-        self._data.append(data)
-
-    def upsert(self, data: list):
-        left_data = [d for d in data if d not in self._data]
-        self._data.extend(left_data)
-        return left_data
-
-    def drop(self):
-        self._data = []

graphgen/models/storage/json_storage.py …aphgen/models/storage/kv/json_storage.pygraphgen/models/storage/json_storage.py renamed to graphgen/models/storage/kv/json_storage.py graphgen/models/storage/json_storage.py renamed to graphgen/models/storage/kv/json_storage.py

@@ -0,0 +1 @@
+from .chunk_service import ChunkService

graphgen/operators/chunk/__init__.py

@@ -0,0 +1,87 @@
+import asyncio
+import os
+from functools import lru_cache
+from typing import Union
+
+import pandas as pd
+from tqdm.asyncio import tqdm as tqdm_async
+
+from graphgen.models import (
+    ChineseRecursiveTextSplitter,
+    RecursiveCharacterSplitter,
+    Tokenizer,
+)
+from graphgen.utils import compute_content_hash, detect_main_language
+
+_MAPPING = {
+    "en": RecursiveCharacterSplitter,
+    "zh": ChineseRecursiveTextSplitter,
+}
+
+SplitterT = Union[RecursiveCharacterSplitter, ChineseRecursiveTextSplitter]
+
+
+@lru_cache(maxsize=None)
+def _get_splitter(language: str, frozen_kwargs: frozenset) -> SplitterT:
+    cls = _MAPPING[language]
+    kwargs = dict(frozen_kwargs)
+    return cls(**kwargs)
+
+
+def split_chunks(text: str, language: str = "en", **kwargs) -> list:
+    if language not in _MAPPING:
+        raise ValueError(
+            f"Unsupported language: {language}. "
+            f"Supported languages are: {list(_MAPPING.keys())}"
+        )
+    frozen_kwargs = frozenset(
+        (k, tuple(v) if isinstance(v, list) else v) for k, v in kwargs.items()
+    )
+    splitter = _get_splitter(language, frozen_kwargs)
+    return splitter.split_text(text)
+
+
+class ChunkService:
+    def __init__(self, **chunk_kwargs):
+        tokenizer_model = os.getenv("TOKENIZER_MODEL", "cl100k_base")
+        self.tokenizer_instance: Tokenizer = Tokenizer(model_name=tokenizer_model)
+        self.chunk_kwargs = chunk_kwargs
+
+    def __call__(self, batch: pd.DataFrame) -> pd.DataFrame:
+        docs = batch.to_dict(orient="records")
+        return pd.DataFrame(self.chunk_documents(docs))
+
+    def chunk_documents(self, new_docs: list) -> list:
+        for doc in new_docs:
+            doc_id = doc.get("_doc_id")
+            doc_type = doc.get("type")
+
+            if doc_type == "text":
+                doc_language = detect_main_language(doc["content"])
+                text_chunks = split_chunks(
+                    doc["content"],
+                    language=doc_language,
+                    **self.chunk_kwargs,
+                )
+
+                return [
+                    {
+                        "_chunk_id": compute_content_hash(chunk_text, prefix="chunk-"),
+                        "content": chunk_text,
+                        "type": "text",
+                        "_doc_id": doc_id,
+                        "length": len(self.tokenizer_instance.encode(chunk_text))
+                        if self.tokenizer_instance
+                        else len(chunk_text),
+                        "language": doc_language,
+                    }
+                    for chunk_text in text_chunks
+                ]
+
+            # other types of documents(images, sequences) are not chunked
+            return [
+                {
+                    "_chunk_id": doc_id.replace("doc-", f"{doc_type}-"),
+                    **doc,
+                }
+            ]

graphgen/operators/chunk/chunk_service.py

@@ -12,7 +12,7 @@
     RDFReader,
     TXTReader,
 )
-from graphgen.utils import logger
+from graphgen.utils import compute_mm_hash, logger
 
 from .parallel_file_scanner import ParallelFileScanner
 
@@ -110,10 +110,16 @@ def read(
             return ray.data.from_items([])
 
         if len(read_tasks) == 1:
-            logger.info("[READ] Successfully read files from %s", input_path)
-            return read_tasks[0]
-        # len(read_tasks) > 1
-        combined_ds = read_tasks[0].union(*read_tasks[1:])
+            combined_ds = read_tasks[0]
+        else:
+            combined_ds = read_tasks[0].union(*read_tasks[1:])
+
+        combined_ds = combined_ds.map(
+            lambda record: {
+                **record,
+                "_doc_id": compute_mm_hash(record),
+            }
+        )
 
         logger.info("[READ] Successfully read files from %s", input_path)
         return combined_ds