Upload train_and_generate.py with huggingface_hub

train_and_generate.py

@@ -1,165 +1,85 @@
-# train_and_generate.py
-# ============================================================================
-# The "main" file: it assembles the model, trains it on log data, generates
-# new log text, and scores log lines for anomalies. Read this alongside
-# TUTORIAL.md. Every block is heavily commented to explain WHAT it does and
-# WHY the values are chosen.
-# ============================================================================
-
+# file: train_and_generate.py
 import os
 import math
 import torch
 import torch.nn as nn
-import torch.nn.functional as F          # functional API (cross_entropy, softmax, etc.)
-from torch.utils.data import DataLoader  # batches + shuffles the dataset for us
-from architecture import TransformerBlock                       # one Transformer layer
+import torch.nn.functional as F
+from torch.utils.data import DataLoader
+from architecture import TransformerBlock
 from dataset import build_or_load_tokenizer, AutoregressiveLogDataset
 
 
-# ============================================================================
-# THE MODEL
-# A decoder-only (GPT-style) language model. It maps a sequence of token IDs
-# to, for every position, a probability distribution over "what token comes
-# next". That single ability powers training, generation, and anomaly scoring.
-# ============================================================================
 class LogSentryLM(nn.Module):
+    """Decoder-only (GPT-style) LM. Defaults are deliberately tiny (~2M params)
+    so it trains in under a minute on a laptop GPU."""
+
     def __init__(self, vocab_size, d_model=128, n_heads=4, n_layers=3, max_seq_len=64):
-        # vocab_size : how many distinct tokens the tokenizer knows. The model
-        #              must output one score per token, so this sets the size of
-        #              both the embedding table and the final output layer.
-        # d_model    : the "width" of the model — the length of the vector used
-        #              to represent each token internally. Bigger = more capacity
-        #              but slower and more memory. 128 is small/fast for a demo.
-        # n_heads    : number of parallel attention "heads" (see architecture.py).
-        #              Must divide d_model evenly.
-        # n_layers   : how many Transformer blocks are stacked. More = deeper
-        #              reasoning, but slower. 3 is tiny by design.
-        # max_seq_len: the longest sequence (in tokens) the model can ever see at
-        #              once — its "context window" / memory span. We need this up
-        #              front because positions are LEARNED (see below), so we must
-        #              allocate one position vector per slot. 64 is just the class
-        #              default; the training script overrides it to 128.
         super().__init__()
+        # max_seq_len is the context window. It's fixed up front because positions
+        # are LEARNED embeddings (one vector per slot), not computed on the fly.
         self.max_seq_len = max_seq_len
-
-        # Token embedding: a lookup table turning each integer token ID into a
-        # learned d_model-length vector. Shape: (vocab_size, d_model).
         self.token_embeddings = nn.Embedding(vocab_size, d_model)
+        self.position_embeddings = nn.Embedding(max_seq_len, d_model)  # attention has no order sense on its own
 
-        # Position embedding: attention has NO built-in sense of order, so we add
-        # a separate learned vector for "this is position 0, 1, 2, ...". This is
-        # why max_seq_len is fixed: there are exactly max_seq_len position slots.
-        self.position_embeddings = nn.Embedding(max_seq_len, d_model)
-
-        # The stack of Transformer blocks — the actual "thinking" layers.
         self.blocks = nn.ModuleList([
             TransformerBlock(d_model=d_model, n_heads=n_heads) for _ in range(n_layers)
         ])
-
-        # A final normalization for training stability before the output layer.
         self.ln_final = nn.LayerNorm(d_model)
-
-        # The "language-model head": projects each d_model vector to vocab_size
-        # logits (one raw score per possible next token). bias=False is the
-        # common GPT convention.
-        self.lm_head = nn.Linear(d_model, vocab_size, bias=False)
+        self.lm_head = nn.Linear(d_model, vocab_size, bias=False)  # project back to vocab logits
 
     def forward(self, idx):
-        # idx: a batch of token-ID sequences, shape (batch, seq_len).
         b, s = idx.shape
-
-        # Build [0, 1, 2, ..., s-1] and shape it (1, s) so it broadcasts over the
-        # batch — these index into the position embedding table.
         positions = torch.arange(0, s, device=idx.device).unsqueeze(0)
-
-        # Input representation = "what the token is" + "where it sits". Adding
-        # them lets the model use both identity and order.
+        # token identity + position — the model needs both
         x = self.token_embeddings(idx) + self.position_embeddings(positions)
 
