| import torch | |
| from torch.utils.data import DataLoader, Dataset | |
| from transformers import AutoTokenizer | |
| from datasets import load_dataset | |
| def setup_tokenizer(num_memory_tokens: int = 32): | |
| tokenizer = AutoTokenizer.from_pretrained("gpt2") | |
| special_tokens = { | |
| "pad_token": "<PAD>", | |
| "bos_token": "<BOS>", | |
| "eos_token": "<EOS>", | |
| "unk_token": "<UNK>", | |
| } | |
| memory_tokens = [f"<MEM_{i}>" for i in range(num_memory_tokens)] | |
| special_tokens["additional_special_tokens"] = memory_tokens | |
| tokenizer.add_special_tokens(special_tokens) | |
| print(f"Tokenizer ready - vocab size: {len(tokenizer)}, memory tokens: {num_memory_tokens}") | |
| return tokenizer | |
| def apply_token_dropout( | |
| input_ids: torch.Tensor, | |
| attention_mask: torch.Tensor, | |
| tokenizer, | |
| dropout_prob: float = 0.15, | |
| num_memory_tokens: int = 32, | |
| ) -> tuple[torch.Tensor, torch.Tensor]: | |
| if dropout_prob <= 0: | |
| return input_ids, attention_mask | |
| dropout_mask = torch.rand_like(input_ids, dtype=torch.float) < dropout_prob | |
| special_token_mask = ( | |
| (input_ids == tokenizer.bos_token_id) | |
| | (input_ids == tokenizer.eos_token_id) | |
| | (input_ids == tokenizer.pad_token_id) | |
| | (input_ids == tokenizer.unk_token_id) | |
| ) | |
| memory_token_mask = torch.zeros_like(input_ids, dtype=torch.bool) | |
| if num_memory_tokens > 0: | |
| memory_token_mask[:, -num_memory_tokens:] = True | |
| dropout_mask = dropout_mask & ~special_token_mask & ~memory_token_mask & attention_mask.bool() | |
| corrupted_ids = input_ids.clone() | |
| corrupted_ids[dropout_mask] = tokenizer.unk_token_id | |
| return corrupted_ids, attention_mask | |
| class PackedTokenDataset(Dataset): | |
| """Fixed-length passages cut from one contiguous token stream.""" | |
| def __init__(self, token_ids: torch.Tensor): | |
| self.token_ids = token_ids | |
| def __len__(self): | |
| return self.token_ids.size(0) | |
| def __getitem__(self, index): | |
| return {"input_ids": self.token_ids[index]} | |
| def create_dataloader( | |
| batch_size: int, | |
| dataset, | |
| tokenizer, | |
| shuffle: bool = True, | |
| num_memory_tokens: int = 32, | |
| max_length: int = 32, | |
| slot_source: str = "appended", | |
| ): | |
| memory_token_ids = torch.tensor( | |
| [tokenizer.convert_tokens_to_ids(f"<MEM_{i}>") for i in range(num_memory_tokens)], | |
| dtype=torch.long, | |
| ) | |
| def collate_fn(batch): | |
| if "input_ids" in batch[0]: | |
| ids_batch = torch.stack([ | |
| torch.as_tensor(item["input_ids"], dtype=torch.long) for item in batch | |
| ]) | |
| mask_batch = torch.ones_like(ids_batch) | |
| else: | |
| texts = [item["text"] for item in batch] | |
| tokenized = tokenizer( | |
| texts, | |
| padding=False, | |
| truncation=True, | |
| max_length=max_length, | |
| return_tensors=None, | |
| ) | |
| max_len = max(len(ids) for ids in tokenized["input_ids"]) | |
| batch_ids = [] | |
| batch_masks = [] | |
| for input_ids, attn_mask in zip(tokenized["input_ids"], tokenized["attention_mask"]): | |
| pad_len = max_len - len(input_ids) | |
| if pad_len > 0: | |
| input_ids = input_ids + [tokenizer.pad_token_id] * pad_len | |
| attn_mask = attn_mask + [0] * pad_len | |
| batch_ids.append(torch.tensor(input_ids, dtype=torch.long)) | |
| batch_masks.append(torch.tensor(attn_mask, dtype=torch.long)) | |
| ids_batch = torch.stack(batch_ids) | |
| mask_batch = torch.stack(batch_masks) | |
| if slot_source == "appended": | |
| B = ids_batch.size(0) | |
| mem_tokens = memory_token_ids.unsqueeze(0).expand(B, -1) | |
| mem_mask = torch.ones(B, num_memory_tokens, dtype=torch.long) | |
| encoder_ids = torch.cat([ids_batch, mem_tokens], dim=1) | |
| encoder_mask = torch.cat([mask_batch, mem_mask], dim=1) | |
| elif slot_source == "token_states": | |
| if ids_batch.size(1) != num_memory_tokens: | |
| raise ValueError( | |
| "token_states requires one memory slot per source position: " | |
| f"source length={ids_batch.size(1)}, memory slots={num_memory_tokens}" | |
| ) | |
| encoder_ids = ids_batch | |
| encoder_mask = mask_batch | |
| else: | |
| raise ValueError(f"Unknown slot_source: {slot_source}") | |
| return { | |
| "input_ids": encoder_ids, | |
| "attention_mask": encoder_mask, | |
| "original_input_ids": ids_batch, | |
| "original_attention_mask": mask_batch, | |
| } | |
| return DataLoader( | |
| dataset, | |
| batch_size=batch_size, | |
| shuffle=shuffle, | |
| collate_fn=collate_fn, | |
| drop_last=shuffle, | |
| num_workers=2, | |
| ) | |
| def load_reconstruction_dataset(cfg: dict, tokenizer=None): | |
| dataset_name = cfg["data"]["dataset"] | |
| if cfg["data"].get("format") == "packed_tokens": | |
| if tokenizer is None: | |
| raise ValueError("packed_tokens format requires a tokenizer") | |
| dataset_config = cfg["data"].get("dataset_config") | |
| split_name = cfg["data"].get("split", "train") | |
| token_budget = cfg["data"].get("token_budget", 500_000) | |
| validation_tokens = cfg["data"].get("validation_tokens", 5_000) | |
| sequence_length = cfg["data"].get("max_length", 256) | |
| print(f"Loading packed dataset: {dataset_name}/{dataset_config or ''}") | |
| ds = load_dataset(dataset_name, dataset_config, split=split_name) | |
| tokens = [] | |
| eos_id = tokenizer.eos_token_id | |
| for row in ds: | |
| text = row["text"].strip() | |
| if not text: | |
| continue | |
| tokens.extend(tokenizer.encode(text, add_special_tokens=False)) | |
| tokens.append(eos_id) | |
| if len(tokens) >= token_budget: | |
| break | |
| if len(tokens) < token_budget: | |
| raise RuntimeError( | |
| f"Dataset yielded {len(tokens):,} tokens, below requested budget {token_budget:,}" | |
| ) | |
| tokens = torch.tensor(tokens[:token_budget], dtype=torch.long) | |
| validation_tokens = min(validation_tokens, len(tokens) // 10) | |
| split = len(tokens) - validation_tokens | |
| split -= split % sequence_length | |
| train_tokens = tokens[:split] | |
| val_tokens = tokens[split:] | |
| val_tokens = val_tokens[: len(val_tokens) - len(val_tokens) % sequence_length] | |
| train_sequences = train_tokens.view(-1, sequence_length) | |
| val_sequences = val_tokens.view(-1, sequence_length) | |
| print( | |
| f"Train: {len(train_sequences)} packed sequences, " | |
| f"Val: {len(val_sequences)}, length: {sequence_length}" | |
| ) | |
| return PackedTokenDataset(train_sequences), PackedTokenDataset(val_sequences) | |
| max_samples = cfg["data"].get("max_samples") | |
| val_ratio = cfg["data"].get("val_ratio", 0.004) | |
| print(f"Loading dataset: {dataset_name}") | |
| ds = load_dataset(dataset_name)["train"] | |
| if max_samples and max_samples < len(ds): | |
| ds = ds.select(range(max_samples)) | |
| if cfg.get("dry_run"): | |
| n = cfg.get("dry_run_samples", 10) | |
| ds = ds.select(range(min(n, len(ds)))) | |
| val_ratio = 0.5 | |
| split = ds.train_test_split(test_size=val_ratio, seed=42) | |
| print(f"Train: {len(split['train'])}, Val: {len(split['test'])}") | |
| return split["train"], split["test"] | |
Xet Storage Details
- Size:
- 7.45 kB
- Xet hash:
- ef896ba650e70c10d22a4b89f1396f4e85b8bcd6518db7b8ea6545d0045b0c29
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.