import copy import logging import os import random import lmdb import numpy as np from bson import ObjectId from torch.utils.data import IterableDataset from tqdm import tqdm from utils.dataset_creation import get_features, load_scalers, scale_item, combine_features from utils.dataset_utils import load_key_stats from utils.lmdb_utils import load_from_lmdb from vsm_datascience_common.cycle_database_connection.cycle_data import get_cycle_by_id from vsm_datascience_common.cycle_database_connection.db_utils import get_cycles_collection logger = logging.getLogger(__name__) def get_prepared_sequence_length(input_sequence: np.ndarray, model_configuration: dict) -> int: """ Returns the length of the input sequence after padding and resampling """ initial_length = len(input_sequence) take_every_nth = model_configuration["preprocessing"]["take_every_nth"] padding_length = get_padding_length(model_configuration, resampled=False) raw_length = initial_length + padding_length adjusted = int(raw_length // take_every_nth) return adjusted def get_padding_length(model_configuration: dict, resampled: bool = True) -> int: """ Returns the length of the input sequence after padding and resampling """ take_every_nth = model_configuration["preprocessing"]["take_every_nth"] padding_length = int(model_configuration["input_window_length"] * ( 1 - model_configuration["preprocessing"]["min_input_length_fraction_for_padding"])) * ( take_every_nth if not resampled else 1) return padding_length def get_number_of_windows(base_length: int, model_configuration: dict) -> int: """ Returns the number of windows for a given sequence length """ window_shift = model_configuration["preprocessing"]["window_shift"] input_window_length = model_configuration["input_window_length"] output_window_length = model_configuration["output_window_length"] output_window_offset = model_configuration["output_window_offset"] if input_window_length > output_window_length + output_window_offset: return (base_length - input_window_length) // window_shift + 1 else: return (base_length - output_window_length - output_window_offset) // window_shift + 1 def get_collated_batch_for_key(sample_key: ObjectId | str, model_configuration: dict, start_cutoff: int = None, end_cutoff: int = None, lmdb_env=None) -> dict: sample_batch_for_key = get_batch_for_key(sample_key, model_configuration, start_cutoff=start_cutoff, end_cutoff=end_cutoff, lmdb_env=lmdb_env) collated = model_configuration["collate_fn"](sample_batch_for_key) return collated def get_batch_for_key(key, model_configuration: dict, start_cutoff: int = None, end_cutoff: int = None, lmdb_env=None) -> np.ndarray: """ Get the batch for a given key from the lmdb database or compute it directly. The return has the same format as the batch produced by the batch_fn in the model configuration. Args: key: key to get the batch for, can also be object id from database model_configuration: model configuration to use start_cutoff: start cutoff for the batch, if None, the whole batch is used, CAUTION: cutoff should not be normalized -> in raw data points end_cutoff: end cutoff for the batch, if None, the whole batch is used, CAUTION: cutoff should not be normalized -> in raw data points lmdb_env: the lmdb environment to use, if None, the features are computed directly from the database Returns: batch: batch for the given key, as returned by the batch_fn in the model configuration """ if lmdb_env is None: # compute and scale features, ignored features are handled internally features = get_scaled_feature_for_key(key, model_configuration) else: # load features from lmdb features = load_from_lmdb(lmdb_env, str(key)) # ignored features are handled internally, thus we need to remove them here ignored_features = model_configuration["feature_config"]["ignored_features"] if "ignored_features" in \ model_configuration[ "feature_config"] else [] # add scaled versions of features ignored_features = ignored_features + [feat + "_scaled" for feat in ignored_features] for feature_set in features: for feature_name in list(features[feature_set].keys()): if feature_name in ignored_features: del features[feature_set][feature_name] processed_chunk = process_chunk([key], [features], model_configuration) # use cutoff, if provided if start_cutoff is not None or end_cutoff is not None: item_length = len(processed_chunk[0]["target_features"]) if start_cutoff is None: start_cutoff = 0 if end_cutoff is None: end_cutoff = item_length # "normalize" cutoff to adjust for added padding and resampling padding_length = get_padding_length(model_configuration, resampled=False) take_every_nth = model_configuration["preprocessing"]["take_every_nth"] # make sure that there is no padding added to the start cutoff, so that it includes the padding added to the sequence start_cutoff_normalized = max(int(start_cutoff // take_every_nth), 0) # end cutoff must be adjusted to include the padding added to the sequence end_cutoff_normalized = min(int((end_cutoff + padding_length) // take_every_nth), item_length) for feature_set in processed_chunk[0]: if feature_set in model_configuration["feature_config"]["feature_sets"]: processed_chunk[0][feature_set] = processed_chunk[0][feature_set][ start_cutoff_normalized:end_cutoff_normalized + 1] if "batch_fn" in model_configuration and model_configuration["batch_fn"] is not None: batch = model_configuration["batch_fn"](processed_chunk, model_configuration) else: batch = processed_chunk if batch is None or len(batch) == 0: raise ValueError(f"Batch is empty for key {key}") return batch def get_scaled_feature_for_key(key: str, model_configuration: dict) -> tuple: feature_config = model_configuration["feature_config"] user_cycles = list( get_cycles_collection().find({"user_id": ObjectId(key)} | feature_config["filter_criteria"]).sort("starts_at", 1)) cycle_features = list() for cycle in user_cycles: features = get_features(cycle, feature_config) cycle_features.append(features) # combine all cycles for the user if len(cycle_features) == 0: raise ValueError(f"No features found for key {key}") features = combine_features(cycle_features, feature_config) scaler_dir = os.path.join(feature_config["dataset_dir"], "scalers") scalers = load_scalers(scaler_dir) # scale features scaled_features = scale_item(features, scalers) return scaled_features def augment_items(identifiers: list, items: list[dict], feature_config: dict, lmdb_env: lmdb.Environment = None) -> list: """ Augments the given sequences by attaching previous sequences Args: identifiers: identifiers of items for finding previous sequences items: actual items feature_config: feature config of the dataset lmdb_env: lmdb environment for loading previous sequences, can be None, in this case the items are computed directly Returns: list of augmented items """ if "augmentation" not in feature_config or \ feature_config["augmentation"]["use_augmentation"] is False: return items augmented_items = list() for i, item in enumerate(items): item_id = identifiers[i] try: item_data = get_cycles_collection().find_one({"_id": ObjectId(item_id)}, {"user_id": 1, "starts_at": 1}) item_user_id = item_data["user_id"] item_starts_at = item_data["starts_at"] except KeyError: print(f"User ID not found for item {item_id}") continue previous_item_identifiers = list(get_cycles_collection().aggregate( # make sure to filter by user first to significantly reduce the number of items [ { "$match": { "user_id": item_user_id, "starts_at": {"$lt": item_starts_at}, } } ] + feature_config["filter_criteria_pipeline"] + [ { "$sort": { "starts_at": -1 } }, { "$project": { "_id": 1, "starts_at": 1 } } ] )) max_lookback = feature_config["augmentation"]["max_lookback"] previous_items = list() for previous_item in previous_item_identifiers: if len(previous_items) >= max_lookback: break previous_item_id = previous_item["_id"] if lmdb_env is None: # compute item cycle = get_cycle_by_id(previous_item_id) features = get_features(cycle, feature_config) scaler_dir = os.path.join(feature_config["dataset_dir"], "scalers") scalers = load_scalers(scaler_dir) previous_item = scale_item(features, scalers) previous_items.append(previous_item) else: # take item from lmdb try: # make sure to parse object id previous_item = load_from_lmdb(lmdb_env, str(previous_item_id)) except KeyError: continue previous_items.append(previous_item) # merge items into one feature_sets = feature_config["feature_sets"] augmented_item = copy.deepcopy(item) for feature_set in feature_sets: if feature_set not in augmented_item: continue if "static" in feature_set: continue for feature_name in augmented_item[feature_set]: if feature_name not in augmented_item[feature_set]: continue augmented_item[feature_set][feature_name] = np.concatenate( [previous_item[feature_set][feature_name] for previous_item in previous_items] + [augmented_item[feature_set][feature_name]]) augmented_items.append(augmented_item) return augmented_items def process_chunk(ids: list, chunk: list, model_configuration: dict, ignored_features: list = None, pad_sequences: bool = True, statics_as_list: bool = True): """ Creates a chunk of data as dataframe for training and inference of a tft model. :param ids: list of ids of the data points in the chunk :param chunk: list of data points :return: dataframe with the data points in the chunk """ if ignored_features is None: ignored_features = list() input_features = list(chunk[0]["target_features"].keys()) feature_config = model_configuration["feature_config"] # variables for padding padding_value = 0 # create dataframe from items in chunk data = [] for i, item in enumerate(chunk): item_length = item["target_features"][input_features[0]].shape[0] take_every_nth = model_configuration["preprocessing"]["take_every_nth"] data_item = dict() for feature_set in feature_config["feature_sets"]: if feature_set in item and len(item[feature_set]) > 0: # fill in the data source = item[feature_set] features_to_use = [feature for feature in source.keys() if "_scaled" in feature] feature_set_data = [source[feature] for feature in features_to_use] padding_length = get_padding_length(model_configuration, resampled=False) if len(feature_set_data) == 0: raise ValueError(f"No features found for key {ids[i]} in feature set {feature_set}") if len(feature_set_data[0]) == 1: # if the feature is constant, we need to repeat it for all time points if statics_as_list: if pad_sequences: desired_length = item_length + padding_length feature_set_data = np.array( [np.full((desired_length,), val) for val in feature_set_data]).T[::take_every_nth] else: feature_set_data = np.array( [np.full((item_length,), val) for val in feature_set_data]).T[::take_every_nth] else: feature_set_data = np.array(feature_set_data).flatten() else: if pad_sequences: padding_values = np.full((padding_length, len(feature_set_data)), [padding_value for vals in feature_set_data]).T feature_set_data = np.concatenate([padding_values, np.array(feature_set_data)], axis=1) # create bins of n values and apply mean # check, if special accumulation function has been specified individual_feature_data = list() for i, feature_to_use in enumerate(features_to_use): try: # get config for feature, make sure to replace scaled to get actual config current_feature_config = \ [x for x in feature_config[feature_set] if x["name"] == feature_to_use.replace("_scaled", "")][0] current_accumulation_fn = current_feature_config["accumulation_fn"] except (KeyError, IndexError): # default to mean current_accumulation_fn = np.mean current_feature_data = apply_fn_to_bins(feature_set_data[i], take_every_nth, current_accumulation_fn) individual_feature_data.append(current_feature_data) # stack individual features feature_set_data = np.stack(individual_feature_data).T if feature_set not in data_item: data_item[feature_set] = feature_set_data else: data_item[feature_set] = np.concatenate([data_item[feature_set], feature_set_data], axis=1) data.append(data_item) return data def apply_fn_to_bins(input_sequence: np.ndarray, bin_size: int, fn: callable) -> np.ndarray: """ Applies a function to the bins of the input sequence along the last dimension and returns the results. Args: input_sequence: sequence to apply the function to (1D or 2D) bin_size: size of the bins fn: function to apply to the bins Returns: np.ndarray: result of the function applied to the bins """ if input_sequence.ndim == 1: input_sequence = input_sequence.reshape(1, -1) dims = 1 else: dims = input_sequence.shape[-1] result = [] for row in input_sequence: row_result = [] for start in range(0, len(row), bin_size): end = min(start + bin_size, len(row)) bin_slice = row[start:end] row_result.append(fn(bin_slice)) result.append(row_result) if dims > 1: return np.array(result) else: return np.array(result[0]) def produce_window_batches(data_chunk: list, model_configuration: dict, offsets: list | np.ndarray = None) -> list: """ Produces batches of windows from the data chunk. Args: data_chunk: data chunk model_configuration: configuration to use offsets: offsets used to determine the length of data to use, uses item[-offset:] of data Returns: data: list of windows """ input_window_length = model_configuration["input_window_length"] output_window_length = model_configuration["output_window_length"] output_window_offset = model_configuration["output_window_offset"] window_shift = model_configuration["preprocessing"]["window_shift"] if offsets is None: offsets = np.zeros(len(data_chunk), dtype=int) data = list() for i, item in enumerate(data_chunk): item_length = item["target_features"][-offsets[i]:].shape[0] total_window_length = input_window_length if input_window_length > output_window_length + output_window_offset \ else output_window_length + output_window_offset if item_length < total_window_length: continue # number of windows is defined by the input window length and output window length with offset num_windows = get_number_of_windows(item_length, model_configuration) for j in range(num_windows): window = dict() for key, value in item.items(): if value is None: window[key] = None else: if key == "target_features": window[key] = value[-offsets[i]:][j * window_shift + output_window_offset: j * window_shift + output_window_length + output_window_offset] else: window[key] = value[-offsets[i]:][j * window_shift:j * window_shift + input_window_length] data.append(window) return data def produce_simple_batches(data_chunk: list, model_configuration: dict, offsets: list | np.ndarray = None) -> list: """ Produces batches without windowing. Args: data_chunk: data chunk model_configuration: configuration to use offsets: offsets used to determine the length of data to use, uses item[-offset:] of data Returns: data: list of items """ pass class LMDBIterableDataset(IterableDataset): def __init__(self, lmdb_env_path: str, lmdb_keys: list[str], model_configuration: dict, batch_size: int = 32): self.lmdb_path = lmdb_env_path self.lmdb_env = None self.lmdb_keys = lmdb_keys self.key_subset = None self.model_configuration = model_configuration self.batch_size = batch_size self.keys_stats = load_key_stats(model_configuration["feature_config"]["dataset_dir"]) self.len = None try: self.model_configuration["batch_fn"] except KeyError: raise KeyError("Batch function not found in model configuration") def set_key_subset(self, key_subset: list[str]): """ Set the key subset to use for the dataset. Args: key_subset: list of keys to use """ self.key_subset = key_subset # reset length self.len = None def get_length_of_data_subset(self, key_set: list[str]): """ Get the length of the data subset. Args: key_set: list of keys to use """ # run simplified version of __iter__ to get the length num_steps = 0 i = 0 num_batch_items = 0 while True: if num_batch_items >= self.batch_size: num_steps += 1 num_batch_items -= self.batch_size else: if i >= len(key_set): if num_batch_items > 0: num_steps += 1 break key = key_set[i] i += 1 # try to fetch stats from key_stats try: current_key_stats = self.keys_stats["by_key"][key] base_length = current_key_stats["item_length"] take_every_nth = self.model_configuration["preprocessing"]["take_every_nth"] padding_length = get_padding_length(self.model_configuration, resampled=False) raw_length = base_length + padding_length item_length = int(raw_length // take_every_nth) except: item = load_from_lmdb(self.lmdb_env, key) item_length = get_prepared_sequence_length( item["target_features"][list(item["target_features"])[0]], self.model_configuration) num_windows = get_number_of_windows(item_length, self.model_configuration) num_batch_items += num_windows return num_steps def __iter__(self): self.init_lmdb_env() # if no key subset is set, use all keys if self.key_subset is None: self.key_subset = self.lmdb_keys random.shuffle(self.key_subset) batch = list() counter = 0 collate_fn = self.model_configuration["collate_fn"] while True: if len(batch) >= self.batch_size: if collate_fn is None: yield batch[:self.batch_size] else: yield collate_fn(batch[:self.batch_size]) batch = batch[self.batch_size:] else: if counter >= len(self.key_subset): if len(batch) > 0: yield collate_fn(batch) break key = self.lmdb_keys[counter] counter += 1 try: current_batch = get_batch_for_key(key, self.model_configuration, lmdb_env=self.lmdb_env) except ValueError: # if betch is empty, try next continue batch.extend(current_batch) def init_lmdb_env(self): if self.lmdb_env is None: self.lmdb_env = lmdb.open(self.lmdb_path, readonly=True, lock=False, readahead=False, meminit=False) def __len__(self): self.init_lmdb_env() # if no key subset is set, use all keys if self.key_subset is None: self.key_subset = self.lmdb_keys if self.len is None: num_steps = self.get_length_of_data_subset(self.key_subset) self.len = num_steps return self.len