Upload infer.py

infer.py

@@ -0,0 +1,434 @@
+"""
+QED-Base-v3 inference library.
+
+Usage:
+    from qed_infer import load_model, load_tokenizer, run
+
+    model = load_model("QED-Base-v3.pt")
+    tokenizer = load_tokenizer("qed-b3-tok.model")
+
+    text = run("Once upon a time", model, tokenizer, max_new_tokens=100)
+
+    # streaming / batched:
+    for texts in generate_stream(model, tokenizer, ["prompt A", "prompt B"]):
+        ...  # texts[i] is the completion-so-far for prompt i
+"""
+
+from __future__ import annotations
+
+import sys
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Iterator, Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import sentencepiece as spm
+
+
+DEVICE = "cuda" if torch.cuda.is_available() else ("mps" if torch.backends.mps.is_available() else "cpu")
+DTYPE = torch.bfloat16 if DEVICE == "cuda" else torch.float32
+
+DEFAULT_CHECKPOINT = "QED-Base-v3.pt"
+DEFAULT_TOKENIZER = "qed-b3-tok.model"
+
+
+@dataclass
+class Config:
+    vocab_size: int = 56000
+    hidden_size: int = 768
+    num_layers: int = 12
+    num_heads: int = 12
+    num_kv_heads: int = 4
+    intermediate_size: int = 1792
+    max_seq_len: int = 2048
+    rope_theta: float = 10000.0
+    rms_eps: float = 1e-6
+
+
+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):
+        variance = x.float().pow(2).mean(dim=-1, keepdim=True)
+        x = x * torch.rsqrt(variance + self.eps)
+        return (self.weight * x).type_as(self.weight)
+
+
+def rotate_half(x):
+    x1, x2 = x.chunk(2, dim=-1)
+    return torch.cat((-x2, x1), dim=-1)
+
+
+class RotaryEmbedding(nn.Module):
+    def __init__(self, head_dim, max_seq_len, theta):
+        super().__init__()
+        inv_freq = 1.0 / (theta ** (torch.arange(0, head_dim, 2).float() / head_dim))
+        positions = torch.arange(max_seq_len).float()
+        freqs = torch.outer(positions, inv_freq)
+        emb = torch.cat([freqs, freqs], dim=-1)
+        self.register_buffer("cos", emb.cos(), persistent=False)
+        self.register_buffer("sin", emb.sin(), persistent=False)
+
+    def forward(self, q, k, offset: int):
+        q_len = q.shape[-2]
+        k_len = k.shape[-2]
+
+        cos_q = self.cos[offset:offset + q_len][None, None, :, :].to(q.dtype)
+        sin_q = self.sin[offset:offset + q_len][None, None, :, :].to(q.dtype)
+
+        k_offset = offset + q_len - k_len
+        cos_k = self.cos[k_offset:k_offset + k_len][None, None, :, :].to(k.dtype)
+        sin_k = self.sin[k_offset:k_offset + k_len][None, None, :, :].to(k.dtype)
+
+        return (
+            q * cos_q + rotate_half(q) * sin_q,
+            k * cos_k + rotate_half(k) * sin_k,
+        )
+
+
+class SwiGLU(nn.Module):
+    def __init__(self, hidden, intermediate):
+        super().__init__()
+        self.gate_proj = nn.Linear(hidden, intermediate, bias=False)
+        self.up_proj = nn.Linear(hidden, intermediate, bias=False)
+        self.down_proj = nn.Linear(intermediate, hidden, bias=False)
+
+    def forward(self, x):
+        return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))
+
+
+class GQAttention(nn.Module):
+    def __init__(self, cfg: Config):
+        super().__init__()
+        self.num_heads = cfg.num_heads
+        self.num_kv_heads = cfg.num_kv_heads
+        self.head_dim = cfg.hidden_size // cfg.num_heads
+
+        self.q_proj = nn.Linear(cfg.hidden_size, cfg.num_heads * self.head_dim, bias=False)
+        self.k_proj = nn.Linear(cfg.hidden_size, cfg.num_kv_heads * self.head_dim, bias=False)
+        self.v_proj = nn.Linear(cfg.hidden_size, cfg.num_kv_heads * self.head_dim, bias=False)
+        self.o_proj = nn.Linear(cfg.hidden_size, cfg.hidden_size, bias=False)
+
+        self.rope = RotaryEmbedding(self.head_dim, cfg.max_seq_len, cfg.rope_theta)
+
+    def forward(self, x, offset: int, past_kv: Optional[tuple] = None):
+        B, T, C = x.shape
+
+        q = self.q_proj(x).view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
+        k = self.k_proj(x).view(B, T, self.num_kv_heads, self.head_dim).transpose(1, 2)
+        v = self.v_proj(x).view(B, T, self.num_kv_heads, self.head_dim).transpose(1, 2)
+
+        if past_kv is not None:
+            past_k, past_v = past_kv
+            k = torch.cat([past_k, k], dim=2)
+            v = torch.cat([past_v, v], dim=2)
+
+        q, k = self.rope(q, k, offset)
+        present = (k, v)
+
+        repeat = self.num_heads // self.num_kv_heads
+        k_rep = k.repeat_interleave(repeat, dim=1)
+        v_rep = v.repeat_interleave(repeat, dim=1)
+
+        y = F.scaled_dot_product_attention(q, k_rep, v_rep, is_causal=T > 1)
+
+        y = y.transpose(1, 2).contiguous().view(B, T, C)
+        return self.o_proj(y), present
+
+
+class QEDBlock(nn.Module):
+    def __init__(self, cfg: Config):
+        super().__init__()
+        self.attn_norm = RMSNorm(cfg.hidden_size, cfg.rms_eps)
+        self.attention = GQAttention(cfg)
+        self.ffn_norm = RMSNorm(cfg.hidden_size, cfg.rms_eps)
+        self.ffn = SwiGLU(cfg.hidden_size, cfg.intermediate_size)
+
+    def forward(self, x, offset: int, past_kv=None):
+        attn_out, present = self.attention(self.attn_norm(x), offset, past_kv)
+        x = x + attn_out
+        x = x + self.ffn(self.ffn_norm(x))
+        return x, present
+
+
+class QEDBaseV3(nn.Module):
+    def __init__(self, cfg: Config):
+        super().__init__()
+        self.cfg = cfg
+        self.embed_tokens = nn.Embedding(cfg.vocab_size, cfg.hidden_size)
+        self.layers = nn.ModuleList([QEDBlock(cfg) for _ in range(cfg.num_layers)])
+        self.final_norm = RMSNorm(cfg.hidden_size, cfg.rms_eps)
+        self.lm_head = nn.Linear(cfg.hidden_size, cfg.vocab_size, bias=False)
+        self.lm_head.weight = self.embed_tokens.weight
+
+    def forward(self, input_ids, offset: int = 0, past_key_values: Optional[list] = None):
+        x = self.embed_tokens(input_ids)
+        new_past = []
+        for i, layer in enumerate(self.layers):
+            past_kv = past_key_values[i] if past_key_values is not None else None
+            x, present = layer(x, offset, past_kv)
+            new_past.append(present)
+        x = self.final_norm(x)
+        return self.lm_head(x), new_past
+
+
+def load_model(model_path: str, cfg: Config = Config()) -> QEDBaseV3:
+    path = Path(model_path)
+    if not path.exists():
+        raise FileNotFoundError(f"Checkpoint not found: {model_path}")
+
+    package = torch.load(path, map_location="cpu", weights_only=True)
+    state_dict = package["state_dict"] if "state_dict" in package else package
+
+    model = QEDBaseV3(cfg)
+    missing, unexpected = model.load_state_dict(state_dict, strict=False)
+    if missing:
+        print(f"[warn] missing keys: {missing}", file=sys.stderr)
+    if unexpected:
+        print(f"[warn] unexpected keys: {unexpected}", file=sys.stderr)
+
+    model.to(DEVICE, dtype=DTYPE)
+    model.eval()
+
+    name = package.get("Name", "QED-Base-v3") if isinstance(package, dict) else "QED-Base-v3"
+    author = package.get("Author", "unknown") if isinstance(package, dict) else "unknown"
+    print(f"Loaded {name} by {author} on {DEVICE} ({DTYPE})", file=sys.stderr)
+    return model
+
+
+def load_tokenizer(tokenizer_path: str) -> spm.SentencePieceProcessor:
+    if not Path(tokenizer_path).exists():
+        raise FileNotFoundError(f"Tokenizer not found: {tokenizer_path}")
+    tok = spm.SentencePieceProcessor()
+    tok.load(tokenizer_path)
+    return tok
+
+
+def _apply_repetition_penalty(logits: torch.Tensor, generated: torch.Tensor, penalty: float):
+    if penalty == 1.0:
+        return logits
+    for b in range(logits.shape[0]):
+        seen = torch.unique(generated[b])
+        vals = logits[b, seen]
+        logits[b, seen] = torch.where(vals > 0, vals / penalty, vals * penalty)
+    return logits
+
+
+def _top_k_top_p_filter(logits: torch.Tensor, top_k: int, top_p: float):
+    if top_k > 0:
+        top_k = min(top_k, logits.size(-1))
+        kth_val = torch.topk(logits, top_k, dim=-1).values[..., -1, None]
+        logits = torch.where(logits < kth_val, torch.full_like(logits, float("-inf")), logits)
+
+    if top_p < 1.0:
+        sorted_logits, sorted_idx = torch.sort(logits, descending=True, dim=-1)
+        probs = F.softmax(sorted_logits, dim=-1)
+        cum_probs = torch.cumsum(probs, dim=-1)
+
+        remove = cum_probs > top_p
+        remove[..., 1:] = remove[..., :-1].clone()
+        remove[..., 0] = False
+
+        sorted_logits[remove] = float("-inf")
+        logits = torch.full_like(logits, float("-inf")).scatter(-1, sorted_idx, sorted_logits)
+
+    return logits
+
+
+@torch.no_grad()
+def generate_stream(
+    model: QEDBaseV3,
+    tokenizer: spm.SentencePieceProcessor,
+    prompts: list[str],
+    max_new_tokens: int = 200,
+    temperature: float = 0.8,
+    top_k: int = 50,
+    top_p: float = 0.95,
+    repetition_penalty: float = 1.15,
+    eos_id: Optional[int] = None,
+    seed: Optional[int] = None,
+) -> Iterator[list[str]]:
+    if seed is not None:
+        torch.manual_seed(seed)
+
+    if eos_id is None:
+        eos_id = tokenizer.eos_id() if tokenizer.eos_id() >= 0 else None
+
+    encoded = [tokenizer.encode(p) for p in prompts]
+    max_len = max(len(e) for e in encoded)
+    pad_id = tokenizer.pad_id() if tokenizer.pad_id() >= 0 else 0
+
+    B = len(prompts)
+    input_ids = torch.full((B, max_len), pad_id, dtype=torch.long, device=DEVICE)
+    for i, e in enumerate(encoded):
+        input_ids[i, max_len - len(e):] = torch.tensor(e, dtype=torch.long, device=DEVICE)
+
+    generated = input_ids.clone()
+    finished = torch.zeros(B, dtype=torch.bool, device=DEVICE)
+    text_so_far = ["" for _ in range(B)]
+
+    logits, past = model(input_ids, offset=0)
+    offset = input_ids.shape[1]
+
+    for _ in range(max_new_tokens):
+        next_logits = logits[:, -1, :].float()
+        next_logits = _apply_repetition_penalty(next_logits, generated, repetition_penalty)
+
+        if temperature <= 0:
+            next_token = next_logits.argmax(dim=-1, keepdim=True)
+        else:
+            next_logits = next_logits / temperature
+            next_logits = _top_k_top_p_filter(next_logits, top_k, top_p)
+            probs = F.softmax(next_logits, dim=-1)
+            next_token = torch.multinomial(probs, num_samples=1)
+
+        next_token = torch.where(
+            finished.unsqueeze(-1), torch.full_like(next_token, pad_id), next_token
+        )
+        generated = torch.cat([generated, next_token], dim=1)
+
+        if eos_id is not None:
+            finished |= next_token.squeeze(-1) == eos_id
+
+        for i in range(B):
+            if not finished[i]:
+                text_so_far[i] = tokenizer.decode(generated[i].tolist())
+
+        yield list(text_so_far)
+
+        if bool(finished.all()):
+            break
+
+        logits, past = model(next_token, offset=offset, past_key_values=past)
+        offset += 1
+
+
+@torch.no_grad()
+def run(
+    prompt: str,
+    model: QEDBaseV3,
+    tokenizer: spm.SentencePieceProcessor,
+    max_new_tokens: int = 200,
+    temperature: float = 0.7,
+    top_k: int = 40,
+    top_p: float = 0.95,
+    repetition_penalty: float = 1.15,
+    seed: Optional[int] = None,
+) -> str:
+    final = ""
+    for texts in generate_stream(
+        model, tokenizer, [prompt],
+        max_new_tokens=max_new_tokens,
+        temperature=temperature,
+        top_k=top_k,
+        top_p=top_p,
+        repetition_penalty=repetition_penalty,
+        seed=seed,
+    ):
+        final = texts[0]
+    return final
+
+
+def _print_help():
+    print(
+        "commands:\n"
+        "  <text>              generate from this prompt\n"
+        "  /temp <float>       set temperature (0 = greedy)\n"
+        "  /top_k <int>        set top-k (0 = disabled)\n"
+        "  /top_p <float>      set top-p\n"
+        "  /rep <float>        set repetition penalty\n"
+        "  /tokens <int>       set max new tokens\n"
+        "  /seed <int|none>    fix or clear the random seed\n"
+        "  /help               show this message\n"
+        "  /quit, /exit        leave"
+    , file=sys.stderr)
+
+
+def _repl():
+    checkpoint_path = Path(DEFAULT_CHECKPOINT)
+    tokenizer_path = Path(DEFAULT_TOKENIZER)
+
+    if not checkpoint_path.exists():
+        print(f"error: checkpoint not found at {checkpoint_path}", file=sys.stderr)
+        sys.exit(1)
+    if not tokenizer_path.exists():
+        print(f"error: tokenizer not found at {tokenizer_path}", file=sys.stderr)
+        sys.exit(1)
+
+    model = load_model(str(checkpoint_path))
+    tokenizer = load_tokenizer(str(tokenizer_path))
+
+    settings = {
+        "max_new_tokens": 200,
+        "temperature": 0.7,
+        "top_k": 40,
+        "top_p": 0.95,
+        "repetition_penalty": 1.15,
+        "seed": None,
+    }
+
+    _print_help()
+    print(file=sys.stderr)
+
+    while True:
+        try:
+            line = input(">>> ")
+        except (EOFError, KeyboardInterrupt):
+            print(file=sys.stderr)
+            break
+
+        line = line.strip()
+        if not line:
+            continue
+
+        if line in ("/quit", "/exit"):
+            break
+        if line == "/help":
+            _print_help()
+            continue
+        if line.startswith("/temp "):
+            settings["temperature"] = float(line.split(maxsplit=1)[1])
+            print(f"temperature set to {settings['temperature']}", file=sys.stderr)
+            continue
+        if line.startswith("/top_k "):
+            settings["top_k"] = int(line.split(maxsplit=1)[1])
+            print(f"top_k set to {settings['top_k']}", file=sys.stderr)
+            continue
+        if line.startswith("/top_p "):
+            settings["top_p"] = float(line.split(maxsplit=1)[1])
+            print(f"top_p set to {settings['top_p']}", file=sys.stderr)
+            continue
+        if line.startswith("/rep "):
+            settings["repetition_penalty"] = float(line.split(maxsplit=1)[1])
+            print(f"repetition_penalty set to {settings['repetition_penalty']}", file=sys.stderr)
+            continue
+        if line.startswith("/tokens "):
+            settings["max_new_tokens"] = int(line.split(maxsplit=1)[1])
+            print(f"max_new_tokens set to {settings['max_new_tokens']}", file=sys.stderr)
+            continue
+        if line.startswith("/seed "):
+            value = line.split(maxsplit=1)[1].strip()
+            settings["seed"] = None if value.lower() == "none" else int(value)
+            print(f"seed set to {settings['seed']}", file=sys.stderr)
+            continue
+        if line.startswith("/"):
+            print(f"unknown command: {line}", file=sys.stderr)
+            continue
+
+        try:
+            output = run(line, model, tokenizer, **settings)
+        except Exception as e:
+            print(f"error during generation: {e}", file=sys.stderr)
+            continue
+
+        print(output)
+        print(file=sys.stderr)
+
+
+if __name__ == "__main__":
+    _repl()