pytensorforge 0.1.0__py3-none-any.whl

src/inference/export.py

@@ -0,0 +1,133 @@
+import hashlib
+import json
+import os
+import shutil
+
+import numpy as np
+
+from src.inference.config import GenerationConfig
+from src.models.gpt.config import GPTConfig
+from src.models.gpt.context import extend_context
+from src.tokenization.base import identity_matches
+from src.training.checkpoint_manager import CheckpointManager
+
+EXPORT_FORMAT_VERSION = 1
+WEIGHTS_FILE = os.path.join("weights", "model.npz")
+
+
+def flatten_state(state, prefix=""):
+    out = {}
+
+    for key, value in state.items():
+        if isinstance(value, dict):
+            out.update(flatten_state(value, f"{prefix}{key}."))
+        elif isinstance(value, list):
+            for i, item in enumerate(value):
+                out.update(flatten_state(item, f"{prefix}{key}.{i}."))
+        else:
+            out[f"{prefix}{key}"] = np.asarray(value)
+
+    return out
+
+
+def _sha256(path):
+    digest = hashlib.sha256()
+
+    with open(path, "rb") as f:
+        for block in iter(lambda: f.read(1 << 20), b""):
+            digest.update(block)
+
+    return digest.hexdigest()
+
+
+def _write_json(path, payload):
+    tmp = path + ".tmp"
+
+    with open(tmp, "w") as f:
+        json.dump(payload, f, indent=2, sort_keys=True)
+
+    os.replace(tmp, path)
+
+
+def export_model(checkpoint_path, output_dir, tokenizer_path, generation=None, dtype="float32", chat_template=None,
+                 context_length=None, context_extension=None):
+    if dtype not in ("float32", "float16"):
+        raise ValueError("dtype must be float32 or float16")
+
+    payload = CheckpointManager(os.path.dirname(os.path.abspath(checkpoint_path))).load(checkpoint_path)
+
+    if payload is None:
+        raise FileNotFoundError(f"checkpoint not found: {checkpoint_path}")
+
+    from src.tokenization.registry import load_tokenizer
+
+    tokenizer = load_tokenizer(tokenizer_path)
+
+    if not identity_matches(payload["tokenizer_identity"], tokenizer.identity):
+        raise ValueError("tokenizer does not match the one this checkpoint was trained with")
+
+    config = GPTConfig.from_dict(payload["model_config"])
+
+    if context_length is not None or context_extension is not None:
+        if context_length is None or context_extension is None:
+            raise ValueError("context extension needs both context_length and context_extension")
+
+        config = extend_context(config, context_length, context_extension)
+
+    from src.inference.chat_template import ChatTemplate
+
+    data_config = payload.get("data_config") or {}
+
+    if chat_template is None and data_config.get("format") == "chat":
+        chat_template = data_config.get("chat_template")
+
+    template = ChatTemplate.resolve(chat_template)
+
+    if template is not None:
+        template.bind(tokenizer)
+    weights = {k: v.astype(dtype) for k, v in flatten_state(payload["model_state"]).items()}
+
+    os.makedirs(os.path.join(output_dir, "weights"), exist_ok=True)
+    os.makedirs(os.path.join(output_dir, "tokenizer"), exist_ok=True)
+
+    weights_path = os.path.join(output_dir, WEIGHTS_FILE)
+    tmp_weights = weights_path + ".tmp"
+
+    with open(tmp_weights, "wb") as f:
+        np.savez(f, **weights)
+
+    os.replace(tmp_weights, weights_path)
+
+    shutil.copyfile(tokenizer_path, os.path.join(output_dir, "tokenizer", "tokenizer.json"))
+
+    generation = generation or GenerationConfig(
+        max_new_tokens=128, temperature=0.8, top_k=40, top_p=0.95,
+    )
+
+    _write_json(os.path.join(output_dir, "generation.json"), generation.to_dict())
+
+    meta = {
+        "format_version": EXPORT_FORMAT_VERSION,
+        "architecture": config.architecture,
+        "model": config.to_dict(),
+        "tokenizer": {"file": "tokenizer/tokenizer.json", "identity": tokenizer.identity},
+        "weights": {
+            "file": WEIGHTS_FILE,
+            "dtype": dtype,
+            "sha256": _sha256(weights_path),
+            "tensors": {k: list(v.shape) for k, v in weights.items()},
+        },
+        "source": {
+            "framework_version": payload.get("framework_version"),
+            "run_id": payload.get("run_id"),
+            "optimizer_steps": payload.get("optimizer_steps"),
+            "tokens_processed": payload.get("tokens_processed"),
+        },
+    }
+
+    if template is not None:
+        meta["chat_template"] = template.to_dict()
+
+    _write_json(os.path.join(output_dir, "config.json"), meta)
+
+    return output_dir

