File size: 6,005 Bytes
56811f1 |
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 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 |
# Copyright Lightning AI. Licensed under the Apache License 2.0, see LICENSE file.
# Derivated from https://github.com/Lightning-AI/litgpt/blob/main/litgpt/generate/base.py
import os
import sys
import time
from pathlib import Path
from typing import Any, Optional
import torch
# support running without installing as a package
wd = Path(__file__).parent.parent.resolve()
sys.path.append(str(wd))
from modello_italia import Italia, ItaliaConfig, Tokenizer
device = 'cuda' if torch.cuda.is_available() else 'cpu'
MI_SYSTEM_PROMPT_SHORT = (
"Tu sei Modello Italia, un modello di linguaggio naturale addestrato da iGenius."
)
def multinomial_num_samples_1(probs: torch.Tensor) -> torch.Tensor:
if torch._dynamo.is_compiling():
# Faster alternative to `torch.multinomial(probs, num_samples=1)` that is also CUDAGraph friendly
distribution = torch.empty_like(probs).exponential_(1)
return torch.argmax(probs / distribution, dim=-1, keepdim=True)
return torch.multinomial(probs, num_samples=1)
def sample(
logits: torch.Tensor, temperature: float = 1.0, top_k: Optional[int] = None
) -> torch.Tensor:
logits = logits[0, -1]
# optionally crop the logits to only the top k options
if top_k is not None:
v, i = torch.topk(logits, min(top_k, logits.size(-1)))
# do not use `torch.where` as in nanogpt because it will repeat top-k collisions
logits = torch.full_like(logits, float("-inf")).scatter_(-1, i, v)
# optionally scale the logits and sample from a probability distribution
if temperature > 0.0:
probs = torch.nn.functional.softmax(logits / temperature, dim=-1)
return multinomial_num_samples_1(probs)
return torch.argmax(logits, dim=-1, keepdim=True)
def next_token(
model: Italia, input_pos: torch.Tensor, x: torch.Tensor, **kwargs: Any
) -> torch.Tensor:
logits = model(x, input_pos)
next = sample(logits, **kwargs)
return next.to(dtype=x.dtype)
@torch.inference_mode()
def generate(
model: Italia,
prompt: torch.Tensor,
tokenizer: Tokenizer,
max_returned_tokens: int,
*,
temperature: float = 1.0,
top_k: Optional[int] = None,
eos_id: Optional[int] = None,
) -> torch.Tensor:
"""Takes a conditioning sequence (prompt) as input and continues to generate as many tokens as requested.
The implementation of this function is modified from A. Karpathy's nanoGPT.
Args:
model: The model to use.
prompt: Tensor of shape (T) with indices of the prompt sequence.
max_returned_tokens: The maximum number of tokens to return (given plus generated).
tokenizer: Tokenizer instance to decode generated tokens
temperature: Scales the predicted logits by 1 / temperature.
top_k: If specified, only sample among the tokens with the k highest probabilities.
"""
T = prompt.size(0)
assert max_returned_tokens > T
device = prompt.device
tokens = [prompt]
input_pos = torch.tensor([T], device=device)
token = next_token(
model,
torch.arange(0, T, device=device),
prompt.view(1, -1),
temperature=temperature,
top_k=top_k,
).clone()
tokens.append(token)
for _ in range(2, max_returned_tokens - T + 1):
token = next_token(
model, input_pos, token.view(1, -1), temperature=temperature, top_k=top_k
).clone()
tokens.append(token)
if token == tokenizer.eos_id:
break
os.system('cls' if os.name == 'nt' else 'clear')
print(tokenizer.decode(torch.cat(tokens)[T:]))
input_pos = input_pos.add_(1)
return torch.cat(tokens)
@torch.inference_mode()
def main(
prompt: str = "Ciao, chi sei?",
*,
num_samples: int = 1,
max_new_tokens: int = 200,
top_k: Optional[int] = 200,
temperature: float = 0.4,
checkpoint_dir: Path = Path("."),
) -> None:
"""Generates text samples based on a pre-trained model and tokenizer.
Args:
prompt: The prompt string to use for generating the samples.
num_samples: The number of text samples to generate.
max_new_tokens: The number of generation steps to take.
top_k: The number of top most probable tokens to consider in the sampling process.
temperature: A value controlling the randomness of the sampling process. Higher values result in more random
samples.
checkpoint_dir: The checkpoint directory to load.
"""
config = ItaliaConfig()
checkpoint_path = checkpoint_dir / "italia.bin"
tokenizer = Tokenizer(checkpoint_dir)
prompt = f"<|system|>{MI_SYSTEM_PROMPT_SHORT}\n<|user|>{prompt}\n<|assistant|>"
encoded = tokenizer.encode(prompt, device=device)
prompt_length = encoded.size(0)
max_returned_tokens = prompt_length + max_new_tokens
print(f"Loading model {str(checkpoint_path)!r}")
t0 = time.perf_counter()
model = Italia(config)
model.load_state_dict(torch.load(checkpoint_path, mmap=True))
model.to(device)
print(
f"Time to instantiate model: {time.perf_counter() - t0:.02f} seconds.",
file=sys.stderr,
)
model.max_seq_length = max_returned_tokens
model.set_kv_cache(batch_size=1, device=device)
model.eval()
for _ in range(num_samples):
t0 = time.perf_counter()
y = generate(
model,
encoded,
tokenizer,
max_returned_tokens,
temperature=temperature,
top_k=top_k,
)
t = time.perf_counter() - t0
for block in model.transformer.h:
block.attn.kv_cache.reset_parameters()
#print(tokenizer.decode(y))
tokens_generated = y.size(0) - prompt_length
print(f"\nTime for inference: {t:.02f} sec total, {tokens_generated / t:.02f} tokens/sec")
if __name__ == "__main__":
from jsonargparse import CLI
torch.set_float32_matmul_precision("high")
CLI(main) |