| """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() |
|
|
