Upload train_and_generate.py with huggingface_hub

train_and_generate.py

@@ -0,0 +1,353 @@
+# 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.
+# ============================================================================
+
+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
+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):
+    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__()
+        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)
+
+        # 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)
+
+    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.
+        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
+
+            # Turn logits into a probability distribution over the vocabulary.
+            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.
+            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.
+
+    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").
+    """
+    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
+    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)
+
+
+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.
+    """
+    scored = []
+    for line in lines:
+        result = score_line(model, tokenizer, line, device, max_seq_len)
+        if result is not None:
+            loss, ppl = result
+            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
+    std = math.sqrt(var)
+    threshold = mean + k * std                                # the cutoff line
+    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):
+    os.makedirs(os.path.dirname(file_path), exist_ok=True)
+    mock_data = """[ERROR] nginx failed. Fix: run systemctl restart nginx
+[INFO] database connection secure.
+[ERROR] out of memory. Fix: upgrade ram capacity
+[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
+    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.
+    if torch.cuda.is_available():
+        device = "cuda"
+    elif torch.backends.mps.is_available():
+        device = "mps"
+    else:
+        device = "cpu"
+    print(f"Using device: {device}")
+
+    LOG_FILE = "data/raw_logs.txt"               # the training corpus
+    TOKENIZER_FILE = "tokenizer/log_tokenizer.json"  # cached BPE tokenizer
+
+    # --- 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.
+    if not os.path.exists(LOG_FILE):
+        print("raw_logs.txt not found - generating synthetic mock logs as fallback.")
+        generate_mock_logs(LOG_FILE)
+
+    # 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 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 = 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.
+    loss_criterion = nn.CrossEntropyLoss()
+
+    # --- 3. Training loop ---
+    EPOCHS = 5                      # one epoch = one full pass over the dataset
+    n_batches = len(dataloader)
+    print("Starting Model Pre-Training Optimization Pass...", flush=True)
+    model.train()                   # enable training behaviors (e.g. dropout)
+    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.
+            loss = loss_criterion(logits.view(-1, logits.size(-1)), y.view(-1))
+            loss.backward()                        # backprop: compute gradients
+            optimizer.step()                       # update the weights
+            total_loss += loss.item()
+
+            # Live, flushed progress so the run is never blind (helps with stdout
+            # buffering when output is piped to a file).
+            if step % 20 == 0 or step == n_batches:
+                print(f"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)
+
+    # --- 4. Generation demo ---
+    print("\n--- Running Inference Engine Generation Phase ---")
+    prompt_text = "Invalid user"  # Seed phrase that appears in the real OpenSSH/Linux logs
+    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
+
+    print(f"User Input Seed Text  : {prompt_text}")
+    print(f"Model Predicted Output: {decoded_output}")
+
+    # --- 5. Anomaly detection ---
+    print("\n--- Anomaly Detection Phase ---", flush=True)
+    with open(LOG_FILE, "r") as f:
+        corpus_lines = [ln.strip() for ln in f if ln.strip()]  # non-empty lines
+
+    # 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}")
+
+    # The lines the model found hardest to predict (usually the rarest log types).
+    print("\nTop 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:")
+    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",
+        "kjshdf!! TOTALLY RANDOM gibberish $$$ 999 not-a-real-log ~~~~",
+    ]
+    for line in test_lines:
+        result = score_line(model, tokenizer, line, device, max_seq_len=max_len)
+        if result is None:
+            print(f"  [skipped: too short] {line}")
+            continue
+        loss, ppl = result
+        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.
+    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
+        "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}")