concat datasets bug fix

vtuber-plan · Jan 2, 2024 · 750375d · 750375d
1 parent 3b73d13
commit 750375d
Show file tree

Hide file tree

Showing 4 changed files with 6 additions and 6 deletions.
diff --git a/hparams/hparams_chat_llama2_13b_lora.json b/hparams/hparams_chat_llama2_13b_lora.json
@@ -111,7 +111,7 @@
         },
         {
             "path": "../../datasets/HanChat/Summary/",
-            "sample": 1.0
+            "sample": 0.1
         },
         {
             "path": "../../datasets/HanChat/TextCorrection/",

diff --git a/katheryne/data/loader/chat.py b/katheryne/data/loader/chat.py
@@ -111,11 +111,11 @@ def create_chat_dataset(hparams: HParams, data_path: List[Union[str, DatasetPath
 
             if isinstance(d_path.sample, int):
                 sample_size = d_path.sample
-                train_dataset = train_dataset[:sample_size]
+                train_dataset = train_dataset.select(list(range(sample_size)))
             elif isinstance(d_path.sample, float):
                 if d_path.sample != 1.0:
                     sample_size = int(d_path.sample * len(train_dataset))
-                    train_dataset = train_dataset[:sample_size]
+                    train_dataset = train_dataset.select(list(range(sample_size)))
             else:
                 raise TypeError("Invalid sample number of dataset path object, need int or float.")
 

diff --git a/katheryne/data/loader/pretrain.py b/katheryne/data/loader/pretrain.py
@@ -125,11 +125,11 @@ def create_pretrain_dataset(hparams: HParams, data_path: List[Union[str, Dataset
 
             if isinstance(d_path.sample, int):
                 sample_size = d_path.sample
-                train_dataset = train_dataset[:sample_size]
+                train_dataset = train_dataset.select(list(range(sample_size)))
             elif isinstance(d_path.sample, float):
                 if d_path.sample != 1.0:
                     sample_size = int(d_path.sample * len(train_dataset))
-                    train_dataset = train_dataset[:sample_size]
+                    train_dataset = train_dataset.select(list(range(sample_size)))
             else:
                 raise TypeError("Invalid sample number of dataset path object, need int or float.")
 

diff --git a/katheryne/train/train.py b/katheryne/train/train.py
@@ -132,7 +132,7 @@ def train(create_dataset, lightning_module_class):
         model.gradient_checkpointing_enable()
 
     # Save Model
-    save_hf_format(model, tokenizer, "./lightning_logs/huggingface_format", sub_folder=f"checkpoint-step-0")
+    # save_hf_format(model, tokenizer, "./lightning_logs/huggingface_format", sub_folder=f"checkpoint-step-0")
 
     # Prepare the data
     print("***** Prepare Dataset *****")