brsx 0.1.0__py3-none-any.whl

brsx/__init__.py

@@ -0,0 +1,8 @@
+"""brsx — AutoForge
+
+Usage:
+    from brsx import automodel
+    automodel.run()
+"""
+
+__version__ = "0.1.0"

brsx/automodel.py

@@ -0,0 +1,95 @@
+"""
+automodel.py — AutoForge single entry point
+
+    from brsx import automodel
+    automodel.run()
+
+Asks 1/2/3/4 and runs the main file inside the matching folder.
+Works both as an installed package and when run directly.
+"""
+
+import os
+import sys
+import importlib
+
+HERE = os.path.dirname(os.path.abspath(__file__))
+
+# key → (label, description, folder, module file (no .py), entry function)
+TASKS = {
+    "1": ("Transformer training",
+          "Train a standard Transformer from scratch",
+          "train_transformer", "full", "main"),
+    "2": ("MTP Transformer training",
+          "Train with Multi-Token Prediction heads",
+          "mtp", "brsx_mtp", "train_mtp"),
+    "3": ("HuggingFace model fine-tune",
+          "Fine-tune a HuggingFace model (local or hub)",
+          "ft_hf_models", "fine_tune", "main"),
+    "4": ("brsx (.pt) model fine-tune",
+          "Fine-tune an existing brsx .pt model",
+          "ft_brsx_model", "brsx_finetune", "main"),
+}
+
+
+def _run_task(folder, module_name, func_name):
+    """Add the task folder to sys.path, import the module, call its entry func.
+
+    Each folder is added to sys.path so the module's own imports
+    (import brsx_calibrate, from brsx_train import GPT, ...) keep working
+    exactly as they do when the file is run directly inside its folder.
+    """
+    folder_path = os.path.join(HERE, folder)
+    if not os.path.isdir(folder_path):
+        print(f"  ! Folder not found: {folder_path}")
+        return
+
+    if folder_path not in sys.path:
+        sys.path.insert(0, folder_path)
+
+    try:
+        module = importlib.import_module(module_name)
+    except Exception as e:
+        print(f"  ! Couldn't load {module_name}.py from {folder}/: {e}")
+        return
+
+    entry = getattr(module, func_name, None)
+    if not callable(entry):
+        print(f"  ! {module_name}.py has no callable {func_name}().")
+        return
+
+    entry()
+
+
+def _print_menu():
+    print("=" * 60)
+    print("                    brsx — AutoModel")
+    print("=" * 60)
+    print("  What would you like to do?")
+    print("-" * 60)
+    for key, (label, desc, *_ ) in TASKS.items():
+        print(f"    {key}) {label}")
+        print(f"       {desc}")
+    print("    q) Quit")
+    print("-" * 60)
+
+
+def run():
+    _print_menu()
+    valid = set(TASKS) | {"q", "quit", "exit"}
+    while True:
+        choice = input("  Choice: ").strip().lower()
+        if choice in valid:
+            break
+        print(f"  ! Pick one of: {', '.join(TASKS)} or q")
+
+    if choice in ("q", "quit", "exit"):
+        print("  See you!")
+        return
+
+    label, _desc, folder, module_name, func_name = TASKS[choice]
+    print(f"\n  → {label}\n")
+    _run_task(folder, module_name, func_name)
+
+
+if __name__ == "__main__":
+    run()

brsx/ft_brsx_model/brsx_finetune.py

