style: improve code style

Author: ChenZiHong-Gavin

graphgen/models/reader/huggingface_reader.py

@@ -191,9 +191,7 @@ def _process_batch(batch: dict[str, "np.ndarray"]) -> dict[str, "np.ndarray"]:
 
             # Rename text_column to 'content' if different
             if self.text_column != "content" and self.text_column in batch:
-                batch["content"] = batch[self.text_column]
-                # Optional: delete old key to avoid duplication
-                # del batch[self.text_column]
+                batch["content"] = batch.pop(self.text_column)
 
             return batch
 

graphgen/operators/read/read.py

@@ -55,11 +55,11 @@ def _build_reader(suffix: str, cache_dir: str | None, **reader_kwargs):
 def _process_huggingface_datasets(hf_uris: List[str], reader_kwargs: dict) -> list:
     """Process HuggingFace datasets and return list of Ray datasets."""
     logger.info("[READ] Processing HuggingFace datasets: %s", hf_uris)
+    hf_reader = HuggingFaceReader(**reader_kwargs)
     read_tasks = []
     for hf_uri in hf_uris:
         # Parse URI format: "huggingface://dataset_name:subset:split"
         uri_part = hf_uri.replace("huggingface://", "")
-        hf_reader = HuggingFaceReader(**reader_kwargs)
         ds = hf_reader.read(uri_part)
         read_tasks.append(ds)
     logger.info("[READ] Successfully loaded %d HuggingFace dataset(s)", len(hf_uris))