fix: stop gap to remove unused colulmns

configs/train/finetune.yaml

@@ -8,6 +8,7 @@ save_name: # CHANGE
 streaming: false
 num_proc: 64
 dataset_path: # update
+revision: null
 max_length: 1024
 batch_size: 32
 

configs/train/finetune_gptj.yaml

@@ -8,6 +8,7 @@ save_name: # CHANGE
 streaming: false
 num_proc: 64
 dataset_path: # CHANGE
+revision: null
 max_length: 1024
 batch_size: 32
 

configs/train/finetune_gptj_lora.yaml

@@ -8,6 +8,7 @@ save_name: # CHANGE
 streaming: false
 num_proc: 64
 dataset_path: # CHANGE
+revision: null
 max_length: 1024
 batch_size: 1 
 

configs/train/finetune_lora.yaml

@@ -8,6 +8,7 @@ save_name: # CHANGE
 streaming: false
 num_proc: 64
 dataset_path: # CHANGE
+revision: null
 max_length: 1024
 batch_size: 4
 

data.py

@@ -61,7 +61,24 @@ def tokenize_inputs(config, tokenizer, examples):
 def load_data(config, tokenizer):
     dataset_path = config["dataset_path"]
 
-    if os.path.exists(dataset_path):
+    if isinstance(dataset_path, list):
+        all_datasets = []
+        for path in dataset_path:
+            dataset = load_dataset(path, split="train")
+
+            current_columns = dataset.column_names
+            columns_to_keep = ["prompt", "response"]
+            to_remove = set(current_columns) - set(columns_to_keep)
+
+            dataset = dataset.remove_columns(to_remove)
+            if "source" not in current_columns:
+                dataset = dataset.add_column("source", [path.split("/")[-1]] * len(dataset))
+            all_datasets.append(dataset)
+
+        dataset = concatenate_datasets(all_datasets)
+
+    # load local json dataset
+    elif os.path.exists(dataset_path):
         if os.path.isdir(dataset_path):
             files = glob.glob(os.path.join(dataset_path, "*_clean.jsonl"))
         else:
@@ -70,9 +87,11 @@ def load_data(config, tokenizer):
         print(f"Reading files {files}")
 
         dataset = load_dataset("json", data_files=files, split="train")
-
+    
+    # read from huggingface
     else:
-        dataset = load_dataset(dataset_path, split="train")
+        revision = config["revision"] 
+        dataset = load_dataset(dataset_path, split="train", revision=revision)
 
     dataset = dataset.train_test_split(test_size=.05, seed=config["seed"])
 
@@ -87,13 +106,13 @@ def load_data(config, tokenizer):
     train_dataset = train_dataset.map(
         lambda ele: tokenize_inputs(config, tokenizer, ele),
         batched=True,
-        remove_columns=["source", "prompt"],
+        remove_columns=["source", "prompt", "id", "response"],
         **kwargs
     )
     val_dataset = val_dataset.map(
         lambda ele: tokenize_inputs(config, tokenizer, ele),
         batched=True,
-        remove_columns=["source", "prompt"],
+        remove_columns=["source", "prompt", "id", "response"],
         **kwargs
     )