| """
|
| 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 |
