pytensorforge 0.1.0__py3-none-any.whl

src/inference/chat_template.py

@@ -0,0 +1,384 @@
+import copy
+import json
+import os
+from dataclasses import dataclass, field
+from typing import List
+
+ROLES = ("system", "user", "assistant")
+
+
+class ChatTemplateError(ValueError):
+    pass
+
+
+class PromptTooLong(ValueError):
+
+    def __init__(self, needed, budget):
+        super().__init__(f"prompt needs {needed} tokens but only {budget} fit in the context window")
+        self.needed = needed
+        self.budget = budget
+
+
+BUILTIN_TEMPLATES = {
+    "plain": {
+        "name": "plain",
+        "add_bos": False,
+        "roles": {
+            "system": {"prefix": ["System: "], "suffix": ["\n"]},
+            "user": {"prefix": ["User: "], "suffix": ["\n"]},
+            "assistant": {"prefix": ["Assistant: "], "suffix": ["\n"]},
+        },
+        "generation_prompt": ["Assistant:"],
+        "stop_tokens": [],
+        "stop_strings": ["\nUser:", "\nSystem:"],
+        "output_lstrip": True,
+    },
+    "ptf-chat": {
+        "name": "ptf-chat",
+        "add_bos": False,
+        "roles": {
+            "system": {"prefix": [{"special": "<|system|>"}, "\n"], "suffix": [{"special": "<|end|>"}, "\n"]},
+            "user": {"prefix": [{"special": "<|user|>"}, "\n"], "suffix": [{"special": "<|end|>"}, "\n"]},
+            "assistant": {"prefix": [{"special": "<|assistant|>"}, "\n"], "suffix": [{"special": "<|end|>"}, "\n"]},
+        },
+        "generation_prompt": [{"special": "<|assistant|>"}, "\n"],
+        "stop_tokens": ["<|end|>"],
+        "stop_strings": [],
+        "output_lstrip": False,
+    },
+}
+
+
+def _normalize_segments(segments, where):
+    if segments is None:
+        return []
+
+    if isinstance(segments, (str, dict)):
+        segments = [segments]
+
+    out = []
+
+    for seg in segments:
+        if isinstance(seg, str):
+            out.append(("text", seg))
+        elif isinstance(seg, dict) and set(seg) == {"text"} and isinstance(seg["text"], str):
+            out.append(("text", seg["text"]))
+        elif isinstance(seg, dict) and set(seg) == {"special"} and isinstance(seg["special"], str) and seg["special"]:
+            out.append(("special", seg["special"]))
+        else:
+            raise ChatTemplateError(f"invalid segment in {where}: {seg!r}")
+
+    return out
+
+
+def _segments_to_json(segments):
+    return [text if kind == "text" else {"special": text} for kind, text in segments]
+
+
+@dataclass
+class ChatTemplate:
+    name: str
+    roles: dict
+    generation_prompt: list
+    add_bos: bool = False
+    stop_tokens: List[str] = field(default_factory=list)
+    stop_strings: List[str] = field(default_factory=list)
+    output_lstrip: bool = False
+    default_system: str = None
+
+    @classmethod
+    def from_dict(cls, data):
+        if not isinstance(data, dict):
+            raise ChatTemplateError("chat template must be a JSON object")
+
+        name = data.get("name") or "custom"
+        roles_in = data.get("roles")
+
+        if not isinstance(roles_in, dict) or not roles_in:
+            raise ChatTemplateError("chat template needs a 'roles' object")
+
+        roles = {}
+
+        for role, spec in roles_in.items():
+            if role not in ROLES:
+                raise ChatTemplateError(f"unsupported role '{role}' in chat template")
+
+            if not isinstance(spec, dict):
+                raise ChatTemplateError(f"role '{role}' must be an object with prefix/suffix")
+
+            roles[role] = {
+                "prefix": _normalize_segments(spec.get("prefix"), f"{role}.prefix"),
+                "suffix": _normalize_segments(spec.get("suffix"), f"{role}.suffix"),
+            }
+
+        for required in ("user", "assistant"):
+            if required not in roles:
+                raise ChatTemplateError(f"chat template must define the '{required}' role")
+
+        stop_tokens = data.get("stop_tokens", [])
+        stop_strings = data.get("stop_strings", [])
+
+        if not all(isinstance(s, str) and s for s in stop_tokens):
+            raise ChatTemplateError("stop_tokens must be non-empty strings")
+
+        if not all(isinstance(s, str) and s for s in stop_strings):
+            raise ChatTemplateError("stop_strings must be non-empty strings")
+
+        default_system = data.get("default_system")
+
+        if default_system is not None and not isinstance(default_system, str):
+            raise ChatTemplateError("default_system must be a string")
+
+        return cls(
+            name=name,
+            roles=roles,
+            generation_prompt=_normalize_segments(data.get("generation_prompt"), "generation_prompt"),
+            add_bos=bool(data.get("add_bos", False)),
+            stop_tokens=list(stop_tokens),
+            stop_strings=list(stop_strings),
+            output_lstrip=bool(data.get("output_lstrip", False)),
+            default_system=default_system,
+        )
+
+    def to_dict(self):
+        return {
+            "name": self.name,
+            "add_bos": self.add_bos,
+            "roles": {
+                role: {"prefix": _segments_to_json(s["prefix"]), "suffix": _segments_to_json(s["suffix"])}
+                for role, s in self.roles.items()
+            },
+            "generation_prompt": _segments_to_json(self.generation_prompt),
+            "stop_tokens": list(self.stop_tokens),
+            "stop_strings": list(self.stop_strings),
+            "output_lstrip": self.output_lstrip,
+            "default_system": self.default_system,
+        }
+
+    @classmethod
+    def builtin(cls, name):
+        if name not in BUILTIN_TEMPLATES:
+            raise ChatTemplateError(f"unknown built-in chat template '{name}'; choose one of {sorted(BUILTIN_TEMPLATES)}")
+
+        return cls.from_dict(copy.deepcopy(BUILTIN_TEMPLATES[name]))
+
+    @classmethod
+    def resolve(cls, spec):
+        if spec is None:
+            return None
+
+        if isinstance(spec, ChatTemplate):
+            return spec
+
+        if isinstance(spec, dict):
+            return cls.from_dict(spec)
+
+        if isinstance(spec, str):
+            if spec in BUILTIN_TEMPLATES:
+                return cls.builtin(spec)
+
+            if os.path.isfile(spec):
+                with open(spec) as f:
+                    return cls.from_dict(json.load(f))
+
+        raise ChatTemplateError(f"chat template '{spec}' is neither a built-in name nor a readable JSON file")
+
+    def special_names(self):
+        names = set(self.stop_tokens)
+
+        for spec in self.roles.values():
+            for kind, text in spec["prefix"] + spec["suffix"]:
+                if kind == "special":
+                    names.add(text)
+
+        for kind, text in self.generation_prompt:
+            if kind == "special":
+                names.add(text)
+
+        return names
+
+    def bind(self, tokenizer):
+        return BoundChatTemplate(self, tokenizer)
+
+
+@dataclass
+class RenderedPrompt:
+    token_ids: List[int]
+    messages_used: int
+    messages_dropped: int
+
+
+class BoundChatTemplate:
+
+    def __init__(self, template, tokenizer):
+        self.template = template
+        self.tokenizer = tokenizer
+        specials = tokenizer.special_tokens or {}
+
+        missing = sorted(n for n in template.special_names() if n not in specials)
+
+        if missing:
+            raise ChatTemplateError(
+                f"chat template '{template.name}' uses special tokens the tokenizer does not have: {missing}. "
+                "Train the tokenizer with them (tokenize --special-token ...) or choose another template."
+            )
+
+        if template.add_bos and tokenizer.bos_id is None:
+            raise ChatTemplateError("chat template requires a BOS token but the tokenizer has none")
+
+        self._specials = dict(specials)
+        self._control_ids = frozenset(
+            i for name, i in specials.items()
+            if name != getattr(tokenizer, "unk_token", None)
+        )
+
+        self.stop_token_ids = [specials[n] for n in template.stop_tokens]
+        self.stop_strings = list(template.stop_strings)
+        self.output_lstrip = template.output_lstrip
+        self._gen_prompt_ids = self._encode_pieces(template.generation_prompt, None)
+
+    @property
+    def name(self):
+        return self.template.name
+
+    def _encode_text(self, text):
+        if not text:
+            return []
+
+        ids = self.tokenizer.encode(text)
+
+        return [i for i in ids if i not in self._control_ids]
+
+    def _encode_pieces(self, segments, content):
+        ids = []
+        run = []
+
+        def flush():
+            if run:
+                ids.extend(self._encode_text("".join(run)))
+                run.clear()
+
+        for kind, text in segments:
+            if kind == "special":
+                flush()
+                ids.append(self._specials[text])
+            elif kind == "content":
+                run.append(content)
+            else:
+                run.append(text)
+
+        flush()
+
+        return ids
+
+    def encode_message(self, role, content):
+        spec = self.template.roles.get(role)
+
+        if spec is None:
+            raise ChatTemplateError(f"this model's chat template does not support the '{role}' role")
+
+        return self._encode_pieces(spec["prefix"] + [("content", None)] + spec["suffix"], content)
+
+    def _continuation_segments(self):
+        prefix = list(self.template.roles.get("assistant", {}).get("prefix", []))
+        gen = list(self.template.generation_prompt)
+
+        if "assistant" not in self.template.roles:
+            raise ChatTemplateError("chat template has no assistant role to train on")
+
+        rest = list(prefix)
+
+        for i, (kind, text) in enumerate(gen):
+            if not rest:
+                raise ChatTemplateError("generation prompt is longer than the assistant prefix")
+
+            head_kind, head_text = rest[0]
+
+            if kind == "special" or head_kind == "special":
+                if (kind, text) != (head_kind, head_text):
+                    raise ChatTemplateError("generation prompt must be a prefix of the assistant prefix to train")
+                rest.pop(0)
+            elif head_text == text:
+                rest.pop(0)
+            elif i == len(gen) - 1 and head_text.startswith(text):
+                rest[0] = (head_kind, head_text[len(text):])
+            else:
+                raise ChatTemplateError("generation prompt must be a prefix of the assistant prefix to train")
+
+        return rest
+
+    def training_tokens(self, messages, eos_after_reply=True):
+        messages = list(messages)
+
+        if self.template.default_system and not any(m["role"] == "system" for m in messages):
+            messages.insert(0, {"role": "system", "content": self.template.default_system})
+
+        continuation = self._continuation_segments()
+        suffix = list(self.template.roles["assistant"]["suffix"])
+        ids = [self.tokenizer.bos_id] if self.template.add_bos else []
+        learn = [0] * len(ids)
+
+        for m in messages:
+            if m["role"] == "assistant":
+                ids.extend(self._gen_prompt_ids)
+                learn.extend([0] * len(self._gen_prompt_ids))
+                reply = self._encode_pieces(continuation + [("content", None)] + suffix, m["content"])
+                ids.extend(reply)
+                learn.extend([1] * len(reply))
+            else:
+                encoded = self.encode_message(m["role"], m["content"])
+                ids.extend(encoded)
+                learn.extend([0] * len(encoded))
+
+        if eos_after_reply and messages and messages[-1]["role"] == "assistant" and self.tokenizer.eos_id is not None:
+            ids.append(self.tokenizer.eos_id)
+            learn.append(1)
+
+        return ids, learn
+
+    def render(self, messages, max_prompt_tokens=None, truncate=True):
+        messages = list(messages)
+
+        if self.template.default_system and not any(m["role"] == "system" for m in messages):
+            messages.insert(0, {"role": "system", "content": self.template.default_system})
+
+        if not messages:
+            raise ChatTemplateError("messages must not be empty")
+
+        encoded = [self.encode_message(m["role"], m["content"]) for m in messages]
+        head = [self.tokenizer.bos_id] if self.template.add_bos else []
+        tail = self._gen_prompt_ids
+
+        pinned = 0
+
+        while pinned < len(messages) and messages[pinned]["role"] == "system":
+            pinned += 1
+
+        fixed = len(head) + len(tail) + sum(len(e) for e in encoded[:pinned])
+        body = encoded[pinned:]
+        total = fixed + sum(len(e) for e in body)
+        dropped = 0
+
+        if max_prompt_tokens is not None and total > max_prompt_tokens:
+            if not truncate:
+                raise PromptTooLong(total, max_prompt_tokens)
+
+            while body and total > max_prompt_tokens and len(body) > 1:
+                total -= len(body[0])
+                body = body[1:]
+                dropped += 1
+
+            if total > max_prompt_tokens:
+                raise PromptTooLong(total, max_prompt_tokens)
+
+        ids = list(head)
+
+        for e in encoded[:pinned]:
+            ids.extend(e)
+
+        for e in body:
+            ids.extend(e)
+
+        ids.extend(tail)
+
+        return RenderedPrompt(ids, len(messages) - dropped, dropped)

