tardits 0.2.0__py3-none-any.whl

tardits/gen.py

@@ -0,0 +1,150 @@
+from argparse import ArgumentParser
+import sys
+import time
+import os
+import yaml
+
+import torch
+from tardits.models import HyperFormer, RealGPT
+from tardits.modules import get_device
+
+
+def main():
+    # --- CLI Arguments ---
+    parser = ArgumentParser(
+        description="Generates text from a trained model using its YAML config"
+    )
+
+    parser.add_argument(
+        "--config",
+        "-c",
+        type=str,
+        default="examples/configs/complex.yaml",
+        help="Path to YAML configuration file",
+    )
+    parser.add_argument(
+        "--checkpoint",
+        type=str,
+        default=None,
+        help="Optional override for model checkpoint path",
+    )
+    parser.add_argument(
+        "--vocab",
+        type=str,
+        default=None,
+        help="Optional override for vocab path",
+    )
+    parser.add_argument(
+        "--prompt",
+        "-p",
+        type=str,
+        default="\n",
+        help="Text prompt to start generation",
+    )
+    parser.add_argument(
+        "--length",
+        "-l",
+        type=int,
+        default=500,
+        help="Number of tokens to generate",
+    )
+    parser.add_argument(
+        "--temp",
+        "-t",
+        type=float,
+        default=1.0,
+        help="Model sampling temperature",
+    )
+    parser.add_argument(
+        "--top_k",
+        "-k",
+        type=int,
+        default=10,
+        help="Top-k sampling threshold",
+    )
+    parser.add_argument(
+        "--follow",
+        "-f",
+        action="store_true",
+        help="Endless stream of tokens",
+    )
+    parser.add_argument(
+        "--delay",
+        type=float,
+        default=0.04,
+        help="Delay between tokens in follow mode (seconds)",
+    )
+
+    args = parser.parse_args()
+
+    # --- Read Config File ---
+    with open(args.config, "r", encoding="utf-8") as f:
+        cfg = yaml.safe_load(f)
+
+    if cfg["model"]["level"] < 1:
+        GPT = RealGPT
+    else:
+        GPT = HyperFormer
+
+    # --- Load Config ---
+    checkpoint_path = args.checkpoint or cfg["trainer"]["checkpoint"]
+    vocab_path = args.vocab or cfg.get("data", {}).get(
+        "vocab_save_path", "examples/saved/chars.txt"
+    )
+    device = get_device(cfg.get("device", "auto"))
+
+    # --- Load Weights ---
+    model = GPT.load_from_checkpoint(checkpoint_path, device=device)
+    model.eval()
+
+    # --- Load Vocabulary ---
+    with open(vocab_path, "r", encoding="utf-8") as ch:
+        chars = list(ch.read())
+
+    stoi = {ch: i for i, ch in enumerate(chars)}
+    itos = {i: ch for i, ch in enumerate(chars)}
+
+    def encode(si):
+        return [stoi[c] for c in si]
+
+    def decode(li):
+        return "".join([itos[i] for i in li])
+
+    prompt = args.prompt.replace("\\n", "\n")
+    idx = torch.tensor([encode(prompt)], dtype=torch.long).to(device)
+
+    # --- Generation ---
+    if args.follow:
+        # 🧹 Clear terminal screen
+        print("\033[2J\033[H", end="")
+        print(
+            f"⏳ Token-Stream (complexity-lvl: {cfg['model']['level']} @ {checkpoint_path}) — Ctrl+C to exit\n"
+        )
+        try:
+            sys.stdout.write(prompt)
+            sys.stdout.flush()
+
+            for token_id in model.generate_stream(
+                idx, temp=args.temp, top_k=args.top_k
+            ):
+                char = itos[int(token_id)]
+                sys.stdout.write(char)
+                sys.stdout.flush()
+
+                if args.delay > 0:
+                    time.sleep(args.delay)
+        except KeyboardInterrupt:
+            print("\n\n🛑 Token stream stopped.")
+
+    else:
+        print(
+            decode(
+                model.generate(
+                    idx, max_new_tokens=args.length, temp=args.temp, top_k=args.top_k
+                )[0].tolist()
+            )
+        )
+
+
+if __name__ == "__main__":
+    main()

tardits/models/__init__.py

@@ -0,0 +1,4 @@
+from .hyper import HyperFormer
+from .real import RealGPT
+
+__all__ = ["HyperFormer", "RealGPT"]