-        # Pass through every Transformer block. is_causal=True forbids a position
-        # from peeking at future tokens (required for next-token prediction).
         for block in self.blocks:
             x = block(x, is_causal=True)
 
-        # Normalize, then project to vocabulary logits.
-        # Output shape: (batch, seq_len, vocab_size).
         return self.lm_head(self.ln_final(x))
 
-    # ------------------------------------------------------------------------
-    # GENERATION: write new log-like text, one token at a time.
-    # @torch.no_grad() disables gradient tracking (faster, less memory) — we are
-    # not training here, just predicting.
-    # ------------------------------------------------------------------------
     @torch.no_grad()
     def generate(self, idx, max_new_tokens, temperature=0.7):
-        # idx           : the seed/prompt token IDs, shape (1, prompt_len).
-        # max_new_tokens: how many tokens to append.
-        # temperature   : randomness dial. <1 sharpens toward the most likely
-        #                 token (safe, repetitive); >1 flattens it (diverse,
-        #                 riskier); 1.0 = the model's true probabilities.
-        self.eval()  # switch off dropout etc. for deterministic inference
-        for _ in range(max_new_tokens):
-            # The model can only handle max_seq_len tokens, so keep just the most
-            # recent window if the sequence has grown beyond it.
-            idx_cond = idx[:, -self.max_seq_len:]
-
-            logits = self(idx_cond)                 # (1, seq_len, vocab_size)
-
-            # We only care about the prediction at the LAST position (the next
-            # token). Divide by temperature to scale confidence before softmax.
-            logits = logits[:, -1, :] / temperature
+        """Autoregressive sampling — predict one token, append, repeat.
 
-            # Turn logits into a probability distribution over the vocabulary.
+        temperature < 1 sharpens toward the likely token (safe/repetitive),
+        > 1 flattens it (more diverse, more mistakes).
+        """
+        self.eval()
+        for _ in range(max_new_tokens):
+            idx_cond = idx[:, -self.max_seq_len:]  # keep only what fits the context window
+            logits = self(idx_cond)
+            logits = logits[:, -1, :] / temperature  # only the last position predicts the next token
             probs = torch.softmax(logits, dim=-1)
-
-            # Randomly SAMPLE the next token from that distribution (vs. always
-            # taking the argmax) so generations have natural variety.
-            next_token = torch.multinomial(probs, num_samples=1)
-
-            # Append the new token and loop — the model now conditions on it too.
+            next_token = torch.multinomial(probs, num_samples=1)  # sample rather than argmax for variety
             idx = torch.cat((idx, next_token), dim=1)
         return idx
 
 
-# ============================================================================
-# ANOMALY SCORING
-# Because the model learned what NORMAL logs look like, a line it cannot
-# predict well is, by definition, unusual. We quantify "surprise" as the
-# average next-token loss; high surprise => likely anomaly.
-# ============================================================================
 @torch.no_grad()
 def score_line(model, tokenizer, text, device, max_seq_len=128):
-    """Score one log line by its average next-token cross-entropy loss.
+    """Anomaly score for one line = average next-token loss.
 
-    The model was trained on normal logs, so a line it can predict well scores
-    LOW (familiar / normal) and a line it cannot predict scores HIGH (surprising
-    / anomalous). Returns (avg_loss, perplexity), or None if the line is too
-    short to score. perplexity = exp(avg_loss) is just a more intuitive scale
-    ("how many tokens the model was effectively choosing between").
+    The model learned normal logs, so a line it predicts well scores LOW and a
+    surprising line scores HIGH. Returns (loss, perplexity); None if too short.
     """
     model.eval()
     ids = tokenizer.encode(text).ids
     if len(ids) < 2:
-        return None  # need at least one (context -> target) pair to score
-    ids = ids[: max_seq_len + 1]                              # respect context window
-    x = torch.tensor([ids[:-1]], dtype=torch.long, device=device)  # inputs  t0..t(n-1)
-    y = torch.tensor([ids[1:]], dtype=torch.long, device=device)   # targets t1..tn
+        return None  # need at least one (context -> target) pair
+    ids = ids[: max_seq_len + 1]
+    x = torch.tensor([ids[:-1]], dtype=torch.long, device=device)
+    y = torch.tensor([ids[1:]], dtype=torch.long, device=device)
     logits = model(x)
-    # cross_entropy compares predicted distributions against the true next tokens
-    # and averages the per-token "surprise". Lower = the model expected this line.
     loss = F.cross_entropy(logits.view(-1, logits.size(-1)), y.view(-1))
-    return loss.item(), math.exp(loss.item())                # (avg loss, perplexity)
+    return loss.item(), math.exp(loss.item())  # perplexity is the intuitive scale
 
 
 def fit_anomaly_threshold(model, tokenizer, lines, device, max_seq_len=128, k=2.0):
-    """Learn a normal-vs-anomalous cutoff from the corpus itself.
-
-    Scores every line, then sets threshold = mean + k * std of the scores. The
-    idea: most lines are normal, so their scores form a "normal" distribution;
-    anything k standard deviations above the average is unusually surprising.
-    Lower k = more sensitive (more flags, more false positives); higher k =
-    stricter. Returns (threshold, mean, std, scored_lines) where scored_lines is
-    a list of (loss, perplexity, line) sorted from most to least anomalous.
+    """Derive the normal-vs-anomalous cutoff from the corpus itself.
+
+    Most lines are normal, so threshold = mean + k*std of the scores. Lower k =
+    more sensitive (more flags/false positives); higher k = stricter.
     """
     scored = []
     for line in lines:
