transformer-toolkit 0.1.0__tar.gz

transformer_toolkit-0.1.0/PKG-INFO

@@ -0,0 +1,28 @@
+Metadata-Version: 2.4
+Name: transformer-toolkit
+Version: 0.1.0
+Summary: Minimal, modular transformer library for training your own LLM
+Author: Govind Barbade
+Project-URL: Homepage, https://github.com/govindbarbade/transformer-toolkit
+Project-URL: Repository, https://github.com/govindbarbade/transformer-toolkit
+Keywords: transformer,llm,deep learning,nlp,pytorch
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+Requires-Dist: torch>=2.0.0
+Requires-Dist: pydantic>=2.0.0
+Provides-Extra: tokenizers
+Requires-Dist: tokenizers>=0.15.0; extra == "tokenizers"
+Provides-Extra: hf
+Requires-Dist: transformers>=4.35.0; extra == "hf"
+Requires-Dist: huggingface_hub>=0.20.0; extra == "hf"
+Requires-Dist: datasets>=2.14.0; extra == "hf"
+Provides-Extra: all
+Requires-Dist: transformer-toolkit[hf,tokenizers]; extra == "all"

transformer_toolkit-0.1.0/pyproject.toml

@@ -0,0 +1,43 @@
+[build-system]
+requires      = ["setuptools>=68", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name            = "transformer-toolkit"
+version         = "0.1.0"
+description     = "Minimal, modular transformer library for training your own LLM"
+readme          = "README.md"
+license         = { file = "LICENSE" }
+requires-python = ">=3.10"
+authors         = [{ name = "Govind Barbade" }]
+keywords        = ["transformer", "llm", "deep learning", "nlp", "pytorch"]
+classifiers     = [
+    "Development Status :: 3 - Alpha",
+    "Intended Audience :: Developers",
+    "Intended Audience :: Science/Research",
+    "License :: OSI Approved :: MIT License",
+    "Programming Language :: Python :: 3.10",
+    "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
+    "Topic :: Scientific/Engineering :: Artificial Intelligence",
+]
+
+dependencies = [
+    "torch>=2.0.0",
+    "pydantic>=2.0.0",
+]
+
+[project.optional-dependencies]
+tokenizers = ["tokenizers>=0.15.0"]
+hf         = ["transformers>=4.35.0",
+              "huggingface_hub>=0.20.0",
+              "datasets>=2.14.0"]
+all        = ["transformer-toolkit[tokenizers,hf]"]
+
+[project.urls]
+Homepage   = "https://github.com/govindbarbade/transformer-toolkit"
+Repository = "https://github.com/govindbarbade/transformer-toolkit"
+
+[tool.setuptools.packages.find]
+where   = ["."]
+include = ["transformer_toolkit*"]

transformer_toolkit-0.1.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

transformer_toolkit-0.1.0/transformer_toolkit/__init__.py

@@ -0,0 +1,16 @@
+from .model      import Transformer, TransformerConfig
+from .trainer    import Trainer, TrainConfig
+from .dataloader import DataConfig, from_files, from_binary, from_hf, from_strings
+from .c_tokenizers import RustBPETokenizer, HFTokenizer, ByteLevelTokenizer, CharLevelTokenizer
+from .hf_hub     import login, push_to_hub, pull_from_hub
+
+__version__ = "0.1.0"
+__author__  = "Govind Barbade"
+
+__all__ = [
+    "Transformer", "TransformerConfig",
+    "Trainer", "TrainConfig",
+    "DataConfig", "from_files", "from_binary", "from_hf", "from_strings",
+    "RustBPETokenizer", "HFTokenizer", "ByteLevelTokenizer", "CharLevelTokenizer",
+    "login", "push_to_hub", "pull_from_hub",
+]

transformer_toolkit-0.1.0/transformer_toolkit/attention.py

@@ -0,0 +1,131 @@
+import torch.nn as nn
+import torch.nn.functional as F
+import math
+
+class MultiHeadAttention(nn.Module):
+    """Classic MHA. Used in original Transformer, BERT, GPT-2."""
+    def __init__(self, dim: int, n_heads: int):
+        super().__init__()
+        self.n_heads = n_heads
+        self.head_dim = dim // n_heads
+        self.qkv = nn.Linear(dim, 3 * dim, bias=False)
+        self.out = nn.Linear(dim, dim, bias=False)
+
+    def forward(self, x, mask=None):
+        B, T, C = x.shape
+        q, k, v = self.qkv(x).split(C, dim=-1)
+        def split(t): return t.view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
+        q, k, v = split(q), split(k), split(v)
+        scores = (q @ k.transpose(-2, -1)) / math.sqrt(self.head_dim)
+        if mask is not None: scores = scores.masked_fill(mask == 0, float('-inf'))
+        return self.out(F.softmax(scores, -1) @ v).transpose(1, 2).reshape(B, T, C) # noqa
+
+
+class GroupedQueryAttention(nn.Module):
+    """Fewer k/v heads than q heads. Used in LLaMA 3, Mistral."""
+    def __init__(self, dim: int, n_heads: int, n_kv_heads: int):
+        super().__init__()
+        self.n_heads    = n_heads
+        self.n_kv_heads = n_kv_heads
+        self.head_dim   = dim // n_heads
+        self.q  = nn.Linear(dim, dim, bias=False)
+        self.kv = nn.Linear(dim, 2 * n_kv_heads * self.head_dim, bias=False)
+        self.out = nn.Linear(dim, dim, bias=False)
+
+    def forward(self, x, mask=None):
+        B, T, C = x.shape
+        q = self.q(x).view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
+        k, v = self.kv(x).split(self.n_kv_heads * self.head_dim, dim=-1)
+        def split_kv(t): return t.view(B, T, self.n_kv_heads, self.head_dim).transpose(1, 2)
+        k, v = split_kv(k), split_kv(v)
+        # repeat k/v to match q heads
+        r = self.n_heads // self.n_kv_heads
+        k = k.repeat_interleave(r, dim=1)
+        v = v.repeat_interleave(r, dim=1)
+        scores = (q @ k.transpose(-2, -1)) / math.sqrt(self.head_dim)
+        if mask is not None: scores = scores.masked_fill(mask == 0, float('-inf'))
+        return self.out((F.softmax(scores, -1) @ v).transpose(1, 2).reshape(B, T, C))
+
+
+class MultiQueryAttention(nn.Module):
+    """Single k/v head shared across all q heads. Used in Falcon, early Gemini."""
+    def __init__(self, dim: int, n_heads: int):
+        super().__init__()
+        self.n_heads  = n_heads
+        self.head_dim = dim // n_heads
+        self.q  = nn.Linear(dim, dim, bias=False)
+        self.k  = nn.Linear(dim, self.head_dim, bias=False)
+        self.v  = nn.Linear(dim, self.head_dim, bias=False)
+        self.out = nn.Linear(dim, dim, bias=False)
+
+    def forward(self, x, mask=None):
+        B, T, C = x.shape
+        q = self.q(x).view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
+        k = self.k(x).view(B, T, 1, self.head_dim).transpose(1, 2).expand(B, self.n_heads, T, self.head_dim)
+        v = self.v(x).view(B, T, 1, self.head_dim).transpose(1, 2).expand(B, self.n_heads, T, self.head_dim)
+        scores = (q @ k.transpose(-2, -1)) / math.sqrt(self.head_dim)
+        if mask is not None: scores = scores.masked_fill(mask == 0, float('-inf'))
+        return self.out((F.softmax(scores, -1) @ v).transpose(1, 2).reshape(B, T, C))
+
+class FlashAttention(nn.Module):
+    """
+    Flash Attention — same result as MHA, far less memory.
+    Uses torch's built-in scaled_dot_product_attention (torch >= 2.0)
+    which calls the CUDA kernel automatically when available.
+    """
+    def __init__(self, dim: int, n_heads: int):
+        super().__init__()
+        self.n_heads  = n_heads
+        self.head_dim = dim // n_heads
+        self.qkv = nn.Linear(dim, 3 * dim, bias=False)
+        self.out = nn.Linear(dim, dim, bias=False)
+
+    def forward(self, x, mask=None, causal=True):
+        B, T, C = x.shape
+        q, k, v = self.qkv(x).split(C, dim=-1)
+        def split(t): return t.view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
+        q, k, v = split(q), split(k), split(v)
+
+        # one line — PyTorch handles the tiling/chunking kernel
+        x = nn.functional.scaled_dot_product_attention(
+            q, k, v,
+            attn_mask=mask,
+            is_causal=causal,       # builds causal mask internally, no extra memory
+            dropout_p=0.0,
+        )
+        return self.out(x.transpose(1, 2).reshape(B, T, C))
+    
+
+class MLAttention(nn.Module):
+    """
+    Multi-Head Latent Attention (DeepSeek-V2/V3).
+    Compresses k/v into a small latent vector instead of caching full k/v.
+    Huge KV cache reduction at inference time.
+    """
+    def __init__(self, dim: int, n_heads: int, latent_dim: int):
+        super().__init__()
+        self.n_heads  = n_heads
+        self.head_dim = dim // n_heads
+        self.kv_down = nn.Linear(dim, latent_dim, bias=False)
+        self.k_up    = nn.Linear(latent_dim, dim, bias=False)
+        self.v_up    = nn.Linear(latent_dim, dim, bias=False)
+
+        self.q       = nn.Linear(dim, dim, bias=False)
+        self.out     = nn.Linear(dim, dim, bias=False)
+
+    def forward(self, x, mask=None):
+        B, T, C = x.shape
+
+        q = self.q(x)
+
+        # compress → this is what you cache at inference, not full k/v
+        latent = self.kv_down(x)           # [B, T, latent_dim]
+
+        k = self.k_up(latent)
+        v = self.v_up(latent)
+
+        def split(t): return t.view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
+        q, k, v = split(q), split(k), split(v)
+
+        x = F.scaled_dot_product_attention(q, k, v, is_causal=True)
+        return self.out(x.transpose(1, 2).reshape(B, T, C))

transformer_toolkit-0.1.0/transformer_toolkit/block.py

@@ -0,0 +1,30 @@
+from feed_forward import FFN
+from normalization import LayerNorm
+from attention import MultiHeadAttention
+from torch import nn
+
+
+class TransformerBlock(nn.Module):
+    """
+    Default: original 'Attention is All You Need' — MHA + FFN + LayerNorm.
+    Swap any component via attn=, ffn=, norm= for modern variants.
+    """
+    def __init__(
+        self,
+        dim:     int,
+        n_heads: int,
+        hidden:  int,
+        norm = None,   # default: LayerNorm (original)
+        attn = None,   # default: MultiHeadAttention (original)
+        ffn  = None,   # default: FFN + ReLU (original)
+    ):
+        super().__init__()
+        self.norm1 = norm or LayerNorm(dim)
+        self.norm2 = norm or LayerNorm(dim)
+        self.attn  = attn or MultiHeadAttention(dim, n_heads)
+        self.ffn   = ffn  or FFN(dim, hidden)
+
+    def forward(self, x, mask=None):
+        x = x + self.attn(self.norm1(x), mask)
+        x = x + self.ffn(self.norm2(x))
+        return x

transformer_toolkit-0.1.0/transformer_toolkit/c_tokenizers.py

@@ -0,0 +1,114 @@
+from abc import ABC, abstractmethod
+
+
+class BaseTokenizer(ABC):
+    @abstractmethod
+    def train(self, texts: list[str], vocab_size: int): ...
+
+    @abstractmethod
+    def encode(self, text: str) -> list[int]: ...
+
+    @abstractmethod
+    def decode(self, ids: list[int]) -> str: ...
+
+    @abstractmethod
+    def save(self, path: str): ...
+
+    @abstractmethod
+    def load(self, path: str): ...
+
+    @property
+    @abstractmethod
+    def vocab_size(self) -> int: ...
+
+class ByteLevelTokenizer(BaseTokenizer):
+    """
+    Zero deps. Every byte is a token (0-255).
+    Works on any text/language out of the box.
+    """
+    def train(self, texts, vocab_size=256): pass   # nothing to train
+
+    def encode(self, text: str) -> list[int]:
+        return list(text.encode("utf-8"))
+
+    def decode(self, ids: list[int]) -> str:
+        return bytes(ids).decode("utf-8", errors="replace")
+
+    def save(self, path: str): pass
+    def load(self, path: str): pass
+
+    @property
+    def vocab_size(self) -> int: return 256
+
+
+class HFTokenizer(BaseTokenizer):
+    """
+    Thin wrapper around any HuggingFace tokenizer.
+    pip install transformers
+    """
+    def __init__(self, model_name: str = "gpt2"):
+        from transformers import AutoTokenizer
+        self._tok = AutoTokenizer.from_pretrained(model_name)
+
+    def train(self, texts, vocab_size=None):
+        raise NotImplementedError("use HF's train_new_from_iterator for custom training")
+
+    def encode(self, text: str) -> list[int]:
+        return self._tok.encode(text)
+
+    def decode(self, ids: list[int]) -> str:
+        return self._tok.decode(ids)
+
+    def save(self, path: str):
+        self._tok.save_pretrained(path)
+
+    def load(self, path: str):
+        from transformers import AutoTokenizer
+        self._tok = AutoTokenizer.from_pretrained(path)
+
+    @property
+    def vocab_size(self) -> int:
+        return len(self._tok)
+    
+
+class RustBPETokenizer(BaseTokenizer):
+    """
+    BPE tokenizer backed by HuggingFace's `tokenizers` Rust crate.
+    Trains ~100x faster than pure Python BPE.
+    pip install tokenizers
+    """
+    def __init__(self):
+        from tokenizers import Tokenizer
+        from tokenizers.models import BPE
+        self._tok = Tokenizer(BPE(unk_token="[UNK]"))
+        self._trained = False
+
+    def train(self, texts: list[str], vocab_size: int = 8000):
+        from tokenizers.trainers import BpeTrainer
+        from tokenizers.pre_tokenizers import Whitespace
+
+        self._tok.pre_tokenizer = Whitespace()
+        trainer = BpeTrainer(
+            vocab_size=vocab_size,
+            special_tokens=["[UNK]", "[PAD]", "[BOS]", "[EOS]"]
+        )
+        self._tok.train_from_iterator(texts, trainer)
+        self._trained = True
+
+    def encode(self, text: str) -> list[int]:
+        return self._tok.encode(text).ids
+
+    def decode(self, ids: list[int]) -> str:
+        return self._tok.decode(ids)
+
+    def save(self, path: str):
+        self._tok.save(path)
+
+    def load(self, path: str):
+        from tokenizers import Tokenizer
+        self._tok = Tokenizer.from_file(path)
+        self._trained = True
+
+    @property
+    def vocab_size(self) -> int:
+        return self._tok.get_vocab_size()

transformer_toolkit-0.1.0/transformer_toolkit/colors.py

@@ -0,0 +1,23 @@
+class C:
+    RESET = "\033[0m";  BOLD  = "\033[1m";  DIM    = "\033[2m"
+    GREEN = "\033[32m"; CYAN  = "\033[36m"; YELLOW = "\033[33m"
+    BLUE  = "\033[34m"; RED   = "\033[31m"; WHITE  = "\033[37m"
+    MAGENTA = "\033[35m"
+
+def _bar(current, total, width=28):
+    filled = int(width * current / max(total, 1))
+    return f"{C.CYAN}{'█' * filled}{'░' * (width - filled)}{C.RESET}"
+
+def _section(title):
+    print(f"\n{C.BOLD}{C.CYAN}{'─' * 52}{C.RESET}")
+    print(f"{C.BOLD}{C.CYAN}  {title}{C.RESET}")
+    print(f"{C.BOLD}{C.CYAN}{'─' * 52}{C.RESET}")
+
+def _info(label, value):
+    print(f"  {C.DIM}{label:<18}{C.RESET} {C.WHITE}{value}{C.RESET}")
+
+def _ok(msg):
+    print(f"  {C.GREEN}✓{C.RESET}  {msg}")
+
+def _err(msg):
+    print(f"  {C.RED}✗{C.RESET}  {msg}")

transformer_toolkit-0.1.0/transformer_toolkit/dataloader.py

@@ -0,0 +1,184 @@
+import torch
+import struct
+from pathlib import Path
+from torch.utils.data import Dataset, DataLoader, IterableDataset
+from c_tokenizers import BaseTokenizer
+from colors import C, _section, _info, _ok, _bar
+
+
+# ─── config ───────────────────────────────────────────────────────────────────
+
+class DataConfig:
+    def __init__(
+        self,
+        seq_len:     int   = 512,
+        batch_size:  int   = 16,
+        shuffle:     bool  = True,
+        num_workers: int   = 0,
+        split:       float = 0.9,
+        streaming:   bool  = False,
+    ):
+        self.seq_len     = seq_len
+        self.batch_size  = batch_size
+        self.shuffle     = shuffle
+        self.num_workers = num_workers
+        self.split       = split
+        self.streaming   = streaming
+
+
+# ─── dataset ──────────────────────────────────────────────────────────────────
+
+class TokenDataset(Dataset):
+    def __init__(self, tokens: torch.Tensor, seq_len: int):
+        self.tokens  = tokens
+        self.seq_len = seq_len
+
+    def __len__(self):
+        return max(0, len(self.tokens) - self.seq_len)
+
+    def __getitem__(self, idx):
+        x = self.tokens[idx:     idx + self.seq_len]
+        y = self.tokens[idx + 1: idx + self.seq_len + 1]
+        return x, y
+
+
+class StreamingDataset(IterableDataset):
+    def __init__(self, paths: list[str], tokenizer: BaseTokenizer, seq_len: int):
+        self.paths     = paths
+        self.tokenizer = tokenizer
+        self.seq_len   = seq_len
+
+    def __iter__(self):
+        buf = []
+        for path in self.paths:
+            for line in open(path, encoding="utf-8", errors="replace"):
+                buf.extend(self.tokenizer.encode(line))
+                while len(buf) >= self.seq_len + 1:
+                    chunk = buf[:self.seq_len + 1]
+                    buf   = buf[self.seq_len:]
+                    yield (
+                        torch.tensor(chunk[:-1], dtype=torch.long),
+                        torch.tensor(chunk[1:],  dtype=torch.long),
+                    )
+
+
+# ─── tokenize + split ─────────────────────────────────────────────────────────
+
+def _tokenize(texts: list[str], tokenizer: BaseTokenizer) -> torch.Tensor:
+    ids = []
+    for i, text in enumerate(texts):
+        print(f"\r  {C.DIM}tokenizing{C.RESET}  {_bar(i+1, len(texts))}  {C.DIM}{i+1}/{len(texts)}{C.RESET}", end="", flush=True)
+        ids.extend(tokenizer.encode(text))
+    print()
+    return torch.tensor(ids, dtype=torch.long)
+
+
+def _split(tokens: torch.Tensor, cfg: DataConfig) -> tuple[TokenDataset, TokenDataset]:
+    n        = len(tokens)
+    split_at = int(n * cfg.split)
+    train_ds = TokenDataset(tokens[:split_at], cfg.seq_len)
+    val_ds   = TokenDataset(tokens[split_at:], cfg.seq_len)
+
+    if len(val_ds) == 0:
+        print(f"  {C.YELLOW}⚠  val set empty — using last 10% of train{C.RESET}")
+        split_at = int(split_at * 0.9)
+        train_ds = TokenDataset(tokens[:split_at], cfg.seq_len)
+        val_ds   = TokenDataset(tokens[split_at:], cfg.seq_len)
+
+    return train_ds, val_ds
+
+
+def _loaders(train_ds, val_ds, cfg, shuffle=None) -> tuple[DataLoader, DataLoader]:
+    s = cfg.shuffle if shuffle is None else shuffle
+    train_dl = DataLoader(train_ds, batch_size=cfg.batch_size, shuffle=s,     num_workers=cfg.num_workers)
+    val_dl   = DataLoader(val_ds,   batch_size=cfg.batch_size, shuffle=False,  num_workers=cfg.num_workers)
+    print(f"  {C.DIM}train{C.RESET}  {C.WHITE}{len(train_ds):>10,}{C.RESET} samples  {C.DIM}│{C.RESET}  {C.YELLOW}{len(train_dl):,}{C.RESET} batches")
+    print(f"  {C.DIM}val  {C.RESET}  {C.WHITE}{len(val_ds):>10,}{C.RESET} samples  {C.DIM}│{C.RESET}  {C.YELLOW}{len(val_dl):,}{C.RESET} batches\n")
+    return train_dl, val_dl
+
+
+# ─── binary ───────────────────────────────────────────────────────────────────
+
+def save_binary(tokens: list[int], path: str):
+    Path(path).write_bytes(struct.pack(f"{len(tokens)}H", *tokens))
+    print(f"  {C.GREEN}✓{C.RESET}  saved {C.CYAN}{len(tokens):,}{C.RESET} tokens → {C.DIM}{path}{C.RESET}")
+
+
+def load_binary(path: str) -> torch.Tensor:
+    raw = Path(path).read_bytes()
+    ids = struct.unpack(f"{len(raw)//2}H", raw)
+    tok = torch.tensor(ids, dtype=torch.long)
+    print(f"  {C.GREEN}✓{C.RESET}  loaded {C.CYAN}{len(tok):,}{C.RESET} tokens ← {C.DIM}{path}{C.RESET}")
+    return tok
+
+
+# ─── public API ───────────────────────────────────────────────────────────────
+
+def from_binary(path: str, cfg: DataConfig) -> tuple[DataLoader, DataLoader]:
+    _section("💾 Binary dataset")
+    _info("path", path)
+    tokens           = load_binary(path)
+    train_ds, val_ds = _split(tokens, cfg)
+    return _loaders(train_ds, val_ds, cfg)
+
+
+def from_files(paths: list[str], tokenizer: BaseTokenizer, cfg: DataConfig) -> tuple[DataLoader, DataLoader]:
+    _section("📂 File dataset")
+    for p in paths: _info("file", p)
+
+    if cfg.streaming:
+        _info("mode", "streaming")
+        split_at = max(1, int(len(paths) * cfg.split))
+        train_ds = StreamingDataset(paths[:split_at], tokenizer, cfg.seq_len)
+        val_ds   = StreamingDataset(paths[split_at:] or paths[-1:], tokenizer, cfg.seq_len)
+        return _loaders(train_ds, val_ds, cfg, shuffle=False)
+
+    _info("mode", "in-memory")
+    texts            = [Path(p).read_text(encoding="utf-8", errors="replace") for p in paths]
+    tokens           = _tokenize(texts, tokenizer)
+    train_ds, val_ds = _split(tokens, cfg)
+    return _loaders(train_ds, val_ds, cfg)
+
+
+def from_hf(dataset_name: str, tokenizer: BaseTokenizer, cfg: DataConfig,
+            split: str = "train", text_col: str = "text") -> tuple[DataLoader, DataLoader]:
+    _section("🤗 HuggingFace dataset")
+    _info("dataset", dataset_name)
+    _info("split", split)
+    _info("streaming", str(cfg.streaming))
+
+    from datasets import load_dataset
+    ds = load_dataset(dataset_name, split=split, streaming=cfg.streaming)
+
+    if cfg.streaming:
+        def _gen():
+            buf = []
+            for row in ds:
+                buf.extend(tokenizer.encode(row[text_col]))
+                while len(buf) >= cfg.seq_len + 1:
+                    chunk = buf[:cfg.seq_len + 1]
+                    buf   = buf[cfg.seq_len:]
+                    yield (
+                        torch.tensor(chunk[:-1], dtype=torch.long),
+                        torch.tensor(chunk[1:],  dtype=torch.long),
+                    )
+        class _HFStream(IterableDataset):
+            def __iter__(self): return _gen()
+        dl = DataLoader(_HFStream(), batch_size=cfg.batch_size, num_workers=cfg.num_workers)
+        _ok("streaming dataloader ready")
+        return dl, dl
+
+    print(f"  {C.YELLOW}⏳ downloading...{C.RESET}", flush=True)
+    texts            = [row[text_col] for row in ds]
+    _ok(f"downloaded {len(texts):,} documents")
+    tokens           = _tokenize(texts, tokenizer)
+    train_ds, val_ds = _split(tokens, cfg)
+    return _loaders(train_ds, val_ds, cfg)
+
+
+def from_strings(texts: list[str], tokenizer: BaseTokenizer, cfg: DataConfig) -> tuple[DataLoader, DataLoader]:
+    _section("📝 String dataset")
+    _info("documents", str(len(texts)))
+    tokens           = _tokenize(texts, tokenizer)
+    train_ds, val_ds = _split(tokens, cfg)
+    return _loaders(train_ds, val_ds, cfg)

transformer_toolkit-0.1.0/transformer_toolkit/feed_forward.py

@@ -0,0 +1,51 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class FFN(nn.Module):
+    """Standard FFN. Used in original Transformer, BERT."""
+    def __init__(self, dim: int, hidden: int):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Linear(dim, hidden),
+            nn.GELU(),
+            nn.Linear(hidden, dim)
+        )
+
+    def forward(self, x):
+        return self.net(x)
+
+
+class SwiGLU(nn.Module):
+    """Gated FFN with Swish. Used in LLaMA, Mistral, PaLM."""
+    def __init__(self, dim: int, hidden: int):
+        super().__init__()
+        self.w1 = nn.Linear(dim, hidden, bias=False)
+        self.w2 = nn.Linear(hidden, dim, bias=False)
+        self.w3 = nn.Linear(dim, hidden, bias=False)
+
+    def forward(self, x):
+        return self.w2(F.silu(self.w1(x)) * self.w3(x))
+
+
+class MoE(nn.Module):
+    """Mixture of Experts. Used in Mixtral, GPT-4 (rumoured)."""
+    def __init__(self, dim: int, hidden: int, n_experts: int, top_k: int = 2):
+        super().__init__()
+        self.top_k   = top_k
+        self.gate    = nn.Linear(dim, n_experts, bias=False)
+        self.experts = nn.ModuleList([SwiGLU(dim, hidden) for _ in range(n_experts)])
+
+    def forward(self, x):
+        B, T, C = x.shape
+        x_flat  = x.view(-1, C)
+        weights, idx = self.gate(x_flat).topk(self.top_k, dim=-1)   # [B*T, top_k]
+        weights = F.softmax(weights, dim=-1)
+        out = torch.zeros_like(x_flat)
+        for i, expert in enumerate(self.experts):
+            mask = (idx == i).any(dim=-1)
+            if mask.any():
+                w = weights[mask][idx[mask] == i].unsqueeze(-1)
+                out[mask] += w * expert(x_flat[mask])
+        return out.view(B, T, C)