modeling.py

from __future__ import annotations
import json
import os
import time
from typing import Optional

os.environ.setdefault("CUDA_VISIBLE_DEVICES", "")
os.environ.setdefault("HIP_VISIBLE_DEVICES", "")
import torch
import torch.nn as nn


def _read_config(weights_dir: str) -> dict:
    path = os.path.join(weights_dir, "config.json")
    if not os.path.exists(path):
        path = weights_dir if weights_dir.endswith(".json") else path
    with open(path) as f:
        return json.load(f)


def _build_config(cfg_json: dict):
    from model_cpu_gpt2 import CPUGPTConfig

    return CPUGPTConfig(
        vocab_size=cfg_json["vocab_size"],
        seq_len=cfg_json.get("seq_len", 1024),
        n_layer=cfg_json["n_layer"],
        n_embd=cfg_json["n_embd"],
        n_head=cfg_json["n_head"],
        ffn_hidden=cfg_json["ffn_hidden"],
        layer_pattern=cfg_json.get("layer_pattern", "SSSL"),
        gla_delta=cfg_json.get("gla_delta", True),
        fno_modes=cfg_json.get("fno_modes", 512),
        gla_chunk=cfg_json.get("gla_chunk", 64),
        landmark_layer_every=cfg_json.get("landmark_layer_every", 0),
        landmark_chunk=cfg_json.get("landmark_chunk", 32),
        landmark_max=cfg_json.get("landmark_max", 64),
        attn_layer_every=cfg_json.get("attn_layer_every", 0),
        dropout=0.0,
    )


class FelaLM:
    def __init__(self, model, cfg, tokenizer, cfg_json):
        self.model = model
        self.cfg = cfg
        self.tok = tokenizer
        self.cfg_json = cfg_json

        def _tid(name):
            i = tokenizer.token_to_id(name)
            return i if i is not None and i >= 0 else None

        self.fim_prefix = _tid("<|fim_prefix|>")
        self.fim_suffix = _tid("<|fim_suffix|>")
        self.fim_middle = _tid("<|fim_middle|>")
        self.fim_pad = _tid("<|fim_pad|>")
        self.eot = _tid("<|endoftext|>")
        self.fim_ok = None not in (self.fim_prefix, self.fim_suffix, self.fim_middle)
        self._stops = {
            t
            for t in (
                self.fim_prefix,
                self.fim_suffix,
                self.fim_middle,
                self.fim_pad,
                self.eot,
            )
            if t is not None
        }

    @torch.no_grad()
    def complete(
        self,
        prefix: str,
        suffix: str = "",
        max_tokens: int = 40,
        temperature: float = 0.0,
        single_line: bool = True,
    ) -> dict:
        prefix = prefix or ""
        suffix = suffix or ""
        used_fim = bool(suffix.strip()) and self.fim_ok
        if used_fim:
            ids = (
                [self.fim_prefix]
                + self.tok.encode(prefix).ids
                + [self.fim_suffix]
                + self.tok.encode(suffix).ids
                + [self.fim_middle]
            )
        else:
            ids = self.tok.encode(prefix).ids
        if not ids:
            ids = [self.eot] if self.eot is not None else [0]
        t0 = time.perf_counter()
        states = self.model.init_state(batch_size=1)
        logits = None
        for tok_id in ids:
            logits, states = self.model.step(
                torch.tensor([tok_id], dtype=torch.long), states
            )
        prefill_ms = (time.perf_counter() - t0) * 1000.0
        out_ids = []
        td = time.perf_counter()
        for _ in range(max_tokens):
            if temperature and temperature > 0:
                probs = torch.softmax(logits.float().reshape(-1) / temperature, -1)
                nxt = int(torch.multinomial(probs, 1).item())
            else:
                nxt = int(logits.float().reshape(-1).argmax().item())
            if nxt in self._stops:
                break
            out_ids.append(nxt)
            piece = self.tok.decode(out_ids)
            if single_line and "\n" in piece:
                break
            logits, states = self.model.step(
                torch.tensor([nxt], dtype=torch.long), states
            )
        decode_ms = (time.perf_counter() - td) * 1000.0
        text = self.tok.decode(out_ids) if out_ids else ""
        if single_line:
            text = text.split("\n", 1)[0]
        n = len(out_ids)
        return {
            "middle": text,
            "n_tokens": n,
            "used_fim": used_fim,
            "prompt_tokens": len(ids),
            "prefill_ms": round(prefill_ms, 1),
            "decode_ms": round(decode_ms, 1),
            "tok_per_s": round(n / (decode_ms / 1000.0), 2)
            if decode_ms > 0 and n
            else 0.0,
        }


