File size: 2,269 Bytes
67ac79e | 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 | """Génération locale avec le checkpoint PyTorch Ys-17M."""
from __future__ import annotations
import argparse
import hashlib
from pathlib import Path
import torch
from tokenizers import Tokenizer
from model import GPT, ModelConfig
def main():
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--prompt", default="Il était une fois")
parser.add_argument("--max-new-tokens", type=int, default=80)
parser.add_argument("--temperature", type=float, default=0.8)
parser.add_argument("--top-k", type=int, default=50)
parser.add_argument("--top-p", type=float, default=0.95)
parser.add_argument("--seed", type=int, default=42)
parser.add_argument("--device", choices=("auto", "cpu", "cuda"), default="auto")
args = parser.parse_args()
selected = args.device
if selected == "auto":
selected = "cuda" if torch.cuda.is_available() else "cpu"
if selected == "cuda" and not torch.cuda.is_available():
parser.error("CUDA indisponible : utilise --device cpu.")
device = torch.device(selected)
torch.set_num_threads(min(8, torch.get_num_threads()))
torch.manual_seed(args.seed)
directory = Path(__file__).resolve().parent
checkpoint = torch.load(directory / "best.pt", map_location="cpu", weights_only=True)
tokenizer_path = directory / "tokenizer.json"
actual_hash = hashlib.sha256(tokenizer_path.read_bytes()).hexdigest()
if actual_hash != checkpoint["tokenizer_hash"]:
raise ValueError("Le tokenizer ne correspond pas au checkpoint Ys-17M.")
tokenizer = Tokenizer.from_file(str(tokenizer_path))
model = GPT(ModelConfig(**checkpoint["model_config"]))
model.load_state_dict(checkpoint["model"], strict=True)
model.to(device).eval()
ids = tokenizer.encode(args.prompt).ids or [tokenizer.token_to_id("<|bos|>")]
inputs = torch.tensor([ids], dtype=torch.long, device=device)
result = model.generate(
inputs,
max_new_tokens=args.max_new_tokens,
temperature=args.temperature,
top_k=args.top_k,
top_p=args.top_p,
eos_id=tokenizer.token_to_id("<|eos|>"),
)
print(tokenizer.decode(result[0].tolist(), skip_special_tokens=True))
if __name__ == "__main__":
main()
|