from QATK.MLDFT import *

model_dir = '/densitymodel/default_params/'

model = DeepDFTModel(
    model_dir=model_dir
)

# Define dataset directories containing HDF5 files with training data
dataset_dirs=[
    "/directory/with/hdf5/files/for/training1/",
    "/directory/with/hdf5/files/for/training2/",
]

# Create the dataset
dataset = GridValuesDataset(
    dataset_dirs=dataset_dirs,
)

# Setup the trainer
trainer = GridValuesModelTrainer(
    grid_values_model=model,
    grid_values_dataset=dataset,
)

# Start the training process
trainer.train()