File size: 2,689 Bytes
d011ff0 127850a d011ff0 | 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 | from __future__ import annotations
import argparse
import json
import sys
from pathlib import Path
import torch
from safetensors.torch import load_model
from tokenizers import Tokenizer
ROOT = Path(__file__).resolve().parent
sys.path.insert(0, str(ROOT / "src"))
sys.dont_write_bytecode = True
from manas_gpt.config import ModelConfig # noqa: E402
from manas_gpt.model import ManasGPT # noqa: E402
def _device(name: str) -> torch.device:
if name == "mps" and not torch.backends.mps.is_available():
raise RuntimeError("MPS was requested but is not available")
if name == "cuda" and not torch.cuda.is_available():
raise RuntimeError("CUDA was requested but is not available")
return torch.device(name)
def main() -> None:
parser = argparse.ArgumentParser(description="Generate a Tiny Manas continuation")
parser.add_argument("--device", choices=("cpu", "cuda", "mps"), required=True)
parser.add_argument("--prompt", required=True)
parser.add_argument("--max-new-tokens", type=int, default=64)
parser.add_argument("--temperature", type=float, default=0.8)
parser.add_argument("--top-k", type=int, default=40)
parser.add_argument("--seed", type=int, default=1337)
args = parser.parse_args()
if args.max_new_tokens < 1:
raise ValueError("--max-new-tokens must be positive")
if args.temperature <= 0:
raise ValueError("--temperature must be positive")
if args.top_k < 1:
raise ValueError("--top-k must be positive")
config_payload = json.loads((ROOT / "config.json").read_text(encoding="utf-8"))
model = ManasGPT(ModelConfig(**config_payload["model_config"]))
missing, unexpected = load_model(model, ROOT / "model.safetensors", strict=True)
if missing or unexpected:
raise RuntimeError(f"Weight mismatch: missing={missing}, unexpected={unexpected}")
device = _device(args.device)
model.eval().to(device)
tokenizer = Tokenizer.from_file(str(ROOT / "tokenizer.json"))
prompt_ids = tokenizer.encode(args.prompt, add_special_tokens=False).ids
if not prompt_ids:
raise ValueError("The prompt encoded to zero tokens")
generator = torch.Generator(device=device).manual_seed(args.seed)
token_ids = torch.tensor([prompt_ids], dtype=torch.long, device=device)
with torch.inference_mode():
generated = model.generate(
token_ids,
max_new_tokens=args.max_new_tokens,
temperature=args.temperature,
top_k=args.top_k,
generator=generator,
use_cache=True,
)
print(tokenizer.decode(generated[0].tolist()))
if __name__ == "__main__":
main()
|