File size: 4,187 Bytes
34979f2 | 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 | """Shard reader for EDT — reads .npy shards directly, no fusion needed.
This replaces the old `torch.load('corpus.pt')` which crashed on files >50GB.
The training reads shards one at a time via mmap (zero RAM cost).
"""
import os, glob, random
import numpy as np
import torch
class ShardCorpusReader:
"""Read tokenized corpus from .npy shards without loading everything in RAM.
Usage:
reader = ShardCorpusReader('data/fractus_1b_shards/')
tokens = reader.get_random_batch(batch_size=8, seq_len=32, device='cuda')
# tokens: (batch_size, seq_len) long tensor
"""
def __init__(self, shard_dir: str):
self.shards = sorted(glob.glob(os.path.join(shard_dir, "*_*.npy")))
if not self.shards:
raise FileNotFoundError(f"No shards found in {shard_dir}")
# Index: for each shard, get its length (mmap, no RAM).
self.shard_lengths = []
for s in self.shards:
arr = np.load(s, mmap_mode="r")
self.shard_lengths.append(len(arr))
self.total_tokens = sum(self.shard_lengths)
print(f" ShardCorpusReader: {len(self.shards)} shards, "
f"{self.total_tokens/1e9:.2f}B tokens", flush=True)
def get_random_batch(self, batch_size: int, seq_len: int, device="cpu"):
"""Get a random batch of token sequences from the shards.
Picks random positions across all shards, reads seq_len tokens from each.
"""
batch_tokens = []
for _ in range(batch_size):
# Pick a random shard (weighted by size).
shard_idx = random.choices(
range(len(self.shards)),
weights=self.shard_lengths
)[0]
shard = np.load(self.shards[shard_idx], mmap_mode="r")
shard_len = self.shard_lengths[shard_idx]
# Pick a random position.
if shard_len <= seq_len + 1:
continue
start = random.randint(0, shard_len - seq_len - 1)
# Read seq_len tokens.
tokens = shard[start:start + seq_len + 1]
batch_tokens.append(tokens)
if not batch_tokens:
# Fallback if all shards too small (shouldn't happen).
shard = np.load(self.shards[0], mmap_mode="r")
tokens = shard[:seq_len + 1]
batch_tokens = [tokens] * batch_size
# Convert to tensor: inp and tgt.
data = np.stack(batch_tokens) # (batch_size, seq_len+1)
data_tensor = torch.from_numpy(data.astype(np.int64)).to(device)
inp = data_tensor[:, :-1] # (batch_size, seq_len)
tgt = data_tensor[:, 1:] # (batch_size, seq_len)
return inp, tgt
def get_random_tokens(self, n: int, device="cpu"):
"""Get n random token IDs from the corpus (for embedding training)."""
tokens = []
for _ in range(n):
shard_idx = random.choices(
range(len(self.shards)),
weights=self.shard_lengths
)[0]
shard = np.load(self.shards[shard_idx], mmap_mode="r")
pos = random.randint(0, self.shard_lengths[shard_idx] - 1)
tokens.append(shard[pos])
return torch.tensor(tokens, dtype=torch.long, device=device)
def stream_tokens(self, batch_size=128, seq_len=64, device="cpu"):
"""Generator that yields batches sequentially through the corpus."""
for shard_path in self.shards:
shard = np.load(shard_path, mmap_mode="r")
shard_len = len(shard)
for start in range(0, shard_len - seq_len - 1, seq_len * batch_size):
batch_tokens = []
for b in range(batch_size):
pos = start + b * seq_len
if pos + seq_len + 1 >= shard_len:
break
batch_tokens.append(shard[pos:pos + seq_len + 1])
if len(batch_tokens) < batch_size:
break
data = np.stack(batch_tokens)
data_tensor = torch.from_numpy(data.astype(np.int64)).to(device)
yield data_tensor[:, :-1], data_tensor[:, 1:]
|