File size: 5,442 Bytes
5ef3875 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 | """Inference: load a trained YCode model and generate answers."""
from __future__ import annotations
import contextlib
import json
import os
from pathlib import Path
from typing import Callable
import torch
from ycode.lm.data import format_chat
from ycode.lm.tokenizer import END, USER
from ycode.lm.train import load_checkpoint, pick_device
# Models larger than this run with 16-bit weights: half the memory (3 GB instead of 6 GB for
# the 1.5B v5) and faster on most CPUs. Smaller models keep full precision.
HALF_PRECISION_PARAMS = 500e6
def stored_params(model_dir: Path) -> int | None:
"""Parameter count from a model folder's info.json, without loading the weights."""
try:
return int(json.loads((Path(model_dir) / "info.json").read_text())["params"])
except (OSError, ValueError, KeyError, TypeError):
return None
class LocalLM:
def __init__(self, model_dir: Path | str, *, device: str = "auto", dtype: str = "auto") -> None:
self.device = pick_device(device)
if self.device == "cpu":
torch.set_num_threads(max(1, os.cpu_count() or 1))
self.dtype = self._pick_dtype(Path(model_dir), dtype)
self.model, self.tokenizer, self.payload = load_checkpoint(Path(model_dir), self.device,
dtype=self.dtype)
self.model.eval()
self.path = Path(model_dir)
def _pick_dtype(self, model_dir: Path, requested: str) -> torch.dtype | None:
if requested == "fp32":
return None
if requested == "auto":
params = stored_params(model_dir)
if params is None or params < HALF_PRECISION_PARAMS:
return None
if self.device.startswith("cuda"):
from ycode.lm.optim import cuda_has_bf16
return torch.bfloat16 if cuda_has_bf16() else torch.float16
return torch.bfloat16
def precision(self):
"""16-bit weights compute under autocast; activations between layers stay fp32."""
if self.dtype is None:
return contextlib.nullcontext()
return torch.autocast(device_type="cuda" if self.device.startswith("cuda") else "cpu", dtype=self.dtype)
@property
def version(self) -> int:
return int(self.payload.get("config", {}).get("arch_version", 1))
@property
def stage(self) -> str:
return str(self.payload.get("stage", "pretrain"))
@property
def num_params(self) -> int:
return self.model.num_params()
def _stream(self, prompt_ids: list[int], *, max_new_tokens: int, temperature: float, top_k: int | None,
top_p: float | None, stop_ids: set[int], on_text: Callable[[str], None] | None,
no_repeat_ngram: int = 0) -> str:
generated: list[int] = []
emitted = 0
def on_token(token: int) -> None:
nonlocal emitted
generated.append(token)
text = self.tokenizer.decode(generated, skip_special=True)
# Hold back a trailing partial UTF-8 character.
if text.endswith("�"):
return
if on_text is not None and len(text) > emitted:
on_text(text[emitted:])
emitted = len(text)
idx = torch.tensor([prompt_ids], dtype=torch.long, device=self.device)
with self.precision():
self.model.generate(idx, max_new_tokens, temperature=temperature, top_k=top_k, top_p=top_p,
stop_ids=stop_ids, on_token=on_token, no_repeat_ngram=no_repeat_ngram)
text = self.tokenizer.decode(generated, skip_special=True)
if on_text is not None and len(text) > emitted:
on_text(text[emitted:])
return text
def complete(self, prompt: str, *, max_new_tokens: int = 200, temperature: float = 0.8,
top_k: int | None = 50, top_p: float | None = 0.95,
on_text: Callable[[str], None] | None = None) -> str:
"""Raw continuation, e.g. complete('def fibonacci(n):')."""
ids = self.tokenizer.encode(prompt, allow_special=False)
return self._stream(ids, max_new_tokens=max_new_tokens, temperature=temperature, top_k=top_k,
top_p=top_p, stop_ids={self.tokenizer.eot_id}, on_text=on_text)
def chat(self, turns: list[tuple[str, str]], *, max_new_tokens: int = 300, temperature: float = 0.7,
top_k: int | None = 40, top_p: float | None = 0.9,
on_text: Callable[[str], None] | None = None, no_repeat_ngram: int = 12) -> str:
"""Answer the last user turn. ``turns`` is a list of (role, text), role in {user, assistant}."""
text = ""
for role, content in turns:
if role == "user":
text += format_chat(content)
else:
text += f"{content.strip()}\n{END}\n"
ids = self.tokenizer.encode(text)
# Keep the prompt within the context window (the model's generate()
# trims from the left, so make sure the newest turn survives).
stop = {self.tokenizer.token_id(END), self.tokenizer.eot_id, self.tokenizer.token_id(USER)}
return self._stream(ids, max_new_tokens=max_new_tokens, temperature=temperature, top_k=top_k,
top_p=top_p, stop_ids=stop, on_text=on_text,
no_repeat_ngram=no_repeat_ngram).strip()
|