tardits/models/hyper.py

@@ -0,0 +1,144 @@
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+
+from tardits.modules import (
+    HyperEmbedding,
+    HyperLMHead,
+    HyperRMSNorm,
+    SAIL,
+    GyRoPE,
+    HyperLinear,
+    LMCoreMixin
+)
+from tardits.modules.conf import ModelConfig
+
+
+class HyperSelfAttention(nn.Module):
+    """
+    Universal multi-headed self-attention for 2^level hypercomplex spaces.
+    """
+
+    def __init__(self, config: ModelConfig):
+        super().__init__()
+        self.config = config
+        self.level = config.level
+        self.head_size = config.n_embd // config.n_head
+        self.n_head = config.n_head
+        self.n_embd = config.n_embd
+        self.hyper_dim = 2**config.level
+
+        # Projections
+        self.q_proj = HyperLinear(self.n_embd, self.n_embd, level=self.level)
+        self.k_proj = HyperLinear(self.n_embd, self.n_embd, level=self.level)
+        self.v_proj = HyperLinear(self.n_embd, self.n_embd, level=self.level)
+        self.out_proj = HyperLinear(self.n_embd, self.n_embd, level=self.level)
+
+        self.rope = GyRoPE(
+            self.head_size, config.block_size, self.n_head, level=self.level
+        )
+
+        self.register_buffer(
+            "tril", torch.tril(torch.ones(config.block_size, config.block_size))
+        )
+        self.tril: torch.Tensor
+
+        self.dropout = nn.Dropout(config.dropout)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # x Shape: [B, T, C, hyper_dim]
+        B, T, C, hyper_dim = x.shape
+
+        q = self.q_proj(x)
+        k = self.k_proj(x)
+        v = self.v_proj(x)
+
+        q = q.view(B, T, self.n_head, self.head_size, hyper_dim).transpose(1, 2)
+        k = k.view(B, T, self.n_head, self.head_size, hyper_dim).transpose(1, 2)
+        v = v.view(B, T, self.n_head, self.head_size, hyper_dim).transpose(1, 2)
+
+        q = self.rope(q)
+        k = self.rope(k)
+
+        # [B, nh, T, hs, hyper_dim] -> [B, nh, T, hs * hyper_dim]
+        q_flat = q.flatten(start_dim=-2)
+        k_flat = k.flatten(start_dim=-2)
+        v_flat = v.flatten(start_dim=-2)
+
+        wei = q_flat @ k_flat.adjoint() * (self.hyper_dim * self.head_size) ** -0.5
+        wei = wei.real
+        wei = wei.masked_fill(self.tril[:T, :T] == 0, float("-inf"))
+        wei = F.softmax(wei, dim=-1)
+        wei = self.dropout(wei)
+
+        wei_complex = wei.to(torch.complex64)
+
+        out_flat = wei_complex @ v_flat  # -> [B, nh, T, hs * hyper_dim]
+
+        # 6. Un-Flatten und Recombine: [B, nh, T, hs * hyper_dim] -> [B, nh, T, hs, hyper_dim]
+        out = out_flat.view(B, self.n_head, T, self.head_size, hyper_dim)
+        
+        # -> [B, T, C, hyper_dim]
+        out = out.transpose(1, 2).contiguous().view(B, T, C, hyper_dim)
+
+        out = self.out_proj(out)
+
+        return out
+
+
+   
+class HyperBlock(nn.Module):
+    """Transformer block: Attention + SAIL with Residuals & RMSNorm"""
+
+    def __init__(self, config: ModelConfig):
+        super().__init__()
+        self.sa = HyperSelfAttention(config)
+        self.sail = SAIL(dim=config.n_embd, level=config.level)
+        self.ln1 = HyperRMSNorm(config.n_embd, level=config.level)
+        self.ln2 = HyperRMSNorm(config.n_embd, level=config.level)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # Attention + Residual
+        x = x + self.ln1(self.sa(x))
+        # FeedForward (SAIL) + Residual
+        x = x + self.ln2(self.sail(x))
+        return x
+
+
+class HyperFormer(nn.Module, LMCoreMixin):
+    """
+    Universal Activation-Free Hypercomplex Transformer (v0.2.0).
+    Dynamically scaled via `config.level` (1=Complex, 2=Quaternion, 3=Octonion, 4=Sedenion...).
+    """
+
+    def __init__(self, config: ModelConfig):
+        super().__init__()
+        self.config = config
+        self.level = config.level
+
+        self.token_embedding = HyperEmbedding(
+            config.vocab_size, config.n_embd, level=self.level
+        )
+        self.blocks = nn.ModuleList([HyperBlock(config) for _ in range(config.n_layer)])
+        self.ln_f = HyperRMSNorm(config.n_embd, level=self.level)
+        self.lm_head = HyperLMHead(config.n_embd, config.vocab_size, level=self.level)
+
+    def forward(self, idx: torch.Tensor, targets: torch.Tensor | None = None):
+        B, T = idx.shape
+        x = self.token_embedding(idx)
+
+        for block in self.blocks:
+            x = block(x)
+
+        x = self.ln_f(x)
+        logits = self.lm_head(x)
+
+        if targets is None:
+            loss = None
+        else:
+            B, T, C = logits.shape
+            logits_flat = logits.view(B * T, C)
+            targets_flat = targets.view(B * T)
+            loss = F.cross_entropy(logits_flat, targets_flat)
+
+        return logits, loss

