File size: 2,677 Bytes
eefb08d | 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 | 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()
|