fix: delete param parallelism for readers

Author: ChenZiHong-Gavin

graphgen/models/reader/csv_reader.py

@@ -14,20 +14,15 @@ class CSVReader(BaseReader):
         - if type is "text", "content" column must be present.
     """
 
-    def read(
-        self,
-        input_path: Union[str, List[str]],
-        parallelism: int = None,
-    ) -> Dataset:
+    def read(self, input_path: Union[str, List[str]]) -> Dataset:
         """
         Read CSV files and return Ray Dataset.
 
         :param input_path: Path to CSV file or list of CSV files.
-        :param parallelism: Number of blocks for Ray Dataset reading.
         :return: Ray Dataset containing validated and filtered data.
         """
 
-        ds = ray.data.read_csv(input_path, override_num_blocks=parallelism)
+        ds = ray.data.read_csv(input_path)
         ds = ds.map_batches(self._validate_batch, batch_format="pandas")
         ds = ds.filter(self._should_keep_item)
         return ds

graphgen/models/reader/json_reader.py

@@ -14,19 +14,14 @@ class JSONReader(BaseReader):
         - if type is "text", "content" column must be present.
     """
 
-    def read(
-        self,
-        input_path: Union[str, List[str]],
-        parallelism: int = 4,
-    ) -> Dataset:
+    def read(self, input_path: Union[str, List[str]]) -> Dataset:
         """
         Read JSON file and return Ray Dataset.
         :param input_path: Path to JSON/JSONL file or list of JSON/JSONL files.
-        :param parallelism: Number of parallel workers for reading files.
         :return: Ray Dataset containing validated and filtered data.
         """
 
-        ds = ray.data.read_json(input_path, override_num_blocks=parallelism)
+        ds = ray.data.read_json(input_path)
         ds = ds.map_batches(self._validate_batch, batch_format="pandas")
         ds = ds.filter(self._should_keep_item)
         return ds

graphgen/models/reader/parquet_reader.py

@@ -14,22 +14,17 @@ class ParquetReader(BaseReader):
     - if type is "text", "content" column must be present.
     """
 
-    def read(
-        self,
-        input_path: Union[str, List[str]],
-        parallelism: int = None,
-    ) -> Dataset:
+    def read(self, input_path: Union[str, List[str]]) -> Dataset:
         """
         Read Parquet files using Ray Data.
 
         :param input_path: Path to Parquet file or list of Parquet files.
-        :param parallelism: Number of blocks for Ray Dataset reading.
         :return: Ray Dataset containing validated documents.
         """
         if not ray.is_initialized():
             ray.init()
 
-        ds = ray.data.read_parquet(input_path, override_num_blocks=parallelism)
+        ds = ray.data.read_parquet(input_path)
         ds = ds.map_batches(self._validate_batch, batch_format="pandas")
         ds = ds.filter(self._should_keep_item)
         return ds

graphgen/models/reader/pdf_reader.py

@@ -68,7 +68,6 @@ def __init__(
     def read(
         self,
         input_path: Union[str, List[str]],
-        parallelism: int = 4,
         **override,
     ) -> Dataset:
 

graphgen/models/reader/pickle_reader.py

@@ -23,22 +23,18 @@ class PickleReader(BaseReader):
     def read(
         self,
         input_path: Union[str, List[str]],
-        parallelism: int = None,
     ) -> Dataset:
         """
         Read Pickle files using Ray Data.
 
         :param input_path: Path to pickle file or list of pickle files.
-        :param parallelism: Number of blocks for Ray Dataset reading.
         :return: Ray Dataset containing validated documents.
         """
         if not ray.is_initialized():
             ray.init()
 
         # Use read_binary_files as a reliable alternative to read_pickle
-        ds = ray.data.read_binary_files(
-            input_path, override_num_blocks=parallelism, include_paths=True
-        )
+        ds = ray.data.read_binary_files(input_path, include_paths=True)
 
         # Deserialize pickle files and flatten into individual records
         def deserialize_batch(batch: pd.DataFrame) -> pd.DataFrame:

graphgen/models/reader/rdf_reader.py

@@ -30,13 +30,11 @@ def __init__(self, *, text_column: str = "content", **kwargs):
     def read(
         self,
         input_path: Union[str, List[str]],
-        parallelism: int = 4,
     ) -> Dataset:
         """
         Read RDF file(s) using Ray Data.
 
         :param input_path: Path to RDF file or list of RDF files.
-        :param parallelism: Number of parallel workers for processing.
         :return: Ray Dataset containing extracted documents.
         """
         if not ray.is_initialized():

graphgen/models/reader/txt_reader.py

@@ -10,17 +10,13 @@ class TXTReader(BaseReader):
     def read(
         self,
         input_path: Union[str, List[str]],
-        parallelism: int = 4,
     ) -> Dataset:
         """
         Read text files from the specified input path.
         :param input_path: Path to the input text file or list of text files.
-        :param parallelism: Number of blocks to override for Ray Dataset reading.
         :return: Ray Dataset containing the read text data.
         """
-        docs_ds = ray.data.read_text(
-            input_path, encoding="utf-8", override_num_blocks=parallelism
-        )
+        docs_ds = ray.data.read_text(input_path, encoding="utf-8")
 
         docs_ds = docs_ds.map(
             lambda row: {

graphgen/operators/__init__.py

@@ -4,6 +4,6 @@
 from .init import init_llm
 from .partition import partition_kg
 from .quiz_and_judge import judge_statement, quiz
-from .read import read_files
+from .read import read
 from .search import search_all
 from .split import chunk_documents

graphgen/operators/evaluate.py

@@ -9,9 +9,13 @@
 from dotenv import load_dotenv
 
 from graphgen.bases.datatypes import QAPair
-
-from .models import LengthEvaluator, MTLDEvaluator, RewardEvaluator, UniEvaluator
-from .utils import logger, set_logger
+from graphgen.models import (
+    LengthEvaluator,
+    MTLDEvaluator,
+    RewardEvaluator,
+    UniEvaluator,
+)
+from graphgen.utils import logger, set_logger
 
 sys_path = os.path.abspath(os.path.dirname(__file__))
 set_logger(os.path.join(sys_path, "cache", "logs", "evaluate.log"))

graphgen/operators/read/__init__.py

@@ -1 +1 @@
-from .read_files import read_files
+from .read import read