Spaces:
Sleeping
Sleeping
File size: 2,933 Bytes
b1e87cf 9e3a5e2 b1e87cf | 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 | import torch
from torch.utils.data import (
DataLoader,
Dataset,
TensorDataset,
WeightedRandomSampler,
random_split,
)
from tqdm.auto import tqdm
from plot2eq.data_utils.augmentation import HandDrawnAugmentation
from plot2eq.data_utils.dataset import SymbolicDataset
class AugmentedSubset(Dataset):
def __init__(self, base_dataset, indices, transform):
self.base_dataset = base_dataset
self.indices = indices
self.transform = transform
def __len__(self):
return len(self.indices)
def __getitem__(self, idx):
points, tokens = self.base_dataset[self.indices[idx]]
if self.transform:
points = self.transform(points)
return points, tokens
def build_dataloaders(
data_dir: str, batch_size: int, val_split: float = 0.05, seed: int = 42
):
print("Loading base dataset...")
base_dataset = SymbolicDataset(
data_dir, drawn_augmentation=False, map_location="cpu"
)
dataset_size = len(base_dataset)
val_size = int(val_split * dataset_size)
train_size = dataset_size - val_size
generator = torch.Generator().manual_seed(seed)
train_indices, val_indices = random_split(
range(dataset_size), # type: ignore
[train_size, val_size],
generator=generator, # type: ignore
)
token_lengths = (
base_dataset.tokens != base_dataset.tokenizer.token_map["<pad>"]
).sum(axis=1) # type: ignore
train_token_lengths = token_lengths[train_indices.indices]
weights = 1.0 / (train_token_lengths.float() + 10.0)
sampler = WeightedRandomSampler(
weights, num_samples=len(train_indices), replacement=True
)
train_transform = HandDrawnAugmentation(p=0.8)
train_dataset = AugmentedSubset(
base_dataset, train_indices, transform=train_transform
)
print("Pre-computing deterministic noisy validation set...")
val_transform = HandDrawnAugmentation(p=1.0)
torch.manual_seed(1337)
val_points_list, val_tokens_list = [], []
for idx in tqdm(val_indices, desc="Caching Val Noise", leave=False): # type: ignore
clean_points, tokens = base_dataset[idx]
noisy_points = val_transform(clean_points)
val_points_list.append(noisy_points)
val_tokens_list.append(tokens)
val_dataset = TensorDataset(
torch.stack(val_points_list), torch.stack(val_tokens_list)
)
train_loader = DataLoader(
train_dataset,
batch_size=batch_size,
shuffle=False,
sampler=sampler,
num_workers=4,
pin_memory=True,
)
val_loader = DataLoader(
val_dataset,
batch_size=batch_size,
shuffle=False,
num_workers=4,
pin_memory=True,
)
return (
train_loader,
val_loader,
len(base_dataset.tokenizer.tokens),
base_dataset.tokenizer.token_map["<pad>"],
)
|