GPT & Decoder-Only Language Models

from transformers import GPT2LMHeadModel, GPT2Tokenizer, pipeline
import torch

tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
model     = GPT2LMHeadModel.from_pretrained('gpt2')
model.eval()

# ── Manual generation: greedy decoding ──
def greedy_generate(prompt: str, max_new: int = 30) -> str:
    inputs = tokenizer(prompt, return_tensors='pt')
    with torch.no_grad():
        output_ids = model.generate(
            **inputs,
            max_new_tokens=max_new,
            do_sample=False,          # Greedy: always pick highest probability token
            pad_token_id=tokenizer.eos_token_id
        )
    return tokenizer.decode(output_ids[0], skip_special_tokens=True)

# ── Sampling strategies ──
def sample_generate(prompt: str, strategy: str = 'top_p') -> str:
    inputs = tokenizer(prompt, return_tensors='pt')
    with torch.no_grad():
        if strategy == 'top_k':
            output_ids = model.generate(
                **inputs, max_new_tokens=50, do_sample=True,
                top_k=50,             # Sample from top-50 tokens
                temperature=0.8,
                pad_token_id=tokenizer.eos_token_id)
        elif strategy == 'top_p':
            output_ids = model.generate(
                **inputs, max_new_tokens=50, do_sample=True,
                top_p=0.92,           # Nucleus sampling: smallest set with 92% probability mass
                temperature=0.8,
                pad_token_id=tokenizer.eos_token_id)
        elif strategy == 'beam':
            output_ids = model.generate(
                **inputs, max_new_tokens=50, num_beams=5,  # Beam search: top-5 sequences
                early_stopping=True,
                pad_token_id=tokenizer.eos_token_id)
    return tokenizer.decode(output_ids[0], skip_special_tokens=True)

prompt = "The history of artificial intelligence began"
print("Greedy:", greedy_generate(prompt, 20))
print("Top-K: ", sample_generate(prompt, 'top_k'))
print("Top-p: ", sample_generate(prompt, 'top_p'))
print("Beam:  ", sample_generate(prompt, 'beam'))

# ── Logits and next-token distribution ──
inputs = tokenizer("The capital of France is", return_tensors='pt')
with torch.no_grad():
    logits = model(**inputs).logits    # (1, seq_len, vocab_size)

# Probability distribution over next token
next_token_logits = logits[0, -1, :]  # Last token's prediction
probs = torch.softmax(next_token_logits, dim=-1)
top5  = probs.topk(5)
print("
Top-5 next token probabilities:")
for prob, idx in zip(top5.values, top5.indices):
    token = tokenizer.decode([idx])
    print(f"  {prob:.3f}: '{token}'")
# ' Paris' should be highest probability