Can't save the model if use the custom constraint

[YiranZhang1014]

Environment details

If you are already running SDV, please indicate the following details about the environment in
which you are running it:

• SDV version: 0.15.0
• Python version: 3.8.13
• Operating System: Windows 11

Problem description

Follow the User Guides, I create a custom constraint and successfully trained the model:

from sdv.constraints import create_custom_constraint

# The constrains defined before. Here I used it
DodUniqueProcess = create_custom_constraint(
    is_valid_fn=dod_data_unique_is_valid
)

dod_data_unique_constraint = DodUniqueProcess(
    column_names=['dod_hosp', 'dod_ssn']
)

constrains = [dod_data_unique_constraint]

model = CTGAN(
    constraints=constrains, 
    batch_size=10000,
    cuda=True, 
    verbose=True, 
    epochs=50)

But when I try to save the model, it showed this error:

---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
f:\MSc Project\Synthetic-Electronic-Health-Records-Data-using-MIMIC-III\SDV_patients_generation.ipynb Cell 42 in <cell line: 1>()
----> 1 model.save(date_set_path + "100_models/" + "patients_model.pkl")

File d:\Environment\anaconda3\envs\msc_project\lib\site-packages\sdv\tabular\base.py:883, in BaseTabularModel.save(self, path)
    880 self._package_versions = get_package_versions(getattr(self, '_model', None))
    882 with open(path, 'wb') as output:
--> 883     pickle.dump(self, output)

AttributeError: Can't pickle local object 'create_custom_constraint.<locals>.CustomConstraint'

What I've tried

I search the problem on Google and it shows:

You can only pickle instances of classes defined at module's top level.

In this case, how should I do for saving my model with the custom constraints?

[npatki]

Hi @Alaskyed, thanks for filing this issue. I can replicate the error.

It seems that pickle is not able to save constraint because it's being created using a factory method now. I'll mark this as a bug and we can use this issue to track its updates.

Workaround

In the meantime, you can unblock your work by using a different package for saving the model. In my (rough) exploration, I found success using the dill package.

import dill

# Save the model
with open('my_model.pkl', 'wb') as file:
  dill.dump(model, file)

# Reload and use the model
with open('my_model.pkl', 'rb') as file:
  model = dill.load(file)

model.sample(num_rows=10)

[npatki]

Hi @Alaskyed, we've submitted the code to fix this issue and it should go out with the next release.

[YiranZhang1014]

Oh! That's great news! Thank you npatki! Following your method I can save the model to file well, and looking forward the next release!

[imrooki]

ImportError: cannot import name 'create_custom_constraint' from 'sdv.constraints'

[imrooki]

Why this error is occurring