import contextlib import copy import gc import glob import inspect import itertools import json import logging import os import pickle import random import shutil import threading import time import warnings from contextlib import contextmanager from dataclasses import dataclass, fields from pathlib import Path from typing import Any, Dict, Optional, Tuple, Union import cupy as cp import dill import numpy as np import pandas as pd import psutil import torch from datasets import Dataset from dython.nominal import associations, cramers_v, identify_nominal_columns from joblib import Parallel, delayed from sklearn.preprocessing import KBinsDiscretizer, OrdinalEncoder from realtabformer import REaLTabFormer as BaseREaLTabFormer, data_utils from realtabformer.data_utils import ( ModelFileName, ModelType, SpecialTokens, TEACHER_FORCING_PRE, encode_partition_numeric_col, encode_processed_column, process_datetime_data ) from realtabformer.realtabformer import _validate_get_device from realtabformer.rtf_datacollator import RelationalDataCollator from realtabformer.rtf_trainer import ResumableTrainer from tqdm import tqdm from transformers import ( EncoderDecoderConfig, EncoderDecoderModel, EarlyStoppingCallback, GPT2Config, GPT2LMHeadModel, LEDConfig, LEDForConditionalGeneration, PretrainedConfig, Seq2SeqTrainer, Seq2SeqTrainingArguments, Trainer, TrainerCallback, TrainingArguments, TrainerState, TrainerControl ) from transformers.models.led import modeling_led data_utils.SPECIAL_COL_SEP = "%" @dataclass(frozen=True) class ColDataType: NUMERIC: str = "NUM" DATETIME: str = "DATETIME" CATEGORICAL: str = "CAT" @staticmethod def types(): return [field.default for field in fields(ColDataType)] data_utils.ColDataType = ColDataType class LEDSeq2SeqModelOutput(modeling_led.LEDSeq2SeqModelOutput): hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None attentions: Optional[Tuple[torch.FloatTensor, ...]] = None modeling_led.LEDSeq2SeqModelOutput = LEDSeq2SeqModelOutput def get_positions(config: PretrainedConfig, encoder=True) -> int: if hasattr(config, "n_positions"): return config.n_positions if hasattr(config, "max_position_embeddings"): return config.max_position_embeddings key = f"max_{'en' if encoder else 'de'}coder_position_embeddings" if hasattr(config, key): return getattr(config, key) raise NotImplementedError(f"Config {config} positions isn't recognized.") def set_positions(config: PretrainedConfig, val: int, encoder=True): try: config.n_positions = val except: pass if hasattr(config, "max_position_embeddings"): config.max_position_embeddings = val led_name = f"max_{'en' if encoder else 'de'}coder_position_embeddings" if hasattr(config, led_name): setattr(config, led_name, val) return config class AdditionalEarlyStopCallback(TrainerCallback): def __init__( self, loss_threshold: float = 1e-6, early_stopping_patience: int = 1, early_stopping_threshold: Optional[float] = 0.0 ): self.loss_threshold = loss_threshold self.early_stopping_patience = early_stopping_patience self.early_stopping_threshold = early_stopping_threshold self.early_stopping_patience_counter = 0 self.best_loss = torch.inf def on_evaluate(self, args, state, control, metrics, **kwargs): loss_value = metrics.get("eval_loss") if loss_value is not None and self.loss_threshold > loss_value: control.should_training_stop = True def on_log(self, args, state, control, logs=None, **kwargs): loss = logs.get("loss") if loss is not None: if not np.isfinite(loss): raise RuntimeError("Training has diverged.") if loss < self.loss_threshold: control.should_training_stop = True elif state.global_step >= min(1000, args.warmup_steps): if loss < self.best_loss - self.early_stopping_threshold: self.early_stopping_patience_counter = 0 self.best_loss = loss else: self.early_stopping_patience_counter += 1 if self.early_stopping_patience_counter >= self.early_stopping_patience: control.should_training_stop = True def on_train_begin(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs): prev_losses = [] if state.log_history: for entry in reversed(state.log_history): if "loss" in entry: prev_losses.append(entry["loss"]) if len(prev_losses) > self.early_stopping_patience: break if prev_losses: self.best_loss = min(prev_losses) if not np.isfinite(self.best_loss): raise RuntimeError("Training has diverged.") if self.best_loss < self.loss_threshold: warnings.warn(f"Stop training due to best loss {self.best_loss}") control.should_training_stop = True elif (state.global_step >= min(1000, args.warmup_steps) and len(prev_losses) > self.early_stopping_patience and all( loss >= prev_losses[-1] - self.early_stopping_threshold for loss in prev_losses[:-1] )): warnings.warn(f"Stop training due to patience reached {prev_losses}") control.should_training_stop = True def make_relational_dataset( in_columns: list, out_columns: list, cache_path: str, vocab: dict, in_out_idx: dict, mask_rate=0, output_max_length: Optional[int] = None, return_token_type_ids: bool = False, ) -> Dataset: with open(os.path.join(cache_path, "in.csv"), 'r') as f: in_rows = sum(1 for _ in f) - 1 encoder_dataset = Dataset.from_csv( os.path.join(cache_path, "in.csv"), num_proc=max(5, min(20, in_rows // 1000)), keep_in_memory=False, cache_dir=os.path.join(cache_path, "in-cache") ) with open(os.path.join(cache_path, "out.csv"), 'r') as f: out_rows = sum(1 for _ in f) - 1 decoder_dataset = Dataset.from_csv( os.path.join(cache_path, "out.csv"), num_proc=max(5, min(20, out_rows // 1000)), keep_in_memory=False, cache_dir=os.path.join(cache_path, "out-cache") ) # Do not add [BOS] and [EOS] here. This will be handled # in the creation of the training_dataset in `get_relational_input_ids`. decoder_dataset = decoder_dataset.map( lambda example: get_input_ids_batched( example, vocab["decoder"], out_columns, mask_rate=mask_rate, return_label_ids=False, return_token_type_ids=return_token_type_ids, affix_bos=False, affix_eos=False, ), remove_columns=decoder_dataset.column_names, num_proc=max(5, min(20, out_rows // 1000)), desc="Decoder dataset", keep_in_memory=False, batched=True, batch_size=min(1000, out_rows // max(5, min(20, out_rows // 1000))) ) training_dataset = encoder_dataset.map( lambda example: get_input_ids_batched( example, vocab["encoder"], in_columns, return_label_ids=False, return_token_type_ids=return_token_type_ids, affix_bos=True, affix_eos=True, ), remove_columns=encoder_dataset.column_names, num_proc=max(5, min(20, in_rows // 1000)), desc="Encoder dataset", keep_in_memory=False, batched=True, batch_size=min(1000, in_rows // max(5, min(20, in_rows // 1000))) ) training_dataset = training_dataset.map( lambda example, idx: get_relational_input_ids_from_ids( example, idx, vocab, decoder_dataset, in_out_idx, output_max_length, ), with_indices=True, num_proc=max(5, min(20, in_rows // 1000)), desc="Encoder dataset connect", keep_in_memory=False ) # If the output_max_length variable is specified, filter # observations that exceed this length. The # `get_relational_input_ids` should have set the # `labels` to None if the output exceeds `output_max_length`. if output_max_length: init_data_length = training_dataset.shape[0] training_dataset = training_dataset.filter( lambda example: example["labels"] is not None ) removed_count = init_data_length - training_dataset.shape[0] if removed_count > 0: warnings.warn( f"A total of {removed_count} out of {init_data_length} has been removed from the training data because they exceeded the `output_max_length` of {output_max_length}." ) return training_dataset def process_numeric_data( series: pd.Series, max_len: int = 10, numeric_precision: int = 4, transform_data: Dict = None, ) -> Tuple[pd.Series, Dict]: is_transform = True if transform_data is None: transform_data = dict() is_transform = False if is_transform: warnings.warn( "Default values will be overridden because transform_data was passed..." ) max_len = transform_data["max_len"] numeric_precision = transform_data["numeric_precision"] else: transform_data["max_len"] = max_len transform_data["numeric_precision"] = numeric_precision has_neg = series.min() < 0 if pd.api.types.is_integer_dtype(series.dtype): is_int = True else: is_int = False # Get the most significant digit if is_transform: mx_sig = transform_data["mx_sig"] max_abs = None else: abs_num_series = series.abs() max_abs = abs_num_series.max() if max_abs == 0: max_abs = 1 if is_int else 1e-6 mx_sig = int(-1 if is_int else (1 if max_abs < 10 else np.ceil(np.log10(max_abs))) + has_neg) transform_data["mx_sig"] = mx_sig if mx_sig <= 0: # The data has no decimal point. # Pad the data with leading zeros if not # aligned to the largest value. # We also don't apply the max_len to integral # valued data because it will basically # remove important information. if is_transform: zfill = transform_data["zfill"] else: zfill = int(max(1, np.ceil(np.log10(max_abs))) + has_neg) # +1 for sign transform_data["zfill"] = zfill series = pd.Series(np.char.zfill(series.to_numpy(dtype='U'), zfill), index=series.index) else: # Make sure that we don't exessively truncate the data. # The max_len should be greater than the mx_sig. # Add a +1 to generate a minimum of tenth place resolution # for this data. assert max_len > ( mx_sig + 1 ), f"The target length {max_len} of the data doesn't include the numeric precision at {mx_sig}. Increase max_len to at least {max_len + (mx_sig + 2 - max_len)}." # Left align first based on the magnitude of the values. # We compute the difference in the most significant digits # of all values with respect to the largest value. # We then pad a leading zero to values with lower most significant # digits. # For example we have the values 1029.61 and 4.269. This will # determine that 1029.61 has the largest magnitude, with most significant # digit of 4. It will pad the value 4.269 with three zeros and convert it # to 0004.269. total_len = int(numeric_precision + 1 + mx_sig) series = pd.Series( np.char.mod(f'%0{total_len}.{numeric_precision}f', series.values), index=series.index ).str[:max_len] # We additionally apply left justify to align based on the trailing precision. # For example, we have 1029.61 and 0004.269 as values. This time we transform the first # value to become 1029.610 to align with the precision of the second value. if is_transform: ljust = transform_data["ljust"] else: ljust = min(max_len, total_len) transform_data["ljust"] = int(ljust) return series, transform_data def tokenize_numeric_col(series: pd.Series, nparts=2, col_zfill=2): # After normalizing the numeric values, we then segment # them based on a fixed partition size (nparts). col = series.name lengths = series.str.len() max_len = lengths.max() if nparts > max_len > 2: # Allow minimum of 0-99 as acceptable singleton range. raise ValueError( f"Partition size {nparts} is greater than the value length {max_len}. Consider reducing the number of partitions..." ) mx = lengths.min() tr = pd.concat([series.str[i : i + nparts] for i in range(0, mx, nparts)], axis=1) tr.columns = encode_partition_numeric_col(col, tr, col_zfill) return tr def process_data( df: pd.DataFrame, numeric_max_len=10, numeric_precision=4, numeric_nparts=2, first_col_type=None, col_transform_data: Dict = None, target_col: str = None, base_idx=None ) -> Tuple[pd.DataFrame, Dict]: # This should receive a dataframe with dtypes that have already been # properly categorized between numeric and categorical. # Date type can be converted as UNIX timestamps. assert first_col_type in [None, ColDataType.CATEGORICAL, ColDataType.NUMERIC] df = df.copy() # Unify the variable for missing data df = df.fillna(pd.NA) # Force cast integral values to Int64Dtype dtype # to save precision if they are represented as float. for c in df: try: if pd.api.types.is_datetime64_any_dtype(df[c].dtype): # Don't cast datetime types. continue if pd.api.types.is_numeric_dtype(df[c].dtype): # Only cast if the column is explicitly numeric type. df[c] = df[c].astype(pd.Int64Dtype()) except TypeError: pass except ValueError: pass if target_col is not None: assert ( first_col_type is None ), "Implicit ordering of columns when teacher-forcing of target is used is not supported yet!" tf_col_name = f"{TEACHER_FORCING_PRE}_{target_col}" assert ( tf_col_name not in df.columns ), f"The column name ({tf_col_name}) must not be in the raw data. Found instead..." target_ser = df[target_col].copy() target_ser.name = tf_col_name df = pd.concat([target_ser, df], axis=1) # Rename the columns to encode the original order by adding a suffix of increasing # integer values. num_cols = len(str(len(df.columns))) if base_idx is None: base_idx = 0 col_idx = {col: f"{str(i + base_idx).zfill(num_cols)}" for i, col in enumerate(df.columns)} # Create a dataframe that will hold the processed data processed_series = [] # Process numerical data numeric_cols = df.select_dtypes(include=np.number).columns if col_transform_data is None: col_transform_data = dict() for c in numeric_cols: col_name = encode_processed_column(col_idx[c], ColDataType.NUMERIC, c) _col_transform_data = col_transform_data.get(c) series, transform_data = process_numeric_data( df[c], max_len=numeric_max_len, numeric_precision=numeric_precision, transform_data=_col_transform_data, ) if _col_transform_data is None: # This means that no transform data is available # before the processing. col_transform_data[c] = transform_data series.name = col_name processed_series.append(series) # Process datetime data datetime_cols = df.select_dtypes(include="datetime").columns for c in datetime_cols: col_name = encode_processed_column(col_idx[c], ColDataType.DATETIME, c) _col_transform_data = col_transform_data.get(c) series, transform_data = process_datetime_data( df[c], transform_data=_col_transform_data, ) if _col_transform_data is None: # This means that no transform data is available # before the processing. col_transform_data[c] = transform_data series.name = col_name processed_series.append(series) processed_df = pd.concat([pd.DataFrame()] + processed_series, axis=1) if not processed_df.empty: # Combine the processed numeric and datetime data. processed_df = pd.concat( [ tokenize_numeric_col(processed_df[col], nparts=numeric_nparts) for col in processed_df.columns ], axis=1, ) # NOTE: The categorical data should be the last to be processed! categorical_cols = df.columns.difference(numeric_cols).difference(datetime_cols) if not categorical_cols.empty: # Process the rest of the data, assumed to be categorical values. processed_cat = df[categorical_cols].astype(str) cat_col_idx = pd.Series([col_idx[c] for c in categorical_cols]) new_cat_col_names = (cat_col_idx + data_utils.SPECIAL_COL_SEP + f"{ColDataType.CATEGORICAL}" + data_utils.SPECIAL_COL_SEP + pd.Series(categorical_cols)) processed_df = pd.concat( [ processed_df, processed_cat.set_axis(new_cat_col_names.tolist(), axis=1) ], axis=1, ) # Get the different sets of column types is_cat = processed_df.columns.str.contains(ColDataType.CATEGORICAL) cat_cols = processed_df.columns[is_cat] numeric_cols = processed_df.columns[~is_cat] if first_col_type == ColDataType.CATEGORICAL: df = processed_df[cat_cols.union(numeric_cols, sort=False)] elif first_col_type == ColDataType.NUMERIC: df = processed_df[numeric_cols.union(cat_cols, sort=False)] else: # Reorder columns to the original order df = processed_df[np.sort(processed_df.columns)] df = df.columns.values.reshape((1, -1)) + data_utils.SPECIAL_COL_SEP + df return df, col_transform_data def build_large_vocab(cache_path, chunk_size, columns, part): id2token = {} curr_id = 0 special_tokens = SpecialTokens.tokens() if special_tokens: id2token.update(dict(enumerate(special_tokens))) curr_id = max(id2token) + 1 column_token_ids = {} def obtain_unique(series): return series.name, series.unique() col_uniques = {col: set() for col in columns} for chunk in pd.read_csv(os.path.join(cache_path, f"{part}.csv"), chunksize=10_000_000 // len(columns)): chunk_uniques = Parallel(n_jobs=20)( delayed(obtain_unique)(chunk[col]) for col in columns ) for col, vals in chunk_uniques: col_uniques[col].update(vals) def collect_unique(col, uniques, dtype): return col, np.sort(np.array([*uniques], dtype=dtype)) unique_values = Parallel(n_jobs=20)( delayed(collect_unique)(col, col_uniques[col], chunk[col].dtype) for col in tqdm(columns, desc="Sorting unique values") ) pbar = tqdm(desc=f"Preparing {part}put vocabulary", total=len(columns)) for col, values in unique_values: new_id2token = dict(enumerate(values, curr_id)) id2token.update(new_id2token) new_curr_id = max(new_id2token) + 1 column_token_ids[col] = list(range(curr_id, new_curr_id)) curr_id = new_curr_id pbar.update() token2id = {v: k for k, v in id2token.items()} return dict( id2token=id2token, token2id=token2id, column_token_ids=column_token_ids, ) def get_input_ids_batched( batch, vocab: Dict, columns: list, mask_rate: float = 0, return_label_ids: Optional[bool] = True, return_token_type_ids: Optional[bool] = False, affix_bos: Optional[bool] = True, affix_eos: Optional[bool] = True, ) -> Dict: # Raise an assertion error while the implementation # is not yet ready. assert return_token_type_ids is False input_ids: list[np.ndarray] = [] token_type_ids: list[int] = [] batch_size = len(batch[columns[0]]) if affix_bos: input_ids.append(np.full(batch_size, vocab["token2id"][SpecialTokens.BOS], dtype=np.int32)) if return_token_type_ids: token_type_ids.append(vocab["token2id"][SpecialTokens.SPTYPE]) vocab_token2id = vocab["token2id"] masked = np.random.random((batch_size, len(columns))) < mask_rate col_names = pd.Series(columns).str.extract( f"([0-9]+{data_utils.SPECIAL_COL_SEP}({'|'.join(ColDataType.types())}))" )[0] for j, k in enumerate(columns): token_ids = [vocab_token2id.get(token, vocab_token2id[SpecialTokens.UNK]) for token in batch[k]] token_ids = np.array(token_ids, dtype=np.int32) if mask_rate > 0: token_ids[masked[:, j]] = vocab_token2id[SpecialTokens.RMASK] input_ids.append(token_ids) if return_token_type_ids: col_name = col_names[j] token_type_ids.append(vocab["token2id"][col_name]) if affix_eos: input_ids.append(np.full(batch_size, vocab["token2id"][SpecialTokens.EOS], dtype=np.int32)) if return_token_type_ids: token_type_ids.append(vocab["token2id"][SpecialTokens.SPTYPE]) input_ids = np.stack(input_ids).T.tolist() data = dict(input_ids=input_ids) if return_label_ids: data["label_ids"] = input_ids if return_token_type_ids: data["token_type_ids"] = np.array(token_type_ids).reshape((1, -1)).repeat(batch_size, axis=0).tolist() return data def get_relational_input_ids_from_ids( example, input_idx, vocab, output_dataset, in_out_idx, output_max_length: Optional[int] = None, ) -> dict: # Start with 2 to take into account the [BOS] and [EOS] tokens sequence_len = 2 # Build the input_ids for the encoder input_ids = example["input_ids"] token_type_ids = example.get("token_type_ids") # Build the label_ids for the decoder output_idx = in_out_idx[input_idx] valid = True label_ids = [vocab["decoder"]["token2id"][SpecialTokens.BOS]] if len(output_idx) > 0: for ids in output_dataset.select(output_idx)["input_ids"]: # Pad each observation with the [BMEM] and [EMEM] tokens tmp_label_ids = [vocab["decoder"]["token2id"][SpecialTokens.BMEM]] tmp_label_ids.extend(ids) tmp_label_ids.append(vocab["decoder"]["token2id"][SpecialTokens.EMEM]) if output_max_length: if (sequence_len + len(tmp_label_ids)) > output_max_length: # This exceeds the expected limit. # Drop this observation. valid = False break label_ids.extend(tmp_label_ids) sequence_len += len(tmp_label_ids) label_ids.append(vocab["decoder"]["token2id"][SpecialTokens.EOS]) payload = dict( input_ids=input_ids, # The variable `labels` is used in the EncoderDecoder model # instead of `label_ids`. labels=label_ids if valid else None, ) if token_type_ids is not None: payload["token_type_ids"] = token_type_ids return payload def relational_to_led_config(relational_config): combined_config: LEDConfig = copy.deepcopy(relational_config.encoder) combined_config.is_encoder_decoder = True combined_config.decoder_layers = relational_config.decoder.decoder_layers combined_config.decoder_attention_heads = relational_config.decoder.decoder_attention_heads combined_config.decoder_ffn_dim = relational_config.decoder.decoder_ffn_dim combined_config.decoder_start_token_id = relational_config.decoder_start_token_id combined_config.max_decoder_position_embeddings = ( relational_config.decoder.max_decoder_position_embeddings) combined_config.add_cross_attention = True combined_config.bos_token_id = relational_config.bos_token_id combined_config.eos_token_id = relational_config.eos_token_id combined_config.pad_token_id = relational_config.pad_token_id combined_config.vocab_size = max(combined_config.vocab_size, relational_config.decoder.vocab_size) return combined_config class REaLTabFormer(BaseREaLTabFormer): # allow LED relational config def __init__( self, *args, early_stopping_threshold: float = 1e-3, n_critic: int = 5, learning_rate: float = 2e-4, lr_scheduler_type: str = "cosine", adam_epsilon: float = 1e-6, **kwargs ): self.n_critic = n_critic super().__init__( *args, **kwargs, early_stopping_threshold=early_stopping_threshold, learning_rate=learning_rate, lr_scheduler_type=lr_scheduler_type, adam_epsilon=adam_epsilon ) def _split_train_eval_dataset(self, dataset: Dataset): test_size = 1 - self.train_size if test_size > 0: dataset = dataset.train_test_split( test_size=test_size, seed=self.random_state ) dataset["train_dataset"] = dataset.pop("train") dataset["eval_dataset"] = dataset.pop("test") # Override `metric_for_best_model` from "loss" to "eval_loss" self.training_args_kwargs["metric_for_best_model"] = "eval_loss" # Make this explicit so that no assumption is made on the # direction of the metric improvement. self.training_args_kwargs["greater_is_better"] = False else: dataset = dict(train_dataset=dataset) self.training_args_kwargs["evaluation_strategy"] = "no" self.training_args_kwargs["load_best_model_at_end"] = False steps_per_epoch = len(dataset["train_dataset"]) // (self.batch_size * torch.cuda.device_count()) self.training_args_kwargs["eval_steps"] = min( max(self.training_args_kwargs["eval_steps"], steps_per_epoch), 1000 ) self.training_args_kwargs["save_steps"] = self.training_args_kwargs["eval_steps"] self.training_args_kwargs["warmup_steps"] = min( 10000, min(5, self.training_args_kwargs.get("num_train_epochs", 5)) * steps_per_epoch ) return dataset def _set_up_relational_coder_configs(self) -> None: def _get_coder(coder_name) -> Union[GPT2Config, LEDConfig]: return getattr(self.relational_config, coder_name) for coder_name in ["encoder", "decoder"]: coder = _get_coder(coder_name) coder.bos_token_id = self.vocab[coder_name]["token2id"][SpecialTokens.BOS] coder.eos_token_id = self.vocab[coder_name]["token2id"][SpecialTokens.EOS] coder.pad_token_id = self.vocab[coder_name]["token2id"][SpecialTokens.PAD] coder.vocab_size = len(self.vocab[coder_name]["id2token"]) if coder_name == "decoder": self.relational_config.bos_token_id = coder.bos_token_id self.relational_config.eos_token_id = coder.eos_token_id self.relational_config.pad_token_id = coder.pad_token_id self.relational_config.decoder_start_token_id = coder.eos_token_id # Make sure that we have at least the number of # columns in the transformed data as positions. # This will prevent runtime error. # `RuntimeError: CUDA error: device-side assert triggered` assert self.relational_max_length if ( coder_name == "decoder" and get_positions(coder, encoder=False) < self.relational_max_length ): coder = set_positions(coder, 128 + self.relational_max_length, encoder=False) elif coder_name == "encoder" and get_positions(coder) < len(self.vocab[coder_name]["column_token_ids"]): positions = 128 + len(self.vocab[coder_name]["column_token_ids"]) coder = set_positions(coder, positions) # This must be set to True for the EncoderDecoderModel to work at least # with GPT2 as the decoder. self.relational_config.decoder.add_cross_attention = True def _fit_relational( self, out_df: pd.DataFrame, in_df: pd.DataFrame, join_on: str, device="cuda", resume_from_checkpoint=None ): # bert2bert.config.decoder_start_token_id = tokenizer.cls_token_id # bert2bert.config.eos_token_id = tokenizer.sep_token_id # bert2bert.config.pad_token_id = tokenizer.pad_token_id # bert2bert.config.vocab_size = bert2bert.config.encoder.vocab_size # All join values in the out_df must be present in the in_df. assert len(set(out_df[join_on].unique()).difference(in_df[join_on])) == 0 # Get the list of index of observations that are related based on # the join_on variable. common_out_idx = ( out_df.reset_index(drop=True) .groupby(join_on) .apply(lambda x: x.index.to_list()) ) # Track the mapping of index from input to the list of output indices. in_out_idx = pd.Series( # Reset the index so that we are sure that the index ids are set properly. dict(in_df[join_on].reset_index(drop=True).items()) ).map(lambda x: common_out_idx.get(x, [])) # Remove the unique id column from the in_df and the out_df in_df = in_df.drop(join_on, axis=1) out_df = out_df.drop(join_on, axis=1) cache_path = str(self.checkpoints_dir).replace("ckpts", "cache-data") os.makedirs(cache_path, exist_ok=True) self._extract_column_info(out_df) max_size = min(10_000_000, out_df.shape[0] * 20) if out_df.size > max_size: self.col_transform_data, chunk_size, self.processed_columns, out_df = self._process_large_data( out_df, cache_path, max_size, "out" ) self.vocab["decoder"] = build_large_vocab(cache_path, chunk_size, self.processed_columns, "out") else: out_df, self.col_transform_data = process_data( out_df, numeric_max_len=self.numeric_max_len, numeric_precision=self.numeric_precision, numeric_nparts=self.numeric_nparts, ) self.processed_columns = out_df.columns.to_list() self.vocab["decoder"] = self._generate_vocab(out_df) out_df.to_csv(os.path.join(cache_path, "out.csv"), index=False) self.relational_col_size = len(self.processed_columns) out_shape = out_df.shape[0], self.relational_col_size del out_df gc.collect() # NOTE: the index starts at zero, but should be adjusted # to account for the special tokens. For relational data, # the index should start at 3 ([[EOS], [BOS], [BMEM]]). self.col_idx_ids = { ix: self.vocab["decoder"]["column_token_ids"][col] for ix, col in enumerate(self.processed_columns) } # Add these special tokens at specific key values # which are used in `REaLSampler._get_relational_col_idx_ids` self.col_idx_ids[-1] = [ self.vocab["decoder"]["token2id"][SpecialTokens.BMEM], self.vocab["decoder"]["token2id"][SpecialTokens.EOS], ] self.col_idx_ids[-2] = [self.vocab["decoder"]["token2id"][SpecialTokens.EMEM]] max_size = min(1_000_000, in_df.shape[0] * 20) if in_df.size > max_size: self.in_col_transform_data, chunk_size, in_columns, in_df = self._process_large_data( in_df, cache_path, max_size, "in" ) if self.parent_vocab is None: self.vocab["encoder"] = build_large_vocab(cache_path, chunk_size, in_columns, "in") else: self.vocab["encoder"] = self.parent_vocab in_shape = in_df.shape[0], len(in_columns) else: in_df, self.in_col_transform_data = process_data( in_df, numeric_max_len=self.numeric_max_len, numeric_precision=self.numeric_precision, numeric_nparts=self.numeric_nparts, col_transform_data=self.parent_col_transform_data, ) if self.parent_vocab is None: self.vocab["encoder"] = self._generate_vocab(in_df) else: self.vocab["encoder"] = self.parent_vocab in_df.to_csv(os.path.join(cache_path, "in.csv"), index=False) in_columns = in_df.columns.tolist() in_shape = in_df.shape del in_df gc.collect() # Load the dataframe into a HuggingFace Dataset # torch.save(dict(vocab=self.vocab, # in_out_idx=in_out_idx, # output_max_length=self.output_max_length, # mask_rate=self.mask_rate, # return_token_type_ids=False,), os.path.join(cache_path, "init-dataset.pkl")) dataset = make_relational_dataset( # in_df=in_df, # out_df=out_df, in_columns=in_columns, out_columns=self.processed_columns, cache_path=cache_path, vocab=self.vocab, in_out_idx=in_out_idx, output_max_length=self.output_max_length, mask_rate=self.mask_rate, return_token_type_ids=False, ) # Compute the longest sequence of labels in the dataset and add a buffer of 1. self.relational_max_length = ( max( dataset.map( lambda example: dict(length=len(example["labels"])), num_proc=min(10, max(1, len(dataset) // 1000)) )[ "length" ] ) + 1 ) # Create train-eval split if specified dataset = self._split_train_eval_dataset(dataset) enc_configs = [ GPT2Config(n_layer=6), LEDConfig(encoder_layers=6), GPT2Config(n_layer=6, n_embd=256, n_inner=1024, n_head=8), LEDConfig( encoder_layers=6, encoder_ffn_dim=1024, encoder_attention_heads=8, d_model=256, attention_window=256 ), GPT2Config(n_layer=4, n_embd=128, n_inner=384, n_head=8), LEDConfig( encoder_layers=4, d_model=128, encoder_ffn_dim=384, encoder_attention_heads=8, attention_window=128 ) ] dec_configs = [ GPT2Config(n_layer=6), LEDConfig(is_decoder=True, add_cross_attention=True, decoder_layers=6), GPT2Config(n_layer=6, n_embd=256, n_inner=1024, n_head=8), LEDConfig( decoder_layers=6, decoder_ffn_dim=1024, decoder_attention_heads=8, d_model=256, is_decoder=True, add_cross_attention=True, attention_window=256 ), GPT2Config(n_layer=4, n_embd=128, n_inner=384, n_head=8), LEDConfig( decoder_layers=4, d_model=128, decoder_ffn_dim=384, decoder_attention_heads=8, is_decoder=True, add_cross_attention=True, attention_window=128 ) ] ee = None for dec_config in dec_configs: for enc_config in enc_configs: if dec_config.model_type != enc_config.model_type: continue if len(in_columns) > 1024 and dec_config.model_type != "led": continue self.relational_config = EncoderDecoderConfig( encoder=enc_config.to_dict().copy(), decoder=dec_config.to_dict().copy() ) try: trainer = self._run_relational_trainer(device, dataset, resume_from_checkpoint) shutil.rmtree(cache_path) return trainer except (torch.cuda.OutOfMemoryError, RuntimeError) as e: ee = e if isinstance(e, torch.cuda.OutOfMemoryError) or "out of memory" in str(e) or "batch size" in str( e): if os.path.exists(self.checkpoints_dir): shutil.rmtree(self.checkpoints_dir) gc.collect() torch.cuda.empty_cache() else: raise e raise ee def _process_large_data(self, df, cache_path, max_size, part): chunk_size = int(np.ceil(max_size / df.shape[0])) csv_paths = [] in_columns = [] chunks = [] for i in range(int(np.ceil(df.shape[-1] / chunk_size))): selected_chunk = df.iloc[:, :chunk_size] df = df.drop(columns=selected_chunk.columns) chunks.append(selected_chunk) all_in_col_transform_data = {} def process_chunk(i, selected_chunk): processed_in_df, this_in_col_transform = process_data( selected_chunk, numeric_max_len=self.numeric_max_len, numeric_precision=self.numeric_precision, numeric_nparts=self.numeric_nparts, col_transform_data=self.parent_col_transform_data, base_idx=i * chunk_size ) csv_path = os.path.join(cache_path, f"temp-{part}-{i}.csv") processed_in_df.to_csv(csv_path, index=False) return this_in_col_transform, csv_path, processed_in_df.columns.tolist() pbar = tqdm( enumerate(chunks), desc=f"Processing {part}put data [chunksize={chunk_size} x {len(chunks)}]", total=len(chunks) ) results = Parallel(n_jobs=min(20, len(chunks)), verbose=10)( delayed(process_chunk)(i, selected_chunk) for i, selected_chunk in pbar ) for t, p, c in results: all_in_col_transform_data.update(t) csv_paths.append(p) in_columns.extend(c) row_chunk_size = 10000 input_files = [open(path, 'r') for path in csv_paths] headers = [f.readline().strip() for f in input_files] merged_header = ",".join(headers) output_file = open(os.path.join(cache_path, f"{part}.csv"), 'w') output_file.write(merged_header + '\n') total_lines = df.shape[0] pbar = tqdm(total=total_lines, desc="Combining chunks") caches = [] for i in range(total_lines): lines = [f.readline() for f in input_files] combined = ",".join(line.strip() for line in lines) caches.append(combined) if (i + 1) % row_chunk_size == 0: output_file.write("\n".join([*caches, ""])) caches = [] pbar.update() if caches: output_file.write("\n".join([*caches, ""])) # Clean up for f in input_files: f.close() output_file.close() for path in csv_paths: os.remove(path) return all_in_col_transform_data, chunk_size, in_columns, df @classmethod def load_from_dir(cls, path: Union[str, Path]): if not os.path.exists(os.path.join(path, "rtf_config.json")): path = glob.glob(f"{path}/*")[0] if isinstance(path, str): path = Path(path) config_file = path / ModelFileName.rtf_config_json model_file = path / ModelFileName.rtf_model_pt assert path.is_dir(), f"Directory {path} does not exist." assert config_file.exists(), f"Config file {config_file} does not exist." assert model_file.exists(), f"Model file {model_file} does not exist." # Load the saved attributes rtf_attrs = json.loads(config_file.read_text()) # Create new REaLTabFormer model instance try: realtf = cls(model_type=rtf_attrs["model_type"]) except KeyError: # Back-compatibility for saved models # before the support for relational data # was implemented. realtf = cls(model_type="tabular") # Set all attributes and handle the # special case for the GPT2Config. for k, v in rtf_attrs.items(): if k == "gpt_config": # Back-compatibility for saved models # before the support for relational data # was implemented. v = GPT2Config.from_dict(v) k = "tabular_config" elif k == "tabular_config": v = GPT2Config.from_dict(v) elif k == "relational_config": v = EncoderDecoderConfig.from_dict(v) elif k in ["checkpoints_dir", "samples_save_dir"]: v = Path(v) elif k == "vocab": if realtf.model_type == ModelType.tabular: # Cast id back to int since JSON converts them to string. v["id2token"] = {int(ii): vv for ii, vv in v["id2token"].items()} elif realtf.model_type == ModelType.relational: v["encoder"]["id2token"] = { int(ii): vv for ii, vv in v["encoder"]["id2token"].items() } v["decoder"]["id2token"] = { int(ii): vv for ii, vv in v["decoder"]["id2token"].items() } else: raise ValueError(f"Invalid model_type: {realtf.model_type}") elif k == "col_idx_ids": v = {int(ii): vv for ii, vv in v.items()} setattr(realtf, k, v) # Implement back-compatibility for REaLTabFormer version < 0.0.1.8.2 # since the attribute `col_idx_ids` is not implemented before. if "col_idx_ids" not in rtf_attrs: if realtf.model_type == ModelType.tabular: realtf.col_idx_ids = { ix: realtf.vocab["column_token_ids"][col] for ix, col in enumerate(realtf.processed_columns) } elif realtf.model_type == ModelType.relational: # NOTE: the index starts at zero, but should be adjusted # to account for the special tokens. For relational data, # the index should start at 3 ([[EOS], [BOS], [BMEM]]). realtf.col_idx_ids = { ix: realtf.vocab["decoder"]["column_token_ids"][col] for ix, col in enumerate(realtf.processed_columns) } # Add these special tokens at specific key values # which are used in `REaLSampler._get_relational_col_idx_ids` realtf.col_idx_ids[-1] = [ realtf.vocab["decoder"]["token2id"][SpecialTokens.BMEM], realtf.vocab["decoder"]["token2id"][SpecialTokens.EOS], ] realtf.col_idx_ids[-2] = [ realtf.vocab["decoder"]["token2id"][SpecialTokens.EMEM] ] # Load model weights if realtf.model_type == ModelType.tabular: realtf.model = GPT2LMHeadModel(realtf.tabular_config) elif realtf.model_type == ModelType.relational: if realtf.relational_config.encoder.model_type == "led": combined_config = relational_to_led_config(realtf.relational_config) realtf.model = LEDForConditionalGeneration(combined_config) else: realtf.model = EncoderDecoderModel(realtf.relational_config) else: raise ValueError(f"Invalid model_type: {realtf.model_type}") realtf.model.load_state_dict( torch.load(model_file.as_posix(), map_location="cpu") ) return realtf def _run_relational_trainer(self, device, dataset, resume_from_checkpoint): # Set up the config and the model self._set_up_relational_coder_configs() # Build the model. if self.relational_config.encoder.model_type == "led" and self.relational_config.decoder.model_type == "led": combined_config = relational_to_led_config(self.relational_config) self.model = LEDForConditionalGeneration(combined_config) else: self.model = EncoderDecoderModel(self.relational_config) if self.parent_gpt2_state_dict is not None: pretrain_load = self.model.encoder.load_state_dict( self.parent_gpt2_state_dict, strict=False ) assert ( not pretrain_load.missing_keys ), "There should be no missing_keys after loading the pretrained GPT2 state!" if self.freeze_parent_model: # We freeze the weights if we use the pretrained # parent table model. for param in self.model.encoder.parameters(): param.requires_grad = False # Tell pytorch to run this model on the GPU. device = torch.device(device) if device == torch.device("cuda"): self.model.cuda() # Set TrainingArguments and the Seq2SeqTrainer training_args_kwargs = dict(self.training_args_kwargs) default_args_kwargs = dict( # predict_with_generate=True, # warmup_steps=2000, fp16=( device == torch.device("cuda") ), # Use fp16 by default if using cuda device auto_find_batch_size=True ) for k, v in default_args_kwargs.items(): if k not in training_args_kwargs: training_args_kwargs[k] = v callbacks = None if training_args_kwargs["load_best_model_at_end"]: callbacks = [ EarlyStoppingCallback( self.early_stopping_patience, self.early_stopping_threshold ), AdditionalEarlyStopCallback( early_stopping_patience=self.early_stopping_patience, early_stopping_threshold=self.early_stopping_threshold ) ] else: callbacks = [AdditionalEarlyStopCallback( early_stopping_patience=self.early_stopping_patience, early_stopping_threshold=self.early_stopping_threshold )] # instantiate trainer trainer = Seq2SeqTrainer( model=self.model, args=Seq2SeqTrainingArguments(**training_args_kwargs), callbacks=callbacks, data_collator=RelationalDataCollator(), **dataset, ) trainer.train(resume_from_checkpoint=resume_from_checkpoint) return trainer def _build_tabular_trainer( self, device="cuda", num_train_epochs: int = None, target_epochs: int = None, ) -> Trainer: device = torch.device(device) # Set TrainingArguments and the Trainer logging.info("Set up the TrainingArguments and the Trainer...") training_args_kwargs: Dict[str, Any] = dict(self.training_args_kwargs) default_args_kwargs = dict( fp16=( device == torch.device("cuda") ), # Use fp16 by default if using cuda device ) for k, v in default_args_kwargs.items(): if k not in training_args_kwargs: training_args_kwargs[k] = v if num_train_epochs is not None: training_args_kwargs["num_train_epochs"] = num_train_epochs # # NOTE: The `ResumableTrainer` will default to its original # # behavior (Trainer) if `target_epochs`` is None. # # Set the `target_epochs` to `num_train_epochs` if not specified. # if target_epochs is None: # target_epochs = training_args_kwargs.get("num_train_epochs") callbacks = None if training_args_kwargs["load_best_model_at_end"]: callbacks = [ EarlyStoppingCallback( self.early_stopping_patience, self.early_stopping_threshold ), AdditionalEarlyStopCallback( early_stopping_patience=self.early_stopping_patience, early_stopping_threshold=self.early_stopping_threshold ) ] else: callbacks = [AdditionalEarlyStopCallback( early_stopping_patience=self.early_stopping_patience, early_stopping_threshold=self.early_stopping_threshold )] assert self.dataset trainer = ResumableTrainer( target_epochs=target_epochs, save_epochs=None, model=self.model, args=TrainingArguments(**training_args_kwargs), data_collator=None, # Use the default_data_collator callbacks=callbacks, **self.dataset, ) return trainer def fit( self, df: pd.DataFrame, in_df: Optional[pd.DataFrame] = None, join_on: Optional[str] = None, resume_from_checkpoint: Union[bool, str] = False, device="cuda", n_critic: Optional[int] = None, **kwargs ): if n_critic is None: n_critic = self.n_critic kwargs["n_critic"] = n_critic device = _validate_get_device(device) if self.model_type == ModelType.relational: assert ( in_df is not None ), "The REaLTabFormer for relational data requires two tables for training." assert join_on is not None, "The column to join the data must not be None." trainer = self._fit_relational( df, in_df, join_on=join_on, device=device, resume_from_checkpoint=resume_from_checkpoint ) try: self.experiment_id = f"id{int((time.time() * 10 ** 10)):024}" torch.cuda.empty_cache() return trainer except Exception as exception: if device == torch.device("cuda"): del self.model torch.cuda.empty_cache() self.model = None raise exception else: return super().fit(df, in_df, join_on, resume_from_checkpoint, device, **kwargs) @torch.no_grad() def sample(self, *args, **kwargs): # Faster sampling return super().sample(*args, **kwargs) def fit_transform_rtf(data: np.ndarray, model_dir: str, prefix: str = None) -> pd.DataFrame: prefix = "" if prefix is None else f"{prefix}-" dim_prefix = prefix if prefix.endswith("-") else "dim" data = pd.DataFrame(data, columns=[f"{dim_prefix}{i:02d}" for i in range(data.shape[1])]) raw_columns = data.columns.tolist() unique_threshold = 500 if prefix == "ctx" else 20 def fast_nunique(col): return col.nunique() results = Parallel(n_jobs=20)( delayed(fast_nunique)(data[col]) for col in data.columns ) uniques = pd.Series(results, index=data.columns) num_columns = data[:10].select_dtypes(include=[np.number]).columns.tolist() unary_columns = { c: data[c].iloc[0] for c, n in uniques.items() if n <= 1 } data = data.drop(columns=[*unary_columns]) num_columns = [ c for c in num_columns if c not in unary_columns and (uniques[c] > unique_threshold or not ( (data[c] - data[c].round()).abs().mean() < 1e-6 and data[c].max() - data[c].min() == uniques[c] - 1)) ] cat_columns = [c for c in data.drop(columns=num_columns).columns] placeholder_columns = [] if data.shape[-1] > 0 else ["placeholder"] if (prefix == "ctx" and data.shape[-1] > 100) or (uniques[num_columns] > unique_threshold).any(): bin_columns = [c for c in num_columns if uniques[c] > unique_threshold] bins = KBinsDiscretizer(n_bins=min(data.shape[0], 200), encode="ordinal", strategy="kmeans") num_columns = [c for c in num_columns if c not in bin_columns] data[bin_columns] = bins.fit_transform((data[bin_columns])).astype(np.int32) torch.save(bins, os.path.join(model_dir, f"{prefix}bins.pkl")) else: bin_columns = [] info = { "num_columns": num_columns, "cat_columns": cat_columns, "bin_columns": bin_columns, "placeholders": placeholder_columns, "unary": unary_columns, "raw_columns": raw_columns } oe = OrdinalEncoder() data[cat_columns] = oe.fit_transform(data[cat_columns]) data[cat_columns] = data[cat_columns].applymap(lambda x: f"C{int(x)}") data[num_columns] = data[num_columns].round(6) data[placeholder_columns] = 0 with open(os.path.join(model_dir, f"{prefix}info.json"), "w") as f: json.dump(info, f, indent=2) torch.save(oe, os.path.join(model_dir, f"{prefix}oe.pkl")) return data def transform_rtf(data: np.ndarray, model_dir: str, prefix: str = None) -> pd.DataFrame: prefix = "" if prefix is None else f"{prefix}-" with open(os.path.join(model_dir, f"{prefix}info.json"), "r") as f: info = json.load(f) oe = torch.load(os.path.join(model_dir, f"{prefix}oe.pkl")) dim_prefix = prefix if prefix.endswith("-") else "dim" data = pd.DataFrame(data, columns=[f"{dim_prefix}{i:02d}" for i in range(data.shape[1])]) num_columns, cat_columns, bin_columns = info["num_columns"], info["cat_columns"], info["bin_columns"] data = data.drop(columns=[*info["unary"]]) data[cat_columns] = oe.transform(data[cat_columns]) data[cat_columns] = data[cat_columns].applymap(lambda x: f"C{int(x)}") data[num_columns] = data[num_columns].round(6) if bin_columns: bins = torch.load(os.path.join(model_dir, f"{prefix}bins.pkl")) data[bin_columns] = bins.transform(data[bin_columns]).astype(np.int32) data[info["placeholders"]] = 0 return data def inverse_transform_rtf(generated: pd.DataFrame, model_dir: str, prefix: str = None) -> np.ndarray: prefix = "" if prefix is None else f"{prefix}-" with open(os.path.join(model_dir, f"{prefix}info.json"), "r") as f: info = json.load(f) oe = torch.load(os.path.join(model_dir, f"{prefix}oe.pkl")) cat_columns = info["cat_columns"] if cat_columns: generated[cat_columns] = oe.inverse_transform( generated[cat_columns].applymap(lambda x: x[len("C"):]).astype(np.int32).values ) bin_columns = info["bin_columns"] if bin_columns: bins = torch.load(os.path.join(model_dir, f"{prefix}bins.pkl")) generated[bin_columns] = bins.inverse_transform(generated[bin_columns].clip( 0, upper={c: max(0, n) for c, n in zip(bin_columns, bins.n_bins_ - 1)}, axis=1).values ) generated = generated.drop(columns=info["placeholders"]) for k, v in info["unary"].items(): generated[k] = v generated = generated[info["raw_columns"]] return generated.values def update_epochs(size, kwargs): if kwargs.get("max_steps", kwargs.get("num_train_epochs", kwargs.get("epochs"))) is None and size < 200: kwargs = kwargs.copy() kwargs["epochs"] = 500 return kwargs def fast_pearson_corr_gpu(df_num: pd.DataFrame) -> pd.DataFrame: data_gpu = cp.asarray(df_num.to_numpy(dtype=cp.float32)) means = cp.mean(data_gpu, axis=0) stds = cp.std(data_gpu, axis=0) standardized = (data_gpu - means) / stds corr_gpu = cp.dot(standardized.T, standardized) / (data_gpu.shape[0] - 1) return pd.DataFrame(cp.asnumpy(cp.abs(corr_gpu)), index=df_num.columns, columns=df_num.columns) def fast_correlation_ratio( categories: pd.Series, measurements: pd.DataFrame, ) -> pd.Series: columns = measurements.columns.tolist() measurements = measurements.values fcat, _ = pd.factorize(categories) # type: ignore cat_num = np.max(fcat) + 1 y_avg_array = np.zeros((cat_num, len(columns))) n_array = np.bincount(fcat) for i in range(0, cat_num): cat_measures = measurements[fcat == i] for j in range(measurements.shape[-1]): y_avg_array[i, j] = np.average(cat_measures[:, j]) n_array = n_array.reshape((-1, 1)) y_total_avg = np.sum(y_avg_array * n_array, axis=0) / np.sum(n_array) numerator = np.sum(n_array * (y_avg_array - y_total_avg) ** 2, axis=0) denominator = np.sum((measurements - y_total_avg) ** 2, axis=0) etas = np.sqrt(np.divide(numerator, denominator, out=np.zeros_like(numerator), where=denominator > 0)) etas = np.clip(etas, 0.0, 1.0) result = pd.Series(etas, index=columns) return result def wide_associations( df: pd.DataFrame, ) -> pd.DataFrame: # identifying categorical columns columns = df.columns nominal_columns = identify_nominal_columns(df) numerical_columns = df.drop(columns=nominal_columns).columns.tolist() # will be used to store associations values pbar = tqdm(total=len(columns) ** 2, desc="Compute corr") corr = pd.DataFrame(np.eye(len(columns)), index=columns, columns=columns, dtype=np.float64) pbar.update(len(columns)) if len(numerical_columns) > 1: pearson_corr = fast_pearson_corr_gpu(df[numerical_columns]) corr.loc[pearson_corr.index, pearson_corr.columns] = pearson_corr pbar.update(len(numerical_columns) ** 2 - len(numerical_columns)) del pearson_corr gc.collect() cat_num_corr = [] pbar = tqdm( desc=f"Cat-num correlations (#cat={len(nominal_columns)}, #num={len(numerical_columns)})", total=len(nominal_columns) ) chunk_size = 20 for i in range(0, len(nominal_columns), chunk_size): this_cat_num_corr = Parallel(n_jobs=chunk_size)( delayed(fast_correlation_ratio)(df[cat], df[numerical_columns]) for cat in nominal_columns[i:i + chunk_size] ) cat_num_corr.extend(this_cat_num_corr) pbar.update(len(nominal_columns[i:i + chunk_size])) for cat, corr_values in zip(nominal_columns, cat_num_corr): corr.loc[cat, corr_values.index] = corr_values.values corr.loc[corr_values.index, cat] = corr_values.values cat_pairs = itertools.combinations(nominal_columns, 2) cat_results = Parallel(n_jobs=20)( delayed(cramers_v)(df[col1], df[col2], False) for col1, col2 in tqdm( cat_pairs, desc=f"Cat-cat correlations ({len(nominal_columns)})", total=len(nominal_columns) * (len(nominal_columns) - 1) // 2 ) ) for (col1, col2), val in zip(cat_pairs, cat_results): corr.loc[col1, col2] = corr.loc[col2, col1] = val pbar.update() return corr.fillna(0) def sort_column_importance(df: pd.DataFrame): columns = [] if df.shape[-1] < 500: corr = associations( df if df.shape[0] < 10_000 else df.sample(n=10_000), plot=False, max_cpu_cores=20, multiprocessing=True, compute_only=True )["corr"].abs() else: corr = wide_associations( df if df.shape[0] < 10_000 else df.sample(n=10_000), ).abs() nuniques = df.nunique() pbar = tqdm(total=df.shape[-1], desc="Sorting column importance") for i, c in enumerate(corr.columns): if c in columns: continue this_round = [c] for c2 in corr.index[i + 1:]: if c2 in columns: continue if corr.loc[c, c2] >= 0.95: this_round.append(c2) if len(this_round) > 1: this_nuniques = nuniques[this_round].sort_values(ascending=False) columns.extend(this_nuniques.index[1:].tolist()) pbar.update(len(this_round) - 1) corr = corr.drop(columns=columns, index=columns) eye = np.eye(corr.shape[0], dtype=np.bool_) corr[eye] = -1 while corr.shape[0] > 0: max_corr = corr.max() max_corr_value = max_corr.max() candidates = max_corr[max_corr == max_corr_value].index new_col = nuniques[candidates].idxmax() columns.append(new_col) corr = corr.drop(columns=[new_col], index=[new_col]) pbar.update() # columns.extend(corr.drop(columns=columns, index=columns).max().sort_values(ascending=False).index.tolist()) return [*reversed(columns)] def save_to(obj, path): with open(path, "wb") as f: pickle.dump(obj, f, protocol=5) def load_from(path): with open(path, "rb") as f: return pickle.load(f) @contextmanager def log_resource_usage(path, descr): process = psutil.Process(os.getpid()) peak_memory = [0] running = [True] def sample(): while running[0]: mem = process.memory_info().rss if mem > peak_memory[0]: peak_memory[0] = mem time.sleep(0.5) thread = threading.Thread(target=sample) thread.start() start_time = time.time() torch.cuda.reset_peak_memory_stats() torch.cuda.synchronize() try: logging.info(f"{descr} ...") yield finally: end_time = time.time() running[0] = False thread.join() peak_memory_mb = peak_memory[0] / (1024 ** 2) torch.cuda.synchronize() peak_gpu = torch.cuda.max_memory_allocated() peak_gpu_memory_mb = peak_gpu / (1024 ** 2) elapsed = end_time - start_time with open(path, 'a') as f: f.write(f"{descr},{peak_memory_mb},{peak_gpu_memory_mb},{elapsed}\n") class CacheBlock(contextlib.AbstractContextManager): def __init__(self, description: str, cache_dir: str): self.description = description self.cache_dir = cache_dir self.cache_file = os.path.join(cache_dir, f"{description}.pkl") def __enter__(self): if os.path.exists(self.cache_file): self._should_skip = True return None else: self._should_skip = False return self def __exit__(self, exc_type, exc_value, traceback): if self._should_skip or exc_type is not None: return False frame = inspect.currentframe().f_back local_vars = frame.f_locals.copy() os.makedirs(self.cache_dir, exist_ok=True) with open(self.cache_file, "wb") as f: dill.dump({ "locals": local_vars, "random_state": random.getstate(), "np_random_state": np.random.get_state(), "torch_state": torch.get_rng_state(), "torch_cuda_state": torch.cuda.get_rng_state_all() if torch.cuda.is_available() else None, "torch_deterministic": torch.backends.cudnn.deterministic, "torch_benchmark": torch.backends.cudnn.benchmark, }, f) return False def resume_from_last(cache_dir): if not os.path.exists(cache_dir): return {} files = sorted( (os.path.join(cache_dir, f) for f in os.listdir(cache_dir) if f.endswith(".pkl")), key=os.path.getmtime ) if not files: return {} last_file = files[-1] logging.info(f"Resuming from {last_file}") print(f"Resuming from {last_file}") with open(last_file, "rb") as f: data = dill.load(f) frame = inspect.currentframe().f_back frame.f_locals.update(data["locals"]) random.setstate(data["random_state"]) np.random.set_state(data["np_random_state"]) torch.set_rng_state(data["torch_state"]) if torch.cuda.is_available() and data["torch_cuda_state"] is not None: torch.cuda.set_rng_state_all(data["torch_cuda_state"]) torch.backends.cudnn.deterministic = data["torch_deterministic"] torch.backends.cudnn.benchmark = data["torch_benchmark"] return data["locals"]