topoexploit 0.1.0__py3-none-any.whl

topogpt3/continuation.py

@@ -0,0 +1,114 @@
+#!/usr/bin/env python3
+"""
+Auto-continuation engine: detects truncated responses and feeds the last
+incomplete lines back so the model can resume where it left off.
+
+Used by both the standard inference pipeline and the HRM "thinking" mode.
+"""
+from __future__ import annotations
+
+import re
+from typing import List, Tuple
+
+
+# ── Fence / block structures ────────────────────────────────────────────────
+
+
+_OPENING = {"{": "}", "(": ")", "[": "]", "<": ">"}
+_CLOSING = {v: k for k, v in _OPENING.items()}
+_OPENERS = frozenset(_OPENING)
+_CLOSERS = frozenset(_CLOSING)
+
+_CODE_FENCE_RE = re.compile(r"^```(\w*)$")
+
+
+def _count_unclosed_brackets(text: str) -> int:
+    stack: List[str] = []
+    for ch in text:
+        if ch in _OPENERS:
+            stack.append(ch)
+        elif ch in _CLOSERS:
+            if stack and _CLOSING.get(ch) == stack[-1]:
+                stack.pop()
+    return len(stack)
+
+
+def _count_unclosed_fences(text: str) -> int:
+    fences = 0
+    for line in text.splitlines():
+        m = _CODE_FENCE_RE.match(line.strip())
+        if m:
+            fences += 1
+    return fences % 2  # 0 = balanced, 1 = unclosed
+
+
+def is_response_complete(text: str, min_chars: int = 10) -> bool:
+    """Heuristic to decide whether a model response looks finished.
+
+    Returns True when the response seems naturally complete (no need to
+    continue), False when it appears truncated and continuation may help.
+    """
+    stripped = text.rstrip()
+
+    if len(stripped) < min_chars:
+        return False
+
+    if _count_unclosed_fences(stripped):
+        return False
+
+    if _count_unclosed_brackets(stripped) > 0:
+        return False
+
+    if not stripped:
+        return False
+
+    last_ch = stripped[-1]
+    if last_ch in {":", ",", ";", "\\"}:
+        return False
+
+    if last_ch in {"\n", ".", "!", "?", ")", "}", "]", "`", '"', "'"}:
+        return True
+
+    return True
+
+
+def extract_tail_for_continuation(text: str, tail_lines: int = 2,
+                                  tail_chars: int = 512) -> str:
+    """Return the last N lines (or up to tail_chars) of `text` as a
+    continuation prefix to feed back into the model.
+
+    The returned string can be prepended as context for the model's next
+    generation call so it continues naturally from that point.
+    """
+    stripped = text.rstrip()
+    if not stripped:
+        return ""
+
+    lines: List[str] = text.splitlines()
+    result: List[str] = []
+    total_chars = 0
+
+    for line in reversed(lines):
+        candidate = line
+        new_total = total_chars + len(candidate) + 1
+        if len(result) >= tail_lines:
+            break
+        if new_total > tail_chars:
+            break
+        result.append(candidate)
+        total_chars = new_total
+
+    result.reverse()
+    return "\n".join(result)
+
+
+def split_at_last_newline(text: str) -> Tuple[str, str]:
+    """Split `text` at the last newline.
+
+    Returns (prefix_without_last_line, last_line).
+    Useful for discarding a trailing incomplete line before continuation.
+    """
+    if "\n" not in text:
+        return "", text
+    idx = text.rfind("\n")
+    return text[:idx], text[idx + 1:]

topogpt3/ewc.py

