Views
No views yet
n_layer: 8n_head: 8n_embd: 512block_size: 1024vocab_size: 50257dropout: 0.11import torch
2import tiktoken
3from model import GPT, GPTConfig # Assuming model.py is available or its classes are defined
4
5# 1. Define model configuration (must match the trained model's config.json)
6# You can load this from config.json if you save it, or define it manually
7config = GPTConfig(
8 vocab_size=50257,
9 block_size=1024,
10 n_layer=8,
11 n_head=8,
12 n_embd=512,
13 dropout=0.1,
14 bias=True
15)
16
17# 2. Initialize the model and load weights
18model = GPT(config)
19state_dict = torch.load("pytorch_model.bin", map_location='cpu') # Replace with path to downloaded model
20model.load_state_dict(state_dict)
21model.eval() # Set to evaluation mode
22device = 'cuda' if torch.cuda.is_available() else 'cpu'
23model.to(device)
24
25# 3. Load the tiktoken tokenizer
26tokenizer = tiktoken.get_encoding("gpt2")
27EOT_TOKEN_ID = tokenizer.eot_token
28
29# 4. Prepare your prompt for text generation
30prompt_text = "Once upon a time there was a pumpkin."
31
32# Encode the prompt
33allowed_special_tokens = 'all'
34input_ids = tokenizer.encode(prompt_text, allowed_special=allowed_special_tokens)
35input_ids_tensor = torch.tensor([input_ids], dtype=torch.long).to(device)
36
37# 5. Generate text
38# Adjust max_new_tokens, temperature, top_k as needed
39generated_output_ids = model.generate(
40 idx=input_ids_tensor,
41 max_new_tokens=100, # Max length for the generated text
42 temperature=0.7,
43 top_k=50
44)
45
46# Decode the generated text (excluding the prompt part)
47generated_text_ids = generated_output_ids[0, len(input_ids):].tolist()
48generated_text = tokenizer.decode(generated_text_ids)
49
50# Clean up any leftover EOT tokens from generation
51generated_text = generated_text.replace(tokenizer.decode([EOT_TOKEN_ID]), "").strip()
52
53print(f"Generated Text: {generated_text}")