simplify code

Former-commit-id: d3731754ab7c28ae81f60784e0e4213f279d93fe

src/llmtuner/dsets/__init__.py

@@ -1,2 +1,3 @@
 from llmtuner.dsets.loader import get_dataset
 from llmtuner.dsets.preprocess import preprocess_dataset
+from llmtuner.dsets.utils import split_dataset

src/llmtuner/dsets/callbacks.py

@@ -1,63 +0,0 @@
-import os
-import json
-import time
-from datetime import timedelta
-
-from transformers import (
-    TrainerCallback,
-    TrainerControl,
-    TrainerState,
-    TrainingArguments
-)
-
-
-class LogCallback(TrainerCallback):
-
-    def __init__(self, runner=None):
-        self.runner = runner
-        self.start_time = time.time()
-        self.tracker = {}
-
-    def on_step_begin(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs):
-        r"""
-        Event called at the beginning of a training step. If using gradient accumulation, one training step
-        might take several inputs.
-        """
-        if self.runner is not None and self.runner.aborted:
-            control.should_epoch_stop = True
-            control.should_training_stop = True
-
-    def on_substep_end(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs):
-        r"""
-        Event called at the end of an substep during gradient accumulation.
-        """
-        if self.runner is not None and self.runner.aborted:
-            control.should_epoch_stop = True
-            control.should_training_stop = True
-
-    def on_log(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs) -> None:
-        r"""
-        Event called after logging the last logs.
-        """
-        if "loss" not in state.log_history[-1]:
-            return
-        cur_time = time.time()
-        cur_steps = state.log_history[-1].get("step")
-        elapsed_time = cur_time - self.start_time
-        avg_time_per_step = elapsed_time / cur_steps if cur_steps != 0 else 0
-        remaining_steps = state.max_steps - cur_steps
-        remaining_time = remaining_steps * avg_time_per_step
-        self.tracker = {
-            "current_steps": cur_steps,
-            "total_steps": state.max_steps,
-            "loss": state.log_history[-1].get("loss", None),
-            "reward": state.log_history[-1].get("reward", None),
-            "learning_rate": state.log_history[-1].get("learning_rate", None),
-            "epoch": state.log_history[-1].get("epoch", None),
-            "percentage": round(cur_steps / state.max_steps * 100, 2) if state.max_steps != 0 else 100,
-            "elapsed_time": str(timedelta(seconds=int(elapsed_time))),
-            "remaining_time": str(timedelta(seconds=int(remaining_time)))
-        }
-        os.makedirs(args.output_dir, exist_ok=True)
-        with open(os.path.join(args.output_dir, "trainer_log.jsonl"), "a", encoding="utf-8") as f:
-            f.write(json.dumps(self.tracker) + "\n")

src/llmtuner/dsets/preprocess.py

@@ -6,7 +6,7 @@ from transformers.tokenization_utils import PreTrainedTokenizer
 from datasets import Dataset
 
 from llmtuner.extras.constants import IGNORE_INDEX
-from llmtuner.extras.template import Template
+from llmtuner.extras.template import get_template
 from llmtuner.hparams import DataArguments
 
 
@@ -19,7 +19,7 @@ def preprocess_dataset(
 ) -> Dataset:
 
     column_names = list(dataset.column_names)
-    prompt_template = Template(data_args.prompt_template)
+    prompt_template = get_template(data_args.prompt_template)
 
     # support question with a single answer or multiple answers
     def get_dialog(examples):

src/llmtuner/dsets/utils.py

@@ -0,0 +1,16 @@
+from typing import Dict
+from datasets import Dataset
+
+
+def split_dataset(
+    dataset: Dataset, dev_ratio: float, do_train: bool
+) -> Dict[str, Dataset]:
+    # Split the dataset
+    if do_train:
+        if dev_ratio > 1e-6:
+            dataset = dataset.train_test_split(test_size=dev_ratio)
+            return {"train_dataset": dataset["train"], "eval_dataset": dataset["test"]}
+        else:
+            return {"train_dataset": dataset}
+    else: # do_eval or do_predict
+        return {"eval_dataset": dataset}