fixes from testing

Author: kimusaku

sapientml_loaddata/generator.py

@@ -45,8 +45,8 @@ def generate_code(self, dataset: Dataset, task: Task):
         code.predict += _render(tpl, dataset=dataset, task=task, script_type="predict")
 
         tpl = template_env.get_template("split.py.jinja")
-        code.validation += _render(tpl, task=task, validation=True)
-        code.test += _render(tpl, task=task, validation=False)
+        code.validation += _render(tpl, dataset=dataset, task=task, validation=True)
+        code.test += _render(tpl, dataset=dataset, task=task, validation=False)
 
         tpl = template_env.get_template("subsample.py.jinja")
         code.validation += _render(tpl, task=task, sample_size=ROW_THRESHOLD_FOR_SAMPLING)

sapientml_loaddata/templates/split.py.jinja

@@ -1,3 +1,5 @@
+{% if (validation and (not dataset.validation_data_path)) or ((not validation) and (not dataset.test_data_path)) %}
+
 # TRAIN-TEST SPLIT
 {% if task.split_method == "random" %}
 {% if task.split_stratify == True %}
@@ -35,10 +37,11 @@ def split_dataset(dataset, split_column_name="{{ task.split_column_name }}", spl
                 train_dataset.loc[:, col] = False
     return train_dataset, test_dataset
 {% endif %}
-{% if not task.test_dataset_path %}
+{% if not dataset.test_data_path %}
 __train_dataset, __test_dataset = split_dataset(__train_dataset)
 {% endif %}
-{% if validation and (not task.validation_dataset_path) %}
+{% endif %}
+{% if validation and (not dataset.validation_data_path) %}
 __train_dataset, __validation_dataset = split_dataset(__train_dataset)
 {% endif %}
 {% if validation %}