src/inference/kv_cache.py

@@ -0,0 +1,65 @@
+import numpy as np
+
+
+class CacheFull(Exception):
+    pass
+
+
+class KVCache:
+
+    def __init__(self, n_layers, n_heads, head_dim, capacity, dtype=np.float32):
+        self.capacity = capacity
+        self.k = np.empty((n_layers, n_heads, capacity, head_dim), dtype=dtype)
+        self.v = np.empty((n_layers, n_heads, capacity, head_dim), dtype=dtype)
+        self.length = 0
+
+    @property
+    def nbytes(self):
+        return self.k.nbytes + self.v.nbytes
+
+    def reset(self):
+        self.length = 0
+
+
+class KVCacheManager:
+
+    def __init__(self, n_layers, n_heads, head_dim, max_bytes=None, dtype=np.float32):
+        self.n_layers = n_layers
+        self.n_heads = n_heads
+        self.head_dim = head_dim
+        self.max_bytes = max_bytes
+        self.dtype = np.dtype(dtype)
+        self._caches = {}
+        self.used_bytes = 0
+
+    def bytes_for(self, capacity):
+        return 2 * self.n_layers * self.n_heads * capacity * self.head_dim * self.dtype.itemsize
+
+    def fits_at_all(self, capacity):
+        return self.max_bytes is None or self.bytes_for(capacity) <= self.max_bytes
+
+    def can_allocate(self, capacity):
+        return self.max_bytes is None or self.used_bytes + self.bytes_for(capacity) <= self.max_bytes
+
+    def allocate(self, key, capacity):
+        if key in self._caches:
+            raise KeyError(f"cache already allocated for {key}")
+
+        if not self.can_allocate(capacity):
+            raise CacheFull(f"kv cache budget of {self.max_bytes} bytes exhausted")
+
+        cache = KVCache(self.n_layers, self.n_heads, self.head_dim, capacity, self.dtype)
+        self._caches[key] = cache
+        self.used_bytes += cache.nbytes
+
+        return cache
+
+    def release(self, key):
+        cache = self._caches.pop(key, None)
+
+        if cache is not None:
+            self.used_bytes -= cache.nbytes
+
+    @property
+    def active(self):
+        return len(self._caches)

src/inference/runtime.py