src/inference/config.py

@@ -0,0 +1,48 @@
+from dataclasses import asdict, dataclass, field, replace
+from typing import List, Optional
+
+
+@dataclass
+class GenerationConfig:
+    max_new_tokens: int = 128
+    temperature: float = 1.0
+    top_k: int = 0
+    top_p: float = 1.0
+    repetition_penalty: float = 1.0
+    do_sample: bool = True
+    seed: Optional[int] = None
+    stop_token_ids: List[int] = field(default_factory=list)
+    stop_strings: List[str] = field(default_factory=list)
+    stop_on_eos: bool = True
+
+    def __post_init__(self):
+        if self.max_new_tokens < 1:
+            raise ValueError("max_new_tokens must be at least 1")
+
+        if self.temperature < 0:
+            raise ValueError("temperature must be >= 0")
+
+        if self.top_k < 0:
+            raise ValueError("top_k must be >= 0")
+
+        if not 0 < self.top_p <= 1:
+            raise ValueError("top_p must be in (0, 1]")
+
+        if self.repetition_penalty <= 0:
+            raise ValueError("repetition_penalty must be > 0")
+
+    @property
+    def greedy(self):
+        return (not self.do_sample) or self.temperature == 0
+
+    def updated(self, **overrides):
+        overrides = {k: v for k, v in overrides.items() if v is not None}
+        return replace(self, **overrides)
+
+    def to_dict(self):
+        return asdict(self)
+
+    @classmethod
+    def from_dict(cls, data):
+        known = set(cls.__dataclass_fields__.keys())
+        return cls(**{k: v for k, v in data.items() if k in known})