@@ -0,0 +1,369 @@
+#!/usr/bin/env python3
+"""
+Elastic Weight Consolidation (EWC) and Experience Replay for TopoGPT3.
+
+Provides EWCRegularizer for regularizing against catastrophic forgetting
+across curriculum tiers, and ReplayBuffer for interleaving data from
+previous tiers during training.
+
+Author: Gris Iscomeback
+License: GPL v3
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Dict, Iterable, List, Optional, Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+logger = logging.getLogger(__name__)
+
+
+class EWCRegularizer:
+    """
+    Elastic Weight Consolidation regularizer.
+
+    Computes the diagonal Fisher information matrix over a dataset and uses
+    it to penalize large deviations from the parameters that were optimal
+    for previous tasks (tiers). The penalty is:
+
+        L_EWC = lambda * sum_i  F_i * (theta_i - theta_opt_i)^2
+
+    where F_i is the Fisher information for parameter i and theta_opt_i
+    is the parameter value at the end of the previous task.
+
+    Usage:
+        ewc = EWCRegularizer(model, lambda_ewc=5000.0)
+        ewc.compute_fisher(dataloader, vocab_size, device)
+        # ... train on new tier ...
+        loss = task_loss + ewc.penalty(model)
+    """
+
+    def __init__(
+        self,
+        model: nn.Module,
+        lambda_ewc: float = 5000.0,
+        fisher_num_samples: int = 50,
+    ):
+        """
+        Args:
+            model: The model to regularize.
+            lambda_ewc: Penalty strength scaling the Fisher-weighted
+                parameter deviation loss.
+            fisher_num_samples: Number of batches used when computing the
+                Fisher information matrix (stored for reference).
+        """
+        self.model = model
+        self.lambda_ewc = lambda_ewc
+        self.fisher_num_samples = fisher_num_samples
+
+        # {param_name: Fisher information tensor}
+        self.fisher: Dict[str, torch.Tensor] = {}
+        # {param_name: parameter snapshot tensor}
+        self.old_params: Dict[str, torch.Tensor] = {}
+
+    def compute_fisher(
+        self,
+        dataloader: Iterable,
+        vocab_size: int,
+        device: torch.device,
+        max_batches: int = 50,
+    ) -> None:
+        """
+        Compute the diagonal Fisher information matrix and snapshot the
+        current parameters as the reference point for the EWC penalty.
+
+        For each batch: forward pass -> per-sample NLL -> backprop ->
+        accumulate squared gradients.  Handles AMP autocast.
+
+        After calling this method, ``self.fisher`` and ``self.old_params``
+        are populated and ready for ``penalty()``.
+
+        Args:
+            dataloader: Iterable yielding ``(bx, by)`` batches.
+            vocab_size: Size of the vocabulary (for reshaping logits).
+            device: Device where the model lives.
+            max_batches: Maximum number of batches to use for estimation.
+        """
+        was_training = self.model.training
+        self.model.train()
+
+        # Ensure all parameters participate in gradient computation
+        prev_requires_grad: Dict[str, bool] = {}
+        for name, p in self.model.named_parameters():
+            prev_requires_grad[name] = p.requires_grad
+            p.requires_grad_(True)
+
+        # Move model to CPU to avoid GPU OOM during backward
+        original_device = device
+        self.model.to("cpu")
+        cpu_device = torch.device("cpu")
+
+        fisher: Dict[str, torch.Tensor] = {}
+
+        batch_count = 0
+        for batch_idx, (bx, by) in enumerate(dataloader):
+            if batch_idx >= max_batches:
+                break
+
+            bx = bx.to(cpu_device, non_blocking=True)
+            by = by.to(cpu_device, non_blocking=True)
+
+            self.model.zero_grad(set_to_none=True)
+
+            # No AMP on CPU — plain float32 forward + backward
+            logits, _recon_loss, _ = self.model(bx)
+            shift_logits = logits[:, :-1, :].contiguous()
+            shift_labels = by[:, 1:].contiguous()
+            loss = F.cross_entropy(
+                shift_logits.view(-1, vocab_size),
+                shift_labels.view(-1),
+            )
+
+            loss.backward()
+
+            for name, p in self.model.named_parameters():
+                if p.grad is not None:
+                    if name not in fisher:
+                        fisher[name] = torch.zeros_like(p.data)
+                    fisher[name] += p.grad.data.pow(2)
+
+            self.model.zero_grad(set_to_none=True)
+            batch_count += 1
+
+        # Average over batches
+        if batch_count > 0:
+            for name in fisher:
+                fisher[name] /= batch_count
+
+        # Move model back to original device
+        self.model.to(original_device)
+
+        # Move fisher tensors to original device for penalty computation
+        self.fisher = {k: v.to(original_device) for k, v in fisher.items()}
+
+        # Snapshot current parameters as the reference point
+        self.old_params = {}
+        for name, p in self.model.named_parameters():
+            if name in self.fisher:
+                self.old_params[name] = p.data.detach().clone()
+
+        # Restore requires_grad flags
+        for name, p in self.model.named_parameters():
+            if name in prev_requires_grad:
+                p.requires_grad_(prev_requires_grad[name])
+
+        # Restore training mode
+        if not was_training:
+            self.model.eval()
+
+        logger.info(
+            f"EWC fisher computed over {batch_count} batches, "
+            f"{len(self.fisher)} parameters tracked."
+        )
+
+    def penalty(self, model: nn.Module) -> torch.Tensor:
+        """
+        Compute the EWC regularization loss.
+
+        Returns:
+            Scalar tensor: lambda * sum_i F_i * (theta_i - theta_opt_i)^2
+        """
+        if not self.fisher or not self.old_params:
+            return torch.tensor(0.0, device=next(model.parameters()).device)
+
+        loss = torch.tensor(0.0, device=next(model.parameters()).device)
+        for name, p in model.named_parameters():
+            if name in self.fisher and name in self.old_params:
+                loss = loss + (
+                    self.fisher[name] * (p - self.old_params[name]).pow(2)
+                ).sum()
+        return self.lambda_ewc * loss
+
+    def save_state(self, path: str) -> None:
+        """
+        Persist the Fisher information matrix and old parameter snapshots
+        to a single ``.pt`` file.
+
+        Args:
+            path: Destination file path (should end in ``.pt``).
+        """
+        state = {
+            "fisher": self.fisher,
+            "old_params": self.old_params,
+            "lambda_ewc": self.lambda_ewc,
+            "fisher_num_samples": self.fisher_num_samples,
+        }
+        torch.save(state, path)
+        logger.info(f"EWC state saved to {path}")
+
+    def load_state(self, path: str) -> None:
+        """
+        Load a previously saved EWC state.
+
+        Args:
+            path: Path to the ``.pt`` file created by ``save_state``.
+        """
+        state = torch.load(path, map_location="cpu", weights_only=True)
+        self.fisher = state["fisher"]
+        self.old_params = state["old_params"]
+        self.lambda_ewc = state.get("lambda_ewc", self.lambda_ewc)
+        self.fisher_num_samples = state.get(
+            "fisher_num_samples", self.fisher_num_samples
+        )
+        logger.info(
+            f"EWC state loaded from {path}: "
+            f"{len(self.fisher)} params, lambda={self.lambda_ewc}"
+        )
+
+
+class ReplayBuffer:
+    """
+    Experience Replay buffer for interleaving data from previous curriculum
+    tiers during training.
+
+    Maintains references to the training dataloaders of previously completed
+    tiers and mixes their batches into the current tier's training loop at a
+    configurable ratio.
+
+    Usage:
+        replay = ReplayBuffer(max_tiers=3, replay_ratio=0.25)
+        replay.add_tier(0, tier0_train_dl)
+        replay.add_tier(1, tier1_train_dl)
+        # During tier-2 training, for each step:
+        bx, by = replay.sample_batch(current_batch_size=4)
+        # bx, by contains replay data from tiers 0,1
+    """
+
+    def __init__(self, max_tiers: int = 3, replay_ratio: float = 0.25):
+        """
+        Args:
+            max_tiers: Maximum number of previous tiers to remember.
+                When exceeded, the oldest tier is dropped.
+            replay_ratio: Fraction of the returned batch that comes from
+                replayed tiers (0.0 = no replay, 1.0 = all replay).
+        """
+        self.max_tiers = max_tiers
+        self.replay_ratio = replay_ratio
+
+        # List of (tier_index, dataloader) tuples
+        self._tiers: List[Tuple[int, Iterable]] = []
+
+        # Cyclic iteration state per tier: stores the iterator and whether
+        # it has been exhausted (we reset it on exhaustion).
+        self._iterators: Dict[int, object] = {}
+
+    def add_tier(self, tier_index: int, dataloader: Iterable) -> None:
+        """
+        Register a tier's training dataloader for future replay.
+
+        If the buffer is full (``max_tiers`` tiers stored), the oldest
+        entry is evicted.
+
+        Args:
+            tier_index: Integer tier identifier (e.g. 0, 1, 2, ...).
+            dataloader: The training dataloader for that tier.
+        """
+        # Evict oldest if at capacity
+        if len(self._tiers) >= self.max_tiers:
+            evicted_idx, _ = self._tiers.pop(0)
+            self._iterators.pop(evicted_idx, None)
+            logger.debug(f"ReplayBuffer: evicted tier {evicted_idx}")
+
+        self._tiers.append((tier_index, dataloader))
+        logger.debug(
+            f"ReplayBuffer: added tier {tier_index}, "
+            f"total stored tiers = {len(self._tiers)}"
+        )
+
+    def _get_next_batch(
+        self, tier_index: int, dataloader: Iterable
+    ) -> Optional[Tuple[torch.Tensor, torch.Tensor]]:
+        """
+        Get the next batch from a tier's dataloader, cycling if necessary.
+
+        Args:
+            tier_index: The tier identifier.
+            dataloader: The tier's dataloader.
+
+        Returns:
+            ``(bx, by)`` tensors or ``None`` if the dataloader is empty.
+        """
+        if tier_index not in self._iterators:
+            self._iterators[tier_index] = iter(dataloader)
+
+        try:
+            batch = next(self._iterators[tier_index])
+        except StopIteration:
+            # Reset iterator and try once more
+            self._iterators[tier_index] = iter(dataloader)
+            try:
+                batch = next(self._iterators[tier_index])
+            except StopIteration:
+                return None
+
+        if isinstance(batch, (list, tuple)) and len(batch) >= 2:
+            return batch[0], batch[1]
+        return None
+
+    def sample_batch(
+        self, current_batch_size: int
+    ) -> Optional[Tuple[torch.Tensor, torch.Tensor]]:
+        """
+        Sample a mixed replay batch from ALL stored previous tiers.
+
+        The returned batch has size approximately
+        ``replay_ratio * current_batch_size`` and is drawn uniformly from
+        all stored tiers, cycling through their dataloaders.
+
+        The caller is responsible for concatenating this with the current
+        tier's batch::
+
+            current_batch = next(current_dataloader)
+            replay = replay_buffer.sample_batch(current_batch_size)
+            if replay is not None:
+                bx = torch.cat([current_batch[0], replay[0]], dim=0)
+                by = torch.cat([current_batch[1], replay[1]], dim=0)
+
+        Args:
+            current_batch_size: The batch size of the current tier.
+                Used to compute the target replay batch size as
+                ``max(1, int(replay_ratio * current_batch_size))``.
+
+        Returns:
+            ``(bx, by)`` tensors on the same device as the dataloader
+            yields them, or ``None`` if no replay tiers are available.
+        """
+        if not self._tiers:
+            return None
+
+        target_replay = max(1, int(self.replay_ratio * current_batch_size))
+        collected_x: List[torch.Tensor] = []
+        collected_y: List[torch.Tensor] = []
+
+        # Round-robin through all stored tiers until we reach target size
+        attempts = 0
+        while len(collected_x) < target_replay and attempts < len(self._tiers) * 3:
+            tier_idx, dl = self._tiers[attempts % len(self._tiers)]
+            batch = self._get_next_batch(tier_idx, dl)
+            if batch is not None:
+                bx, by = batch
+                collected_x.append(bx)
+                collected_y.append(by)
+            attempts += 1
+
+        if not collected_x:
+            return None
+
+        # Concatenate and trim/pad to target size
+        all_x = torch.cat(collected_x, dim=0)
+        all_y = torch.cat(collected_y, dim=0)
+
+        if all_x.shape[0] > target_replay:
+            all_x = all_x[:target_replay]
+            all_y = all_y[:target_replay]
+
+        return all_x, all_y

topogpt3/exploitgym_config.py

@@ -0,0 +1,130 @@
+#!/usr/bin/env python3
+"""
+TopoExploit Training Configuration
+
+122M parameter model trained on ExploitGym vulnerability/exploit data.
+
+Tier structure:
+  Tier 1: vulnerability_analysis (256 seq, 3 epochs)
+  Tier 2: patch_analysis        (512 seq, 4 epochs)
+  Tier 3: exploit_development   (1024 seq, 5 epochs)
+
+Author: Gris Iscomeback
+License: GPL v3
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Optional, Tuple
+
+from .model import TopoGPT2Config
+from .exploitgym_loader import ExploitGymLoaderConfig
+
+
+@dataclass
+class TopoExploitConfig:
+    """Configuration for TopoExploit: 122M params on ExploitGym data."""
+
+    # --- Model: 122M parameters ---
+    SCALE: str = "large"
+    DEVICE: str = "cuda"
+    RANDOM_SEED: int = 42
+    USE_AMP: bool = True
+    GRADIENT_CHECKPOINTING: bool = True
+
+    # --- Training hyperparams (tuned for 122M) ---
+    BATCH_SIZE: int = 2
+    GRAD_ACCUM_STEPS: int = 16
+    LEARNING_RATE: float = 2e-4
+    WEIGHT_DECAY: float = 0.1
+    GRADIENT_CLIP_NORM: float = 1.0
+    WARMUP_RATIO: float = 0.05
+
+    # --- Regularizacion (anti-overfitting) ---
+    LABEL_SMOOTHING: float = 0.05
+    EARLY_STOP_PATIENCE: int = 1
+    EARLY_STOP_MIN_DELTA: float = 0.0
+    BEST_CKPT_RESTORE: bool = True
+
+    # --- Curriculum: 3 tiers over ExploitGym ---
+    # Tier 1 gets more epochs (weak tier, small data), Tier 2 capped (easy tier, huge data)
+    TIER_EPOCHS: Tuple[int, int, int] = (6, 3, 4)
+    TIER_SEQ_LEN: Tuple[int, int, int] = (256, 512, 768)
+    TIER_NAMES: Tuple[str, str, str] = (
+        "vulnerability_analysis",
+        "patch_analysis",
+        "exploit_development",
+    )
+
+    # --- Progressive window sizes ---
+    PROGRESSIVE_WINDOW: Tuple[int, int, int] = (128, 256, 256)
+
+    # --- Grassmannian tracking ---
+    GRASS_TRACK_EVERY: int = 200
+    GRASS_MAX_RANK: int = 16
+    GRASS_ELBOW_RATIO: float = 0.05
+    GRASS_FISHER_BATCHES: int = 40
+    GRASS_FISHER_GRADS: int = 64
+
+    # --- ExploitGym data ---
+    DATA_CACHE: str = "data_exploitgym"
+    REPO_CACHE: str = ".exploitgym_repo"
+    MAX_TOKENS_PER_TIER: Tuple[int, int, int] = (50_000_000, 80_000_000, 100_000_000)
+    VAL_FRACTION: float = 0.10
+    HOLDOUT_FRACTION: float = 0.10
+
+    # --- Checkpointing ---
+    CHECKPOINT_DIR: str = "checkpoints_topoexploit"
+    LOG_FILE: str = "topoexploit_history.jsonl"
+    LOG_INTERVAL_STEPS: int = 50
+    CHECKPOINT_EVERY_MIN: int = 15
+    MAX_CHECKPOINTS: int = 5
+
+    # --- Gauss optimization ---
+    USE_GAUSS_COMPLEX_MULT: bool = True
+
+    # --- Limits ---
+    MAX_TRAIN_BATCHES_PER_EPOCH: int = 200     # cap per-tier to balance 18:1 data ratio
+    MAX_EVAL_BATCHES: int = 200
+    EFFICIENCY_THROUGHPUT_BATCHES: int = 20
+
+    # --- Anti-forgetting: EWC + Replay ---
+    USE_EWC: bool = True
+    EWC_LAMBDA: float = 5000.0
+    EWC_FISHER_BATCHES: int = 50
+    USE_REPLAY: bool = True
+    REPLAY_RATIO: float = 0.25
+    REPLAY_MAX_TIERS: int = 3
+
+    def build_topogpt2_config(self, max_seq_len: int,
+                              attn_window: Optional[int] = None) -> TopoGPT2Config:
+        cfg = TopoGPT2Config(
+            DEVICE=self.DEVICE,
+            RANDOM_SEED=self.RANDOM_SEED,
+            USE_AMP=self.USE_AMP,
+            SCALE=self.SCALE,
+            BATCH_SIZE=self.BATCH_SIZE,
+            GRAD_ACCUM_STEPS=self.GRAD_ACCUM_STEPS,
+            LEARNING_RATE=self.LEARNING_RATE,
+            WEIGHT_DECAY=self.WEIGHT_DECAY,
+            GRADIENT_CLIP_NORM=self.GRADIENT_CLIP_NORM,
+            WARMUP_RATIO=self.WARMUP_RATIO,
+            GRADIENT_CHECKPOINTING=self.GRADIENT_CHECKPOINTING,
+            CHECKPOINT_DIR=self.CHECKPOINT_DIR,
+            DATA_DIR=self.DATA_CACHE,
+        )
+        cfg.MAX_SEQ_LEN = max_seq_len
+        if attn_window is not None and attn_window > 0:
+            cfg.ATTN_WINDOW = attn_window
+        return cfg
+
+    def build_loader_config(self) -> ExploitGymLoaderConfig:
+        return ExploitGymLoaderConfig(
+            data_cache=self.DATA_CACHE,
+            repo_cache=self.REPO_CACHE,
+            max_tokens_per_tier=self.MAX_TOKENS_PER_TIER,
+            val_fraction=self.VAL_FRACTION,
+            holdout_fraction=self.HOLDOUT_FRACTION,
+            split_seed=self.RANDOM_SEED,
+        )