code: """
  This module allows automatic splitting of a DataFrame
  into smaller DataFrames (by clusters of columns) and doing
  model training and text generation on each sub-DF independently.
  
  Then we can concat each sub-DF back into one final synthetic dataset.
  
  For example usage, please see our Jupyter Notebook.
  """
  import abc
  import glob
  import gzip
  import io
  import json
  import logging
  import shutil
  import tempfile
  import threading
  import time
  
  from copy import deepcopy
  from dataclasses import dataclass, field
  from itertools import zip_longest
  from math import ceil
  from pathlib import Path
  from typing import Callable, Dict
  from typing import Iterator as IteratorType
  from typing import List, Optional, Type, Union
  
  import cloudpickle
  import gretel_synthetics.const as const
  import numpy as np
  import pandas as pd
  
  from gretel_synthetics.config import (
      BaseConfig,
      config_from_model_dir,
      CONFIG_MAP,
      LocalConfig,
  )
  from gretel_synthetics.errors import TooManyInvalidError
  from gretel_synthetics.generate import generate_text, GenText, SeedingGenerator
  from gretel_synthetics.tokenizers import BaseTokenizerTrainer
  from gretel_synthetics.train import train
  from tqdm.auto import tqdm
  
  logging.basicConfig()
  logger = logging.getLogger(__name__)
  logger.setLevel(logging.INFO)
  
  MAX_INVALID = 1000
  BATCH_SIZE = 15
  FIELD_DELIM = "field_delimiter"
  GEN_LINES = "gen_lines"
  READ = "read"
  WRITE = "write"
  HEADER_FILE = "headers.json"
  ORIG_HEADERS = "original_headers.json"
  CHECKPOINT_DIR = "checkpoint_dir"
  CONFIG_FILE = "model_params.json"
  TRAIN_FILE = "train.csv"
  PATH_HOLDER = "___path_holder___"
  FILE = "file"
  MEMORY = "memory"
  
  
  @dataclass
  class GenerationSummary:
      """A class to capture the summary data after synthetic data is generated."""
  
      valid_lines: int = 0
      invalid_lines: int = 0
      is_valid: bool = False
  
  
  class _BatchEpochCallback:
      """
      Wrapper class to take a user supplied callback and inject the batch number.  The batch number
      is then available in the EpochState object when it is supplied to the callback.
      """
  
      def __init__(self, user_callback: callable, batch_number: int):
          self._batch_number = batch_number
          self._user_callback = user_callback
  
      def callback(self, epoch_state):
          epoch_state.batch = self._batch_number
          self._user_callback(epoch_state)
  
  
  @dataclass
  class Batch:
      """A representation of a synthetic data workflow.  It should not be used
      directly. This object is created automatically by the primary batch handler,
      such as ``DataFrameBatch``.  This class holds all of the necessary information
      for training, data generation and DataFrame re-assembly.
      """
  
      checkpoint_dir: str
      input_data_path: str
      headers: List[str]
      config: LocalConfig
      gen_data_count: int = 0
  
      training_df: Type[pd.DataFrame] = field(default_factory=lambda: None, init=False)
      gen_data_stream: io.StringIO = field(default_factory=io.StringIO, init=False)
      gen_data_invalid: List[GenText] = field(default_factory=list, init=False)
      validator: Callable = field(default_factory=lambda: None, init=False)
  
      def __post_init__(self):
          self.reset_gen_data()
  
      @property
      def synthetic_df(self) -> pd.DataFrame:
          """Get a DataFrame constructed from the generated lines"""
          if not self.gen_data_stream.getvalue():  # pragma: no cover
              return pd.DataFrame()
          self.gen_data_stream.seek(0)
          return pd.read_csv(self.gen_data_stream, sep=self.config.field_delimiter)
  
      def set_validator(self, fn: Callable, save=True):
          """Assign a validation callable to this batch. Optionally
          pickling and saving the validator for loading later
          """
          self.validator = fn
          if save:
              p = Path(self.checkpoint_dir) / "validator.p.gz"
              with gzip.open(p, "w") as fout:
                  fout.write(cloudpickle.dumps(fn))
  
      def load_validator_from_file(self):
          """Load a saved validation object if it exists"""
          p = Path(self.checkpoint_dir) / "validator.p.gz"
          if p.exists():
              with gzip.open(p, "r") as fin:
                  self.validator = cloudpickle.loads(fin.read())
  
      def reset_gen_data(self):
          """Reset all objects that accumulate or track synthetic
          data generation
          """
          self.gen_data_invalid = []
          self.gen_data_stream = io.StringIO()
          self.gen_data_stream.write(
              self.config.field_delimiter.join(self.headers) + "\n"
          )
          self.gen_data_count = 0
  
      def add_valid_data(self, data: GenText):
          """Take a ``gen_text`` object and add the generated
          line to the generated data stream
          """
          self.gen_data_stream.write(data.text + "\n")
          self.gen_data_count += 1
  
      def _basic_validator(self, raw_line: str):  # pragma: no cover
          return len(raw_line.split(self.config.field_delimiter)) == len(self.headers)
  
      def get_validator(self):
          """If a custom validator is set, we return that. Otherwise,
          we return the built-in validator, which simply checks if a generated
          line has the right number of values based on the number of headers
          for this batch.
  
          This at least makes sure the resulting DataFrame will be the right
          shape
          """
          if self.validator is not None:
              return self.validator
  
          return self._basic_validator
  
  
  def _create_batch_from_dir(batch_dir: str):
      path = Path(batch_dir)
      if not path.is_dir():  # pragma: no cover
          raise ValueError("%s is not a directory" % batch_dir)
  
      if not (path / HEADER_FILE).is_file():  # pragma: no cover
          raise ValueError("missing headers")
      headers = json.loads(open(path / HEADER_FILE).read())
  
      if not (path / CONFIG_FILE).is_file():  # pragma: no cover
          raise ValueError("missing model param file")
  
      config = config_from_model_dir(batch_dir)
  
      # training path can be empty, since we will not need access
      # to training data simply for read-only data generation
      train_path = ""
  
      # Wrap the user supplied callback with a _BatchEpochCallback so we have the batch number too.
      if config.epoch_callback is not None:
          batch_count = int(Path(batch_dir).name.split("_")[-1])
          config.epoch_callback = _BatchEpochCallback(
              config.epoch_callback, batch_count
          ).callback
  
      batch = Batch(
          checkpoint_dir=batch_dir,
          input_data_path=train_path,
          headers=headers,
          config=config,
      )
  
      batch.load_validator_from_file()
  
      return batch
  
  
  def _crawl_checkpoint_for_batches(checkpoint_dir: str):
      logger.info("Looking for and loading batch data...")
      matching_dirs = glob.glob(str(Path(checkpoint_dir) / "batch_*"))
      if not matching_dirs:
          raise ValueError(
              "checkpoint directory does not exist or does not contain batch data"
          )
  
      batches = []
      for batch_dir in matching_dirs:
          idx = int(Path(batch_dir).name.split("_")[-1])
          batches.append((idx, _create_batch_from_dir(batch_dir)))
  
      logger.info("Found and loaded %d batches", len(batches))
      return dict(sorted(batches, key=lambda b: b[0]))
  
  
  def _build_batch_dirs(
      base_ckpoint: str, headers: List[List[str]], config: dict
  ) -> dict:
      """Return a mapping of batch number => ``Batch`` object"""
      out = {}
      logger.info("Creating directory structure for batch jobs...")
      base_path = Path(config["checkpoint_dir"])
      if not base_path.is_dir():
          base_path.mkdir()
      for i, headers in enumerate(headers):
          ckpoint = Path(base_ckpoint) / f"batch_{i}"
          if not ckpoint.is_dir():
              ckpoint.mkdir()
          checkpoint_dir = str(ckpoint)
          input_data_path = str(ckpoint / "train.csv")
          new_config = deepcopy(config)
          new_config.update(
              {"checkpoint_dir": checkpoint_dir, "input_data_path": input_data_path}
          )
  
          # Determine what BaseConfig subclass to use, if the config template does
          # not have a model type then we'll default to using a LocalConfig which gives
          # us backwards compat to 0.14.0
          config_class_str = new_config.get(const.MODEL_TYPE, None)
          if config_class_str is None:
              config_class = LocalConfig
          else:
              config_class = CONFIG_MAP[config_class_str]
  
          # Wrap the user supplied callback with a _BatchEpochCallback so we have the batch number too.
          if new_config.get("epoch_callback") is not None:
              new_config["epoch_callback"] = _BatchEpochCallback(
                  new_config.get("epoch_callback"), i
              ).callback
  
          out[i] = Batch(
              checkpoint_dir=checkpoint_dir,
              input_data_path=input_data_path,
              headers=headers,
              config=config_class(**new_config),
          )
          # try and load any previously saved validators
          out[i].load_validator_from_file()
  
          # we write the headers out as well incase we load these
          # batches back in via "read" mode only later...it's the only
          # way to get the header names back
          with open(ckpoint / HEADER_FILE, "w") as fout:
              fout.write(json.dumps(headers))
  
      return out
  
  
  def _validate_batch_seed_values(
      batch: Batch, seed_values: Union[dict, List[dict]]
  ) -> Union[str, List[str]]:
      """Validate that seed values line up with the first N columns in a batch. Also construct
      an appropiate seed string based on the values in the batch
      """
      ret_str = True
      if isinstance(seed_values, dict):
          seed_values = [seed_values]
      elif isinstance(seed_values, list):
          ret_str = False
      else:
          raise TypeError("seed_values should be a dict or list of dicts")
  
      seed_strings = []
  
      for seed in seed_values:
          if len(seed) > len(batch.headers):
              raise RuntimeError(
                  "The number of seed fields is greater than the number of columns in the first batch"
              )
  
          headers_to_seed = batch.headers[: len(seed)]
          tmp = []
          for header in headers_to_seed:
              value = seed.get(header)
              if value is None:
                  raise RuntimeError(
                      f"The header: {header} is not in the seed values mapping"
                  )  # noqa
              tmp.append(str(value))
  
          seed_strings.append(
              batch.config.field_delimiter.join(tmp) + batch.config.field_delimiter
          )
  
      if ret_str:
          return seed_strings[0]
      else:
          return seed_strings
  
  
  class _BufferedRecords(abc.ABC):
      """Base class for all buffers used when
      generating records
      """
  
      @abc.abstractmethod
      def add(self, record: dict):
          ...
  
      @abc.abstractmethod
      def get_records(self):
          ...
  
      def cleanup(self):
          pass
  
  
  class _BufferedDicts(_BufferedRecords):
  
      _records: List[dict]
  
      def __init__(self):
          self._records = []
  
      def add(self, record: dict):
          self._records.append(record)
  
      def get_records(self):
          return self._records
  
  
  class _BufferedDataFrame(_BufferedRecords):
      """Buffer dictionaries into a memory or file, then
      load it as a DataFrame and set the column order
      based on the provided list. This allows
      datatypes to be inferred as if the values were
      being read from a CSV on disk.
  
      NOTE: The cleanup() method must be called when done
      with this class.
      """
  
      def __init__(self, delim: str, columns: List[str], method: str = FILE):
          self.delim = delim
          self.columns = columns
          self.headers_set = False
          self.method = method
  
          # Create our actual buffer file-like object
          if self.method == FILE:
              self.buffer = tempfile.TemporaryFile(mode="w+")
          elif self.method == MEMORY:
              self.buffer = io.StringIO()
          else:
              raise ValueError("Invalid method")
  
      def add(self, record: dict):
          # write the columns names into the buffer, we
          # use the first dict to specify the order and
          # assume subsequent dicts have the same order
          if not self.headers_set:
              _columns = self.delim.join(record.keys())
              self.buffer.write(_columns + "\n")
              self.headers_set = True
          _row = self.delim.join(record.values())
          self.buffer.write(_row + "\n")
  
      @property
      def df(self) -> pd.DataFrame:
          self.buffer.seek(0)
          return pd.read_csv(self.buffer, sep=self.delim)[self.columns]
  
      def get_records(self) -> pd.DataFrame:
          return self.df
  
      def cleanup(self):
          if self.method == FILE:
              self.buffer.close()
  
  
  @dataclass
  class GenerationProgress:
      """
      This class should not have to be used directly.
  
      It is used to communicate the current progress of record generation.
  
      When a callback function is passed to the ``RecordFactory.generate_all()`` method,
      each time the callback is called an instance of this class will be passed
      as the single argument::
  
          def my_callback(data: GenerationProgress):
              ...
  
          factory: RecordFactory
          df = factory.generate_all(output="df", callback=my_callback)
  
      This class is used to periodically communicate progress of generation to the user,
      through a callback that can be passed to ``RecordFactory.generate_all()`` method.
      """
  
      current_valid_count: int = 0
      """The number of valid lines/records that
      were generated so far.
      """
  
      current_invalid_count: int = 0
      """The number of invalid lines/records that
      were generated so far.
      """
  
      new_valid_count: int = 0
      """The number of new valid lines/records that
      were generated since the last progress callback.
      """
  
      new_invalid_count: int = 0
      """The number of new valid lines/records that
      were generated since the last progress callback.
      """
  
      completion_percent: float = 0.0
      """The percentage of valid lines/records that have been generated."""
  
      timestamp: float = field(default_factory=time.time)
      """The timestamp from when the information in this object has been captured."""
  
  
  class _GenerationCallback:
      """
      Wrapper around a callback function that is sending progress updates only once
      per configured time period (``update_interval``).
  
      Args:
          callback_fn: Callback function to be invoked with current progress.
          update_interval: Number of seconds to wait between sending progress update.
      """
  
      def __init__(self, callback_fn: callable, update_interval: int = 30):
          self._callback_fn = callback_fn
          self._update_interval = update_interval
  
          self._last_update_time = int(time.monotonic())
          self._last_progress = GenerationProgress()
  
      def update_progress(
          self,
          num_lines: int,
          valid_count: int,
          invalid_count: int,
          *,
          force_update=False,
      ):
  
          """
          Method that's being called from the generator with a progress update.
  
          Args:
              num_lines: Total number of lines to be generated.
              valid_count: Number of valid lines that were generated so far.
              invalid_count: Number of invalid lines that were generated so far.
              final_update:
                  Is this the final update? It is ``True`` when sending last update, after the
                  whole generation was complete.
          """
          now = int(time.monotonic())
  
          if now - self._last_update_time >= self._update_interval or force_update:
              current_progress = GenerationProgress(
                  current_valid_count=valid_count,
                  current_invalid_count=invalid_count,
                  new_valid_count=valid_count - self._last_progress.current_valid_count,
                  new_invalid_count=invalid_count
                  - self._last_progress.current_invalid_count,
                  completion_percent=0
                  if num_lines == 0
                  else round(valid_count / num_lines * 100, 2),
              )
  
              self._callback_fn(current_progress)
              self._last_update_time = now
              self._last_progress = current_progress
  
  
  @dataclass
  class _FactoryCounter:
      num_lines: int = 0
      """The target number of lines to generate when
      iterating or generating all records.
      """
  
      max_invalid: int = MAX_INVALID
      """The number of max invalid lines to tolerate before
      stopping generation and raising a ``RunTimeError.``
      """
  
      valid_count: int = 0
      """The number of valid records / lines that have been generated
      """
  
      invalid_count: int = 0
      """The number of invalid records / lines that were generated
      """
  
  
  def _threading_generation_callback(
      counter: _FactoryCounter, callback: _GenerationCallback, event: threading.Event
  ):
      while not event.is_set():
          try:
              callback.update_progress(
                  counter.num_lines, counter.valid_count, counter.invalid_count
              )
          except Exception:
              event.set()
              break
          time.sleep(1)
  
  
  class RecordFactory:
      """A stateful factory that can be used to generate and validate entire
      records, regardless of the number of underlying header clusters that were
      used to build multiple sub-models.
  
      Instances of this class should be created by calling the appropiate method
      of the ``DataFrameBatch`` instance. This class should not have to
      be used directly. You should be able to create an instance like so::
  
          factory = batcher.create_record_factory(num_lines=50)
  
      The class is init'd with default capacity and limits as specified
      by the ``num_lines`` and ``max_invalid`` attributes.  At any time,
      you can inspect the state of the instance by doing::
  
          factory.summary
  
      The factory instance can be used one of two ways: buffered or unbuffered.
  
      For unbuffered mode, the entire instance can be used as an iterator to
      create synthetic records. Each record will be a dictionary.
  
      NOTE:
          All values in the generated dictionaries will be strings.
  
      The ``valid_count`` and ``invalid_count`` counters will update as
      records are generated.
  
      When creating the record factory, you may also provide an entire
      record validator::
  
          def validator(rec: dict):
              ...
  
          factory = batcher.create_record_factory(num_lines=50, validator=validator)
  
      Each generated record dict will be passed to the validator. This validator may either
      return False or raise an exception to mark a record as invalid.
  
      At any point, you may reset the state of the factory by calling::
  
          factory.reset()
  
      This will reset all counters and allow you to keep generating records.
  
      Finally, you can generate records in buffered mode, where generated records
      will be buffered in memory and returned as one collection. By default, a list
      of dicts will be returned::
  
          factory.generate_all()
  
      You may request the records to be returned as a DataFrame.  The dtypes will
      be inferred as if you were reading the data from a CSV::
  
          factory.generate_all(output="df")
  
      NOTE:
          When using ``generate_all``, the factory states will be reset automatically.
      """
  
      validator: Callable
      """An optional callable that will receive a fully constructed record for one
      final validation before returning or yielding a single record. Records that
      do not pass this validation will also increment the ``invalid_count.``
      """
  
      _batches: Dict[int, Batch]
      _header_list: List[str]
      _seed_fields: Union[str, List[str]]
      _record_generator: IteratorType[dict]
      _delimiter: str
      _parallelism: int
      _counter = _FactoryCounter
      _invalid_cache_size: int
      _thread_event: threading.Event = None
  
      invalid_cache: List[dict]
  
      def __init__(
          self,
          *,
          num_lines: int,
          batches: dict,
          header_list: list,
          delimiter: str,
          seed_fields: Union[dict, list] = None,
          max_invalid=MAX_INVALID,
          validator: Optional[Callable] = None,
          parallelism: int = 4,
          invalid_cache_size: int = 100,
      ):
          self._counter = _FactoryCounter()
          self._counter.num_lines = num_lines
          self.max_invalid = max_invalid
          self._batches = batches
          self._header_list = header_list
          self._seed_fields = seed_fields
          self._delimiter = delimiter
          self._parallelism = parallelism
          self.validator = validator
          self._invalid_cache_size = invalid_cache_size
          self.reset()
  
          if self._seed_fields is not None:
              self._seed_fields = _validate_batch_seed_values(
                  self._batches[0], self._seed_fields
              )
  
          if isinstance(self._seed_fields, list):
              logger.info(
                  "Adjusting num_lines and parallelism because seed_fields is a list, will only target %d lines",
                  len(self._seed_fields),
              )  # noqa
              self._parallelism = 1
              self._counter.num_lines = len(self._seed_fields)
  
      def _cache_invalid(self, line: GenText):
          self.invalid_cache.append(line.as_dict())
          self.invalid_cache = self.invalid_cache[: self._invalid_cache_size]
  
      def _get_record(self) -> IteratorType[dict]:
          # our actual batch line generators
          generators = []
  
          # if we have a list of seed fields, we do special
          # handling to create the proper generator
          seed_generator = None  # assume no seeds to start
          if isinstance(self._seed_fields, list):
              seed_generator = SeedingGenerator(
                  self._batches[0].config,
                  seed_list=self._seed_fields,
                  line_validator=self._batches[0].get_validator(),
                  max_invalid=self.max_invalid * 10000,
              )
              generators.append((self._batches[0], seed_generator))
  
          for idx, batch in self._batches.items():
              start_string = None
              if idx == 0 and seed_generator:
                  # We've already added the first batch's generator to the list
                  # so we just continue on to the next one
                  continue
              if idx == 0:
                  # In the event we have seeds that aren't a list, (i.e. static seeds)
                  start_string = self._seed_fields
              generators.append(
                  (
                      batch,
                      # We seed the low level API with much higher limits on
                      # valid / invalid generation because we will enforce
                      # those limits in this high level instance.
                      generate_text(
                          batch.config,
                          line_validator=batch.get_validator(),
                          max_invalid=self.max_invalid * 10000,
                          num_lines=self._counter.num_lines * 10000,
                          start_string=start_string,
                          parallelism=self._parallelism,
                      ),
                  )
              )
  
          # At this point, we've created our list of generators. Below here
          # is what gets run on every next() call, which tries to construct
          # a full record from all the underlying batches.
  
          # keep looping as long as our target line count is less than
          # our total line count
          while self._counter.valid_count < self._counter.num_lines:
              # loop over each batch line generater and attempt
              # to construct a full line, we'll only count a
              # full line once we get through each generator
  
              # if we are using a watchdog thread to monitor generation
              # and it throws an exception, a threading event will be set
              # that signals generation should stop
              if self._thread_event and self._thread_event.is_set():
                  break
  
              if self._counter.invalid_count >= self.max_invalid:
                  raise RuntimeError("Invalid record count exceeded during generation")
  
              seed_cache = None
              if seed_generator:
                  # If we're using a seeding generator (from a list of seeds)
                  # we cache the next seed we are about to use to generate
                  # the next record.
                  seed_cache = seed_generator.settings.start_string[0]
  
              record = {}
              batch: Batch
              for batch, gen in generators:
                  while True:
  
                      # see above usage for watchdog thread exception handling
                      if self._thread_event and self._thread_event.is_set():
                          break
  
                      line = next(gen)  # type:  GenText
                      if line.valid is False:
                          self._cache_invalid(line)
                          self._counter.invalid_count += 1
                          if self._counter.invalid_count > self.max_invalid:
                              raise RuntimeError(
                                  "Invalid record count exceeded during generation"
                              )
                          continue
                      partial_rec = dict(
                          zip_longest(batch.headers, line.values_as_list(), fillvalue="")
                      )
                      record.update(partial_rec)
                      break
  
              # Do a final validation, if configured, on the fully constructed
              # record, if this validation fails, we'll still increment our
              # invalid count.
  
              valid = True  # assume we have a valid record
  
              if self.validator is not None:
                  try:
                      _valid = self.validator(record)
                      if _valid is False:
                          valid = False
                  except Exception:
                      valid = False
  
              if not valid:
                  self._counter.invalid_count += 1
                  if seed_cache:
                      seed_generator.settings.start_string.insert(0, seed_cache)
                  continue  # back to the while start
  
              self._counter.valid_count += 1
              yield record
  
      def __iter__(self):
          return self
  
      def __next__(self):
          return next(self._record_generator)
  
      def reset(self):
          self._counter.valid_count = 0
          self._counter.invalid_count = 0
          self._record_generator = self._get_record()
          self._thread_event = None
          self.invalid_cache = []
  
      def generate_all(
          self,
          output: Optional[str] = None,
          callback: Optional[callable] = None,
          callback_interval: int = 30,
          callback_threading: bool = False,
      ):
          """Attempt to generate the full number of records that was set when
          creating the ``RecordFactory.``  This method will create a buffer
          that holds all records and then returns the the buffer once
          generation is complete.
  
          Args:
              output: How the records should be returned. If ``None``, which is the
                  default, then a list of record dicts will be returned. Other options
                  that are supported are: 'df' for a DataFrame.
              callback: An optional callable that will periodically be called with
                  a ``GenerationProgress`` instance as the single argument while
                  records are being generated.
              callback_interval: If using a callback, the minimum number of seconds that
                  should occur between callbacks.
              callback_threading: If enabled, a watchdog thread will be used to execute
                  the callback. This will ensure that the callback is called regardless
                  of invalid or valid counts. If callback threading is disabled, the callback
                  will only be called after valid records are generated. If the callback
                  raises and exception, then a threading event will be set which will trigger
                  the stopping of generation.
  
          Returns:
              Generated records in an object that is dependent on the ``output`` param.  By default
              this will be a list of dicts.
          """
          progress_callback = None
          if callback:
              progress_callback = _GenerationCallback(callback, callback_interval)
  
          self.reset()
          if output is not None and output not in ("df",):
              raise ValueError("invalid output type")
  
          _iter = tqdm(self._record_generator, total=self._counter.num_lines)
  
          buffer = None  # type: _BufferedRecords
  
          if output == "df":
              buffer = _BufferedDataFrame(self._delimiter, self._header_list)
  
          if not buffer:
              buffer = _BufferedDicts()
  
          callback_thread = None
          if callback_threading:
              if not progress_callback:
                  raise ValueError(
                      "Cannot use callback_threading without a progress callback"
                  )
              self._thread_event = threading.Event()
              callback_thread = threading.Thread(
                  target=_threading_generation_callback,
                  args=(self._counter, progress_callback, self._thread_event),
              )
              callback_thread.start()
  
          try:
              for rec in _iter:
                  # NOTE: This iterator will block while no records are being
                  # succesfully generated. If callbacks need to occur in this
                  # situation, ensure the callback threading option is enabled
                  #
                  # If threading is enabled, and the callback encounters an exception,
                  # a threading event will be set and the generator will break out of its
                  # loop and generation will cease.
                  buffer.add(rec)
  
                  if progress_callback and not callback_threading:
                      progress_callback.update_progress(
                          self._counter.num_lines,
                          self._counter.valid_count,
                          self._counter.invalid_count,
                      )
  
          except (RuntimeError, StopIteration) as err:
              logger.warning(
                  f"Runtime error on iteration, returning current buffer, {str(err)}"
              )
          finally:
              if callback_threading:
                  self._thread_event.set()
                  callback_thread.join()
  
          # send final progress update
          if progress_callback:
              progress_callback.update_progress(
                  self._counter.num_lines,
                  self._counter.valid_count,
                  self._counter.invalid_count,
                  force_update=True,
              )
  
          out_records = buffer.get_records()
          buffer.cleanup()
          return out_records
  
      @property
      def summary(self):
          return {
              "num_lines": self._counter.num_lines,
              "max_invalid": self._counter.max_invalid,
              "valid_count": self._counter.valid_count,
              "invalid_count": self._counter.invalid_count,
          }
  
  
  class DataFrameBatch:
      """Create a multi-batch trainer / generator. When created, the directory
      structure to store models and training data will automatically be created.
      The directory structure will be created under the "checkpoint_dir" location
      provided in the ``config`` template. There will be one directory per batch,
      where each directory will be called "batch_N" where N is the batch number, starting
      from 0.
  
      Training and generating can happen per-batch or we can loop over all batches to
      do both train / generation functions.
  
      Example:
          When creating this object, you must explicitly create the training data
          from the input DataFrame before training models::
  
              my_batch = DataFrameBatch(df=my_df, config=my_config)
              my_batch.create_training_data()
              my_batch.train_all_batches()
  
      Args:
          df: The input, source DataFrame
          batch_size:  If ``batch_headers`` is not provided we automatically break up
              the number of columns in the source DataFrame into batches of N columns.
          batch_headers:  A list of lists of strings can be provided which will control
              the number of batches. The number of inner lists is the number of batches, and each
              inner list represents the columns that belong to that batch
          config: A template training config to use, this will be used as kwargs for each Batch's
              synthetic configuration. This may also be a sucblass of ``BaseConfig``. If this is used,
              you can set the ``input_data_path`` param to the constant ``PATH_HOLDER`` as it does not
              really matter
          tokenizer_class:  An optional ``BaseTokenizerTrainer`` subclass. If not provided the default
              tokenizer will be used for the underlying ML engine.
  
      NOTE:
          When providing a config, the source of training data is not necessary, only the
          ``checkpoint_dir`` is needed. Each batch will control its input training data path
          after it creates the training dataset.
      """
  
      batches: Dict[int, Batch]
      """A mapping of ``Batch`` objects to a batch number. The batch number (key)
      increments from 0..N where N is the number of batches being used.
      """
  
      batch_size: int
      """The max number of columns allowed for a single DF batch
      """
  
      # NOTE: Allowing a dict is for backwards compat
      config: Union[dict, BaseConfig]
      """The template config that will be used for all batches. If a dict
      is provided we default to a TensorFlowConfig.
      """
  
      mode: Union[WRITE, READ]
  
      master_header_list: List[str]
      """During training, this is the original column order. When reading from
      disk, we concatenate all headers from all batches together. This list is not
      guaranteed to preserve the original header order.
      """
  
      original_headers: List[str]
      """Stores the original header list / order from the original training data that was used.
      This is written out to the model directory during training and loaded back in when
      using read-only mode.
      """
  
      def __init__(
          self,
          *,
          df: pd.DataFrame = None,
          batch_size: int = BATCH_SIZE,
          batch_headers: List[List[str]] = None,
          config: Union[dict, BaseConfig] = None,
          tokenizer: BaseTokenizerTrainer = None,
          mode: str = WRITE,
          checkpoint_dir: str = None,
      ):
  
          if mode not in (WRITE, READ):  # pragma: no cover
              raise ValueError("mode must be read or write")
  
          self.mode = mode
  
          # If the config was a subclass of BaseConfig, then we convert
          # it to a dict and utilize that dict as our template. We do this
          # because when we re-create the batches we want to utilize the
          # Config constructors to set some attrs for us
          if isinstance(config, BaseConfig):
              config = config.as_dict()
  
          self.tokenizer = tokenizer
  
          self.original_headers = None
  
          if self.mode == READ:
              if isinstance(config, dict):
                  _ckpoint_dir = config.get("checkpoint_dir")
              else:
                  _ckpoint_dir = checkpoint_dir
  
              if _ckpoint_dir is None:
                  raise ValueError("checkpoint_dir required for read mode")
              else:
                  self._read_checkpoint_dir = _ckpoint_dir
  
          if self.mode == WRITE:
              if not config:
                  raise ValueError("config is required!")
  
              checkpoint_path = Path(config[CHECKPOINT_DIR])
              overwrite = config.get("overwrite", False)
              if (
                  not overwrite
                  and checkpoint_path.is_dir()
                  and any(checkpoint_path.iterdir())
              ):
                  raise RuntimeError(
                      "checkpoint_dir already exists and is non-empty, set overwrite on config or remove model directory!"
                  )  # noqa
  
              if overwrite and checkpoint_path.is_dir():
                  shutil.rmtree(checkpoint_path)
  
              if not isinstance(df, pd.DataFrame):
                  raise ValueError("df must be a DataFrame in write mode")
  
              if FIELD_DELIM not in config:
                  raise ValueError("field_delimiter must be in config")
  
              if GEN_LINES not in config:
                  config[GEN_LINES] = df.shape[0]
  
              self._source_df = df
              self.batch_size = batch_size
              self.config = config
              self._source_df.fillna("", inplace=True)
              self.master_header_list = list(self._source_df.columns)
  
              if not batch_headers:
                  self.batch_headers = self._create_header_batches()
              else:  # pragma: no cover
                  self.batch_headers = batch_headers
  
              self.batches = _build_batch_dirs(
                  self.config["checkpoint_dir"], self.batch_headers, self.config
              )
  
              # Preserve the original order of the DF headers
              self.original_headers = list(self._source_df)
              with open(Path(self.config[CHECKPOINT_DIR]) / ORIG_HEADERS, "w") as fout:
                  fout.write(json.dumps(list(self.original_headers)))
          else:
              self.batches = _crawl_checkpoint_for_batches(self._read_checkpoint_dir)
              self.master_header_list = []
              for batch in self.batches.values():
                  self.master_header_list.extend(batch.headers)
  
              try:
                  self.original_headers = json.loads(
                      open(Path(self._read_checkpoint_dir) / ORIG_HEADERS).read()
                  )
              except FileNotFoundError:
                  self.original_headers = None
  
              logger.info("Validating underlying models exist via generation test...")
              try:
                  self.generate_all_batch_lines(parallelism=1, num_lines=1)
              except Exception as err:
                  raise RuntimeError(
                      "Error testing generation during model load"
                  ) from err
  
      def _create_header_batches(self):
          num_batches = ceil(len(self._source_df.columns) / self.batch_size)
          tmp = np.array_split(list(self._source_df.columns), num_batches)
          return [list(row) for row in tmp]
  
      def create_training_data(self):
          """Split the original DataFrame into N smaller DataFrames. Each
          smaller DataFrame will have the same number of rows, but a subset
          of the columns from the original DataFrame.
  
          This method iterates over each ``Batch`` object and assigns
          a smaller training DataFrame to the ``training_df`` attribute
          of the object.
  
          Finally, a training CSV is written to disk in the specific
          batch directory
          """
          if self.mode == READ:  # pragma: no cover
              raise RuntimeError("Method cannot be used in read-only mode")
          for i, batch in self.batches.items():
              logger.info(f"Generating training DF and CSV for batch {i}")
              out_df = self._source_df[batch.headers]
              batch.training_df = out_df.copy(deep=True)
              out_df.to_csv(
                  batch.input_data_path,
                  header=False,
                  index=False,
                  sep=self.config[FIELD_DELIM],
              )
  
      def train_batch(self, batch_idx: int):
          """Train a model for a single batch. All model information will
          be written into that batch's directory.
  
          Args:
              batch_idx: The index of the batch, from the ``batches`` dictionary
          """
          if self.tokenizer is not None:
              _tokenizer = deepcopy(self.tokenizer)
              _tokenizer.config = self.batches[batch_idx].config
          else:
              _tokenizer = None
  
          if self.mode == READ:  # pragma: no cover
              raise RuntimeError("Method cannot be used in read-only mode")
          try:
              train(self.batches[batch_idx].config, _tokenizer)
          except KeyError:
              raise ValueError("batch_idx is invalid")
  
      def train_all_batches(self):
          """Train a model for each batch."""
          if self.mode == READ:  # pragma: no cover
              raise RuntimeError("Method cannot be used in read-only mode")
          for idx in self.batches.keys():
              self.train_batch(idx)
  
      def set_batch_validator(self, batch_idx: int, validator: Callable):
          """Set a validator for a specific batch. If a validator is configured
          for a batch, each generated record from that batch will be sent
          to the validator.
  
          Args:
              batch_idx: The batch number .
              validator: A callable that should take exactly one argument,
                  which will be the raw line generated from the ``generate_text``
                  function.
          """
          if self.mode == READ:  # pragma: no cover
              raise RuntimeError("Method cannot be used in read-only mode")
          if not callable(validator):
              raise ValueError("validator must be callable!")
          try:
              self.batches[batch_idx].set_validator(validator)
          except KeyError:
              raise ValueError("invalid batch number!")
  
      def generate_batch_lines(
          self,
          batch_idx: int,
          max_invalid=MAX_INVALID,
          raise_on_exceed_invalid: bool = False,
          num_lines: int = None,
          seed_fields: Union[dict, List[dict]] = None,
          parallelism: int = 0,
      ) -> GenerationSummary:
          """Generate lines for a single batch. Lines generated are added
          to the underlying ``Batch`` object for each batch. The lines
          can be accessed after generation and re-assembled into a DataFrame.
  
          Args:
              batch_idx: The batch number
              max_invalid: The max number of invalid lines that can be generated, if
                  this is exceeded, generation will stop
              raise_on_exceed_invalid: If true and if the number of lines generated exceeds the ``max_invalid``
                  amount, we will re-raise the error thrown by the generation module which will interrupt
                  the running process. Otherwise, we will not raise the caught exception and just return ``False``
                  indicating that the batch failed to generate all lines.
              num_lines: The number of lines to generate, if ``None``, then we use the number from the
                  batch's config
              seed_fields: A dictionary that maps field/column names to initial seed values for those columns. This seed
                  will only apply to the first batch that gets trained and generated. Additionally, the fields provided
                  in the mapping MUST exist at the front of the first batch.
  
                  NOTE:
                      This param may also be a list of dicts. If this is the case, then ``num_lines`` will automatically
                      be set to the list length downstream, and a 1:1 ratio will be used for generating valid lines for
                      each prefix.
              parallelism: The number of concurrent workers to use. ``1`` (the default) disables parallelization,
                  while a non-positive value means "number of CPUs + x" (i.e., use ``0`` for using as many workers
                  as there are CPUs). A floating-point value is interpreted as a fraction of the available CPUs,
                  rounded down.
          """
          try:
              batch = self.batches[batch_idx]
          except KeyError:  # pragma: no cover
              raise ValueError("invalid batch index")
  
          seed_string = None
  
          # If we are on batch 0 and we have seed values, we want to validate that
          # the seed values line up properly with the first N columns.
          if batch_idx == 0 and seed_fields is not None:
              seed_string = _validate_batch_seed_values(batch, seed_fields)
  
          batch: Batch
          batch.reset_gen_data()
          validator = batch.get_validator()
          if num_lines is None:
              num_lines = batch.config.gen_lines
  
          if isinstance(seed_fields, list):
              num_lines = len(seed_fields)
  
          t = tqdm(total=num_lines, desc="Valid record count ")
          t2 = tqdm(total=max_invalid, desc="Invalid record count ")
          line: GenText
          summary = GenerationSummary()
          try:
              for line in generate_text(
                  batch.config,
                  line_validator=validator,
                  max_invalid=max_invalid,
                  num_lines=num_lines,
                  start_string=seed_string,
                  parallelism=parallelism,
              ):
                  if line.valid is None or line.valid is True:
                      batch.add_valid_data(line)
                      t.update(1)
                      summary.valid_lines += 1
                  else:
                      t2.update(1)
                      batch.gen_data_invalid.append(line)
                      summary.invalid_lines += 1
          except TooManyInvalidError:
              if raise_on_exceed_invalid:
                  raise
              else:
                  return summary
          t.close()
          t2.close()
          summary.is_valid = batch.gen_data_count >= num_lines
          return summary
  
      def create_record_factory(
          self,
          *,
          num_lines: int,
          max_invalid: int = MAX_INVALID,
          validator: Callable = None,
          seed_fields: Union[dict, List[dict]] = None,
          parallellism: int = 4,
          **kwargs,
      ) -> RecordFactory:
          if validator is not None:
              if not callable(validator):
                  raise ValueError("validator must be callable")
          return RecordFactory(
              num_lines=num_lines,
              batches=self.batches,
              delimiter=self.batches[0].config.field_delimiter,
              header_list=self.original_headers or self.master_header_list,
              seed_fields=seed_fields,
              max_invalid=max_invalid,
              validator=validator,
              parallelism=parallellism,
              **kwargs,
          )
  
      def generate_all_batch_lines(
          self,
          max_invalid=MAX_INVALID,
          raise_on_failed_batch: bool = False,
          num_lines: int = None,
          seed_fields: Union[dict, List[dict]] = None,
          parallelism: int = 0,
      ) -> Dict[int, GenerationSummary]:
          """Generate synthetic lines for all batches. Lines for each batch
          are added to the individual ``Batch`` objects. Once generateion is
          done, you may re-assemble the dataset into a DataFrame.
  
          Example::
  
              my_batch.generate_all_batch_lines()
              # Wait for all generation to complete
              synthetic_df = my_batch.batches_to_df()
  
          Args:
              max_invalid: The number of invalid lines, per batch. If this number
                  is exceeded for any batch, generation will stop.
              raise_on_failed_batch: If True, then an exception will be raised if any single batch
                  fails to generate the requested number of lines. If False, then the failed batch
                  will be set to ``False`` in the result dictionary from this method.
              num_lines: The number of lines to create from each batch.  If ``None`` then the value
                  from the config template will be used.
  
                  NOTE:
                      Will be overridden / ignored if ``seed_fields`` is a list. Will be set to the len of the list.
              seed_fields: A dictionary that maps field/column names to initial seed values for those columns. This seed
                  will only apply to the first batch that gets trained and generated. Additionally, the fields provided
                  in the mapping MUST exist at the front of the first batch.
  
                  NOTE:
                      This param may also be a list of dicts. If this is the case, then ``num_lines`` will automatically
                      be set to the list length downstream, and a 1:1 ratio will be used for generating valid lines for
                      each prefix.
              parallelism: The number of concurrent workers to use. ``1`` (the default) disables parallelization,
                  while a non-positive value means "number of CPUs + x" (i.e., use ``0`` for using as many workers
                  as there are CPUs). A floating-point value is interpreted as a fraction of the available CPUs,
                  rounded down.
  
          Returns:
              A dictionary of batch number to a dictionary that reports the number of valid, invalid lines and bool value
              that shows if each batch was able to generate the full number of requested lines::
  
                  {
                      0: GenerationSummary(valid_lines=1000, invalid_lines=10, is_valid=True),
                      1: GenerationSummary(valid_lines=500, invalid_lines=5, is_valid=True)
                  }
          """
          batch_status = {}
          for idx in self.batches.keys():
              batch_status[idx] = self.generate_batch_lines(
                  idx,
                  max_invalid=max_invalid,
                  raise_on_exceed_invalid=raise_on_failed_batch,
                  num_lines=num_lines,
                  seed_fields=seed_fields,
                  parallelism=parallelism,
              )
          return batch_status
  
      def batch_to_df(self, batch_idx: int) -> pd.DataFrame:  # pragma: no cover
          """Extract a synthetic data DataFrame from a single batch.
  
          Args:
              batch_idx: The batch number
  
          Returns:
              A DataFrame with synthetic data
          """
          try:
              return self.batches[batch_idx].synthetic_df
          except KeyError:
              raise ValueError("batch_idx is invalid!")
  
      def batches_to_df(self) -> pd.DataFrame:
          """Convert all batches to a single synthetic data DataFrame.
  
          Returns:
              A single DataFrame that is the concatenation of all the
              batch DataFrames.
          """
          batch_iter = iter(self.batches.values())
          base_batch = next(batch_iter)
          accum_df = base_batch.synthetic_df
  
          for batch in batch_iter:
              accum_df = pd.concat([accum_df, batch.synthetic_df], axis=1)
  
          return accum_df[self.original_headers or self.master_header_list]
  
apis: pandas.read_csv pandas.DataFrame pandas.concat