src/inference/engine.py

@@ -0,0 +1,241 @@
+import numpy as np
+
+from src.inference.kv_cache import KVCache
+
+
+def _layer_norm(x, gamma, beta, eps):
+    mean = x.mean(axis=-1, keepdims=True)
+    var = ((x - mean) ** 2).mean(axis=-1, keepdims=True)
+    return gamma * ((x - mean) / np.sqrt(var + eps)) + beta
+
+
+def _softmax(x):
+    x = x - x.max(axis=-1, keepdims=True)
+    np.exp(x, out=x)
+    x /= x.sum(axis=-1, keepdims=True)
+    return x
+
+
+def _gelu(x):
+    return 0.5 * x * (1.0 + np.tanh(0.7978845608028654 * (x + 0.044715 * (x * x * x))))
+
+
+def _sigmoid(x):
+    return 1.0 / (1.0 + np.exp(-x))
+
+
+ACTIVATIONS = {
+    "gelu": _gelu,
+    "relu": lambda x: np.maximum(0, x),
+    "tanh": np.tanh,
+    "sigmoid": _sigmoid,
+}
+
+
+class _Layer:
+
+    def __init__(self, w, prefix):
+        g = lambda name: np.ascontiguousarray(w[f"{prefix}.{name}"], dtype=np.float32)
+
+        self.ln1_g = g("norm1.gamma")
+        self.ln1_b = g("norm1.beta")
+        self.wqkv = np.concatenate([g("attn.Wq"), g("attn.Wk"), g("attn.Wv")], axis=1)
+        self.wo = g("attn.Wo")
+        self.ln2_g = g("norm2.gamma")
+        self.ln2_b = g("norm2.beta")
+        self.fc1_k = g("fc1.kernel")
+        self.fc1_b = g("fc1.bias")
+        self.fc2_k = g("fc2.kernel")
+        self.fc2_b = g("fc2.bias")
+
+
+class InferenceModel:
+
+    def __init__(self, config, weights):
+        if config.activation not in ACTIVATIONS:
+            raise ValueError(f"unsupported activation '{config.activation}'")
+
+        self.config = config
+        self.n_layers = config.n_layers
+        self.n_heads = config.n_heads
+        self.d_model = config.d_model
+        self.head_dim = config.d_model // config.n_heads
+        self.context_length = config.context_length
+        self.eps = config.norm_eps
+        self.scale = np.float32(1.0 / np.sqrt(self.head_dim))
+        self.act = ACTIVATIONS[config.activation]
+
+        f32 = lambda name: np.ascontiguousarray(weights[name], dtype=np.float32)
+
+        self.tok_emb = f32("token_emb.embedding")
+        self.rope = None
+        self.pos_emb = None
+
+        if getattr(config, "position_encoding", "learned") == "rope":
+            self.rope = config.rotary_tables()
+
+            if "pos_emb.embedding" in weights:
+                raise ValueError("rope model must not carry learned positional embeddings")
+        else:
+            self.pos_emb = f32("pos_emb.embedding")
+        self.ln_f_g = f32("final_norm.gamma")
+        self.ln_f_b = f32("final_norm.beta")
+        self.layers = [_Layer(weights, f"blocks.{i}") for i in range(config.n_layers)]
+
+        if config.tie_weights:
+            self.head_w = self.tok_emb.T
+            self.head_b = None
+        else:
+            self.head_w = f32("lm_head.kernel")
+            self.head_b = f32("lm_head.bias")
+
+        if self.tok_emb.shape != (config.vocab_size, config.d_model):
+            raise ValueError("token embedding shape does not match config")
+
+        if self.pos_emb is not None and self.pos_emb.shape != (config.context_length, config.d_model):
+            raise ValueError("positional embedding shape does not match config")
+
+    @property
+    def nbytes(self):
+        total = self.tok_emb.nbytes + self.ln_f_g.nbytes + self.ln_f_b.nbytes
+
+        if self.pos_emb is not None:
+            total += self.pos_emb.nbytes
+
+        if self.head_b is not None:
+            total += self.head_w.nbytes + self.head_b.nbytes
+
+        for layer in self.layers:
+            total += sum(v.nbytes for v in vars(layer).values())
+
+        return total
+
+    def new_cache(self, capacity):
+        return KVCache(self.n_layers, self.n_heads, self.head_dim, min(capacity, self.context_length))
+
+    def _logits(self, x):
+        out = _layer_norm(x, self.ln_f_g, self.ln_f_b, self.eps) @ self.head_w
+
+        if self.head_b is not None:
+            out = out + self.head_b
+
+        return out
+
+    def _ffn(self, layer, x):
+        h = _layer_norm(x, layer.ln2_g, layer.ln2_b, self.eps)
+        return self.act(h @ layer.fc1_k + layer.fc1_b) @ layer.fc2_k + layer.fc2_b
+
+    def forward_chunk(self, token_ids, cache, want_logits=True):
+        ids = np.asarray(token_ids, dtype=np.int64)
+        T = ids.size
+        p0 = cache.length
+
+        if T == 0:
+            raise ValueError("empty chunk")
+
+        if p0 + T > cache.capacity:
+            raise ValueError("chunk exceeds kv cache capacity")
+
+        x = self.tok_emb[ids]
+        L = p0 + T
+
+        if self.pos_emb is not None:
+            x = x + self.pos_emb[p0:L]
+        else:
+            cos, sin = self.rope.tables(np.arange(p0, L))
+
+        mask = None
+        if T > 1:
+            mask = np.arange(L)[None, :] > (p0 + np.arange(T))[:, None]
+
+        H, hd, d = self.n_heads, self.head_dim, self.d_model
+
+        for li, layer in enumerate(self.layers):
+            h = _layer_norm(x, layer.ln1_g, layer.ln1_b, self.eps)
+            qkv = h @ layer.wqkv
+
+            q = qkv[:, :d].reshape(T, H, hd).transpose(1, 0, 2)
+            k = qkv[:, d:2 * d].reshape(T, H, hd).transpose(1, 0, 2)
+
+            if self.rope is not None:
+                q = self.rope.rotate(q, cos, sin)
+                k = self.rope.rotate(k, cos, sin)
+
+            cache.k[li, :, p0:L] = k
+            cache.v[li, :, p0:L] = qkv[:, 2 * d:].reshape(T, H, hd).transpose(1, 0, 2)
+
+            scores = (q @ cache.k[li, :, :L].transpose(0, 2, 1)) * self.scale
+
+            if mask is not None:
+                scores[:, mask] = -1e9
+
+            ctx = _softmax(scores) @ cache.v[li, :, :L]
+            x = x + ctx.transpose(1, 0, 2).reshape(T, d) @ layer.wo
+            x = x + self._ffn(layer, x)
+
+        cache.length = L
+
+        if not want_logits:
+            return None
+
+        return self._logits(x[-1:])[0]
+
+    def prefill(self, token_ids, cache, chunk_size=256):
+        ids = np.asarray(token_ids, dtype=np.int64)
+        logits = None
+
+        for start in range(0, ids.size, chunk_size):
+            end = min(start + chunk_size, ids.size)
+            logits = self.forward_chunk(ids[start:end], cache, want_logits=(end == ids.size))
+
+        return logits
+
+    def decode_batch(self, token_ids, caches):
+        B = len(caches)
+        ids = np.asarray(token_ids, dtype=np.int64)
+        pos = np.array([c.length for c in caches], dtype=np.int64)
+
+        for c in caches:
+            if c.length + 1 > c.capacity:
+                raise ValueError("kv cache capacity exceeded")
+
+        x = self.tok_emb[ids]
+
+        if self.pos_emb is not None:
+            x = x + self.pos_emb[pos]
+        else:
+            cos, sin = self.rope.tables(pos)
+            q_cos, q_sin = cos[:, None, None, :], sin[:, None, None, :]
+            k_cos, k_sin = cos[:, None, :], sin[:, None, :]
+
+        H, hd, d = self.n_heads, self.head_dim, self.d_model
+
+        for li, layer in enumerate(self.layers):
+            h = _layer_norm(x, layer.ln1_g, layer.ln1_b, self.eps)
+            qkv = h @ layer.wqkv
+
+            q = qkv[:, :d].reshape(B, H, 1, hd)
+            k = qkv[:, d:2 * d].reshape(B, H, hd)
+            v = qkv[:, 2 * d:].reshape(B, H, hd)
+
+            if self.rope is not None:
+                q = self.rope.rotate(q, q_cos, q_sin)
+                k = self.rope.rotate(k, k_cos, k_sin)
+
+            ctx = np.empty((B, d), dtype=np.float32)
+
+            for b, c in enumerate(caches):
+                n = int(pos[b])
+                c.k[li, :, n] = k[b]
+                c.v[li, :, n] = v[b]
+
+                scores = (q[b] @ c.k[li, :, :n + 1].transpose(0, 2, 1)) * self.scale
+                ctx[b] = (_softmax(scores) @ c.v[li, :, :n + 1]).reshape(d)
+
+            x = x + ctx @ layer.wo
+            x = x + self._ffn(layer, x)
+
+        for c in caches:
+            c.length += 1
+
+        return self._logits(x)