fix: allow to set eval datasets, fix shutdown of processes

Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

Author: mudler

backend/python/trl/backend.py

@@ -256,6 +256,39 @@ def _do_training(self, request, job):
         else:
             dataset = load_dataset(request.dataset_source, split=dataset_split)
 
+        # Evaluation dataset setup
+        save_steps = request.save_steps if request.save_steps > 0 else 500
+        eval_dataset = None
+        eval_strategy = extra.get("eval_strategy", "steps")
+        eval_steps = int(extra.get("eval_steps", str(save_steps)))
+
+        if eval_strategy != "no":
+            eval_split = extra.get("eval_split")
+            eval_dataset_source = extra.get("eval_dataset_source")
+            if eval_split:
+                # Load a specific split as eval dataset
+                if os.path.exists(request.dataset_source):
+                    if request.dataset_source.endswith('.json') or request.dataset_source.endswith('.jsonl'):
+                        eval_dataset = load_dataset("json", data_files=request.dataset_source, split=eval_split)
+                    elif request.dataset_source.endswith('.csv'):
+                        eval_dataset = load_dataset("csv", data_files=request.dataset_source, split=eval_split)
+                    else:
+                        eval_dataset = load_dataset(request.dataset_source, split=eval_split)
+                else:
+                    eval_dataset = load_dataset(request.dataset_source, split=eval_split)
+            elif eval_dataset_source:
+                # Load eval dataset from a separate source
+                eval_dataset = load_dataset(eval_dataset_source, split="train")
+            else:
+                # Auto-split the training set
+                eval_split_ratio = float(extra.get("eval_split_ratio", "0.1"))
+                split = dataset.train_test_split(test_size=eval_split_ratio)
+                dataset = split["train"]
+                eval_dataset = split["test"]
+
+        if eval_strategy == "no":
+            eval_dataset = None
+
         # Training config
         output_dir = request.output_dir or f"./output-{job.job_id}"
         num_epochs = request.num_epochs if request.num_epochs > 0 else 3
@@ -265,7 +298,6 @@ def _do_training(self, request, job):
         warmup_steps = request.warmup_steps if request.warmup_steps > 0 else 5
         weight_decay = request.weight_decay if request.weight_decay > 0 else 0.01
         max_steps = request.max_steps if request.max_steps > 0 else -1
-        save_steps = request.save_steps if request.save_steps > 0 else 500
         seed = request.seed if request.seed > 0 else 3407
         optimizer = request.optimizer or "adamw_torch"
 
@@ -308,6 +340,12 @@ def _do_training(self, request, job):
         if save_total_limit:
             _save_kwargs["save_total_limit"] = save_total_limit
 
+        # Eval arguments
+        _eval_kwargs = {}
+        if eval_dataset is not None:
+            _eval_kwargs["eval_strategy"] = eval_strategy
+            _eval_kwargs["eval_steps"] = eval_steps
+
         # Common training arguments shared by all methods
         _common_args = dict(
             output_dir=output_dir,
@@ -323,6 +361,7 @@ def _do_training(self, request, job):
             logging_steps=1,
             report_to="none",
             **_save_kwargs,
+            **_eval_kwargs,
             **common_train_kwargs,
         )
 
@@ -343,6 +382,7 @@ def _do_training(self, request, job):
                 model=model,
                 args=training_args,
                 train_dataset=dataset,
+                eval_dataset=eval_dataset,
                 processing_class=tokenizer,
                 callbacks=[progress_cb.get_callback()],
             )
@@ -365,6 +405,7 @@ def _do_training(self, request, job):
                 model=model,
                 args=training_args,
                 train_dataset=dataset,
+                eval_dataset=eval_dataset,
                 processing_class=tokenizer,
                 callbacks=[progress_cb.get_callback()],
             )
@@ -399,6 +440,7 @@ def _do_training(self, request, job):
                 model=model,
                 args=training_args,
                 train_dataset=dataset,
+                eval_dataset=eval_dataset,
                 processing_class=tokenizer,
                 reward_funcs=reward_funcs,
                 callbacks=[progress_cb.get_callback()],
@@ -420,6 +462,7 @@ def _do_training(self, request, job):
                 model=model,
                 args=training_args,
                 train_dataset=dataset,
+                eval_dataset=eval_dataset,
                 processing_class=tokenizer,
                 callbacks=[progress_cb.get_callback()],
             )
@@ -440,6 +483,7 @@ def _do_training(self, request, job):
                 model=model,
                 args=training_args,
                 train_dataset=dataset,
+                eval_dataset=eval_dataset,
                 processing_class=tokenizer,
                 callbacks=[progress_cb.get_callback()],
             )
@@ -460,6 +504,7 @@ def _do_training(self, request, job):
                 model=model,
                 args=training_args,
                 train_dataset=dataset,
+                eval_dataset=eval_dataset,
                 processing_class=tokenizer,
                 callbacks=[progress_cb.get_callback()],
             )
@@ -478,6 +523,7 @@ def _do_training(self, request, job):
                 model=model,
                 args=training_args,
                 train_dataset=dataset,
+                eval_dataset=eval_dataset,
                 processing_class=tokenizer,
                 callbacks=[progress_cb.get_callback()],
             )