Add exact vocab.json (95 tokens); align model param names; add guard

bot_model.py

@@ -1,3 +1,4 @@
+# bot_model.py
 import math
 import torch
 import torch.nn as nn
@@ -9,51 +10,50 @@ class CausalSelfAttention(nn.Module):
         assert d_model % num_heads == 0
         self.num_heads = num_heads
         self.d_k = d_model // num_heads
-        self.q = nn.Linear(d_model, d_model)
-        self.k = nn.Linear(d_model, d_model)
-        self.v = nn.Linear(d_model, d_model)
-        self.o = nn.Linear(d_model, d_model)
-        self.drop = nn.Dropout(dropout)
+        # NOTE: match checkpoint names exactly
+        self.W_q = nn.Linear(d_model, d_model)
+        self.W_k = nn.Linear(d_model, d_model)
+        self.W_v = nn.Linear(d_model, d_model)
+        self.W_o = nn.Linear(d_model, d_model)
+        self.dropout = nn.Dropout(dropout)
 
     def forward(self, x):
         B, T, C = x.size()
-        q = self.q(x).view(B, T, self.num_heads, C//self.num_heads).transpose(1,2)
-        k = self.k(x).view(B, T, self.num_heads, C//self.num_heads).transpose(1,2)
-        v = self.v(x).view(B, T, self.num_heads, C//self.num_heads).transpose(1,2)
-        att = (q @ k.transpose(-2,-1)) / math.sqrt(self.d_k)
+        q = self.W_q(x).view(B, T, self.num_heads, C//self.num_heads).transpose(1, 2)
+        k = self.W_k(x).view(B, T, self.num_heads, C//self.num_heads).transpose(1, 2)
+        v = self.W_v(x).view(B, T, self.num_heads, C//self.num_heads).transpose(1, 2)
+        att = (q @ k.transpose(-2, -1)) / math.sqrt(self.d_k)
         mask = torch.tril(torch.ones(T, T, device=x.device))
-        att = att.masked_fill(mask==0, float('-inf'))
+        att = att.masked_fill(mask == 0, float('-inf'))
         att = F.softmax(att, dim=-1)
-        att = self.drop(att)
+        att = self.dropout(att)
         y = att @ v
-        y = y.transpose(1,2).contiguous().view(B, T, C)
-        return self.o(y)
+        y = y.transpose(1, 2).contiguous().view(B, T, C)
+        return self.W_o(y)
 
 class GPTBlock(nn.Module):
-    def __init__(self, d_model, num_heads, d_ff, dropout=0.1):
+    def __init__(self, d_model, num_heads, d_ff=768, dropout=0.1):
         super().__init__()
         self.ln1 = nn.LayerNorm(d_model)
         self.attn = CausalSelfAttention(d_model, num_heads, dropout)
         self.ln2 = nn.LayerNorm(d_model)
-        self.ff = nn.Sequential(
-            nn.Linear(d_model, d_ff),
-            nn.ReLU(),
-            nn.Linear(d_ff, d_model),
-            nn.Dropout(dropout)
+        # NOTE: match checkpoint module name "ffn"
+        self.ffn = nn.Sequential(
+            nn.Linear(d_model, d_ff), nn.ReLU(), nn.Linear(d_ff, d_model), nn.Dropout(dropout)
         )
     def forward(self, x):
         x = x + self.attn(self.ln1(x))
-        x = x + self.ff(self.ln2(x))
+        x = x + self.ffn(self.ln2(x))
         return x
 
 class GPT(nn.Module):
     def __init__(self, vocab_size, d_model=192, num_heads=6, num_layers=4, max_len=320, dropout=0.1):
         super().__init__()
         self.token_emb = nn.Embedding(vocab_size, d_model)
-        self.pos_emb = nn.Embedding(max_len, d_model)
-        self.blocks = nn.ModuleList([GPTBlock(d_model, num_heads, 4*d_model, dropout) for _ in range(num_layers)])
-        self.ln_f = nn.LayerNorm(d_model)
-        self.head = nn.Linear(d_model, vocab_size, bias=False)
+        self.pos_emb   = nn.Embedding(max_len, d_model)
+        self.blocks    = nn.ModuleList([GPTBlock(d_model, num_heads, 4*d_model, dropout) for _ in range(num_layers)])
+        self.ln_f      = nn.LayerNorm(d_model)
+        self.head      = nn.Linear(d_model, vocab_size, bias=False)
         self.cfg = dict(vocab_size=vocab_size, d_model=d_model, num_heads=num_heads, num_layers=num_layers, max_len=max_len, dropout=dropout)
 
     def forward(self, idx):

export_vocab.py

@@ -0,0 +1,137 @@
+# export_vocab.py
+import json, random
+
+# ---- Gen config (must mirror your training script that produced this checkpoint) ----
+EXTRA_CHARSET = (
+    "ABCDEFGHIJKLMNOPQRSTUVWXYZ"
+    "abcdefghijklmnopqrstuvwxyz"
+    "0123456789"
+    " +-*/=?:.,'\"()[]"
+    "$£€₹¥₩₦"
+)
+
+ITEMS = ["apple","mango","banana","orange","pear","kiwi","pen","notebook","ticket","cookie","product"]
+ASKS = [
+    "how much is {n} {item_plural}?",
+    "what's the cost of {n} {item_plural}?",
+    "price for {n} {item_plural}?",
+    "how much would {n} {item_plural} cost?",
+    "what do I pay for {n} {item_plural}?",
+    "total price of {n} {item_plural}?",
+    "what's {n} {item_plural} going to cost?"
+]
+COPULA = ["is", "costs", "priced at", "= ", "comes to", "sells for"]
+
+def s(x, w): return w if x == 1 else (w + "s")
+
+def money(v: float) -> str:
+    return str(int(round(v))) if abs(v - round(v)) < 1e-9 else f"{v:.2f}"
+
+CURRENCIES = [
+    {"prefix":"$",  "name_sing":"dollar",  "name_plur":"dollars"},
+    {"prefix":"£",  "name_sing":"pound",   "name_plur":"pounds"},
+    {"prefix":"€",  "name_sing":"euro",    "name_plur":"euros"},
+    {"prefix":"₹",  "name_sing":"rupee",   "name_plur":"rupees"},
+    {"prefix":"¥",  "name_sing":"yen",     "name_plur":"yen"},
+    {"prefix":"₩",  "name_sing":"won",     "name_plur":"won"},
+    {"prefix":"₦",  "name_sing":"naira",   "name_plur":"naira"},
+    {"prefix":"A$", "name_sing":"dollar",  "name_plur":"dollars"},
+    {"prefix":"C$", "name_sing":"dollar",  "name_plur":"dollars"},
+    {"prefix":"CHF","name_sing":"franc",   "name_plur":"francs"},
+    {"prefix":"د.إ","name_sing":"dirham",  "name_plur":"dirhams"},
+    {"prefix":"",   "name_sing":"dollar",  "name_plur":"dollars"},
+    {"prefix":"",   "name_sing":"pound",   "name_plur":"pounds"},
+    {"prefix":"",   "name_sing":"rupee",   "name_plur":"rupees"},
+    {"prefix":"",   "name_sing":"euro",    "name_plur":"euros"}
+]
+EXTRA_CURRENCY_WORDS = [
+    ("peso","pesos"), ("dinar","dinars"), ("dirham","dirhams"), ("riyal","riyals"),
+    ("rial","rials"), ("krona","kronor"), ("krone","kroner"), ("koruna","koruny"),
+    ("rand","rand"), ("shekel","shekels"), ("lira","lire"), ("ruble","rubles"),
+    ("rouble","roubles"), ("yuan","yuan"), ("renminbi","renminbi"), ("ringgit","ringgit"),
+    ("baht","baht"), ("dong","dong"), ("taka","taka"), ("rupiah","rupiah"),
+    ("shilling","shillings"), ("kuna","kuna"), ("zloty","zlotys"), ("forint","forints"),
+    ("leu","lei"), ("riel","riel"), ("kyat","kyat"), ("quetzal","quetzales"),
+    ("bolivar","bolivares"), ("colon","colones"), ("sol","soles")
+]
+SYNTHETIC_CURRENCY_WORDS = [
+    ("credit","credits"), ("token","tokens"), ("coin","coins"), ("unit","units"),
+    ("galactic credit","galactic credits"), ("space buck","space bucks")
+]
+WORD_POOL = [{"prefix":"", "name_sing":a, "name_plur":b} for (a,b) in (EXTRA_CURRENCY_WORDS + SYNTHETIC_CURRENCY_WORDS)]
+
+# We only need enough samples to surface all characters that occurred in training text
+def cur_str(cur: dict, v: float) -> str:
+    if cur["prefix"]:
+        return f"{cur['prefix']}{money(v)}"
+    name = cur["name_sing"] if abs(v - 1.0) < 1e-9 else cur["name_plur"]
+    return f"{money(v)} {name}"
+
+def make_addition():
+    k = random.randint(2, 6)
+    nums = [random.randint(0, 200) for _ in range(k)]
+    expr = " + ".join(map(str, nums))
+    s_ = sum(nums)
+    tail = random.choice(["", " = ?", " = ", "?"])
+    return f"q: {expr}{tail} a: step 1: take {nums[0]}. step 2: add {nums[1]}. step 3: {expr} = {s_}. final answer: {s_} <EOS>"
+
+def make_subtraction():
+    a = random.randint(0,200); b = random.randint(0,a)
+    s_ = a - b
+    tail = random.choice(["", " = ?", " = ", "?"])
+    return f"q: {a} - {b}{tail} a: step 1: take {a}. step 2: subtract {b}. step 3: {a} - {b} = {s_}. final answer: {s_} <EOS>"
+
+def make_multiplication():
+    a = random.randint(0,50); b = random.randint(0,20)
+    s_ = a*b
+    tail = random.choice(["", " = ?", " = ", "?"])
+    return f"q: {a} × {b}{tail} a: step 1: take {a}. step 2: multiply by {b}. step 3: {a} × {b} = {s_}. final answer: {s_} <EOS>"
+
+def make_division():
+    b = random.randint(1, 20); qv = random.randint(0, 50); a = b*qv
+    tail = random.choice(["", " = ?", " = ", "?"])
+    return f"q: {a} ÷ {b}{tail} a: step 1: take {a}. step 2: divide by {b}. step 3: {a} ÷ {b} = {qv}. final answer: {qv} <EOS>"
+
+def make_unit_price():
+    item = random.choice(ITEMS); qty = random.randint(1, 15)
+    unit_price = random.choice([0.39,0.5,0.75,1.2,1.5,2.25,3.4,4.99]) if (random.random()<0.5) else random.randint(1,30)
+    cost = unit_price * qty; target = random.randint(1, 30)
+    cur = random.choice(CURRENCIES) if random.random() < 0.7 else random.choice(WORD_POOL)
+    cop = random.choice(COPULA); ask = random.choice(ASKS)
+    if random.random() < 0.5:
+        q = f"{qty} {s(qty,item)} {cop} {cur_str(cur,cost)}, {ask.format(n=target, item_plural=s(target,item))}"
+    else:
+        q = f"{qty} {s(qty,item)} cost {cur_str(cur,cost)}. {ask.format(n=target, item_plural=s(target,item))}"
+    unit = cost/qty; total = unit*target
+    return f"q: {q} a: step 1: find unit price: {cur_str(cur,cost)} / {qty} = {cur_str(cur,unit)} per {item}. step 2: multiply by {target}: {cur_str(cur,unit)} × {target} = {cur_str(cur,total)}. final answer: {cur_str(cur,total)} <EOS>"
+
+def make_group_price():
+    item = random.choice(ITEMS); group_qty = random.randint(2, 20)
+    price = random.choice([9.0,10.0,12.5,7.2,15.75,19.99,3.3,4.5,21.95]) if (random.random()<0.7) else float(random.randint(2,60))
+    target = random.randint(1,40); cur = random.choice(CURRENCIES)
+    q = f"{group_qty} {s(group_qty,item)} for {cur_str(cur,price)}, what's the price for {target} {s(target,item)}?"
+    unit = price/group_qty; total = unit*target
+    return f"q: {q} a: step 1: unit price = {cur_str(cur,price)} / {group_qty} = {cur_str(cur,unit)} per {item}. step 2: total = {cur_str(cur,unit)} × {target} = {cur_str(cur,total)}. final answer: {cur_str(cur,total)} <EOS>"
+
+def build_seed_pool():
+    pool = []
+    for _ in range(200): pool.append(make_addition())
+    for _ in range(150): pool.append(make_subtraction())
+    for _ in range(150): pool.append(make_multiplication())
+    for _ in range(150): pool.append(make_division())
+    for _ in range(150): pool.append(make_unit_price())
+    for _ in range(100): pool.append(make_group_price())
+    random.shuffle(pool)
+    return pool
+
+if __name__ == "__main__":
+    SEED_POOL = build_seed_pool()
+    seed_text = "\n".join(SEED_POOL)
+    base_chars = sorted(set(seed_text) | set(EXTRA_CHARSET))
+    # IMPORTANT: the older checkpoint DID NOT include <EOS> in the vocab
+    SPECIAL = ["<PAD>", "<UNK>"]
+    itos_list = SPECIAL + base_chars
+    print("len(itos) =", len(itos_list))
+    with open("vocab.json", "w", encoding="utf-8") as f:
+        json.dump({"itos": itos_list}, f, ensure_ascii=False, indent=2)
+    print("Wrote vocab.json")