| import argparse |
| import json |
| from pathlib import Path |
|
|
| import torch |
| import tiktoken |
| from safetensors.torch import load_file |
|
|
| from modeling_ys import GPTModel |
|
|
|
|
| TOKENIZER_VOCAB_SIZE = 50257 |
| END_OF_TEXT_ID = 50256 |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser(description="Génère du texte avec Ys.") |
| parser.add_argument("--prompt", default="Il était une fois") |
| parser.add_argument("--max-new-tokens", type=int, default=120) |
| parser.add_argument("--temperature", type=float, default=0.8, |
| help="0 pour choisir le token le plus probable") |
| parser.add_argument("--top-k", type=int, default=40, |
| help="0 pour désactiver le filtrage top-k") |
| parser.add_argument("--device", choices=("auto", "cpu", "cuda"), default="auto") |
| args = parser.parse_args() |
|
|
| repo_dir = Path(__file__).resolve().parent |
| config = json.loads((repo_dir / "config.json").read_text(encoding="utf-8")) |
| device = args.device |
| if device == "auto": |
| device = "cuda" if torch.cuda.is_available() else "cpu" |
| if device == "cuda" and not torch.cuda.is_available(): |
| raise SystemExit("CUDA n'est pas disponible. Relance avec --device cpu.") |
| if args.temperature < 0: |
| raise SystemExit("--temperature doit être positif ou égal à 0.") |
|
|
| model = GPTModel(config) |
| state_dict = load_file(str(repo_dir / "model.safetensors"), device="cpu") |
| model.load_state_dict(state_dict, strict=True) |
| model.to(device).eval() |
|
|
| tokenizer = tiktoken.get_encoding("gpt2") |
| token_ids = torch.tensor([tokenizer.encode(args.prompt)], dtype=torch.long, device=device) |
| if token_ids.shape[1] == 0: |
| token_ids = torch.tensor([[END_OF_TEXT_ID]], dtype=torch.long, device=device) |
|
|
| for _ in range(args.max_new_tokens): |
| context = token_ids[:, -config["context_length"]:] |
| with torch.inference_mode(): |
| logits = model(context)[:, -1, :TOKENIZER_VOCAB_SIZE] |
|
|
| if args.temperature == 0: |
| next_id = torch.argmax(logits, dim=-1, keepdim=True) |
| else: |
| logits = logits / args.temperature |
| if 0 < args.top_k < logits.shape[-1]: |
| threshold = torch.topk(logits, args.top_k, dim=-1).values[:, -1, None] |
| logits = logits.masked_fill(logits < threshold, float("-inf")) |
| probabilities = torch.softmax(logits, dim=-1) |
| next_id = torch.multinomial(probabilities, num_samples=1) |
|
|
| if next_id.item() == END_OF_TEXT_ID: |
| break |
| token_ids = torch.cat((token_ids, next_id), dim=1) |
|
|
| print(tokenizer.decode(token_ids[0].tolist())) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
