generator_model.py

from ctgan import CTGANSynthesizer
import configparser
import ast
import os
import logging


class GeneratorModel:
    def __init__(self, data):
        self.data = data

        config = configparser.ConfigParser()
        config.read("attributes.ini")
        self.discrete_attributes = ast.literal_eval(config.get("ATTRIBUTES", "discrete_attributes"))
        print(self.discrete_attributes)

        if len(self.discrete_attributes) == 0:
            logging.warning("No Discrete Attribute defined")

    def train_and_save_ctgan(self, epochs=100):
        ctgan = CTGANSynthesizer(epochs=epochs)
        ctgan.fit(train_data=self.data, discrete_columns=self.discrete_attributes)
        logging.info("CTGAN model successfully trained")

        try:
            if not os.path.exists("trained_model"):
                os.mkdir("trained_model")
            ctgan.save("trained_model/trained_generator.pth")

            logging.info("Generator model successfully saved")
        except:
            logging.error("Generator model could not be saved")

        return ctgan