@@ -0,0 +1,161 @@
+import json
+import os
+
+import numpy as np
+
+from src.inference.chat_template import ChatTemplate
+from src.inference.config import GenerationConfig
+from src.inference.engine import InferenceModel
+from src.inference.export import EXPORT_FORMAT_VERSION, _sha256
+from src.inference.scheduler import BatchScheduler, GenerationRequest
+from src.models.gpt.config import GPTConfig
+from src.models.gpt.context import describe_context, extend_context
+from src.tokenization.base import identity_matches
+from src.tokenization.registry import load_tokenizer
+
+
+class TextGenerator:
+
+    def __init__(
+        self,
+        engine,
+        tokenizer,
+        default_config=None,
+        max_batch_size=8,
+        cache_budget_bytes=None,
+        prefill_chunk=256,
+        chat_template=None,
+        metadata=None,
+    ):
+        self.engine = engine
+        self.tokenizer = tokenizer
+        self.default_config = default_config or GenerationConfig()
+        self.chat_template = chat_template
+        self.metadata = metadata or {}
+        self.scheduler = BatchScheduler(
+            engine,
+            tokenizer,
+            max_batch_size=max_batch_size,
+            cache_budget_bytes=cache_budget_bytes,
+            prefill_chunk=prefill_chunk,
+        )
+
+    @property
+    def config(self):
+        return self.engine.config
+
+    @property
+    def weights_nbytes(self):
+        return self.engine.nbytes
+
+    def bound_chat_template(self, fallback=None):
+        template = self.chat_template or ChatTemplate.resolve(fallback)
+
+        if template is None:
+            raise ValueError("model has no chat template and no fallback was given")
+
+        return template.bind(self.tokenizer)
+
+    def _request(self, prompt, overrides, timeout_s=None, truncate_prompt=False, request_id=None):
+        kwargs = {}
+
+        if request_id:
+            kwargs["request_id"] = request_id
+
+        return GenerationRequest(
+            prompt=prompt,
+            config=self.default_config.updated(**overrides),
+            timeout_s=timeout_s,
+            truncate_prompt=truncate_prompt,
+            **kwargs,
+        )
+
+    def submit(self, request):
+        return self.scheduler.submit(request)
+
+    def generate(self, prompt, timeout_s=None, truncate_prompt=False, **overrides):
+        handle = self.submit(self._request(prompt, overrides, timeout_s, truncate_prompt))
+        return handle.result(drive=not self.scheduler.background)
+
+    def generate_stream(self, prompt, timeout_s=None, truncate_prompt=False, events=False, **overrides):
+        handle = self.submit(self._request(prompt, overrides, timeout_s, truncate_prompt))
+
+        stream = handle.events(drive=not self.scheduler.background)
+
+        try:
+            for ev in stream:
+                if events:
+                    yield ev
+                elif ev.text:
+                    yield ev.text
+        finally:
+            stream.close()
+            handle.cancel()
+
+    def start(self):
+        self.scheduler.start()
+
+    def stop(self):
+        self.scheduler.stop()
+
+
+def load_model(path, max_batch_size=8, cache_budget_bytes=None, verify=True, prefill_chunk=256, chat_template=None,
+               context_length=None, context_extension=None):
+    with open(os.path.join(path, "config.json")) as f:
+        meta = json.load(f)
+
+    if meta["format_version"] != EXPORT_FORMAT_VERSION:
+        raise ValueError(f"unsupported export format version {meta['format_version']}")
+
+    weights_path = os.path.join(path, meta["weights"]["file"])
+
+    if verify and _sha256(weights_path) != meta["weights"]["sha256"]:
+        raise ValueError("model weights failed checksum verification")
+
+    config = GPTConfig.from_dict(meta["model"])
+
+    if context_length is not None or context_extension is not None:
+        if context_length is None or context_extension is None:
+            raise ValueError("context extension needs both context_length and context_extension")
+
+        config = extend_context(config, context_length, context_extension)
+
+    with np.load(weights_path, allow_pickle=False) as archive:
+        weights = {k: archive[k] for k in archive.files}
+
+    tokenizer = load_tokenizer(os.path.join(path, meta["tokenizer"]["file"]))
+
+    if not identity_matches(meta["tokenizer"]["identity"], tokenizer.identity):
+        raise ValueError("bundled tokenizer does not match the exported model")
+
+    gen_path = os.path.join(path, "generation.json")
+    default = GenerationConfig()
+
+    if os.path.isfile(gen_path):
+        with open(gen_path) as f:
+            default = GenerationConfig.from_dict(json.load(f))
+
+    template = ChatTemplate.resolve(chat_template if chat_template is not None else meta.get("chat_template"))
+
+    if template is not None:
+        template.bind(tokenizer)
+
+    engine = InferenceModel(config, weights)
+    del weights
+
+    return TextGenerator(
+        engine,
+        tokenizer,
+        default,
+        max_batch_size=max_batch_size,
+        cache_budget_bytes=cache_budget_bytes,
+        prefill_chunk=prefill_chunk,
+        chat_template=template,
+        metadata={
+            "weights_sha256": meta["weights"]["sha256"],
+            "weights_dtype": meta["weights"]["dtype"],
+            "source": meta.get("source", {}),
+            "path": os.path.abspath(path),
+            "context": describe_context(config),
+        },
+    )

src/inference/sampling.py

@@ -0,0 +1,42 @@
+import numpy as np
+
+
+def _softmax(x):
+    x = x - x.max()
+    e = np.exp(x)
+    return e / e.sum()
+
+
+def sample_token(logits, config, rng, seen_ids=None):
+    logits = np.array(logits, dtype=np.float64, copy=True)
+
+    penalty = config.repetition_penalty
+
+    if penalty != 1.0 and seen_ids:
+        idx = np.fromiter(seen_ids, dtype=np.int64, count=len(seen_ids))
+        vals = logits[idx]
+        logits[idx] = np.where(vals > 0, vals / penalty, vals * penalty)
+
+    if config.greedy:
+        return int(np.argmax(logits))
+
+    logits /= config.temperature
+
+    candidates = None
+
+    if 0 < config.top_k < logits.size:
+        candidates = np.argpartition(logits, -config.top_k)[-config.top_k:]
+        logits = logits[candidates]
+
+    if config.top_p < 1.0:
+        order = np.argsort(-logits, kind="stable")
+        probs = _softmax(logits[order])
+        cutoff = int(np.searchsorted(np.cumsum(probs), config.top_p)) + 1
+        keep = order[:cutoff]
+        candidates = keep if candidates is None else candidates[keep]
+        logits = logits[keep]
+
+    probs = _softmax(logits)
+    choice = int(rng.choice(probs.size, p=probs))
+
+    return int(candidates[choice]) if candidates is not None else choice