File size: 4,327 Bytes
ad68b7f | 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 | """
Post-training validation script.
Tests the trained model by generating text from prompts.
Usage:
python Base/scripts/validate_model.py --checkpoint_dir Base/out/pretrain/custom-100m-10m-test/final
python Base/scripts/validate_model.py --checkpoint_dir Base/out/pretrain/custom-100m-3b/final
"""
import argparse
from pathlib import Path
import torch
from litgpt import Tokenizer
from litgpt.config import Config
from litgpt.model import GPT
def load_model(checkpoint_dir: str):
"""Load a pretrained model from a litgpt checkpoint directory."""
checkpoint_dir = Path(checkpoint_dir)
if not checkpoint_dir.exists():
raise FileNotFoundError(f"Checkpoint directory not found: {checkpoint_dir}")
config = Config.from_checkpoint(checkpoint_dir)
model_path = checkpoint_dir / "lit_model.pth"
if not model_path.exists():
raise FileNotFoundError(f"Model file not found: {model_path}")
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")
with torch.device("meta"):
model = GPT(config)
checkpoint = torch.load(str(model_path), map_location="cpu", weights_only=False)
# Handle both direct state_dict and wrapped checkpoint formats
if "model" in checkpoint:
state_dict = checkpoint["model"]
else:
state_dict = checkpoint
model.load_state_dict(state_dict, assign=True)
model = model.to(device)
model.eval()
return model, config, device
def generate_text(model, tokenizer, device, prompt, max_new_tokens=200, temperature=0.8, top_k=50):
"""Generate text from a prompt."""
input_ids = tokenizer.encode(prompt, device=device).unsqueeze(0)
with torch.no_grad():
for _ in range(max_new_tokens):
# Crop to block_size if needed
idx_cond = input_ids[:, -model.max_seq_length:]
logits = model(idx_cond)
logits = logits[:, -1, :] / temperature
# Top-k filtering
if top_k is not None:
v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
logits[logits < v[:, [-1]]] = float("-inf")
probs = torch.softmax(logits, dim=-1)
next_id = torch.multinomial(probs, num_samples=1)
input_ids = torch.cat([input_ids, next_id], dim=1)
# Stop on EOS
if next_id.item() == tokenizer.eos_id:
break
return tokenizer.decode(input_ids[0])
def main():
parser = argparse.ArgumentParser(description="Validate a pretrained LitGPT model")
parser.add_argument(
"--checkpoint_dir",
type=str,
required=True,
help="Path to the checkpoint directory (e.g., Base/out/pretrain/custom-100m-10m-test/final)",
)
parser.add_argument(
"--tokenizer_dir",
type=str,
default="Base/checkpoints/EleutherAI/pythia-160m",
help="Path to the tokenizer directory",
)
parser.add_argument(
"--max_new_tokens",
type=int,
default=200,
help="Maximum tokens to generate",
)
args = parser.parse_args()
print(f"Loading model from: {args.checkpoint_dir}")
model, config, device = load_model(args.checkpoint_dir)
print(f"\nModel config: {config.name}")
print(f"Parameters: {sum(p.numel() for p in model.parameters()):,}")
print(f"Block size: {config.block_size}")
tokenizer = Tokenizer(Path(args.tokenizer_dir))
test_prompts = [
"The future of artificial intelligence",
"In a groundbreaking study, researchers found that",
"The most important thing about education is",
"Once upon a time, in a land far away,",
]
print(f"\n{'='*60}")
print("GENERATION TESTS")
print(f"{'='*60}")
for i, prompt in enumerate(test_prompts, 1):
print(f"\n--- Prompt {i}: \"{prompt}\" ---")
output = generate_text(
model, tokenizer, device, prompt,
max_new_tokens=args.max_new_tokens,
)
print(f"Generated:\n{output}\n")
print(f"{'='*60}")
print("Validation complete!")
print("Note: A freshly pretrained model will produce semi-coherent text.")
print("Quality improves with more training data and compute.")
if __name__ == "__main__":
main()
|