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