tardits/models/real.py

@@ -0,0 +1,167 @@
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+
+from tardits.modules import LMCoreMixin
+from tardits.modules.conf import ModelConfig
+
+
+class RotaryPositionalEmbedding(nn.Module):
+    """🌀 Rotates query/key pairwise in the complex plane based on position."""
+
+    def __init__(self, head_size, block_size):
+        super().__init__()
+        self.head_size = head_size
+
+        inv_freq = 1.0 / (
+            10000.0 ** (torch.arange(0, head_size, 2).float() / head_size)
+        )
+        t = torch.arange(block_size, dtype=torch.float32)
+        freqs = torch.outer(t, inv_freq)
+        self.emb = torch.cat((freqs, freqs), dim=-1)
+
+        self.register_buffer("cos_cached", self.emb.cos(), persistent=False)
+        self.register_buffer("sin_cached", self.emb.sin(), persistent=False)
+        self.cos_cached: torch.Tensor
+        self.sin_cached: torch.Tensor
+
+    def _rotate_half(self, x):
+        x1 = x[..., : self.head_size // 2]
+        x2 = x[..., self.head_size // 2 :]
+        return torch.cat((-x2, x1), dim=-1)
+
+    def forward(self, x):
+        # x.shape: (B, nh, T, hs)
+        T = x.shape[2]
+        cos = self.cos_cached[:T, :].unsqueeze(0).unsqueeze(1)  # (1, 1, T, hs)
+        sin = self.sin_cached[:T, :].unsqueeze(0).unsqueeze(1)  # (1, 1, T, hs)
+
+        return (x * cos) + (self._rotate_half(x) * sin)
+
+
+RoPE = RotaryPositionalEmbedding
+
+
+class SwiGLU(nn.Module):
+    """🚀 SwiGLU Feed-Forward Network."""
+
+    def __init__(self, n_embd, hidden_dim, dropout):
+        super().__init__()
+        self.w = nn.Linear(n_embd, hidden_dim, bias=False)
+        self.v = nn.Linear(n_embd, hidden_dim, bias=False)
+        self.out_proj = nn.Linear(hidden_dim, n_embd, bias=False)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x):
+        gate = F.silu(self.w(x))
+        activated = gate * self.v(x)
+        out = self.out_proj(activated)
+
+        return self.dropout(out)
+
+
+class CasualSelfAttention(nn.Module):
+    """casual (and causal) multi-headed self-attention"""
+
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.head_size = config.n_embd // config.n_head
+        self.n_head = config.n_head
+        self.n_embd = config.n_embd
+
+        self.c_attn = nn.Linear(self.n_embd, 3 * self.n_embd, bias=False)
+
+        self.rope = RoPE(self.head_size, config.block_size)
+
+        self.register_buffer(
+            "tril", torch.tril(torch.ones(config.block_size, config.block_size))
+        )
+        self.tril: torch.Tensor
+
+        self.dropout = nn.Dropout(config.dropout)
+
+        self.proj = nn.Linear(config.n_embd, config.n_embd)
+        self.proj_dropout = nn.Dropout(config.dropout)
+
+    def forward(self, x):
+        B, T, C = x.shape
+
+        qkv = self.c_attn(x)  # (B,T,3*C)
+        q, k, v = qkv.split(C, dim=2)  # (B,T,C)
+
+        # view(B,T,nh,hs): (B,T,C) == (B,T,nh*hs) --> (B,T,nh,hs)
+        # transpose(1, 2): (B,T,nh,hs) --> (B,nh,T,hs)
+        q = q.view(B, T, self.n_head, self.head_size).transpose(1, 2)
+        k = k.view(B, T, self.n_head, self.head_size).transpose(1, 2)
+        v = v.view(B, T, self.n_head, self.head_size).transpose(1, 2)
+
+        # fancy RotaryPositionalEmbedding
+        q = self.rope(q)
+        k = self.rope(k)
+
+        wei = q @ k.transpose(-2, -1) * (self.head_size**-0.5)  # (B,nh,T,T)
+        wei = wei.masked_fill(self.tril[:T, :T] == 0, float("-inf"))  # (B,nh,T,T)
+        wei = F.softmax(wei, dim=-1)  # (B,nh,T,T)
+        wei = self.dropout(wei)  # (B,nh,T,T)
+
+        # weighted aggregation of value
+        out = wei @ v  # (B,nh,T,T) @ (B,nh,T,hs) --> (B,nh,T,hs)
+
+        out = out.transpose(1, 2).reshape(B, T, C)  # -> (B,T,nh,hs) -> (B,T,C)
+
+        out = self.proj(out)
+        out = self.proj_dropout(out)
+
+        return out  # (B,T,C)
+
+
+class Block(nn.Module):
+    """Transformer block: communication followed by computation"""
+
+    def __init__(self, config):
+        # n_embd: embedding dimension, n_head: number of heads
+        super().__init__()
+        self.sa = CasualSelfAttention(config)
+        self.ffwd = SwiGLU(config.n_embd, int(config.n_embd * (8 / 3)), config.dropout)
+        self.ln1 = nn.RMSNorm(config.n_embd)
+        self.ln2 = nn.RMSNorm(config.n_embd)
+
+    def forward(self, x):
+        x = x + self.ln1(self.sa(x))  # x + because "residual connections"...
+        x = x + self.ln2(self.ffwd(x))
+        return x
+
+
+class RealGPT(nn.Module, LMCoreMixin):
+    def __init__(
+        self,
+        config: ModelConfig,
+    ):
+        super().__init__()
+
+        self.config = config
+        self.token_embedding_table = nn.Embedding(config.vocab_size, config.n_embd)
+
+        self.blocks = nn.Sequential(*[Block(config) for _ in range(config.n_layer)])
+
+        self.ln_f = nn.RMSNorm(config.n_embd)
+        self.lm_head = nn.Linear(config.n_embd, config.vocab_size)
+
+    def forward(self, idx, targets=None):
+        B, T = idx.shape
+        x = self.token_embedding_table(idx)  # (B,T,C)
+        x = self.blocks(x)
+        x = self.ln_f(x)
+        logits = self.lm_head(x)
+
+        if targets is None:
+            loss = None
+        else:
+            B, T, C = logits.shape
+            logits = logits.view(B * T, C)
+            targets = targets.view(B * T)
+
+            loss = F.cross_entropy(logits, targets)
+
+        return logits, loss

tardits/modules/__init__.py

@@ -0,0 +1,27 @@
+from .trainer import Trainer
+from .base import get_device
+from .conf import ModelConfig, TrainerConfig, MLFlowConfig
+from .core import (
+    SAIL,
+    GyRoPE,
+    HyperRMSNorm,
+    HyperEmbedding,
+    HyperLMHead,
+    HyperLinear,
+)
+from .base import LMCoreMixin
+
+__all__ = [
+    "Trainer",
+    "ModelConfig",
+    "TrainerConfig",
+    "MLFlowConfig",
+    "SAIL",
+    "GyRoPE",
+    "LMCoreMixin"
+    "HyperRMSNorm",
+    "HyperEmbedding",
+    "HyperLMHead",
+    "HyperLinear",
+    "get_device"
+]

tardits/modules/base.py

@@ -0,0 +1,99 @@
+# pyright: reportAttributeAccessIssue=false
+import torch
+from torch.nn import functional as F
+import os
+
+
+def get_device(requested_device: str = "auto") -> torch.device:
+    if requested_device == "auto" or not requested_device:
+        if torch.cuda.is_available():
+            return torch.device("cuda")
+        elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
+            return torch.device("mps")  # Apple Silicon GPU
+        return torch.device("cpu")
+    return torch.device(requested_device)
+
+
+class LMCoreMixin:
+    @torch.no_grad()
+    def generate(self, idx, max_new_tokens, temp=1.0, top_k=10):
+        device = next(self.parameters()).device
+        idx = idx.to(device)
+        # idx is (B,T) array of indices in the current context
+        for _ in range(max_new_tokens):
+            # crop idx to the last block_size max_new_tokens
+            idx_cond = idx[:, -self.config.block_size :]
+            # get predictions
+            logits, _ = self.forward(idx_cond)
+            # focus only on the last timestep
+            logits = logits[:, -1, :]  # becomes (B,C)
+            # temp
+            logits = logits / temp
+            # top_k
+            v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
+            logits[logits < v[:, [-1]]] = -float("Inf")
+            # apply softmax to get probabilities
+            probs = F.softmax(logits, dim=-1)  # (B,C)
+            # sample from distribution
+            idx_next = torch.multinomial(probs, num_samples=1)  # (B,1)
+            # append sampled index to the running sequence
+            idx = torch.cat((idx, idx_next), dim=1)  # (B,T+1)
+        return idx
+
+    @torch.no_grad()
+    def generate_stream(self, idx, temp=1.0, top_k=10):
+        while True:
+            # slicing:
+            idx_cond = idx[:, -self.config.block_size :]
+            logits, _ = self.forward(idx_cond)
+            logits = logits[:, -1, :]  # becomes (B,C)
+
+            # temp/top_k
+            logits = logits / temp
+            v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
+            logits[logits < v[:, [-1]]] = -float("Inf")
+
+            # sampling
+            probs = F.softmax(logits, dim=-1)  # (B,C)
+            idx_next = torch.multinomial(probs, num_samples=1)  # (B,1)
+
+            # append
+            idx = torch.cat((idx, idx_next), dim=1)  # (B,T+1)
+
+            # yield single token
+            yield idx_next.item()
+
+    def save_checkpoint(self, filepath="checkpoint.pt"):
+        """saves model weights and config to disk"""
+        dirname = os.path.dirname(filepath)
+        if dirname:
+            os.makedirs(dirname, exist_ok=True)
+        checkpoint = {"model_state_dict": self.state_dict(), "config": self.config}
+        torch.save(checkpoint, filepath)
+        # print(f"💾 checkpoint successfully saved: {filepath}")
+
+    def get_num_params(self):
+        """returns number of trainable weights"""
+        total_params = 0
+        for p in self.parameters():
+            if p.requires_grad:
+                if p.is_complex():
+                    total_params += p.numel() * 2
+                else:
+                    total_params += p.numel()
+        return total_params
+
+    @classmethod
+    def load_from_checkpoint(
+        cls, filepath="checkpoint.pt", device: str | torch.device = "cpu"
+    ):
+        """loads model state from checkpoint file"""
+        if not os.path.exists(filepath):
+            raise FileNotFoundError(f"❌ file not found:\n{filepath}")
+
+        checkpoint = torch.load(filepath, map_location=device, weights_only=False)
+        model = cls(checkpoint["config"])  # type: ignore
+        model.load_state_dict(checkpoint["model_state_dict"])
+        model.to(device)
+        # print(f"🚀 checkpoint successfully restored from:\n'{filepath}'")
+        return model

tardits/modules/conf.py

@@ -0,0 +1,31 @@
+from dataclasses import dataclass
+
+
+@dataclass
+class ModelConfig:
+    level: int = 1
+    block_size: int = 8  # maximum context length
+    vocab_size: int = 74
+    n_head: int = 4
+    n_embd: int = 32
+    n_layer: int = 3
+    dropout: float = 0.2
+
+
+@dataclass
+class TrainerConfig:
+    batch_size: int = 32  # how many independent sequences processed in parallel?
+    max_iters: int = 5000
+    eval_interval: int = 500
+    learning_rate: float = 3e-4
+    eval_iters: int = 200
+    checkpoint: str = ""
+    save_interval: int = 0
+
+
+@dataclass
+class MLFlowConfig:
+    experiment: str = "sail_experiment"
+    log_interval: int = 10
+    deep_log_interval: int = 100
+    system_metrics: bool = False