@@ -169,19 +89,17 @@ def fit_anomaly_threshold(model, tokenizer, lines, device, max_seq_len=128, k=2.
             scored.append((loss, ppl, line))
 
     losses = [s[0] for s in scored]
-    mean = sum(losses) / len(losses)                          # average surprise
-    var = sum((l - mean) ** 2 for l in losses) / len(losses)  # spread of surprise
+    mean = sum(losses) / len(losses)
+    var = sum((l - mean) ** 2 for l in losses) / len(losses)
     std = math.sqrt(var)
-    threshold = mean + k * std                                # the cutoff line
-    scored.sort(key=lambda s: s[0], reverse=True)             # most anomalous first
+    threshold = mean + k * std
+    scored.sort(key=lambda s: s[0], reverse=True)  # most anomalous first
     return threshold, mean, std, scored
 
 
-# ============================================================================
-# HELPER: synthetic logs, used only as a fallback if the real corpus is missing.
-# The lines are duplicated so the tiny model has an obvious pattern to learn.
-# ============================================================================
 def generate_mock_logs(file_path):
+    """Fallback synthetic corpus. Lines are duplicated so the tiny model has an
+    obvious pattern to latch onto if the real logs are missing."""
     os.makedirs(os.path.dirname(file_path), exist_ok=True)
     mock_data = """[ERROR] nginx failed. Fix: run systemctl restart nginx
 [INFO] database connection secure.
@@ -189,139 +107,93 @@ def generate_mock_logs(file_path):
 [ERROR] nginx failed. Fix: run systemctl restart nginx
 [INFO] disk space clean.
 [ERROR] out of memory. Fix: upgrade ram capacity
-""" * 50  # Duplicate entries to create a recognizable token distribution pattern
+""" * 50
     with open(file_path, "w") as f:
         f.write(mock_data)
 
 
-# ============================================================================
-# MAIN PROGRAM: data -> model -> train -> generate -> detect anomalies -> save.
-# Only runs when you execute `python train_and_generate.py` directly.
-# ============================================================================
 if __name__ == "__main__":
-    # --- Pick the fastest available compute device ---
-    # CUDA = NVIDIA GPU; MPS = Apple Silicon GPU (Metal); else CPU. We train on
-    # whichever is present, preferring the GPU for speed.
+    # Prefer GPU: cuda -> apple mps -> cpu
     if torch.cuda.is_available():
         device = "cuda"
     elif torch.backends.mps.is_available():
         device = "mps"
     else:
         device = "cpu"
-    print(f"Using device: {device}")
+    print(f"[SETUP] Using device: {device}")
 
-    LOG_FILE = "data/raw_logs.txt"               # the training corpus
-    TOKENIZER_FILE = "tokenizer/log_tokenizer.json"  # cached BPE tokenizer
+    LOG_FILE = "data/raw_logs.txt"
+    TOKENIZER_FILE = "tokenizer/log_tokenizer.json"
 
-    # --- 1. Setup data ---
-    # raw_logs.txt is the combined real-world corpus (Linux + Apache + OpenSSH
-    # samples from Loghub). Fall back to synthetic mock logs if it is missing.
+    # --- Data ---
     if not os.path.exists(LOG_FILE):
-        print("raw_logs.txt not found - generating synthetic mock logs as fallback.")
+        print("[DATA] raw_logs.txt not found — generating synthetic mock logs as fallback.")
         generate_mock_logs(LOG_FILE)
+    tokenizer = build_or_load_tokenizer(LOG_FILE, TOKENIZER_FILE, vocab_size=10000)  # ceiling; small corpus fills ~5.5k
 
-    # Train (or load) a Byte-Pair Encoding tokenizer. vocab_size=10000 is a
-    # ceiling; the small corpus only fills ~5,500 of it, which is fine.
-    tokenizer = build_or_load_tokenizer(LOG_FILE, TOKENIZER_FILE, vocab_size=10000)
-
-    # max_len is THE context window: the model reads up to 128 tokens at a time
-    # and learns dependencies within that span. Larger = sees more context but
-    # costs more compute (attention grows with the square of this length);
-    # smaller = cheaper but shorter "memory". 128 comfortably covers a full log
-    # line. (This value is passed to the model as max_seq_len.)
     max_len = 128
-
-    # stride=64 (half the window) slides the window 64 tokens at a time. stride=1
-    # would create ~187k near-duplicate windows (very slow); 64 gives ~2,900
-    # windows that still cover the corpus. Smaller stride = more/overlapping data.
     dataset = AutoregressiveLogDataset(LOG_FILE, tokenizer, max_seq_len=max_len, stride=64)
+    dataloader = DataLoader(dataset, batch_size=64, shuffle=True)  # bigger batch = better GPU use
+    print(f"[DATA] {len(dataset)} windows | {len(dataloader)} batches/epoch")
 
-    # DataLoader batches windows together and shuffles them each epoch. batch_size
-    # =64 means 64 windows are processed per step — bigger batches use the GPU
-    # more efficiently and give smoother gradients, but use more memory.
-    dataloader = DataLoader(dataset, batch_size=64, shuffle=True)
-    print(f"Dataset: {len(dataset)} windows | {len(dataloader)} batches/epoch")
-
-    # --- 2. Build the model ---
-    # vocab_size comes from the trained tokenizer; max_seq_len must match the data
-    # window so the position embeddings line up. .to(device) moves weights to GPU.
+    # --- Model ---
     model = LogSentryLM(vocab_size=tokenizer.get_vocab_size(), max_seq_len=max_len).to(device)
-
-    # AdamW: a robust optimizer. lr=1e-3 (the learning rate) is the step size for
-    # each weight update — too high diverges (loss -> nan), too low learns slowly.
-    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)
-
-    # Cross-entropy is the standard loss for predicting one class (here: the next
-    # token) out of many. Training minimizes it.
+    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)  # 1e-3 is fine at this size; drop to 3e-4 if it diverges
     loss_criterion = nn.CrossEntropyLoss()
 
-    # --- 3. Training loop ---
-    EPOCHS = 5                      # one epoch = one full pass over the dataset
+    # --- Train ---
+    EPOCHS = 5
     n_batches = len(dataloader)
-    print("Starting Model Pre-Training Optimization Pass...", flush=True)
-    model.train()                   # enable training behaviors (e.g. dropout)
+    print("[TRAIN] Starting pre-training pass...", flush=True)
+    model.train()
     for epoch in range(EPOCHS):
         total_loss = 0
         for step, (x, y) in enumerate(dataloader, start=1):
-            # x = input tokens, y = the same tokens shifted by one (the targets).
             x, y = x.to(device), y.to(device)
-
-            optimizer.zero_grad()                  # clear gradients from last step
-            logits = model(x)                      # forward pass -> predictions
-            # Flatten (batch, seq, vocab) -> (batch*seq, vocab) and targets to
-            # (batch*seq,) so loss is averaged over every predicted position.
+            optimizer.zero_grad()
+            logits = model(x)
+            # flatten (B, T, vocab) -> (B*T, vocab) so loss covers every position
             loss = loss_criterion(logits.view(-1, logits.size(-1)), y.view(-1))
-            loss.backward()                        # backprop: compute gradients
-            optimizer.step()                       # update the weights
+            loss.backward()
+            optimizer.step()
             total_loss += loss.item()
 
-            # Live, flushed progress so the run is never blind (helps with stdout
-            # buffering when output is piped to a file).
+            # flush so progress shows even when piped to a file (stdout buffering)
             if step % 20 == 0 or step == n_batches:
-                print(f"Epoch {epoch+1}/{EPOCHS} | batch {step}/{n_batches} "
+                print(f"[TRAIN] Epoch {epoch+1}/{EPOCHS} | batch {step}/{n_batches} "
                       f"| loss {loss.item():.4f}", flush=True)
 
-        # Average loss per epoch should fall steadily — that means it's learning.
-        print(f"== Epoch {epoch+1}/{EPOCHS} done | Avg Cross-Entropy Loss: "
-              f"{total_loss/n_batches:.4f} ==", flush=True)
+        print(f"[TRAIN] Epoch {epoch+1}/{EPOCHS} done | avg CE loss "
+              f"{total_loss/n_batches:.4f}", flush=True)
 
-    # --- 4. Generation demo ---
-    print("\n--- Running Inference Engine Generation Phase ---")
-    prompt_text = "Invalid user"  # Seed phrase that appears in the real OpenSSH/Linux logs
+    # --- Generate (quick sanity check) ---
+    prompt_text = "Invalid user"  # a phrase that actually appears in the corpus
     prompt_tokens = tokenizer.encode(prompt_text).ids
-    # Wrap in a batch of 1 and move to the device. dtype=long because token IDs
-    # are integers used to index the embedding table.
     input_tensor = torch.tensor([prompt_tokens], dtype=torch.long, device=device)
 
-    # Generate 6 new tokens following the seed. temperature=0.5 keeps it crisp.
     generated_ids = model.generate(input_tensor, max_new_tokens=6, temperature=0.5)
-    decoded_output = tokenizer.decode(generated_ids[0].tolist())  # IDs -> text
+    decoded_output = tokenizer.decode(generated_ids[0].tolist())
 
-    print(f"User Input Seed Text  : {prompt_text}")
-    print(f"Model Predicted Output: {decoded_output}")
+    print(f"\n[GEN] Seed  : {prompt_text}")
+    print(f"[GEN] Output: {decoded_output}")
 
-    # --- 5. Anomaly detection ---
-    print("\n--- Anomaly Detection Phase ---", flush=True)
+    # --- Anomaly detection ---
     with open(LOG_FILE, "r") as f:
-        corpus_lines = [ln.strip() for ln in f if ln.strip()]  # non-empty lines
+        corpus_lines = [ln.strip() for ln in f if ln.strip()]
 
-    # Learn the normal-score distribution from the training corpus and derive the
-    # anomaly threshold (mean + 2*std of per-line scores).
     threshold, mean, std, scored = fit_anomaly_threshold(
         model, tokenizer, corpus_lines, device, max_seq_len=max_len, k=2.0
     )
-    print(f"Score stats: mean={mean:.3f} std={std:.3f} | "
-          f"anomaly threshold (mean+2*std)={threshold:.3f}")
+    print(f"\n[ANOMALY] Score stats: mean={mean:.3f} std={std:.3f} | "
+          f"threshold (mean+2*std)={threshold:.3f}", flush=True)
 
-    # The lines the model found hardest to predict (usually the rarest log types).
-    print("\nTop 5 most anomalous lines in the corpus:")
+    print("[ANOMALY] Top 5 most anomalous lines in the corpus:")
     for loss, ppl, line in scored[:5]:
         flag = "ANOMALY" if loss > threshold else "normal "
         print(f"  [{flag}] loss={loss:.3f} ppl={ppl:8.1f} | {line[:90]}")
 
-    # Predict on fresh lines: a couple of real-looking logs plus an injected
-    # off-distribution line that the model has never seen (should score huge).
-    print("\nPredicting anomalies on new lines:")
+    # A couple of real lines + one injected off-distribution line (should score huge)
+    print("[ANOMALY] Predicting on new lines:")
     test_lines = [
         "Dec 10 06:55:46 LabSZ sshd[24200]: Invalid user webmaster from 173.234.31.186",
         "Failed password for root from 112.95.230.3 port 49204 ssh2",
@@ -336,18 +208,15 @@ if __name__ == "__main__":
         flag = "ANOMALY" if loss > threshold else "normal "
         print(f"  [{flag}] loss={loss:.3f} ppl={ppl:8.1f} | {line[:90]}")
 
-    # --- 6. Save everything the standalone predictor needs ---
-    # We store the weights AND the learned threshold + config, so
-    # detect_anomalies.py can rebuild the exact model and judge new lines without
-    # retraining.
+    # --- Save weights + threshold so detect_anomalies.py can run standalone ---
     MODEL_FILE = "model/logsentry_lm.pt"
     os.makedirs(os.path.dirname(MODEL_FILE), exist_ok=True)
     torch.save({
-        "model_state": model.state_dict(),   # the learned weights
+        "model_state": model.state_dict(),
         "vocab_size": tokenizer.get_vocab_size(),
         "max_seq_len": max_len,
         "threshold": threshold,
         "score_mean": mean,
         "score_std": std,
     }, MODEL_FILE)
-    print(f"\nSaved model + anomaly threshold to {MODEL_FILE}")
+    print(f"\n[SAVE] Saved model + anomaly threshold to {MODEL_FILE}")