@@ -0,0 +1,307 @@
+"""
+==============================================================
+  brsx_finetune.py  -  Fine-Tune for brsx's own models
+==============================================================
+NOT HF. Fine-tunes your own GPT model (model.pt + config.json + tokenizer.json).
+
+- Model architecture is IDENTICAL to the GPT in brsx_train.py (MultiheadAttention).
+- Injects LoRA manually; only into the MLP Linear layers (not into the
+  attention out_proj - PyTorch looks it up internally via .weight, so
+  wrapping it breaks things). MLP LoRA is enough for fine-tuning.
+- 'full' mode trains all weights.
+- Interactive: everything is asked at runtime, no flags.
+- If no data is found, downloads it from brsx-labs/datasets.
+
+Must be in the same folder:  brsx_train.py  (for GPT + CFG)
+Optional:                     brsx_tokenizer.py (falls back to byte-level if missing)
+
+Run:  python brsx_finetune.py
+==============================================================
+"""
+
+import os, sys, json, time, math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+try:
+    from brsx_train import GPT, CFG
+except ImportError:
+    print("[ERROR] brsx_train.py must be in the same folder (for GPT, CFG).")
+    sys.exit(1)
+
+HF_REPO = "brsx-labs/datasets"
+DOMAINS = {
+    "1": ("Code (code.txt)",       "code.txt"),
+    "2": ("General (general.txt)", "general.txt"),
+    "3": ("Math (math.txt)",       "math.txt"),
+    "4": ("Chat (chat.txt)",       "chat.txt"),
+}
+
+
+# --------------------------------------------------------------------------
+# Input helpers
+# --------------------------------------------------------------------------
+def ask(msg, default=None, cast=str, hint=None):
+    if hint: print(f"   ({hint})")
+    sfx = f" [{default}]" if default is not None else ""
+    while True:
+        raw = input(f" > {msg}{sfx}: ").strip()
+        if not raw and default is not None: return default
+        if not raw: print("   Enter a value."); continue
+        try: return cast(raw)
+        except ValueError: print("   Invalid, try again.")
+
+def ask_choice(title, options, hint=None):
+    print(f"\n{title}")
+    if hint: print(f"   ({hint})")
+    for k, v in options.items(): print(f"   {k}) {v}")
+    while True:
+        c = input(" > Choice: ").strip()
+        if c in options: return c
+        print("   Pick from the list.")
+
+
+# --------------------------------------------------------------------------
+# LoRA
+# --------------------------------------------------------------------------
+class LoRALinear(nn.Module):
+    """Wraps an nn.Linear. Original weight is frozen; only A,B are trained."""
+    def __init__(self, base: nn.Linear, r=8, alpha=16, dropout=0.0):
+        super().__init__()
+        self.base = base
+        for p in self.base.parameters():
+            p.requires_grad = False
+        self.r = r
+        self.scale = alpha / r
+        self.drop = nn.Dropout(dropout)
+        self.A = nn.Linear(base.in_features, r, bias=False)
+        self.B = nn.Linear(r, base.out_features, bias=False)
+        nn.init.kaiming_uniform_(self.A.weight, a=math.sqrt(5))
+        nn.init.zeros_(self.B.weight)
+
+    def forward(self, x):
+        return self.base(x) + self.scale * self.B(self.A(self.drop(x)))
+
+
+def inject_lora(model, r=8, alpha=16, dropout=0.05):
+    """
+    Injects LoRA ONLY into the MLP Linear layers.
+    Not into the MultiheadAttention out_proj (PyTorch looks for .weight -> breaks).
+    """
+    n = 0
+    for blk in model.blocks:
+        for i, layer in enumerate(blk.mlp):
+            if isinstance(layer, nn.Linear):
+                blk.mlp[i] = LoRALinear(layer, r, alpha, dropout)
+                n += 1
+    print(f"   LoRA injected: {n} layers (MLP)")
+    if n == 0:
+        print("   [!] No LoRA layers could be injected - blk.mlp structure may differ.")
+        print("       To inspect layers: for n,_ in model.named_modules(): print(n)")
+    return model
+
+
+# --------------------------------------------------------------------------
+# Load model + tokenizer
+# --------------------------------------------------------------------------
+def load_brsx(model_dir, device):
+    cfgp = os.path.join(model_dir, "config.json")
+    mp   = os.path.join(model_dir, "model.pt")
+    if not os.path.exists(cfgp) or not os.path.exists(mp):
+        print(f"[ERROR] {model_dir} must contain config.json and model.pt.")
+        sys.exit(1)
+
+    with open(cfgp, encoding="utf-8") as f:
+        conf = json.load(f)
+
+    cfg = CFG()
+    cfg.n_layer = conf["n_layer"]; cfg.n_head = conf["n_head"]
+    cfg.n_embd  = conf["n_embd"];  cfg.block_size = conf["block_size"]
+    cfg.vocab_size = conf["vocab_size"]; cfg.device = device
+
+    model = GPT(cfg).to(device)
+    model.load_state_dict(torch.load(mp, map_location=device))
+    try:
+        pc = model.param_count()/1e6
+    except Exception:
+        pc = sum(p.numel() for p in model.parameters())/1e6
+    print(f"   Model loaded: {pc:.1f}M params "
+          f"(n_layer={cfg.n_layer}, n_embd={cfg.n_embd}, vocab={cfg.vocab_size})")
+
+    tok = None
+    if conf.get("has_tokenizer"):
+        try:
+            from brsx_tokenizer import BPETokenizer
+            tok = BPETokenizer()
+            tok.load(os.path.join(model_dir, "tokenizer.json"))
+            print("   Tokenizer loaded (BPE).")
+        except Exception as e:
+            print(f"   [!] Could not load tokenizer ({e}) -> byte-level.")
+    return model, cfg, tok
+
+
+# --------------------------------------------------------------------------
+# Data
+# --------------------------------------------------------------------------
+def get_data_text(tok):
+    print("\n--- DATA ---")
+    path = input(" > Data .txt path (EMPTY = download a ready-made dataset): ").strip()
+    if path and os.path.exists(path):
+        print(f"   Found: {path}")
+    else:
+        if path: print(f"   '{path}' not found, switching to download.")
+        dom = ask_choice("Which dataset?", {k: v[0] for k, v in DOMAINS.items()})
+        fname = DOMAINS[dom][1]
+        from huggingface_hub import hf_hub_download
+        print(f"   Downloading: {HF_REPO}/{fname} ...")
+        path = hf_hub_download(repo_id=HF_REPO, filename=fname, repo_type="dataset")
+        print(f"   Downloaded: {path}")
+    with open(path, "r", encoding="utf-8") as f:
+        return f.read()
+
+
+def encode_data(text, tok, cfg):
+    if tok is not None:
+        ids = tok.encode(text)
+    else:
+        ids = list(text.encode("utf-8"))
+    data = torch.tensor(ids, dtype=torch.long)
+    if len(data) < cfg.block_size + 2:
+        print("[ERROR] Data too short.")
+        sys.exit(1)
+    return data
+
+
+# --------------------------------------------------------------------------
+# Training
+# --------------------------------------------------------------------------
+def train_ft(model, cfg, data, mode, steps, lr, batch, out_dir, save_every, tok):
+    device = cfg.device
+    if mode == "lora":
+        params = [p for p in model.parameters() if p.requires_grad]
+    else:
+        for p in model.parameters(): p.requires_grad = True
+        params = list(model.parameters())
+    ntr = sum(p.numel() for p in params)
+    print(f"   Trainable params: {ntr/1e6:.2f}M")
+
+    opt = torch.optim.AdamW(params, lr=lr)
+    use_amp = (device == "cuda")
+    scaler = torch.amp.GradScaler('cuda', enabled=use_amp)
+
+    def get_batch():
+        ix = torch.randint(0, len(data) - cfg.block_size - 1, (batch,))
+        x = torch.stack([data[i:i+cfg.block_size] for i in ix]).to(device)
+        y = torch.stack([data[i+1:i+cfg.block_size+1] for i in ix]).to(device)
+        return x, y
+
+    model.train()
+    t0 = time.time()
+    for step in range(1, steps+1):
+        x, y = get_batch()
+        opt.zero_grad(set_to_none=True)
+        if use_amp:
+            with torch.amp.autocast('cuda'):
+                _, loss = model(x, y)
+            scaler.scale(loss).backward(); scaler.step(opt); scaler.update()
+        else:
+            _, loss = model(x, y)
+            loss.backward(); opt.step()
+
+        if step % 20 == 0 or step == 1:
+            ips = step / (time.time() - t0)
+            v = torch.cuda.memory_allocated()/1e6 if device == "cuda" else 0
+            print(f" step {step}/{steps} | loss {loss.item():.4f} | {ips:.1f} it/s"
+                  + (f" | VRAM {v:.0f}MB" if device=='cuda' else ""))
+
+        if save_every and step % save_every == 0 and step < steps:
+            _save(model, cfg, tok, out_dir)
+            print(f"   [checkpoint] {out_dir}")
+
+    _save(model, cfg, tok, out_dir)
+    print(f"\n Training finished ({time.time()-t0:.1f}s). Saved -> {out_dir}")
+
+
+def _save(model, cfg, tok, out_dir):
+    os.makedirs(out_dir, exist_ok=True)
+    torch.save(model.state_dict(), os.path.join(out_dir, "model.pt"))
+    conf = {"n_layer": cfg.n_layer, "n_head": cfg.n_head, "n_embd": cfg.n_embd,
+            "block_size": cfg.block_size, "vocab_size": cfg.vocab_size,
+            "has_tokenizer": tok is not None, "finetuned": True}
+    with open(os.path.join(out_dir, "config.json"), "w", encoding="utf-8") as f:
+        json.dump(conf, f, indent=2)
+    if tok is not None:
+        try:
+            tok.save(os.path.join(out_dir, "tokenizer.json"))
+        except Exception:
+            pass
+
+
+# --------------------------------------------------------------------------
+# main
+# --------------------------------------------------------------------------
+def main():
+    print("=" * 60)
+    print("        brsx Fine-Tune  (for your own models)")
+    print("=" * 60)
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    print(f" Device: {device}"
+          + (f" ({torch.cuda.get_device_name(0)})" if device=='cuda' else ""))
+
+    print("\n--- MODEL ---")
+    print("Folder of YOUR OWN model to fine-tune")
+    print("(folder containing model.pt + config.json + tokenizer.json).")
+    model_dir = ask("Model folder", default="brsx_model")
+    model_dir = os.path.abspath(os.path.expanduser(model_dir))
+
+    model, cfg, tok = load_brsx(model_dir, device)
+
+    mode_key = ask_choice("FT type:", {
+        "1": "lora - Only small adapters are trained (fast, low memory, RECOMMENDED)",
+        "2": "full - Whole model is trained (stronger but more memory)",
+    }, hint="If unsure, pick 1 (lora)")
+    mode = {"1": "lora", "2": "full"}[mode_key]
+
+    if mode == "lora":
+        r = ask("LoRA rank r (typically 4-16)", default=8, cast=int,
+                hint="Larger r = more capacity, more params")
+        model = inject_lora(model, r=r, alpha=2*r, dropout=0.05)
+        model.to(device)
+
+    text = get_data_text(tok)
+    data = encode_data(text, tok, cfg)
+    print(f"   Data: {len(data)} tokens")
+
+    print("\n--- TRAINING SETTINGS ---")
+    lr = ask("Learning rate", default=(2e-4 if mode=="lora" else 5e-5), cast=float,
+             hint="2e-4 for LoRA, lower (5e-5) is better for full FT")
+    steps = ask("Max steps", default=1000, cast=int)
+    batch = ask("batch_size", default=8, cast=int,
+                hint="On 6GB try 8-16; lower it if you hit OOM")
+
+    print("\n--- CHECKPOINT ---")
+    if steps >= 1000:
+        sug = max(1, steps // 5)
+        save_every = ask("Save every N steps (0=off)", default=sug, cast=int)
+    else:
+        save_every = 0
+        print("   Few steps -> checkpointing off.")
+
+    out_dir = ask("Output folder", default="brsx_model_ft")
+    out_dir = os.path.abspath(os.path.expanduser(out_dir))
+
+    print("\n" + "=" * 60)
+    print(f" model={model_dir}")
+    print(f" mode={mode} | lr={lr} | steps={steps} | batch={batch}")
+    print(f" checkpoint={save_every} | output={out_dir}")
+    print("=" * 60)
+    if input(" Start? (y/n) [y]: ").strip().lower() not in ("","y","yes","e"):
+        print("Cancelled."); return
+
+    train_ft(model, cfg, data, mode, steps, lr, batch, out_dir, save_every, tok)
+    print("\nDone. Use the model.pt in the output folder with your own chat.py.")
+
+
+if __name__ == "__main__":
+    main()