File size: 2,528 Bytes
511fd40 83c84a4 511fd40 83c84a4 511fd40 83c84a4 511fd40 83c84a4 511fd40 83c84a4 511fd40 83c84a4 511fd40 83c84a4 511fd40 83c84a4 511fd40 83c84a4 511fd40 83c84a4 | 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 | # file: dataset.py
import os
import torch
from torch.utils.data import Dataset
from tokenizers import Tokenizer
from tokenizers.models import BPE
from tokenizers.trainers import BpeTrainer
from tokenizers.pre_tokenizers import Whitespace
def build_or_load_tokenizer(corpus_path, tokenizer_save_path, vocab_size=1000):
"""Train a BPE tokenizer on the corpus, or load the cached one if present.
NOTE: the cache is keyed only by path — if you change vocab_size, delete the
saved json first or you'll silently keep the old vocab.
"""
if os.path.exists(tokenizer_save_path):
return Tokenizer.from_file(tokenizer_save_path)
tokenizer = Tokenizer(BPE(unk_token="<unk>"))
tokenizer.pre_tokenizer = Whitespace() # split on spaces first so merges stay within words
# "\n" is reserved as a token because each log line is a record boundary
trainer = BpeTrainer(vocab_size=vocab_size, special_tokens=["<pad>", "<unk>", "\n"])
tokenizer.train([corpus_path], trainer)
tokenizer.save(tokenizer_save_path)
return tokenizer
class AutoregressiveLogDataset(Dataset):
"""Sliding windows over the token stream for next-token training."""
def __init__(self, corpus_path, tokenizer, max_seq_len=64, stride=None):
self.tokenizer = tokenizer
self.max_seq_len = max_seq_len
# stride = how far the window advances each step. Default half-window is a
# balance: stride=1 gives ~187k near-duplicate windows (slow), stride=max_seq_len
# gives no overlap (less coverage).
self.stride = stride if stride is not None else max_seq_len // 2
# Small corpus, so just read + encode it once. Would stream for a big one.
with open(corpus_path, "r") as f:
raw_text = f.read()
self.encoded_tokens = tokenizer.encode(raw_text).ids
# Precompute window starts; each needs max_seq_len + 1 tokens (the +1 is the last target)
last_start = len(self.encoded_tokens) - self.max_seq_len - 1
self.start_positions = list(range(0, max(0, last_start) + 1, self.stride))
def __len__(self):
return len(self.start_positions)
def __getitem__(self, idx):
start = self.start_positions[idx]
chunk = self.encoded_tokens[start : start + self.max_seq_len + 1]
# y is x shifted by one — at every position the model predicts the next token
x = torch.tensor(chunk[:-1], dtype=torch.long)
y = torch.tensor(chunk[1:], dtype=torch.long)
return x, y
|