Initialization code for the OpenSoftware-World-OSW1 AI model. (This code was written by Claude and edited by OpenSoftware-World.)

Download one of our OpenSoftware-World-OSW1:5m, OpenSoftware-World-OSW1:10m, or OpenSoftware-World-OSW1:100m AI models and place it in the same folder as model_init.py. You can then start chatting with the OpenSoftware-World-OSW1 AI model.

model_init.py

@@ -0,0 +1,237 @@
+import os
+import re
+import sys
+import glob
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+NUM_THREADS = os.cpu_count() or 4
+torch.set_num_threads(NUM_THREADS)
+try:
+    torch.set_num_interop_threads(max(1, NUM_THREADS // 2))
+except RuntimeError:
+    pass
+
+DEVICE = torch.device("cpu")
+print(f"🧵 Number of CPU threads : {NUM_THREADS}")
+TOKEN_RE = re.compile(r"\w+|[^\w\s]", re.UNICODE)
+
+def tokenize(text: str):
+    return TOKEN_RE.findall(text.lower())
+
+class Vocab:
+    PAD, UNK, BOS, EOS = "<pad>", "<unk>", "<bos>", "<eos>"
+
+    def __init__(self):
+        self.stoi = {}
+        self.itos = []
+
+    def encode(self, text, add_bos=False, add_eos=False):
+        ids = [self.stoi.get(t, self.stoi[Vocab.UNK]) for t in tokenize(text)]
+        if add_bos:
+            ids = [self.stoi[Vocab.BOS]] + ids
+        if add_eos:
+            ids = ids + [self.stoi[Vocab.EOS]]
+        return ids
+
+    def decode(self, ids):
+        toks = [self.itos[i] for i in ids if 0 <= i < len(self.itos)]
+        toks = [t for t in toks if t != Vocab.PAD and t != Vocab.BOS]
+        out = []
+        for t in toks:
+            if t == Vocab.EOS:
+                break
+            out.append(t)
+        text = " ".join(out)
+        text = re.sub(r"\s+([.,!?;:])", r"\1", text)
+        return text
+
+    def __len__(self):
+        return len(self.itos)
+
+class CausalSelfAttention(nn.Module):
+    def __init__(self, d_model, n_head, dropout):
+        super().__init__()
+        assert d_model % n_head == 0, "d_model must be evenly divisible by n_head"
+        self.n_head = n_head
+        self.head_dim = d_model // n_head
+        self.qkv = nn.Linear(d_model, 3 * d_model)
+        self.proj = nn.Linear(d_model, d_model)
+        self.attn_drop = nn.Dropout(dropout)
+        self.resid_drop = nn.Dropout(dropout)
+
+    def forward(self, x, attn_mask):
+        B, T, C = x.shape
+        qkv = self.qkv(x)
+        q, k, v = qkv.split(C, dim=2)
+        q = q.view(B, T, self.n_head, self.head_dim).transpose(1, 2)
+        k = k.view(B, T, self.n_head, self.head_dim).transpose(1, 2)
+        v = v.view(B, T, self.n_head, self.head_dim).transpose(1, 2)
+
+        att = (q @ k.transpose(-2, -1)) / math.sqrt(self.head_dim)
+        att = att.masked_fill(attn_mask, float("-inf"))
+        att = F.softmax(att, dim=-1)
+        att = self.attn_drop(att)
+        out = att @ v
+        out = out.transpose(1, 2).contiguous().view(B, T, C)
+        return self.resid_drop(self.proj(out))
+
+class TransformerBlock(nn.Module):
+    def __init__(self, d_model, n_head, d_ff, dropout):
+        super().__init__()
+        self.ln1 = nn.LayerNorm(d_model)
+        self.attn = CausalSelfAttention(d_model, n_head, dropout)
+        self.ln2 = nn.LayerNorm(d_model)
+        self.mlp = nn.Sequential(
+            nn.Linear(d_model, d_ff),
+            nn.GELU(),
+            nn.Linear(d_ff, d_model),
+            nn.Dropout(dropout),
+        )
+
+    def forward(self, x, attn_mask):
+        x = x + self.attn(self.ln1(x), attn_mask)
+        x = x + self.mlp(self.ln2(x))
+        return x
+
+class OSW1Model(nn.Module):
+    def __init__(self, vocab_size, cfg: dict, pad_id: int):
+        super().__init__()
+        self.cfg = cfg
+        self.pad_id = pad_id
+        self.block_size = cfg["block_size"]
+
+        self.tok_emb = nn.Embedding(vocab_size, cfg["d_model"])
+        self.pos_emb = nn.Embedding(cfg["block_size"], cfg["d_model"])
+        self.drop = nn.Dropout(cfg["dropout"])
+        self.blocks = nn.ModuleList([
+            TransformerBlock(cfg["d_model"], cfg["n_head"], cfg["d_ff"], cfg["dropout"])
+            for _ in range(cfg["n_layer"])
+        ])
+        self.ln_f = nn.LayerNorm(cfg["d_model"])
+        self.head = nn.Linear(cfg["d_model"], vocab_size, bias=False)
+        self.head.weight = self.tok_emb.weight  # weight tying
+
+    def forward(self, idx):
+        B, T = idx.shape
+        pos = torch.arange(T, device=idx.device).unsqueeze(0)
+        x = self.drop(self.tok_emb(idx) + self.pos_emb(pos))
+
+        mask = torch.triu(torch.ones(T, T, dtype=torch.bool, device=idx.device), diagonal=1)
+        for block in self.blocks:
+            x = block(x, mask)
+        x = self.ln_f(x)
+        return self.head(x)
+
+    @torch.no_grad()
+    def generate(self, idx, max_new_tokens, temperature=0.85, top_k=40, eos_id=None):
+        self.eval()
+        for _ in range(max_new_tokens):
+            idx_cond = idx[:, -self.block_size:]
+            logits = self(idx_cond)
+            logits = logits[:, -1, :] / max(temperature, 1e-5)
+            if top_k is not None:
+                v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
+                logits[logits < v[:, [-1]]] = float("-inf")
+            probs = F.softmax(logits, dim=-1)
+            next_id = torch.multinomial(probs, num_samples=1)
+            idx = torch.cat([idx, next_id], dim=1)
+            if eos_id is not None and next_id.item() == eos_id:
+                break
+        return idx
+
+def find_checkpoint():
+    candidates = glob.glob("opensoftware_world_osw1_*.pth")
+    if not candidates:
+        return None
+    candidates.sort(key=os.path.getmtime, reverse=True)
+    return candidates[0]
+
+def load_checkpoint(path: str):
+    print(f"📦 Loading: {path}")
+    ckpt = torch.load(path, map_location="cpu")
+
+    cfg = ckpt["config"]
+    vocab = Vocab()
+    vocab.stoi = ckpt["vocab_stoi"]
+    vocab.itos = ckpt["vocab_itos"]
+    pad_id = ckpt["pad_id"]
+
+    model = OSW1Model(len(vocab), cfg, pad_id=pad_id).to(DEVICE)
+    model.load_state_dict(ckpt["model_state_dict"])
+    model.eval()
+
+    param_count = ckpt.get("param_count", sum(p.numel() for p in model.parameters()))
+    training_time = ckpt.get("training_time_sec", None)
+    final_loss = ckpt.get("final_loss", None)
+
+    print("\n" + "=" * 64)
+    print("🧠  OpenSoftware-World OSW1 — LOADED MODEL INFORMATION")
+    print("=" * 64)
+    print(f"  File                 : {path}")
+    print(f"  Vocab size           : {len(vocab):,}")
+    print(f"  Number of parameters : {param_count:,}")
+    print(f"  d_model / n_layer    : {cfg['d_model']} / {cfg['n_layer']}")
+    print(f"  n_head / d_ff        : {cfg['n_head']} / {cfg['d_ff']}")
+    print(f"  Context window       : {cfg['block_size']}")
+    if training_time is not None:
+        print(f"  Training time        : {training_time/60:.2f} minutes")
+    if final_loss is not None:
+        print(f"  Final training loss  : {final_loss:.4f}")
+    print("=" * 64 + "\n")
+
+    return model, vocab, cfg
+
+def chat_loop(model: OSW1Model, vocab: Vocab):
+    print("=" * 64)
+    print("💬 OSW1 ready! You can start chatting. Type 'exit' to quit.")
+    print("=" * 64)
+
+    eos_id = vocab.stoi[Vocab.EOS]
+    bos_id = vocab.stoi[Vocab.BOS]
+
+    while True:
+        try:
+            user_in = input("\nYou: ").strip()
+        except (EOFError, KeyboardInterrupt):
+            print("\n👋 Goodbye!")
+            break
+
+        if user_in.lower() in ("exit", "quit"):
+            print("👋 Goodbye!")
+            break
+        if not user_in:
+            continue
+
+        ids = [bos_id] + vocab.encode(user_in)
+        x = torch.tensor([ids], dtype=torch.long)
+        out = model.generate(x, max_new_tokens=60, temperature=0.85, top_k=40, eos_id=eos_id)
+        answer_ids = out[0, len(ids):].tolist()
+        answer = vocab.decode(answer_ids)
+        print(f"OSW1: {answer if answer else '(...silence...)'}")
+
+def main():
+    if len(sys.argv) > 1:
+        ckpt_path = sys.argv[1]
+        if not os.path.isfile(ckpt_path):
+            print(f"❌ File not found: {ckpt_path}")
+            sys.exit(1)
+    else:
+        ckpt_path = find_checkpoint()
+        if ckpt_path is None:
+            print(
+                "❌ No checkpoint files found in the directory.\n"
+                "   Please train a model using 'python train_osw1.py' or\n"
+                "   specify a checkpoint file using 'python model_init.py <file_path>'."
+            )
+            sys.exit(1)
+
+    model, vocab, cfg = load_checkpoint(ckpt_path)
+    chat_loop(model, vocab)
+
+
+if __name__ == "__main__":
+    main()