File size: 3,902 Bytes
4012ebc | 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 | #!/usr/bin/env python3
"""Run inference from the Hugging Face model repo without cloning it.
Usage:
pip install torch tokenizers huggingface-hub
python examples/inference_from_hf.py "Polska jest" 80
"""
from __future__ import annotations
import importlib.util
import sys
import time
from pathlib import Path
import torch
import torch.nn.functional as F
from huggingface_hub import hf_hub_download
from tokenizers import Tokenizer
REPO_ID = "SlayerLab/slayer-gpt-tokenizer-model"
TEMP = 0.7
TOP_K = 40
TOP_P = 0.92
REP_PEN = 1.15
NGRAM = 3
EOT = 0
def load_model_module(path: str):
spec = importlib.util.spec_from_file_location("slayer_gpt_model", path)
if spec is None or spec.loader is None:
raise RuntimeError(f"Could not load model module from {path}")
module = importlib.util.module_from_spec(spec)
sys.modules[spec.name] = module
spec.loader.exec_module(module)
return module
def banned_next_tokens(seq: list[int], n: int) -> set[int]:
if len(seq) < n - 1:
return set()
prefix = tuple(seq[-(n - 1):])
banned: set[int] = set()
for i in range(len(seq) - n + 1):
if tuple(seq[i:i + n - 1]) == prefix:
banned.add(seq[i + n - 1])
return banned
@torch.no_grad()
def generate(model, tokenizer: Tokenizer, prompt: str, max_new_tokens: int, block_size: int, device: str) -> tuple[str, float]:
idx = torch.tensor(tokenizer.encode(prompt).ids, dtype=torch.long, device=device)[None]
start = time.time()
generated = 0
for _ in range(max_new_tokens):
logits, _ = model(idx[:, -block_size:])
logits = logits[:, -1, :].float()
for token_id in set(idx[0].tolist()):
logits[0, token_id] /= REP_PEN if logits[0, token_id] > 0 else 1 / REP_PEN
for token_id in banned_next_tokens(idx[0].tolist(), NGRAM):
logits[0, token_id] = -float("inf")
logits /= TEMP
kth = torch.topk(logits, TOP_K)[0][..., -1, None]
logits[logits < kth] = -float("inf")
sorted_logits, sorted_indices = torch.sort(logits, descending=True)
cumulative = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
remove = cumulative > TOP_P
remove[..., 1:] = remove[..., :-1].clone()
remove[..., 0] = False
logits[0, sorted_indices[0][remove[0]]] = -float("inf")
next_id = torch.multinomial(F.softmax(logits, dim=-1), 1)
generated += 1
if next_id.item() == EOT:
break
idx = torch.cat([idx, next_id], dim=1)
tokens_per_second = generated / max(time.time() - start, 1e-6)
return tokenizer.decode(idx[0].tolist()), tokens_per_second
def main() -> None:
prompt = sys.argv[1] if len(sys.argv) > 1 else "Polska jest"
max_new_tokens = int(sys.argv[2]) if len(sys.argv) > 2 else 80
device = "mps" if torch.backends.mps.is_available() else "cuda" if torch.cuda.is_available() else "cpu"
model_py = hf_hub_download(REPO_ID, "scripts/model.py")
ckpt_path = hf_hub_download(REPO_ID, "model/ckpt.pt")
tokenizer_path = hf_hub_download(REPO_ID, "tokenizers/polish_bpe_32k.json")
model_module = load_model_module(model_py)
ckpt = torch.load(ckpt_path, map_location="cpu")
model = model_module.GPT(model_module.GPTConfig(**ckpt["model_args"]))
state_dict = ckpt["model"]
for key in list(state_dict):
if key.startswith("_orig_mod."):
state_dict[key[len("_orig_mod."):]] = state_dict.pop(key)
model.load_state_dict(state_dict)
model.eval().to(device)
tokenizer = Tokenizer.from_file(tokenizer_path)
text, tps = generate(
model,
tokenizer,
prompt,
max_new_tokens,
ckpt["model_args"]["block_size"],
device,
)
print(f"[repo={REPO_ID} device={device} {tps:.1f} tok/s]\n")
print(text)
if __name__ == "__main__":
main()
|