more ckpt

Author: kddubey

src/grouping_trainer/pretrain.py

@@ -131,7 +131,7 @@ class PretrainingConfig(BaseModel):
     # Logging / checkpointing
     wandb_project: str = "grouping-trainer"
     num_logs: int = 2000
-    num_checkpoints: int = 10
+    num_checkpoints: int = 50
 
 
 def _load_model_and_tokenizer(base_model: str) -> tuple[Any, PreTrainedTokenizerBase]:
@@ -159,9 +159,13 @@ def _get_train_sampler(self, train_dataset=None) -> Sampler | None:
 
 
 def _sort_dataset_by_length_desc(dataset: Dataset) -> Dataset:
-    lengths = np.array([len(input_ids) for input_ids in dataset["input_ids"]])
-    sorted_indices = np.argsort(-lengths)
-    sorted_dataset = dataset.select(sorted_indices)
+    dataset_with_length = dataset.map(
+        lambda batch: {"length": [len(input_ids) for input_ids in batch["input_ids"]]},
+        batched=True,
+    )
+    assert isinstance(dataset_with_length, Dataset)
+    sorted_dataset = dataset_with_length.sort("length", reverse=True).remove_columns("length")
+    assert isinstance(sorted_dataset, Dataset)
     return sorted_dataset  # type: ignore[bad-return]