File size: 7,974 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 | from typing import Optional, Tuple, Union
import torch
from src.utils.file_utils import open_local_or_remote
def locations_to_index_tuple(locations: torch.Tensor, num_dims: int = 2) -> Tuple:
"""
Convert a tensor of locations to a tuple of index tensors for advanced indexing.
Args:
locations (torch.Tensor): A tensor of shape `[L, D]` where `L` is the number of
locations and `D >= num_dims`.
num_dims (int): The number of dimensions to extract. The first num_dims columns of
the locations tensor are used. We explicitly specify this to make the
function call traceable.
Returns:
Tuple: A tuple of `num_dims` tensors, each of shape `[L]` representing the
indices for one dimension.
Example:
>>> locations = torch.tensor([[0, 10], [1, 20], [2, 5]])
>>> locations_to_index_tuple(locations, num_dims=2)
(tensor([0, 1, 2]), tensor([10, 20, 5]))
>>> locations = torch.tensor([[0, 10], [1, 20], [2, 5]])
>>> locations_to_index_tuple(locations, num_dims=1)
(tensor([0, 1, 2]))
"""
return tuple(locations[:, i] for i in range(num_dims))
def extract_locations(
data: torch.tensor, locations: torch.tensor, num_dims: int = 2
) -> torch.tensor:
"""
Extracts the elements from a tensor at the specified indices.
Args:
data (torch.tensor): The input tensor of N dimensions from which to extract elements.
locations (torch.tensor): Tensor of shape [L, D] where L is the number of
elements where each D dimensional row reprecents the first D dimensions
of the data tensor to extract.
num_dims (int): The number of dimensions to extract. The first num_dims columns of
the locations tensors are used. We need to specify to make this function call traceable.
Returns:
torch.tensor: A tensor of shape [L,...] with total N-num_dims+1 dimensions
containing the extracted elements.
Example:
>>> data = torch.tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
>>> locations = torch.tensor([[0, 1], [1, 2]])
>>> extract_locations(data, locations, num_dims=2)
tensor([2, 6]) # (Index 0,1 gives 2 (First row, second column); Index 1,2 gives 6 (Second row, third column))
>>> locations = torch.tensor([[0, 1], [2, 0]])
>>> extract_locations(data, locations, num_dims=1)
tensor([[1, 2, 3], [7, 8, 9]])
# (num_dims = 1 implies we are extracting based on the first dimension only.
# Thus, we get the first row (from [0,1] as 1 is ignored) and the third row
# (from [2,0] as 0 is ignored) of the data tensor.
"""
# Separate the locations for each of the first D dimensions
index_tuple = locations_to_index_tuple(locations=locations, num_dims=num_dims)
# Use indexing with a tuple of index tensors
extracted_values = data[index_tuple]
return extracted_values
def merge_list_of_keyed_tensors_to_single_tensor(
data: list[dict[str, torch.Tensor]],
index_key: str,
value_key: str,
) -> torch.Tensor:
"""
Converts a list of dictionaries of id to tensors into a single tensor by squeezing
the tensors along the specified index key.
e.g.,
data = [
[
{'user_id': 123,
{'semantic_id': torch.tensor([21, 32, 124]),
other features.....}
],
[
{'user_id': 456,
{'semantic_id': torch.tensor([11, 22, 33]),
other features.....}
]
output:
tensor([...,
[21, 32, 124], # row 123
...,
[11, 22, 33], # row 456
...
])
Args:
data (list[dict[str, torch.Tensor]]): A list of dictionaries where each dictionary
contains an index key and a value key.
index_key (str): The key in the dictionary that contains the index for each row.
value_key (str): The key in the dictionary that contains the tensor to be merged.
"""
batch_size = len(data)
dimensions = torch.tensor(data[0][value_key]).size()
output_tensor = torch.zeros((batch_size, *dimensions))
for row in data:
index = row[index_key]
value = row[value_key]
if index < batch_size:
output_tensor[index] = torch.tensor(value)
else:
raise IndexError(
f"Index {index} out of bounds for batch size {batch_size}."
)
return output_tensor
def deduplicate_rows_in_tensor(
file_path: Optional[str] = None, return_tensor: bool = False
) -> Union[None, torch.Tensor]:
"""
Identifies and de-duplicate repeated rows in a PyTorch tensor.
Rows that are not duplicated will have a new column with value 0,
while rows that are duplicated will have a new column indicating the number of duplicates from 1 to N-1
where N is the number of duplicates for that row.
Args:
file_path: Optional; Path to a file containing the tensor data.
return_tensor: If True, returns the modified tensor; otherwise, saves it to the file.
Returns:
If return_tensor is True, returns the modified tensor with a new column indicating
the number of duplicates for each row. If False, saves the modified tensor to the file.
"""
if not file_path.endswith(".pt"):
return None
data = torch.load(open_local_or_remote(file_path, mode="rb"))
assert len(data.size()) == 2, "Input data must be a 2D PyTorch tensor."
# Use torch.unique to get unique rows and their inverse indices
unique_rows, inverse_indices, counts = torch.unique(
data, dim=0, return_inverse=True, return_counts=True
)
output_indices = torch.zeros_like(inverse_indices)
# Find indices where counts > 1 (meaning duplicates exist)
duplicate_indices = torch.where(counts > 1)[0]
for i in range(len(duplicate_indices)):
# Calculate number of collisions
num_of_collisions = counts[duplicate_indices[i]]
# Gather the indices where the collisions occur
indices_to_change = torch.where(inverse_indices == duplicate_indices[i])[0]
# Create a range based on the number of collision, starting from 1
range_to_add = torch.arange(1, num_of_collisions + 1)
# Scatter to those specific indices
output_indices = output_indices.scatter(0, indices_to_change, range_to_add)
# Concatenate the duplicate indicator column to the original data
result = torch.cat((data, output_indices.unsqueeze(1)), dim=1).long()
if return_tensor:
return result
else:
# Save the result to a file
torch.save(result, file_path)
return None
def transpose_tensor_from_file(
file_path: Optional[str] = None,
return_tensor: bool = False,
dim1: int = -2,
dim2: int = -1,
) -> Union[None, torch.Tensor]:
"""
Transposes a PyTorch tensor from a file accoridng to designated dimensions.
Args:
file_path: Optional; Path to a file containing the tensor data.
return_tensor: If True, returns the modified tensor; otherwise, saves it to the file.
dim1: The first dimension to transpose (default: -2).
dim2: The second dimension to transpose (default: -1).
Returns:
If return_tensor is True, returns the modified tensor. If False, saves the modified tensor to the file.
"""
if not file_path.endswith(".pt"):
return None
data = torch.load(open_local_or_remote(file_path, mode="rb"))
# Transpose the tensor
result = data.transpose(dim1, dim2)
if return_tensor:
return result
else:
# Save the result to a file
torch.save(result, file_path)
return None
|