def _read_bf16_state(weights_dir: str) -> dict:
    from safetensors import safe_open

    st = {}
    path = os.path.join(weights_dir, "model.safetensors")
    with safe_open(path, framework="pt", device="cpu") as f:
        for k in f.keys():
            st[k] = f.get_tensor(k).float()
    return st


def _read_int8_state(weights_dir: str) -> dict:
    from safetensors import safe_open

    st = {}
    path = os.path.join(weights_dir, "model_int8.safetensors")
    with safe_open(path, framework="pt", device="cpu") as f:
        keys = list(f.keys())
        for k in keys:
            if k.startswith("keep."):
                st[k[len("keep.") :]] = f.get_tensor(k).float()
        for k in keys:
            if k.startswith("int8."):
                base = k[len("int8.") :]
                w = f.get_tensor(k).float()
                s = f.get_tensor("scale." + base).float()
                st[base] = w * s.reshape([-1] + [1] * (w.dim() - 1))
    return st


def _apply_state(model, st: dict) -> None:
    params = dict(model.named_parameters())
    params.update(dict(model.named_buffers()))
    keys = set(st)
    for k in keys:
        dst = params.get(k)
        if dst is None:
            raise KeyError(f"Checkpoint key {k!r} has no home in the model")
        with torch.no_grad():
            dst.copy_(st[k])
    missing = set(params) - keys
    if missing:
        raise KeyError(f"Missing {len(missing)} params, e.g. {sorted(missing)[:5]}")


def _resolve_quant(weights_dir: str, quant: str) -> str:
    if quant == "auto":
        has_int8 = os.path.exists(os.path.join(weights_dir, "model_int8.safetensors"))
        has_bf16 = os.path.exists(os.path.join(weights_dir, "model.safetensors"))
        return "bf16" if has_bf16 else ("int8" if has_int8 else "bf16")
    return quant


def load_model(
    weights_dir: str = ".", threads: Optional[int] = None, quant: str = "bf16"
) -> FelaLM:
    from model_cpu_gpt2 import CPUGPT
    from cpu_patch import enable_cpu_delta
    from tokenizers import Tokenizer

    if threads:
        torch.set_num_threads(threads)
    if weights_dir.endswith(".safetensors"):
        weights_dir = os.path.dirname(os.path.abspath(weights_dir)) or "."
    cfg_json = _read_config(weights_dir)
    cfg = _build_config(cfg_json)
    model = CPUGPT(cfg)
    model.lm_head = nn.Linear(cfg.n_embd, cfg.vocab_size, bias=False)
    quant = _resolve_quant(weights_dir, quant)
    if quant == "int8":
        st = _read_int8_state(weights_dir)
    else:
        st = _read_bf16_state(weights_dir)
    _apply_state(model, st)
    model.eval()
    enable_cpu_delta(model)
    model.prepare_inference()
    tok_path = os.path.join(weights_dir, "tokenizer.json")
    tokenizer = Tokenizer.from_file(tok_path)
    return FelaLM(model, cfg, tokenizer, cfg_json)


def from_pretrained(repo_id: str = "lowdown-labs/FELA-autocomplete") -> FelaLM:
    from huggingface_hub import hf_hub_download

    d = os.path.dirname(hf_hub_download(repo_id, "config.json"))
    hf_hub_download(repo_id, "model.safetensors")
    hf_hub_download(repo_id, "tokenizer.json")
    hf_hub_download(repo_id, "model_cpu_gpt2.py")
    for f in ("cpu_delta.py", "cpu_landmark.py", "cpu_swa.py", "cpu_patch.py"):
        hf_hub_download(repo_id, f)
    return load_model(d)