| import argparse |
| import json |
| import math |
| from pathlib import Path |
| from contextlib import nullcontext |
|
|
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| import sentencepiece as spm |
|
|
| torch.set_num_threads(12) |
| torch.set_num_interop_threads(12) |
|
|
| |
| |
| |
| |
|
|
| class RMSNorm(nn.Module): |
| def __init__(self, dim, eps=1e-6): |
| super().__init__() |
| self.weight = nn.Parameter(torch.ones(dim)) |
| self.eps = eps |
| def forward(self, x): |
| return x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps) * self.weight |
|
|
| def rope_cache(seq_len, head_dim, device, dtype): |
| inv_freq = 1.0 / (10000 ** (torch.arange(0, head_dim, 2, device=device, dtype=torch.float32) / head_dim)) |
| pos = torch.arange(seq_len, device=device, dtype=torch.float32) |
| freqs = torch.outer(pos, inv_freq) |
| cos = freqs.cos().to(dtype=dtype)[None, None, :, :] |
| sin = freqs.sin().to(dtype=dtype)[None, None, :, :] |
| return cos, sin |
|
|
| def apply_rope(x, cos, sin): |
| x1 = x[..., ::2] |
| x2 = x[..., 1::2] |
| out = torch.empty_like(x) |
| out[..., ::2] = x1 * cos - x2 * sin |
| out[..., 1::2] = x1 * sin + x2 * cos |
| return out |
|
|
| class CausalSelfAttention(nn.Module): |
| def __init__(self, dim, n_head, dropout): |
| super().__init__() |
| assert dim % n_head == 0 |
| self.n_head = n_head |
| self.head_dim = dim // n_head |
| assert self.head_dim % 2 == 0 |
| self.qkv = nn.Linear(dim, 3 * dim) |
| self.proj = nn.Linear(dim, dim) |
| self.dropout = dropout |
| def forward(self, x): |
| B, T, C = x.shape |
| q, k, v = self.qkv(x).chunk(3, dim=-1) |
| 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) |
| cos, sin = rope_cache(T, self.head_dim, x.device, x.dtype) |
| q = apply_rope(q, cos, sin) |
| k = apply_rope(k, cos, sin) |
| if hasattr(F, "scaled_dot_product_attention"): |
| a = F.scaled_dot_product_attention(q, k, v, dropout_p=0.0, is_causal=True) |
| else: |
| att = (q @ k.transpose(-2, -1)) / math.sqrt(self.head_dim) |
| mask = torch.triu(torch.ones(T, T, device=x.device, dtype=torch.bool), diagonal=1) |
| att = att.masked_fill(mask, float("-inf")) |
| att = F.softmax(att, dim=-1) |
| a = att @ v |
| a = a.transpose(1, 2).contiguous().view(B, T, C) |
| return self.proj(a) |
|
|
| class SwiGLU(nn.Module): |
| def __init__(self, dim, dropout): |
| super().__init__() |
| hidden = 4 * dim |
| self.fc = nn.Linear(dim, hidden * 2) |
| self.proj = nn.Linear(hidden, dim) |
| self.drop = nn.Dropout(dropout) |
| def forward(self, x): |
| x1, x2 = self.fc(x).chunk(2, dim=-1) |
| return self.drop(self.proj(F.silu(x1) * x2)) |
|
|
| class Block(nn.Module): |
| def __init__(self, dim, n_head, dropout): |
| super().__init__() |
| self.n1 = RMSNorm(dim) |
| self.attn = CausalSelfAttention(dim, n_head, dropout) |
| self.n2 = RMSNorm(dim) |
| self.mlp = SwiGLU(dim, dropout) |
| def forward(self, x): |
| x = x + self.attn(self.n1(x)) |
| x = x + self.mlp(self.n2(x)) |
| return x |
|
|
| class GPT(nn.Module): |
| def __init__(self, vocab_size, block_size, n_layer, n_head, n_embd, dropout=0.0): |
| super().__init__() |
| self.block_size = block_size |
| self.tok_emb = nn.Embedding(vocab_size, n_embd) |
| self.drop = nn.Dropout(dropout) |
| self.blocks = nn.ModuleList([Block(n_embd, n_head, dropout) for _ in range(n_layer)]) |
| self.norm_f = RMSNorm(n_embd) |
| self.lm_head = nn.Linear(n_embd, vocab_size, bias=False) |
| self.lm_head.weight = self.tok_emb.weight |
| def forward(self, idx): |
| B, T = idx.shape |
| if T > self.block_size: |
| idx = idx[:, -self.block_size:] |
| x = self.tok_emb(idx) |
| x = self.drop(x) |
| for block in self.blocks: |
| x = block(x) |
| x = self.norm_f(x) |
| logits = self.lm_head(x) |
| return logits |
|
|
| def get_stop_ids(sp): |
| ids = set() |
| for piece in ("<|user|>", "<|system|>"): |
| pid = sp.piece_to_id(piece) |
| if pid != sp.unk_id(): |
| ids.add(pid) |
| return ids |
|
|
| DEFAULT_CONFIG = { |
| "vocab_size": 32000, |
| "block_size": 512, |
| "n_layer": 10, |
| "n_head": 8, |
| "n_embd": 576, |
| } |
|
|
|
|
| def pick_device(): |
| if torch.cuda.is_available(): |
| return torch.device("cuda") |
| if torch.backends.mps.is_available(): |
| return torch.device("mps") |
| return torch.device("cpu") |
|
|
| _MODULE_DIR = Path(__file__).resolve().parent |
|
|
| DEFAULT_CKPT = _MODULE_DIR / "qWisp-base-v1.pt" |
| DEFAULT_TOKENIZER = _MODULE_DIR / "qWisp.model" |
|
|
| class Wisp: |
| def __init__(self): |
| self.device = pick_device() |
| self.model = None |
| self.tokenizer = None |
| self.stop_ids = None |
| self.config = DEFAULT_CONFIG.copy() |
|
|
| def load( |
| self, |
| ckpt=str(DEFAULT_CKPT), |
| tokenizer=str(DEFAULT_TOKENIZER), |
| device=None, |
| ): |
| if device is not None: |
| self.device = torch.device(device) |
|
|
| self.tokenizer = spm.SentencePieceProcessor() |
| self.tokenizer.load(tokenizer) |
|
|
| self.model = GPT( |
| vocab_size=self.config["vocab_size"], |
| block_size=self.config["block_size"], |
| n_layer=self.config["n_layer"], |
| n_head=self.config["n_head"], |
| n_embd=self.config["n_embd"], |
| dropout=0.0, |
| ) |
|
|
| obj = torch.load(ckpt, map_location=self.device) |
|
|
| state_dict = obj["model"] if isinstance(obj, dict) and "model" in obj else obj |
|
|
| self.model.load_state_dict(state_dict, strict=True) |
| self.model.to(self.device) |
| self.model.eval() |
|
|
| self.stop_ids = get_stop_ids(self.tokenizer) |
|
|
| return self |
|
|
| def unload(self): |
| self.model = None |
| self.tokenizer = None |
| self.stop_ids = None |
|
|
| if torch.cuda.is_available(): |
| torch.cuda.empty_cache() |
|
|
| def encode(self, text): |
| if self.tokenizer is None: |
| raise RuntimeError("Model not loaded.") |
| return self.tokenizer.encode(text, out_type=int) |
|
|
| def decode(self, ids): |
| if self.tokenizer is None: |
| raise RuntimeError("Model not loaded.") |
| return self.tokenizer.decode(ids) |
|
|
| @torch.no_grad() |
| def generate( |
| self, |
| prompt, |
| max_new_tokens=200, |
| temperature=0.8, |
| top_k=50, |
| ): |
| if self.model is None: |
| raise RuntimeError("Model not loaded.") |
|
|
| ids = self.tokenizer.encode(prompt, out_type=int) |
|
|
| if len(ids) == 0: |
| ids = [self.tokenizer.bos_id()] |
|
|
| x = torch.tensor([ids], dtype=torch.long, device=self.device) |
|
|
| self.model.eval() |
|
|
| for _ in range(max_new_tokens): |
| x_cond = x[:, -self.model.block_size :] |
| logits = self.model(x_cond) |
| logits = logits[:, -1] / max(temperature, 1e-6) |
|
|
| if top_k is not None and top_k > 0: |
| values, _ = torch.topk(logits, min(top_k, logits.size(-1))) |
| logits = torch.where( |
| logits < values[:, [-1]], |
| torch.full_like(logits, float("-inf")), |
| logits, |
| ) |
|
|
| probs = F.softmax(logits, dim=-1) |
|
|
| next_token = torch.multinomial(probs, 1) |
| token = int(next_token.item()) |
|
|
| if token == self.tokenizer.eos_id() or token in self.stop_ids: |
| break |
|
|
| x = torch.cat((x, next_token), dim=1) |
|
|
| return self.tokenizer.decode(x[0].tolist()) |
|
|
| def chat( |
| self, |
| message, |
| max_new_tokens=200, |
| temperature=0.8, |
| top_k=50, |
| syspr="Answer to what user have said." |
| ): |
| prompt = f"<|system|> {syspr}\n<|user|> {message}\n<|assistant|>" |
| return self.generate( |
| prompt, |
| max_new_tokens=max_new_tokens, |
| temperature=temperature, |
| top_k=top_k, |
| ) |
|
|
| def __call__( |
| self, |
| message, |
| max_new_tokens=200, |
| temperature=0.8, |
| top_k=50, |
| ): |
| return self.chat( |
| message, |
| max_new_tokens=max_new_tokens, |
| temperature=temperature, |
| top_k=top_k, |
| ) |
|
|
| def __repr__(self): |
| loaded = self.model is not None |
| return ( |
| f"Wisp(" |
| f"loaded={loaded}, " |
| f"device='{self.device}', " |
| f"layers={self.config['n_layer']}, " |
| f"hidden={self.config['n_embd']}, " |
| f"vocab={self.config['vocab_size']}" |
| f")" |
| ) |
