File size: 13,214 Bytes
d9115cb | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 | from typing import Any, Dict, List, Optional, Union
import numpy as np
import tensorflow as tf
import torch
from src.utils.file_utils import load_json
from src.data.loading.components.interfaces import (
BaseDatasetConfig,
SemanticIDDatasetConfig,
)
from src.data.loading.components.interfaces import TokenizerConfig
from src.utils.utils import load_tokenize
# support functions
def convert_bytes_to_string(
batch_or_row: Dict[str, np.ndarray],
dataset_config: BaseDatasetConfig,
features_to_apply: Optional[List[str]] = [],
**kwargs,
) -> Dict[str, np.ndarray]:
# For each feature to apply, cast its np.ndarray of bytes to string.
for k in batch_or_row:
if is_feature_in_features_to_apply(features_to_apply, k):
batch_or_row[k] = batch_or_row[k].astype(str)
return batch_or_row
def is_feature_in_features_to_apply(features_to_apply: List[str], k: str) -> bool:
if len(features_to_apply) > 0 and k not in features_to_apply:
return False
return True
def filter_features_to_consider(
batch_or_row: Dict[str, tf.Tensor],
dataset_config: BaseDatasetConfig,
features_to_apply: Optional[List[str]] = [],
**kwargs,
) -> Dict[str, tf.Tensor]:
batch_or_row = map_feature_names(batch_or_row, dataset_config)
features_to_consider = set(dataset_config.features_to_consider)
if hasattr(dataset_config, "keep_user_id") and dataset_config.keep_user_id:
if dataset_config.user_id_field not in features_to_consider:
features_to_consider.add(dataset_config.user_id_field)
if hasattr(dataset_config, "keep_item_id") and dataset_config.keep_item_id:
if dataset_config.item_id_field not in features_to_consider:
features_to_consider.add(dataset_config.item_id_field)
if len(dataset_config.features_to_consider):
# Given a batch or row, filter the features to consider.
return {k: v for k, v in batch_or_row.items() if k in features_to_consider}
# if not specified, we consider all features
return batch_or_row
def convert_to_dense_numpy_array(
batch_or_row: Dict[str, tf.Tensor],
dataset_config: BaseDatasetConfig,
features_to_apply: Optional[List[str]] = [],
**kwargs,
) -> Dict[str, np.ndarray]:
# Transform a tfrecord example to a dictionary of numpy arrays, converting sparse tensors to dense numpy arrays.
for k in batch_or_row:
if is_feature_in_features_to_apply(features_to_apply, k):
batch_or_row[k] = tf.sparse.to_dense(batch_or_row[k]).numpy()
return batch_or_row
def map_feature_names(
batch_or_row: Dict[str, np.ndarray],
dataset_config: BaseDatasetConfig,
features_to_apply: Optional[List[str]] = [],
**kwargs,
) -> Dict[str, np.ndarray]:
# Given a batch or row, map the feature names to the desired feature names.
if dataset_config.feature_map:
batch_or_row = {
v: batch_or_row[k]
for k, v in dataset_config.feature_map.items()
if k in batch_or_row
}
return batch_or_row
def convert_fields_to_tensors(
batch_or_row: Dict[str, np.ndarray],
dataset_config: BaseDatasetConfig,
features_to_apply: Optional[List[str]] = [],
**kwargs,
) -> Dict[str, torch.Tensor]:
# Given a batch or row, convert all fields to torch tensors. Uses the field type map to determine the dtype, defaulting to torch.long
# if no dtype is specified.
for k, v in batch_or_row.items():
if is_feature_in_features_to_apply(features_to_apply, k):
if isinstance(v, int) or isinstance(v, float):
v = [int(v)]
batch_or_row[k] = torch.tensor(v, dtype=dataset_config.field_type_map.get(k, torch.long)) # type: ignore
return batch_or_row
def filter_sequence_length_row(row: Dict[str, torch.Tensor], dataset_config: BaseDatasetConfig, features_to_apply: Optional[List[str]] = [], **kwargs) -> Dict[str, np.ndarray]: # type: ignore
# Only works for a row right now. This filters out rows that have fields with sequence length smaller than the min threshold.
# TODO(lneves): Make this work for a batch as well without creating batches of different sizes.
for _, tensor in row.items():
if len(tensor) < dataset_config.min_sequence_length:
return None
return row
def filter_empty_feature(row: Dict[str, torch.Tensor], dataset_config: BaseDatasetConfig, features_to_apply: Optional[List[str]] = [], **kwargs) -> Dict[str, np.ndarray]: # type: ignore
# Only works for a row right now. This filters out rows that have fields with empty tensors.
for k, v in row.items():
if is_feature_in_features_to_apply(features_to_apply, k):
if len(v) == 0:
return None
return row
def map_sparse_id_to_semantic_id(
row: Dict[str, torch.Tensor],
dataset_config: SemanticIDDatasetConfig,
features_to_apply: Optional[List[str]] = [],
num_hierarchies: Optional[int] = None,
**kwargs,
) -> Dict[str, torch.Tensor]:
"""
Given a row of data, maps the sparse ids to semantic ids
based on the id_map in the dataset config.
"""
for k, v in row.items():
if is_feature_in_features_to_apply(features_to_apply, k):
id_map: torch.Tensor = dataset_config.semantic_id_map.get(k, None)
# id_map is a D x N tensor
# where N is the number of unique items in the dataset
# and D is the number of hierarchies (semantic id digits)
if id_map is not None:
# flatten the semantic id sequence
if num_hierarchies is None:
row[k] = id_map.t()[v].view(-1)
else:
assert num_hierarchies <= id_map.size(
0
), "num_hierarchies must be less than or equal to the number of hierarchies in the semantic id map."
row[k] = id_map[:num_hierarchies].t()[v].view(-1)
else:
raise ValueError(f"Semantic id map not found for feature {k}")
return row
def trim_sequence_row(
row: Dict[str, Any],
dataset_config: BaseDatasetConfig,
sequence_length: int,
should_trim_left: bool,
features_to_apply: Optional[List[str]] = [],
**kwargs,
) -> Dict[str, Any]:
"""
Trim the sequences in the row to the sequence_length.
This function handles only rows (not batches) and assumes that the sequences are not
padded in the first dimension (the dimension to truncate).
Args:
row (Dict[str, Any]): A dictionary representing a row of data where each key
is a feature name and, if the feature is being trimmed, the corresponding
value is a sequential object to be truncated. The value will be trimmed on
the side determined by should_trim_left to the specified sequence_length in
the first dimension.
dataset_config (BaseDatasetConfig): The dataset configuration object.
sequence_length (int): The desired length to trim the sequences to.
should_trim_left (bool): If True, trim the left side of the sequence.
If False, trim the right side of the sequence.
features_to_apply (Optional[List[str]]): A list of feature names to apply the
trimming to. If empty, all features in the row will be trimmed.
Returns:
Dict[str, Any]:
A dictionary identical to the input row, but with the sequences trimmed to
the specified sequence_length on the specified side. Sequences that are
shorter than sequence_length will remain unchanged.
"""
if should_trim_left:
for k, v in row.items():
if is_feature_in_features_to_apply(features_to_apply, k):
v = v[-sequence_length:]
row[k] = v
else:
for k, v in row.items():
if is_feature_in_features_to_apply(features_to_apply, k):
v = v[:sequence_length]
row[k] = v
return row
def tokenize_text_features(
batch_or_row: Dict[str, Any],
features_to_apply: Optional[List[str]] = [],
tokenizer_config: Optional[TokenizerConfig] = None,
**kwargs,
) -> Dict[str, Any]:
# Tokenize text features. features_to_apply must contain only text features.
# This works for both rows and batches.
tokenize = load_tokenize(config=tokenizer_config)
batch_or_row_masks = {}
for k, v in batch_or_row.items():
if is_feature_in_features_to_apply(features_to_apply, k):
k_mask = "_".join([k, "mask"])
if isinstance(v, np.ndarray) or isinstance(v, list):
tokenized_seq_list = [tokenize(s) for s in v]
batch_or_row[k] = torch.stack(
[seq["input_ids"].flatten() for seq in tokenized_seq_list]
) # seq_length x token_seq_length x 1
batch_or_row_masks[k_mask] = torch.stack(
[seq["attention_mask"].flatten() for seq in tokenized_seq_list]
) # seq_length x token_seq_length x 1
else:
tokenized_seq = tokenize(v)
batch_or_row[k] = tokenized_seq[
"input_ids"
].flatten() # token_seq_length
batch_or_row_masks[k_mask] = tokenized_seq[
"attention_mask"
].flatten() # token_seq_length
batch_or_row.update(batch_or_row_masks)
return batch_or_row
def preprocess_categorical_feature_to_idx(
batch_or_row: Dict[str, Any],
features_to_apply: Optional[List[str]] = [],
mapping_file: Optional[str] = "",
**kwargs,
) -> Dict[str, Any]:
# Translate categorical features to indices by looking at the mapping provided.
# features_to_apply must contain name of the categorical features whose mapping is available in the mapping_file.
# This works for both rows and batches.
# Load the mapping if a mapping file is provided
if mapping_file:
category_to_idx = load_json(mapping_file)
else:
raise ValueError("A valid path to the mapping file must be provided.")
# Helper function to translate feature values to index
def translate_to_index(value: Union[str, List[str]]) -> Union[int, List[int]]:
if isinstance(value, list):
return [
category_to_idx.get(v, 0) for v in value
] # Translate each element in the list
else:
return category_to_idx.get(
value, 0
) # Default to 0 (e.g., '<OOV>') if not found
# Determine if we are handling a single row or a batch of rows
is_batch = isinstance(batch_or_row, list)
# Apply the mapping to the appropriate features
if is_batch:
for row in batch_or_row:
for feature in features_to_apply:
if feature in row:
row[feature] = translate_to_index(row[feature])
else:
for feature in features_to_apply:
if feature in batch_or_row:
# if it's a sequence feature then process the entire sequence
batch_or_row[feature] = translate_to_index(batch_or_row[feature])
return batch_or_row
def map_sparse_id_to_embedding(
row: Dict[str, Any],
dataset_config = None,
features_to_apply: Optional[List[str]] = [],
sparse_id_field: str = "id",
embedding_field_to_add: str = "embedding",
**kwargs,
) -> Dict[str, Any]:
# Map sparse id to pre-computed embedding
embedding_map: torch.Tensor = dataset_config.embedding_map.get(
sparse_id_field, None
)
# embedding_map is an N x d tensor
# where N is the number of unique items in the dataset
# and d is the dimension of the embedding
if embedding_map is not None:
row[embedding_field_to_add] = embedding_map[row[sparse_id_field]].squeeze()
else:
raise ValueError(f"Embedding map not found")
return row
def squeeze_tensor_in_place(
batch_or_row: Dict[str, Any],
features_to_apply: Optional[List[str]] = [],
**kwargs,
) -> Dict[str, Any]:
# Squeeze the dimensions of the features to apply
# This squeeze is done in place, it does not create a new tensor
for k, v in batch_or_row.items():
if is_feature_in_features_to_apply(features_to_apply, k):
if isinstance(v, torch.Tensor):
batch_or_row[k] = v.squeeze_()
elif isinstance(v, list):
batch_or_row[k] = [
item.squeeze_() if isinstance(item, torch.Tensor) else item
for item in v
]
else:
raise ValueError(
f"Unsupported type for feature {k}: {type(v)}. Expected torch.Tensor or list."
)
return batch_or_row |