Upload 4 files

QED-Base-v2.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:71c4e5291a4f0bc77b16abc77014635d1b43d55b7e774347cc7d5ecd0ef01a8c
+size 421245563

infer.py

@@ -0,0 +1,6 @@
+from model import load_model, load_tokenizer, run
+
+model = load_model("QED-Base-v2.pt")
+tokenizer = load_tokenizer("tok.model")
+text = run("The future of ai will be", model, tokenizer, max_new_tokens=100)
+print(text)

model.py

@@ -0,0 +1,325 @@
+"""
+QED-Base-v2 inference library.
+
+Usage:
+    from qed_infer import load_model, load_tokenizer, run
+
+    model = load_model("QED-Base-v2.pt")
+    tokenizer = load_tokenizer("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 "cpu"
+DTYPE = torch.bfloat16 if torch.cuda.is_available() else torch.float32
+
+
+@dataclass
+class Config:
+    vocab_size: int = 48000
+    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 QEDBaseV2(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()) -> QEDBaseV2:
+    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 = QEDBaseV2(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-v2") if isinstance(package, dict) else "QED-Base-v2"
+    author = package.get("Author", "unknown") if isinstance(package, dict) else "unknown"
+    print(f"Loaded {name} by {author} on {DEVICE} ({DTYPE})")
+    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: QEDBaseV2,
+    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,
+) -> Iterator[list[str]]:
+    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: QEDBaseV2,
+    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,
+) -> 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,
+    ):
+        final = texts[0]
+    return final

tok.model

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5507c05b9eb158901b28fa8b9f2462b87c24d0dfc87b046ad5c0659ec953fe2f
+size 1051850