| """Portable native WiggleGPT loader and exact GPT-2 tokenizer.
|
|
|
| The model class is the byte-preserved source in native/model_bio.py.
|
| Model construction and all forwards are CUDA only. Tokenization is CPU only.
|
| """
|
| from __future__ import annotations
|
| import base64
|
| import hashlib
|
| import os
|
| from pathlib import Path
|
| import subprocess
|
| import sys
|
| ROOT = Path(__file__).resolve().parent
|
| sys.path.insert(0, str(ROOT / 'native'))
|
|
|
| DESKTOP_GPU_APPS = {
|
| "shellhost.exe", "crossdeviceresume.exe", "lghub_system_tray.exe", "dcv2.exe",
|
| "snippingtool.exe", "steelseriesggclient.exe", "explorer.exe", "textinputhost.exe",
|
| "startmenuexperiencehost.exe", "searchhost.exe", "shellexperiencehost.exe",
|
| "msedgewebview2.exe", "windowsterminal.exe", "chrome.exe", "msedge.exe", "firefox.exe",
|
| "chatgpt.exe", "codex-computer-use-swift.exe", "razerappengine.exe", "notepad++.exe",
|
| "obsidian.exe", "dwm.exe",
|
| }
|
|
|
| def get_encoder():
|
| import tiktoken
|
| from tiktoken_ext.openai_public import r50k_pat_str
|
|
|
| raw = (ROOT / "tokenizer" / "gpt2.tiktoken").read_bytes()
|
|
|
| if hashlib.sha256(raw).hexdigest() != "306cd27f03c1a714eca7108e03d66b7dc042abe8c258b44c199a7ed9838dd930":
|
| raise ValueError("Bundled GPT-2 merge table is damaged.")
|
| ranks = {base64.b64decode(token): int(rank) for token, rank in
|
| (line.split() for line in raw.splitlines() if line)}
|
| return tiktoken.Encoding(name="gpt2", pat_str=r50k_pat_str, mergeable_ranks=ranks,
|
| special_tokens={"<|endoftext|>": 50256}, explicit_n_vocab=50257)
|
|
|
| def check_gpu():
|
| result = subprocess.run(
|
| ["nvidia-smi", "--query-compute-apps=pid,process_name", "--format=csv,noheader,nounits"],
|
| capture_output=True, text=True, check=True,
|
| )
|
| busy = []
|
| for line in result.stdout.splitlines():
|
| fields = line.split(",", 1)
|
|
|
|
|
| desktop = len(fields) == 2 and Path(fields[1].strip()).name.lower() in DESKTOP_GPU_APPS
|
| if fields[0].strip().isdigit() and int(fields[0]) != os.getpid() and not desktop:
|
| busy.append(line.strip())
|
| if busy:
|
| raise RuntimeError("Another GPU job is running. Finish it before loading chat: " + "; ".join(busy))
|
|
|
| def load_model(checkpoint_path):
|
| check_gpu()
|
| import torch
|
| from model_bio import GPT, GPTConfig
|
|
|
| if not torch.cuda.is_available():
|
| raise RuntimeError("WiggleGPT needs CUDA here; CPU model execution is disabled.")
|
| torch.set_num_threads(2)
|
| torch.cuda.set_device(0)
|
| free, total = torch.cuda.mem_get_info()
|
| budget = min(total - int(2.5 * 1024**3), free - int(0.5 * 1024**3))
|
| if budget < 2 * 1024**3:
|
| raise RuntimeError("Less than 2 GiB is available after leaving desktop room.")
|
| torch.cuda.set_per_process_memory_fraction(budget / total, 0)
|
| torch.set_float32_matmul_precision("high")
|
| checkpoint_path = Path(checkpoint_path)
|
| checkpoint = torch.load(checkpoint_path, map_location="cpu", mmap=True, weights_only=False)
|
| config = GPTConfig(**checkpoint["model_args"])
|
|
|
| with torch.device("cuda:0"):
|
| model = GPT(config)
|
| state = {key.removeprefix("_orig_mod."): value for key, value in checkpoint["model"].items()}
|
| model.load_state_dict(state, strict=True)
|
| model.eval()
|
| metadata = {"checkpoint": str(checkpoint_path), "saved_step": checkpoint.get("iter_num", checkpoint.get("step")),
|
| "parameters": sum(parameter.numel() for parameter in model.parameters()),
|
| "model_args": checkpoint["model_args"], "torch": torch.__version__}
|
| del checkpoint, state
|
| encoder = get_encoder()
|
| print(f"CUDA: {torch.cuda.get_device_name(0)}; {metadata['parameters']:,} parameters; "
|
| f"{budget / 1024**3:.2f} GiB memory cap", flush=True)
|
| print(f"Checkpoint: {checkpoint_path}", flush=True)
|
| return model, config, encoder, metadata
|
|
|
