[pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

applications/ColossalChat/benchmarks/benchmark_dpo.py

@@ -5,10 +5,11 @@ import resource
 from contextlib import nullcontext
 
 import torch
-from coati.dataset import DataCollatorForPreferenceDataset, StatefulDistributedSampler, load_tokenized_dataset
+from coati.dataset import DataCollatorForPreferenceDataset, StatefulDistributedSampler
 from coati.models import convert_to_lora_module, disable_dropout
 from coati.trainer import DPOTrainer
 from coati.utils import load_checkpoint
+from dummy_dataset import DummyLLMDataset
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
 import colossalai
@@ -18,7 +19,6 @@ from colossalai.cluster import DistCoordinator
 from colossalai.logging import get_dist_logger
 from colossalai.nn.lr_scheduler import CosineAnnealingWarmupLR
 from colossalai.nn.optimizer import HybridAdam
-from dummy_dataset import DummyLLMDataset
 
 logger = get_dist_logger()
 
@@ -165,9 +165,11 @@ def train(args):
 
     # configure dataset
     mode_map = {"train": "train", "valid": "validation", "test": "test"}
-    train_dataset = DummyLLMDataset(["chosen_input_ids", "chosen_loss_mask", "rejected_input_ids",
-                                     "rejected_loss_mask"], 
-                                     args.max_length, args.dataset_size)
+    train_dataset = DummyLLMDataset(
+        ["chosen_input_ids", "chosen_loss_mask", "rejected_input_ids", "rejected_loss_mask"],
+        args.max_length,
+        args.dataset_size,
+    )
     data_collator = DataCollatorForPreferenceDataset(tokenizer=tokenizer, max_length=args.max_length)
 
     train_dataloader = plugin.prepare_dataloader(

applications/ColossalChat/benchmarks/benchmark_orpo.py

@@ -5,10 +5,11 @@ import resource
 from contextlib import nullcontext
 
 import torch
-from coati.dataset import DataCollatorForPreferenceDataset, StatefulDistributedSampler, load_tokenized_dataset
+from coati.dataset import DataCollatorForPreferenceDataset, StatefulDistributedSampler
 from coati.models import convert_to_lora_module, disable_dropout
 from coati.trainer import ORPOTrainer
 from coati.utils import load_checkpoint
+from dummy_dataset import DummyLLMDataset
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
 import colossalai
@@ -18,7 +19,7 @@ from colossalai.cluster import DistCoordinator
 from colossalai.logging import get_dist_logger
 from colossalai.nn.lr_scheduler import CosineAnnealingWarmupLR
 from colossalai.nn.optimizer import HybridAdam
-from dummy_dataset import DummyLLMDataset
+
 logger = get_dist_logger()
 
 
@@ -152,9 +153,11 @@ def train(args):
     # configure dataset
     coordinator.print_on_master(f"Load dataset: {args.dataset}")
     mode_map = {"train": "train", "valid": "validation", "test": "test"}
-    train_dataset = DummyLLMDataset(["chosen_input_ids", "chosen_loss_mask", "rejected_input_ids",
-                                     "rejected_loss_mask"], 
-                                     args.max_length, args.dataset_size)
+    train_dataset = DummyLLMDataset(
+        ["chosen_input_ids", "chosen_loss_mask", "rejected_input_ids", "rejected_loss_mask"],
+        args.max_length,
+        args.dataset_size,
+    )
     data_collator = DataCollatorForPreferenceDataset(tokenizer=tokenizer, max_length=args.max_length)
 
     train_dataloader = plugin.prepare_dataloader(

applications/ColossalChat/benchmarks/benchmark_sft.py

@@ -6,10 +6,11 @@ import resource
 from contextlib import nullcontext
 
 import torch
-from coati.dataset import DataCollatorForSupervisedDataset, StatefulDistributedSampler, load_tokenized_dataset
+from coati.dataset import DataCollatorForSupervisedDataset, StatefulDistributedSampler
 from coati.models import convert_to_lora_module
 from coati.trainer import SFTTrainer
 from coati.utils import load_checkpoint
+from dummy_dataset import DummyLLMDataset
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
 import colossalai
@@ -19,7 +20,6 @@ from colossalai.cluster import DistCoordinator
 from colossalai.logging import get_dist_logger
 from colossalai.nn.lr_scheduler import CosineAnnealingWarmupLR
 from colossalai.nn.optimizer import HybridAdam
-from dummy_dataset import DummyLLMDataset
 
 logger = get_dist_logger()
 

applications/ColossalChat/benchmarks/dummy_dataset.py

@@ -1,5 +1,6 @@
 import torch
-from torch.utils.data import Dataset, DataLoader
+from torch.utils.data import Dataset
+
 
 class DummyLLMDataset(Dataset):
     def __init__(self, keys, seq_len, size=500):
@@ -11,7 +12,7 @@ class DummyLLMDataset(Dataset):
     def _generate_data(self):
         data = {}
         for key in self.keys:
-            data[key] = torch.ones(self.seq_len, dtype = torch.long)
+            data[key] = torch.ones(self.seq_len, dtype=torch.long)
         return data
 
     def __len__(self):