File size: 5,025 Bytes
d993048 | 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 | """
Step 5: Generate Armenian text from a pretrained ArmGPT checkpoint.
Quick sampling tool for the Stage 1 (pretrained) model produced by
4_train.py. For interactive chat with the fine-tuned model, use 8_chat.py.
Usage:
python 5_generate.py
python 5_generate.py --prompt "Հայաստանի"
python 5_generate.py --temperature 0.5 --length 500
python 5_generate.py --checkpoint checkpoints/step_5000.pt
"""
import argparse
import os
import sys
# Force UTF-8 stdout/stderr on Windows so Armenian text can be printed
if sys.platform == "win32":
sys.stdout.reconfigure(encoding="utf-8", errors="replace")
sys.stderr.reconfigure(encoding="utf-8", errors="replace")
import torch
from core.model import GPT
from core import detect_tokenizer_type, load_tokenizer as _load_tokenizer
def load_tokenizer(data_dir, tokenizer_type=None):
"""Load the tokenizer used during training.
If tokenizer_type is None, auto-detects from data_dir.
"""
if tokenizer_type is None:
tokenizer_type = detect_tokenizer_type(data_dir)
return _load_tokenizer(data_dir, tokenizer_type)
def main():
parser = argparse.ArgumentParser(description="Generate Armenian text with ArmGPT")
parser.add_argument("--checkpoint", type=str, default="checkpoints/final.pt",
help="Path to model checkpoint")
parser.add_argument("--prompt", type=str, default="Հայաստան",
help="Starting text (Armenian)")
parser.add_argument("--length", type=int, default=200,
help="Number of tokens/characters to generate")
parser.add_argument("--temperature", type=float, default=0.6,
help="Randomness: 0.1=safe, 0.8=balanced, 1.5=creative")
parser.add_argument("--top_k", type=int, default=20,
help="Only sample from top k tokens (0=all)")
parser.add_argument("--repetition_penalty", type=float, default=1.15,
help="Penalty for repeating tokens already in context. "
"1.0=off, 1.1-1.3 typical, helps escape repetition loops")
parser.add_argument("--num_samples", type=int, default=1,
help="How many samples to generate")
parser.add_argument("--data_dir", type=str, default="data",
help="Directory containing the tokenizer file")
parser.add_argument("--tokenizer", type=str, default=None,
choices=["char", "bpe"],
help="Tokenizer type. If omitted, auto-detects from data_dir.")
args = parser.parse_args()
# Load checkpoint
if not os.path.exists(args.checkpoint):
print(f"Error: checkpoint not found at {args.checkpoint}")
print("Train a model first with: python 4_train.py")
return
print(f"Loading model from {args.checkpoint}...")
checkpoint = torch.load(args.checkpoint, map_location="cpu", weights_only=False)
cfg = checkpoint["config"]
# Determine device
if torch.cuda.is_available():
device = "cuda"
elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
device = "mps"
else:
device = "cpu"
# Load tokenizer
tokenizer = load_tokenizer(args.data_dir, args.tokenizer)
# Create model and load weights
model = GPT(
vocab_size=tokenizer.vocab_size,
n_layer=cfg["n_layer"],
n_head=cfg["n_head"],
n_embd=cfg["n_embd"],
block_size=cfg["block_size"],
dropout=0.0, # no dropout during generation
).to(device)
state_dict = checkpoint["model"]
# Strip torch.compile() prefix if present
if any(k.startswith("_orig_mod.") for k in state_dict):
state_dict = {k.removeprefix("_orig_mod."): v for k, v in state_dict.items()}
model.load_state_dict(state_dict)
model.eval()
# Encode the prompt
prompt_ids = tokenizer.encode(args.prompt)
if len(prompt_ids) == 0:
print("Warning: prompt produced no tokens. Using default seed.")
prompt_ids = [0]
print(f"\nDevice: {device}")
print(f"Prompt: {args.prompt}")
print(f"Temperature: {args.temperature}")
print(f"Generating {args.length} tokens...\n")
# Generate
top_k = args.top_k if args.top_k > 0 else None
for i in range(args.num_samples):
context = torch.tensor([prompt_ids], dtype=torch.long, device=device)
output = model.generate(context, max_new_tokens=args.length,
temperature=args.temperature, top_k=top_k,
repetition_penalty=args.repetition_penalty)
text = tokenizer.decode(output[0].tolist())
if args.num_samples > 1:
print(f"--- Sample {i+1} ---")
print(text)
if args.num_samples > 1:
print()
if __name__ == "__main__":
main()
|