File size: 2,931 Bytes
81c887b |
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 |
import torch
from transformers import AutoTokenizer
from model import SmallLanguageModel, ModelConfig
def create_model_config(vocab_size):
"""Create model configuration matching training"""
return ModelConfig(
vocab_size=vocab_size,
block_size=512,
n_layer=12,
n_head=12,
n_embd=768,
dropout=0.1,
bias=True
)
def generate_text(prompt, model, tokenizer, max_length=100, temperature=0.8, top_k=50):
model.eval()
input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(device)
with torch.no_grad():
for _ in range(max_length):
# Get model predictions
outputs = model(input_ids)
next_token_logits = outputs[:, -1, :] / temperature
# Apply top-k filtering
top_k_logits, top_k_indices = torch.topk(next_token_logits, top_k, dim=-1)
next_token_logits[0, :] = float('-inf')
next_token_logits[0, top_k_indices[0]] = top_k_logits[0]
# Sample from the filtered distribution
probs = torch.softmax(next_token_logits, dim=-1)
next_token = torch.multinomial(probs, num_samples=1)
# Append to input_ids
input_ids = torch.cat([input_ids, next_token], dim=-1)
# Stop if we generate the EOS token
if next_token[0].item() == tokenizer.eos_token_id:
break
return tokenizer.decode(input_ids[0], skip_special_tokens=True)
if __name__ == "__main__":
# Load tokenizer
tokenizer = AutoTokenizer.from_pretrained("gpt2")
tokenizer.pad_token = tokenizer.eos_token
# Setup device
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Using device: {device}")
# Create and load model
config = create_model_config(tokenizer.vocab_size)
model = SmallLanguageModel(config).to(device)
# Load trained weights
try:
checkpoint = torch.load("small_language_model.pt", map_location=device)
model.load_state_dict(checkpoint)
print("Loaded model from small_language_model.pt")
except FileNotFoundError:
print("No saved model found. Please train the model first.")
exit(1)
# Generate some example texts
prompts = [
"Once upon a time",
"The meaning of life is",
"In the distant future",
"The best way to learn programming is"
]
print("\nGenerating text samples:\n")
for prompt in prompts:
print(f"Prompt: {prompt}")
generated_text = generate_text(
prompt,
model,
tokenizer,
max_length=100,
temperature=0.8,
top_k=50
)
print(f"Generated: {generated_text}\n")
|