fix: Offsets

This commit is contained in:
Tibo De Peuter 2025-12-09 22:08:52 +01:00
parent 926cde17d3
commit 291be01069
Signed by: tdpeuter
GPG key ID: 38297DE43F75FFE2

View file

@ -49,12 +49,12 @@ class Dataset(TorchDataset, ABC):
return len(self.dataset) return len(self.dataset)
def process_data(self): def process_data(self):
self.chunk_offsets = self.get_offsets()
if self.size == -1: if self.size == -1:
# Just use the whole dataset # Just use the whole dataset
self.bytes = ''.join(tqdm(self.data, desc="Encoding data")).encode('utf-8', errors='replace') self.bytes = ''.join(tqdm(self.data, desc="Encoding data")).encode('utf-8', errors='replace')
else: else:
# Use only partition, calculate offsets # Use only partition, calculate offsets
self.chunk_offsets = self.get_offsets()
self.bytes = ''.join(tqdm(self.data[:len(self.chunk_offsets)], desc="Encoding data")).encode('utf-8', errors='replace') self.bytes = ''.join(tqdm(self.data[:len(self.chunk_offsets)], desc="Encoding data")).encode('utf-8', errors='replace')
self.tensor = torch.tensor(list(self.bytes), dtype=torch.long) self.tensor = torch.tensor(list(self.bytes), dtype=torch.long)