Upload crazy.py

idk dont mind the name this is the script i use to train it but also to say! if you wanna run it Ctrl + F find all "Eclipsed" and replace it with your PC user

crazy.py

@@ -0,0 +1,597 @@
+import torch
+import torch.nn as nn
+from torch.utils.data import Dataset, DataLoader
+import os
+import time
+import threading
+import glob
+import re
+import json
+import numpy as np
+import wikipediaapi
+
+# ============================================================
+#  5x2T — Word Level Model with Disk Offloaded Optimizer
+#  Uses numpy for disk saves — much more memory efficient
+# ============================================================
+
+# ---------------- CONFIG ----------------
+MAX_TRAIN_MIN   = 60
+BATCH_SIZE      = 2
+SEQ_LENGTH      = 64
+EMBED_SIZE      = 192
+HIDDEN_SIZE     = 384
+NUM_LAYERS      = 1
+DROPOUT         = 0.2
+LEARNING_RATE   = 0.01
+GRAD_CLIP       = 1.0
+AUTOSAVE_MIN    = 5
+MAX_VOCAB       = 995600
+TEMPERATURE     = 0.8
+RESPONSE_LENGTH = 40
+MAX_LENGTH      = 200
+UNK_TOKEN       = "<UNK>"
+PAD_TOKEN       = "<PAD>"
+
+BASE_DIR      = r"C:\Users\Eclipsed\Downloads\5x2T"
+DATASET_DIR   = os.path.join(BASE_DIR, "datasets")
+MODEL_DIR     = os.path.join(BASE_DIR, "5x2T-2")
+MODEL_PATH    = os.path.join(MODEL_DIR, "model.pth")
+VOCAB_PATH    = os.path.join(MODEL_DIR, "vocab.json")
+OFFLOAD_DIR   = os.path.join(MODEL_DIR, "offload")
+
+DATASET_FOLDERS = [
+    os.path.join(DATASET_DIR, "chat_dataset"),
+    os.path.join(DATASET_DIR, "python_data"),
+    os.path.join(DATASET_DIR, "lua_dataset"),
+    os.path.join(DATASET_DIR, "dic_dataset"),
+    os.path.join(DATASET_DIR, "Wiki_dataset"),
+    r"E:\5x2T",
+]
+
+DEVICE = torch.device("cpu")
+
+
+# ---------------- DISK OFFLOAD OPTIMIZER ----------------
+class DiskOffloadSGD:
+    """
+    SGD optimizer that stores momentum buffers on disk as numpy files.
+    Large buffers are processed in chunks to avoid RAM spikes.
+    """
+    CHUNK = 4_000_000  # process 4 million elements at a time
+
+    def __init__(self, params, lr=0.01, momentum=0.9, offload_dir=OFFLOAD_DIR):
+        self.params      = list(params)
+        self.lr          = lr
+        self.momentum    = momentum
+        self.offload_dir = offload_dir
+        os.makedirs(offload_dir, exist_ok=True)
+
+        print(f"  Initialising {len(self.params)} momentum buffers on disk...")
+        for i, p in enumerate(self.params):
+            path = os.path.join(offload_dir, f"m_{i}.npy")
+            if not os.path.exists(path):
+                # Save in chunks to avoid allocating the full array at once
+                shape = p.data.shape
+                total = p.data.numel()
+                flat  = np.zeros(total, dtype=np.float32)
+                np.save(path, flat.reshape(shape))
+                del flat
+        print(f"  Momentum buffers ready in {offload_dir}\n")
+
+    def zero_grad(self):
+        for p in self.params:
+            if p.grad is not None:
+                p.grad.detach_()
+                p.grad.zero_()
+
+    def step(self):
+        for i, p in enumerate(self.params):
+            if p.grad is None:
+                continue
+
+            path  = os.path.join(self.offload_dir, f"m_{i}.npy")
+            shape = p.data.shape
+            total = p.data.numel()
+
+            # Use memory mapped file — only the chunk we touch is in RAM
+            buf_mm  = np.load(path, mmap_mode="r+")
+            buf_flat = buf_mm.reshape(-1)
+            grad_flat = p.grad.data.reshape(-1).numpy()
+            data_flat = p.data.reshape(-1).numpy()
+
+            # Process in chunks so RAM never spikes
+            for start in range(0, total, self.CHUNK):
+                end = min(start + self.CHUNK, total)
+                buf_flat[start:end]  = (self.momentum * buf_flat[start:end]
+                                        + grad_flat[start:end])
+                data_flat[start:end] -= self.lr * buf_flat[start:end]
+
+            # Write updated data back to param tensor
+            p.data.copy_(torch.from_numpy(data_flat.reshape(shape)))
+
+            # Flush mmap and free
+            buf_mm.flush()
+            del buf_mm, buf_flat, grad_flat, data_flat
+
+    def state_dict(self):
+        return {"lr": self.lr, "momentum": self.momentum}
+
+    def load_state_dict(self, state):
+        self.lr       = state.get("lr", self.lr)
+        self.momentum = state.get("momentum", self.momentum)
+
+
+# ---------------- DATASET DISCOVERY ----------------
+def find_all_txt_files(folders):
+    all_files = []
+    for folder in folders:
+        if os.path.exists(folder):
+            found = glob.glob(os.path.join(folder, "**", "*.txt"), recursive=True)
+            all_files.extend(found)
+            print(f"  [{os.path.basename(folder)}] -> {len(found)} file(s)")
+        else:
+            print(f"  [SKIP] Not found: {folder}")
+    if not all_files:
+        raise FileNotFoundError("No .txt files found. Check your dataset paths.")
+    print(f"\n  Total files: {len(all_files)}\n")
+    return all_files
+
+
+# ---------------- TOKENISER ----------------
+def tokenise(text):
+    return re.findall(r"\b\w+\b|[\"'.,!?;:\-\n]", text.lower())
+
+
+def build_vocab(files, max_vocab=MAX_VOCAB):
+    print("  Building vocabulary...")
+    freq = {}
+    total_tokens = 0
+    for f in files:
+        try:
+            with open(f, "r", encoding="utf-8", errors="ignore") as file:
+                tokens = tokenise(file.read())
+                for t in tokens:
+                    freq[t] = freq.get(t, 0) + 1
+                total_tokens += len(tokens)
+        except Exception as e:
+            print(f"  [WARNING] Could not read {f}: {e}")
+
+    sorted_vocab = sorted(freq.items(), key=lambda x: x[1], reverse=True)
+    vocab_words  = [PAD_TOKEN, UNK_TOKEN] + [w for w, _ in sorted_vocab[:max_vocab - 2]]
+    word2idx     = {w: i for i, w in enumerate(vocab_words)}
+    idx2word     = {i: w for i, w in enumerate(vocab_words)}
+
+    print(f"  Total tokens    : {total_tokens:,}")
+    print(f"  Unique words    : {len(freq):,}")
+    print(f"  Vocab size      : {len(vocab_words):,}\n")
+
+    return vocab_words, word2idx, idx2word
+
+
+def save_vocab(vocab_words, path):
+    with open(path, "w", encoding="utf-8") as f:
+        json.dump(vocab_words, f)
+
+
+def load_vocab(path):
+    with open(path, "r", encoding="utf-8") as f:
+        vocab_words = json.load(f)
+    word2idx = {w: i for i, w in enumerate(vocab_words)}
+    idx2word = {i: w for i, w in enumerate(vocab_words)}
+    return vocab_words, word2idx, idx2word
+
+
+# ---------------- DATASET ----------------
+class WordDataset(Dataset):
+    def __init__(self, files, word2idx):
+        self.data = []
+        unk_idx   = word2idx.get(UNK_TOKEN, 1)
+        for f in files:
+            try:
+                with open(f, "r", encoding="utf-8", errors="ignore") as file:
+                    tokens     = tokenise(file.read())
+                    self.data += [word2idx.get(t, unk_idx) for t in tokens]
+            except Exception as e:
+                print(f"  [WARNING] Could not read {f}: {e}")
+
+        if not self.data:
+            raise ValueError("Dataset is empty after tokenisation.")
+        print(f"  Dataset tokens: {len(self.data):,}\n")
+
+    def __len__(self):
+        return len(self.data) - SEQ_LENGTH
+
+    def __getitem__(self, idx):
+        x = torch.tensor(self.data[idx:idx + SEQ_LENGTH],         dtype=torch.long)
+        y = torch.tensor(self.data[idx + 1:idx + SEQ_LENGTH + 1], dtype=torch.long)
+        return x, y
+
+
+# ---------------- MODEL ----------------
+class Model(nn.Module):
+    def __init__(self, vocab_size):
+        super().__init__()
+        self.embed   = nn.Embedding(vocab_size, EMBED_SIZE, padding_idx=0)
+        self.dropout = nn.Dropout(DROPOUT)
+        self.lstm    = nn.LSTM(
+            EMBED_SIZE, HIDDEN_SIZE,
+            num_layers=NUM_LAYERS,
+            batch_first=True,
+            dropout=0
+        )
+        self.norm = nn.LayerNorm(HIDDEN_SIZE)
+        self.fc   = nn.Linear(HIDDEN_SIZE, vocab_size)
+
+    def forward(self, x, hc=None):
+        x = self.dropout(self.embed(x))
+        x, hc = self.lstm(x, hc)
+        x = self.norm(x)
+        x = self.fc(x)
+        return x, hc
+
+
+# ---------------- SETUP ----------------
+def setup_dirs():
+    os.makedirs(MODEL_DIR,   exist_ok=True)
+    os.makedirs(OFFLOAD_DIR, exist_ok=True)
+    os.makedirs(os.path.join(MODEL_DIR, "questions"), exist_ok=True)
+
+
+# ---------------- GENERATE ----------------
+def generate(model, word2idx, idx2word, seed_text, length=RESPONSE_LENGTH, temperature=TEMPERATURE):
+    model.eval()
+    tokens  = tokenise(seed_text)
+    unk_idx = word2idx.get(UNK_TOKEN, 1)
+    indices = [word2idx.get(t, unk_idx) for t in tokens]
+    hc      = None
+
+    with torch.no_grad():
+        for _ in range(min(length, MAX_LENGTH)):
+            x        = torch.tensor([indices[-SEQ_LENGTH:]], dtype=torch.long)
+            out, hc  = model(x, hc)
+            logits   = out[0, -1] / temperature
+            probs    = torch.softmax(logits, dim=0)
+            next_idx = torch.multinomial(probs, 1).item()
+            indices.append(next_idx)
+
+    generated = indices[len(tokens):]
+    words     = [idx2word.get(i, UNK_TOKEN) for i in generated]
+    return " ".join(words)
+
+
+def format_response(text):
+    text = re.sub(r' ([.,!?;:])', r'\1', text)
+    text = re.sub(r'\n ', '\n', text)
+    if text:
+        text = text[0].upper() + text[1:]
+    return text
+
+
+# ---------------- WIKIPEDIA ----------------
+wiki_api = wikipediaapi.Wikipedia(
+    language='en',
+    extract_format=wikipediaapi.ExtractFormat.WIKI,
+    user_agent="5x2T-AI/1.0"
+)
+
+def search_wikipedia(query):
+    try:
+        search_term = query.lower()
+        for prefix in ["what is ", "what are ", "who is ", "who was ",
+                        "tell me about ", "explain ", "define ",
+                        "what was ", "how does ", "how do "]:
+            search_term = search_term.replace(prefix, "")
+        search_term = search_term.replace("?", "").strip()
+        page = wiki_api.page(search_term)
+        if page.exists():
+            return f"[Wikipedia: {page.title}]\n{page.summary[:600]}"
+        return None
+    except Exception as e:
+        print(f"  [WARNING] Wikipedia lookup failed: {e}")
+        return None
+
+
+def should_search_wiki(text):
+    triggers = [
+        "what is", "what are", "who is", "who was",
+        "tell me about", "explain", "define", "what was",
+        "how does", "how do"
+    ]
+    return any(text.lower().strip().startswith(t) for t in triggers)
+
+
+# ---------------- TRAINING ----------------
+def train():
+    setup_dirs()
+    print("=" * 55)
+    print("  5x2T — Word Level Training (Disk Offload)")
+    print(f"  Device      : {DEVICE}")
+    print(f"  Offload dir : {OFFLOAD_DIR}")
+    print(f"  Target      : {MAX_TRAIN_MIN} minutes")
+    print("=" * 55 + "\n")
+
+    print("Scanning dataset folders...")
+    files = find_all_txt_files(DATASET_FOLDERS)
+
+    if os.path.exists(VOCAB_PATH):
+        print("  Found existing vocab — loading...")
+        vocab_words, word2idx, idx2word = load_vocab(VOCAB_PATH)
+        print(f"  Vocab size: {len(vocab_words):,}\n")
+    else:
+        vocab_words, word2idx, idx2word = build_vocab(files)
+        save_vocab(vocab_words, VOCAB_PATH)
+        print(f"  Vocab saved to {VOCAB_PATH}\n")
+
+    print("Loading dataset...")
+    dataset = WordDataset(files, word2idx)
+    loader  = DataLoader(
+        dataset, batch_size=BATCH_SIZE,
+        shuffle=True, num_workers=0
+    )
+
+    vocab_size = len(vocab_words)
+    model      = Model(vocab_size)
+    criterion  = nn.CrossEntropyLoss(ignore_index=0)
+    optimizer  = DiskOffloadSGD(
+        model.parameters(),
+        lr=LEARNING_RATE,
+        momentum=0.9,
+        offload_dir=OFFLOAD_DIR
+    )
+
+    param_count = sum(p.numel() for p in model.parameters())
+    print(f"  Model parameters : {param_count:,}")
+    print(f"  Vocab size       : {vocab_size:,}")
+    print(f"  Optimizer        : DiskOffloadSGD (numpy on disk)")
+    print(f"  Offload folder   : {OFFLOAD_DIR}\n")
+
+    if os.path.exists(MODEL_PATH + ".npz"):
+        load_path = MODEL_PATH + ".npz"
+    elif os.path.exists(MODEL_PATH):
+        load_path = MODEL_PATH
+    else:
+        load_path = None
+
+    if load_path:
+        try:
+            if load_path.endswith(".npz"):
+                raw         = np.load(load_path)
+                checkpoint  = {k: torch.from_numpy(raw[k]) for k in raw.files}
+            else:
+                checkpoint  = torch.load(load_path, map_location="cpu")
+            model_state = model.state_dict()
+            loaded = 0
+            for k in checkpoint.keys():
+                if k in model_state and checkpoint[k].shape == model_state[k].shape:
+                    model_state[k] = checkpoint[k]
+                    loaded += 1
+            model.load_state_dict(model_state)
+            print(f"  Resumed from checkpoint ({loaded} layers matched)\n")
+        except Exception as e:
+            print(f"  Could not load checkpoint: {e} — starting fresh\n")
+
+    print("-" * 55)
+    print("  Training started...\n")
+
+    start_time    = time.time()
+    epoch         = 0
+    best_loss     = float("inf")
+    total_tokens  = 0
+    last_autosave = 0
+    epoch_loss    = 0
+    batches       = 0
+    loss          = None
+
+    def print_progress():
+        while True:
+            elapsed_sec = time.time() - start_time
+            elapsed_min = elapsed_sec / 60
+            speed = total_tokens / (elapsed_sec + 1e-5)
+            avg_loss = epoch_loss / max(batches, 1) if batches > 0 else 0
+            mins = int(elapsed_sec // 60)
+            secs = int(elapsed_sec % 60)
+            current_loss = loss.item() if loss is not None else 0.0
+            print(
+                f"  Epoch {epoch+1:>3} | "
+                f"Batch {batches:>5} | "
+                f"Loss: {current_loss:.4f} | "
+                f"Avg: {avg_loss:.4f} | "
+                f"Speed: {speed:.0f} tok/s | "
+                f"Time: {mins:02d}:{secs:02d}/{MAX_TRAIN_MIN:02d}:00",
+                end="\r"
+            )
+            time.sleep(1)
+
+    # Start progress printing thread
+    progress_thread = threading.Thread(target=print_progress, daemon=True)
+    progress_thread.start()
+
+    while (time.time() - start_time) / 60 < MAX_TRAIN_MIN:
+        epoch_loss = 0
+        batches    = 0
+
+        for x, y in loader:
+            optimizer.zero_grad()
+            out, _ = model(x)
+            out    = out.view(-1, vocab_size)
+            y      = y.view(-1)
+            loss   = criterion(out, y)
+            loss.backward()
+            nn.utils.clip_grad_norm_(model.parameters(), GRAD_CLIP)
+            optimizer.step()
+
+            epoch_loss   += loss.item()
+            batches      += 1
+            total_tokens += x.numel()
+
+            elapsed_min = (time.time() - start_time) / 60
+            mins = int(elapsed_min)
+            secs = int((elapsed_min - mins) * 60)
+            if elapsed_min - last_autosave >= AUTOSAVE_MIN:
+                try:
+                    torch.save(model.state_dict(), MODEL_PATH)
+                    last_autosave = elapsed_min
+                    print(f"\n  [Autosave] {mins:02d}:{secs:02d} -> {MODEL_PATH}")
+                except MemoryError:
+                    try:
+                        print(f"\n  [Autosave] RAM full - saving directly to disk...")
+                        tmp_path = MODEL_PATH + ".tmp"
+                        with open(tmp_path, "wb") as f:
+                            state = {k: v.numpy() for k, v in model.state_dict().items()}
+                            np.savez_compressed(f, **state)
+                        os.replace(tmp_path, MODEL_PATH + ".npz")
+                        last_autosave = elapsed_min
+                        print(f"\n  [Autosave] {mins:02d}:{secs:02d} -> {MODEL_PATH}.npz")
+                    except OSError:
+                        print(f"\n  [5xSc-404] Low storage or memory - autosave skipped")
+                    except Exception as e:
+                        print(f"\n  [5xSc-9512] Unknown autosave error: {e}")
+                except OSError:
+                    print(f"\n  [5xSc-404] Low storage or memory - autosave skipped")
+                except KeyboardInterrupt:
+                    print(f"\n  [5xSc-80082] Training stopped early - saving...")
+                    try:
+                        torch.save(model.state_dict(), MODEL_PATH)
+                    except Exception:
+                        state = {k: v.numpy() for k, v in model.state_dict().items()}
+                        np.savez_compressed(MODEL_PATH + ".npz", **state)
+                    print(f"  Model saved. Exiting.")
+                    raise
+                except Exception as e:
+                    err = str(e).lower()
+                    if "corrupt" in err or "invalid" in err:
+                        print(f"\n  [5xSc-312] Corruption detected: {e}")
+                    elif "allocat" in err or "memory" in err:
+                        print(f"\n  [5xSc-500] Memory allocation failed: {e}")
+                    else:
+                        print(f"\n  [5xSc-9512] Unknown error: {e}")
+            if elapsed_min >= MAX_TRAIN_MIN:
+                break
+
+        print()
+        epoch    += 1
+        avg_loss  = epoch_loss / max(batches, 1)
+
+        if avg_loss < best_loss:
+            best_loss = avg_loss
+            try:
+                torch.save(model.state_dict(), MODEL_PATH)
+            except MemoryError:
+                print(f"  [Save] RAM full — saving directly to disk...")
+                tmp_path = MODEL_PATH + ".tmp"
+                with open(tmp_path, "wb") as f:
+                    state = {k: v.numpy() for k, v in model.state_dict().items()}
+                    np.savez_compressed(f, **state)
+                os.replace(tmp_path, MODEL_PATH + ".npz")
+            print(f"  [Saved] Best loss: {best_loss:.4f}\n")
+
+        if (time.time() - start_time) / 60 >= MAX_TRAIN_MIN:
+            break
+
+    print("-" * 55)
+    print(f"  Done! Epochs: {epoch} | Best loss: {best_loss:.4f}")
+    print(f"  Model saved to: {MODEL_PATH}\n")
+
+    print("  Sample generation:")
+    seed   = '"what is marxism"\n"'
+    sample = generate(model, word2idx, idx2word, seed_text=seed, length=40)
+    print(f"  {format_response(sample)}\n")
+
+    return model, word2idx, idx2word
+
+
+# ---------------- CHAT ----------------
+def chat(model=None, word2idx=None, idx2word=None):
+    print("=" * 55)
+    print("  5x2T — Chat")
+    print("  Commands:")
+    print("    quit       — exit")
+    print("    temp X     — temperature e.g. temp 0.7")
+    print("    length X   — response length e.g. length 60")
+    print("    maxlen X   — max length cap e.g. maxlen 300")
+    print("    wiki X     — force Wikipedia lookup e.g. wiki Python")
+    print("=" * 55 + "\n")
+
+    if model is None:
+        if not os.path.exists(VOCAB_PATH):
+            print("[ERROR] No vocab found. Run training first.")
+            return
+        if not os.path.exists(MODEL_PATH):
+            print("[ERROR] No model found. Run training first.")
+            return
+        vocab_words, word2idx, idx2word = load_vocab(VOCAB_PATH)
+        model = Model(len(vocab_words))
+        model.load_state_dict(torch.load(MODEL_PATH, map_location="cpu"))
+        model.eval()
+        param_count = sum(p.numel() for p in model.parameters())
+        print(f"  Vocab size       : {len(vocab_words):,}")
+        print(f"  Model parameters : {param_count:,}")
+        print(f"  Device           : {DEVICE}\n")
+
+    temperature     = TEMPERATURE
+    response_length = RESPONSE_LENGTH
+    max_length      = MAX_LENGTH
+
+    while True:
+        user_input = input("You: ").strip()
+
+        if not user_input:
+            continue
+        if user_input.lower() in ("quit", "exit", "q"):
+            print("Goodbye.")
+            break
+        if user_input.lower().startswith("temp "):
+            try:
+                temperature = float(user_input.split()[1])
+                print(f"  Temperature -> {temperature}\n")
+            except:
+                print("  Usage: temp 0.8\n")
+            continue
+        if user_input.lower().startswith("length "):
+            try:
+                response_length = int(user_input.split()[1])
+                print(f"  Length -> {response_length}\n")
+            except:
+                print("  Usage: length 50\n")
+            continue
+        if user_input.lower().startswith("maxlen "):
+            try:
+                max_length = int(user_input.split()[1])
+                print(f"  Max length -> {max_length}\n")
+            except:
+                print("  Usage: maxlen 300\n")
+            continue
+
+        if user_input.lower().startswith("wiki "):
+            query  = user_input[5:].strip()
+            result = search_wikipedia(query)
+            reply  = result if result else f"No Wikipedia page found for '{query}'"
+            print(f"5x2T: {reply}\n")
+            continue
+
+        wiki_result = None
+        if should_search_wiki(user_input):
+            wiki_result = search_wikipedia(user_input)
+
+        if wiki_result:
+            print(f"5x2T: {wiki_result}\n")
+        else:
+            seed  = f'"{user_input.lower()}"\n"'
+            raw   = generate(model, word2idx, idx2word,
+                             seed_text=seed,
+                             length=min(response_length, max_length),
+                             temperature=temperature)
+            reply = format_response(raw)
+            print(f"5x2T: {reply}\n")
+
+
+# ---------------- ENTRY POINT ----------------
+if __name__ == "__main__":
+    import sys
+    if len(sys.argv) > 1 and sys.argv[1] == "chat":
+        chat()
+    else:
+        model, word2idx, idx2word = train()
+        print("\nStarting chat...\n")
+        chat(model, word2idx, idx2word)