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)