File size: 1,406 Bytes
54a634f | 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 | import os
import sys
import torch
import torch.nn.functional as F
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
from model.model import ModelConfig, Transformer
def generate(model: Transformer, prompt_tokens: torch.Tensor, max_new_tokens: int = 50, temperature: float = 0.8, top_k: int = 40):
model.eval()
tokens = prompt_tokens.clone()
with torch.no_grad():
for _ in range(max_new_tokens):
logits = model(tokens)
logits = logits[:, -1, :] / temperature
if top_k is not None:
v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
logits[logits < v[:, [-1]]] = -float('Inf')
probs = F.softmax(logits, dim=-1)
next_token = torch.multinomial(probs, num_samples=1)
tokens = torch.cat((tokens, next_token), dim=1)
return tokens
if __name__ == "__main__":
device = "cuda" if torch.cuda.is_available() else "cpu"
config = ModelConfig.get_125m()
model = Transformer(config).to(device)
dummy_input = torch.randint(0, config.vocab_size, (1, 10), device=device)
print("Generating sample sequence from model...")
output = generate(model, dummy_input, max_new_tokens=20)
print(f"Generated Tokens Shape: {output.shape}")
print(f"Generated Token IDs: {output[0].tolist()}")
|