Upload dataset.py with huggingface_hub

dataset.py

@@ -1,10 +1,4 @@
-# dataset.py
-# ============================================================================
-# Turns raw log TEXT into model-ready NUMBERS. Two responsibilities:
-#   1. build_or_load_tokenizer  -> trains/loads a BPE tokenizer (text <-> IDs)
-#   2. AutoregressiveLogDataset -> serves (input, target) token windows for
-#      next-token-prediction training.
-# ============================================================================
+# file: dataset.py
 import os
 import torch
 from torch.utils.data import Dataset
@@ -15,73 +9,50 @@ from tokenizers.pre_tokenizers import Whitespace
 
 
 def build_or_load_tokenizer(corpus_path, tokenizer_save_path, vocab_size=1000):
-    """Return a BPE tokenizer, training a new one if none is cached on disk.
+    """Train a BPE tokenizer on the corpus, or load the cached one if present.
 
-    Byte-Pair Encoding (BPE) starts from individual characters and repeatedly
-    merges the most frequent adjacent pair into a new token. Frequent log
-    substrings (e.g. "sshd", "Failed", "port") become single tokens, while rare
-    text stays split into smaller pieces — keeping the vocabulary compact yet
-    able to represent ANY input.
+    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.
     """
-    # Reuse the cached tokenizer if it exists. IMPORTANT: if you change
-    # vocab_size, delete this file so a new tokenizer is trained.
     if os.path.exists(tokenizer_save_path):
         return Tokenizer.from_file(tokenizer_save_path)
 
-    # unk_token: the symbol emitted for anything the tokenizer can't represent.
     tokenizer = Tokenizer(BPE(unk_token="<unk>"))
-    # Pre-tokenize on whitespace first, so merges happen within words, not across
-    # spaces.
-    tokenizer.pre_tokenizer = Whitespace()
-    # vocab_size is a CEILING on how many tokens to learn. special_tokens are
-    # reserved IDs: <pad> (padding), <unk> (unknown), "\n" (newline — meaningful
-    # in logs, where each line is a record).
+    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)   # learn merges from the corpus
-    tokenizer.save(tokenizer_save_path)        # cache for next time
+    tokenizer.train([corpus_path], trainer)
+    tokenizer.save(tokenizer_save_path)
     return tokenizer
 
 
 class AutoregressiveLogDataset(Dataset):
-    """Slices the tokenized corpus into overlapping windows for LM training.
+    """Sliding windows over the token stream for next-token training."""
 
-    The core trick of language modeling: the target is the input shifted by one
-    position, so at every step the model learns to predict the NEXT token.
-    """
     def __init__(self, corpus_path, tokenizer, max_seq_len=64, stride=None):
         self.tokenizer = tokenizer
-        # max_seq_len: window length in tokens = how much context the model sees
-        # per training example. Must match the model's max_seq_len.
         self.max_seq_len = max_seq_len
-
-        # stride controls how far the sliding window advances between samples:
-        #   stride=1            -> maximally overlapping windows. Huge, highly
-        #                          redundant dataset (slow to train).
-        #   stride=max_seq_len  -> non-overlapping chunks (least data, fastest).
-        #   default = half-window -> a good balance of coverage vs. redundancy.
+        # 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
 
-        # Read the whole corpus and encode it once into a single long list of
-        # token IDs. (Fine for a small corpus; a large one would stream instead.)
+        # 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 every valid window START index. Each window needs
-        # max_seq_len + 1 tokens (the +1 is the final target). We step by stride.
+        # 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):
-        # Number of training windows = number of start positions.
         return len(self.start_positions)
 
     def __getitem__(self, idx):
-        # Grab a window of max_seq_len + 1 tokens starting at this position.
         start = self.start_positions[idx]
         chunk = self.encoded_tokens[start : start + self.max_seq_len + 1]
-        # x = all but the last token; y = all but the first (i.e. x shifted by 1).
-        # So for every position i, the model sees x[i] and must predict y[i]=x[i+1].
-        x = torch.tensor(chunk[:-1], dtype=torch.long)   # inputs
-        y = torch.tensor(chunk[1:], dtype=torch.long)    # targets (next tokens)
+        # 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