add keep column feature

src/text2dataset/main.py

@@ -82,6 +82,12 @@
     default="ja",
     help="Target language for translation. This is used only for DeepL API.",
 )
+@click.option(
+    "--keep_columns",
+    type=str,
+    default="txt",
+    help="Columns to keep in the output dataset. Specify the column names separated by comma.",
+)
 def main(
     model_id: str,
     batch_size: int,
@@ -105,6 +111,7 @@ def main(
     top_p: float,
     max_tokens: int,
     target_lang: str,
+    keep_columns: str,
 ):
     # Text in source_column of the Dataset will be translated into Japanese.
     state = State(0, 0, 0)
@@ -127,8 +134,10 @@ def main(
     os.makedirs(output_dir, exist_ok=True)
     state_path = os.path.join(output_dir, "state.jsonl")
     ds = create_dataset(input_path, state)
+    # keep only the specified columns
+    ds = ds.select_columns(keep_columns.split(","))
     # batch dataloader
-    data_loader = ds["train"].batch(batch_size=batch_size)
+    data_loader = ds.batch(batch_size=batch_size)
 
     if use_wandb:
         config_parameters = dict(locals())

src/text2dataset/reader.py

@@ -2,33 +2,47 @@
 from text2dataset.utils import State
 
 
-def create_dataset(input_path: str, state: State) -> IterableDataset:
+def create_dataset(
+    input_path: str, state: State, split: str = "train"
+) -> IterableDataset:
     """
     Create a Iterabledataset from the input path.
     The input path can be a local file path or a HuggingFace Datasets dataset name.
     Use the state to skip already processed examples.
     """
     match input_path.split(".")[-1]:
         case "csv":
-            ds = load_dataset("csv", data_files=input_path, streaming=True)
+            ds = load_dataset("csv", data_files=input_path, streaming=True, split=split)
         case "json":
-            ds = load_dataset("json", data_files=input_path, streaming=True)
+            ds = load_dataset(
+                "json", data_files=input_path, streaming=True, split=split
+            )
         case "jsonl":
-            ds = load_dataset("json", data_files=input_path, streaming=True)
+            ds = load_dataset(
+                "json", data_files=input_path, streaming=True, split=split
+            )
         case "parquet":
-            ds = load_dataset("parquet", data_files=input_path, streaming=True)
+            ds = load_dataset(
+                "parquet", data_files=input_path, streaming=True, split=split
+            )
         case "tar":
-            ds = load_dataset("webdataset", data_files=input_path, streaming=True)
+            ds = load_dataset(
+                "webdataset", data_files=input_path, streaming=True, split=split
+            )
         case "arrow":
-            ds = load_dataset("arrow", data_files=input_path, streaming=True)
+            ds = load_dataset(
+                "arrow", data_files=input_path, streaming=True, split=split
+            )
         case "txt":
-            ds = load_dataset("text", data_files=input_path, streaming=True)
+            ds = load_dataset(
+                "text", data_files=input_path, streaming=True, split=split
+            )
         case _:
-            ds = load_dataset(input_path, streaming=True)
+            ds = load_dataset(input_path, streaming=True, split=split)
 
     # skip already processed examples
     if state.last_saved_example_num > 0:
-        ds["train"] = ds["train"].skip(state.last_saved_example_num)
+        ds = ds.skip(state.last_saved_example_num)
 
     return ds
 
@@ -39,7 +53,7 @@ def test_create_dataset():
         current_shard_id=0, last_saved_example_num=0, total_processed_examples=0
     )
     ds = create_dataset(input_path, state)
-    assert ds["train"] is not None
+    assert ds is not None
     # error happing case
     try:
         create_dataset("hogehuga", state)
@@ -55,10 +69,11 @@ def test_create_dataset():
         current_shard_id=0, last_saved_example_num=0, total_processed_examples=0
     )
     ds = create_dataset(input_path, state)
-    print(ds["train"])
-    # iterabledatasetdict to datasetdict
-    input_path = "data/english_quotes.json"
+    print(ds)
+    print(next(iter(ds)))
+    state = State(
+        current_shard_id=0, last_saved_example_num=10, total_processed_examples=10
+    )
     ds = create_dataset(input_path, state)
-    print(ds["train"])
-    print(next(iter(ds["train"])))
-    ds.map(lambda x: x["txt"], batched=True)
+    print(ds)
+    print(next(iter(ds)))