Spaces:
Sleeping
Sleeping
| import math, re, torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| # ========================= | |
| # Model | |
| # ========================= | |
| class CausalSelfAttention(nn.Module): | |
| def __init__(self, d_model, num_heads, dropout=0.1): | |
| super().__init__() | |
| assert d_model % num_heads == 0 | |
| self.num_heads = num_heads | |
| self.d_k = d_model // num_heads | |
| 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.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 = F.softmax(att, dim=-1) | |
| att = self.dropout(att) | |
| y = att @ v | |
| 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=256, 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.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.ffn(self.ln2(x)) | |
| return x | |
| class GPT(nn.Module): | |
| def __init__(self, vocab_size, d_model=128, num_heads=4, num_layers=2, max_len=128, 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, 256, dropout) for _ in range(num_layers)]) | |
| self.ln_f = nn.LayerNorm(d_model) | |
| self.head = nn.Linear(d_model, vocab_size, bias=False) | |
| def forward(self, idx): | |
| B, T = idx.shape | |
| pos = torch.arange(0, T, device=idx.device).unsqueeze(0) | |
| x = self.token_emb(idx) + self.pos_emb(pos) | |
| for blk in self.blocks: | |
| x = blk(x) | |
| x = self.ln_f(x) | |
| return self.head(x) | |
| # ========================= | |
| # Tokenization helpers | |
| # ========================= | |
| def encode(s: str, stoi: dict, space_id: int): | |
| return [stoi.get(c, space_id) for c in s] | |
| def decode(ids, itos: dict): | |
| return "".join(itos[i] for i in ids) | |
| # ========================= | |
| # Greedy generation | |
| # ========================= | |
| def greedy_until_eos(model, idx, itos, max_new_tokens=256, stop_text="<EOS>"): | |
| gen = idx | |
| for _ in range(max_new_tokens): | |
| idx_cond = gen[:, -model.pos_emb.num_embeddings:] | |
| logits = model(idx_cond)[:, -1, :] | |
| next_id = torch.argmax(logits, dim=-1, keepdim=True) | |
| gen = torch.cat([gen, next_id], dim=1) | |
| if stop_text in decode(gen[0].tolist(), itos): | |
| break | |
| return gen | |
| # ========================= | |
| # Number parsing | |
| # ========================= | |
| _UNITS = { | |
| "zero":0,"one":1,"two":2,"three":3,"four":4,"five":5,"six":6,"seven":7,"eight":8,"nine":9, | |
| "ten":10,"eleven":11,"twelve":12,"thirteen":13,"fourteen":14,"fifteen":15,"sixteen":16, | |
| "seventeen":17,"eighteen":18,"nineteen":19 | |
| } | |
| _TENS = {"twenty":20,"thirty":30,"forty":40,"fifty":50,"sixty":60,"seventy":70,"eighty":80,"ninety":90} | |
| _SCALES = {"hundred":100,"thousand":1000,"million":1_000_000} | |
| def words_to_int(text: str): | |
| t = text.lower().replace("-", " ") | |
| toks = re.findall(r"[a-z]+|\d+", t) | |
| for tok in toks: | |
| if tok.isdigit(): | |
| return int(tok) | |
| total, current, sign = 0, 0, 1 | |
| for tok in toks: | |
| if tok in _UNITS: current += _UNITS[tok] | |
| elif tok in _TENS: current += _TENS[tok] | |
| elif tok == "hundred": current = (current or 1) * 100 | |
| elif tok == "thousand": total += (current or 1) * 1000; current = 0 | |
| elif tok == "million": total += (current or 1) * 1_000_000; current = 0 | |
| elif tok == "minus" or tok == "negative": sign = -1 | |
| return sign * (total + current) | |
| def parse_number(piece: str): | |
| m = re.search(r"[-+]?\d[\d,]*\.?\d*", piece.replace(" ", "")) | |
| if m: | |
| s = m.group(0).replace(",", "") | |
| try: | |
| return int(s) if "." not in s else int(round(float(s))) | |
| except ValueError: | |
| return None | |
| return words_to_int(piece) | |
| def extract_numbers(text: str): | |
| t = text.lower().strip() | |
| parts = re.split(r"[+]|(?:and)|,| ", t) | |
| nums = [parse_number(p) for p in parts if p.strip()] | |
| return [n for n in nums if n is not None] | |
| # ========================= | |
| # Helpers for final answer | |
| # ========================= | |
| def patch_final_answer(text: str, nums=None) -> str: | |
| try: | |
| s = sum(nums) | |
| text = re.sub(r"(=\s*)(-?\d+)", fr"\g<1>{s}", text, count=1) | |
| if "final answer:" in text: | |
| text = re.sub(r"(final answer:\s*)(-?\d+)", fr"\g<1>{s}", text) | |
| return text | |
| except Exception: | |
| return text | |
| def render_add_steps(nums: list[int]) -> str: | |
| steps = [] | |
| total = nums[0] | |
| steps.append(f"step 1: take {nums[0]}.") | |
| for i, n in enumerate(nums[1:], start=2): | |
| prev = total | |
| total += n | |
| steps.append(f"step {i}: add {n}. {prev} + {n} = {total}.") | |
| steps.append(f"final answer: {total}") | |
| return " ".join(steps) | |