from transformers import GPT2Model, GPT2Config, AutoModelForCausalLM, AutoTokenizer
import torch

# Load GPT-2 (smallest version: 117M parameters)
config = GPT2Config()
print("── GPT-2 Small Configuration ──")
print(f"Layers (n_layer):   {config.n_layer}")       # 12
print(f"Heads (n_head):     {config.n_head}")        # 12
print(f"Hidden dim (n_embd):{config.n_embd}")        # 768
print(f"FFN inner dim:      {config.n_embd * 4}")    # 3072
print(f"Vocab size:         {config.vocab_size}")    # 50257 (BPE)
print(f"Max sequence len:   {config.n_positions}")  # 1024
print(f"Activation:         {config.activation_function}") # gelu_new

# Compare GPT-2 sizes
variants = {
    'GPT-2 Small':  'gpt2',
    'GPT-2 Medium': 'gpt2-medium',
    'GPT-2 Large':  'gpt2-large',
    'GPT-2 XL':     'gpt2-xl',
}
sizes = {
    'GPT-2 Small':  {'layers': 12, 'heads': 12, 'dim': 768,  'params': 117e6},
    'GPT-2 Medium': {'layers': 24, 'heads': 16, 'dim': 1024, 'params': 345e6},
    'GPT-2 Large':  {'layers': 36, 'heads': 20, 'dim': 1280, 'params': 762e6},
    'GPT-2 XL':     {'layers': 48, 'heads': 25, 'dim': 1600, 'params': 1542e6},
}
print(f"\n{'Variant':<14} {'Layers':>7} {'Heads':>6} {'Dim':>5} {'Params':>10}")
print("-" * 46)
for name, s in sizes.items():
    print(f"{name:<14} {s['layers']:>7} {s['heads']:>6} {s['dim']:>5} {s['params']/1e6:>8.0f}M")

# Load the actual model and count parameters
model = AutoModelForCausalLM.from_pretrained('gpt2')
total = sum(p.numel() for p in model.parameters())
print(f"\nActual loaded parameters: {total:,}")

# Inspect the first transformer block
print(f"\nFirst block sub-modules:")
for name, module in model.transformer.h[0].named_children():
    print(f"  {name}: {type(module).__name__}")

import torch
import numpy as np

# Illustrate sampling strategies on a toy distribution
# Suppose the model's logits for the next token are:
logits = torch.tensor([2.0, 1.5, 0.5, -0.5, -1.0, -2.0, -3.0, -4.0])
vocab  = ['cat', 'dog', 'bird', 'fish', 'the', 'and', 'xyz', 'qqq']

print("=== Decoding Strategy Comparison ===\n")

# --- Greedy ---
greedy_token = vocab[logits.argmax().item()]
print(f"Greedy: always picks '{greedy_token}'")

# --- Temperature ---
print("\nTemperature effect on probabilities:")
for T in [0.1, 0.5, 1.0, 1.5, 2.0]:
    scaled = logits / T
    probs = torch.softmax(scaled, dim=0)
    top3 = [(vocab[i], f"{probs[i]:.3f}") for i in probs.topk(3).indices.tolist()]
    print(f"  T={T}: top-3 = {top3}")

# --- Top-k ---
print("\nTop-k sampling (k=3):")
probs = torch.softmax(logits, dim=0)
topk_vals, topk_idx = probs.topk(3)
topk_probs_renorm = topk_vals / topk_vals.sum()
for i, (idx, p) in enumerate(zip(topk_idx.tolist(), topk_probs_renorm.tolist())):
    print(f"  {vocab[idx]:8s}: {p:.3f}")

# --- Top-p (nucleus) ---
print("\nTop-p sampling (p=0.9):")
sorted_probs, sorted_idx = probs.sort(descending=True)
cumsum = torch.cumsum(sorted_probs, dim=0)
# Keep tokens until cumulative probability exceeds p=0.9
nucleus_mask = cumsum <= 0.9
# Always include at least the first token
nucleus_mask[0] = True
nucleus_tokens = sorted_idx[nucleus_mask]
nucleus_probs  = sorted_probs[nucleus_mask]
nucleus_renorm = nucleus_probs / nucleus_probs.sum()
for idx, p in zip(nucleus_tokens.tolist(), nucleus_renorm.tolist()):
    print(f"  {vocab[idx]:8s}: {p:.3f}")

from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

tokenizer = AutoTokenizer.from_pretrained('gpt2')
model     = AutoModelForCausalLM.from_pretrained('gpt2')
model.eval()

# GPT-2's tokenizer doesn't have a pad token — set it to eos
tokenizer.pad_token = tokenizer.eos_token

prompt = "In the future, artificial intelligence will"

def generate_text(prompt, strategy_name, **generate_kwargs):
    inputs = tokenizer(prompt, return_tensors='pt')
    with torch.no_grad():
        output_ids = model.generate(
            inputs['input_ids'],
            max_new_tokens=80,
            pad_token_id=tokenizer.eos_token_id,
            **generate_kwargs
        )
    # Decode only the newly generated tokens (skip the prompt)
    new_tokens = output_ids[0, inputs['input_ids'].shape[1]:]
    generated  = tokenizer.decode(new_tokens, skip_special_tokens=True)
    print(f"\n── {strategy_name} ──")
    print(f"Prompt: {prompt}")
    print(f"Output: {generated}")

# Strategy 1: Greedy (deterministic, repetitive)
generate_text(prompt, "Greedy Decoding",
    do_sample=False)

# Strategy 2: Temperature only (creative but can be incoherent)
generate_text(prompt, "Temperature T=0.7",
    do_sample=True, temperature=0.7)

# Strategy 3: Top-k sampling
generate_text(prompt, "Top-k (k=50)",
    do_sample=True, top_k=50)

# Strategy 4: Top-p (nucleus) sampling — RECOMMENDED DEFAULT
generate_text(prompt, "Top-p (p=0.92, T=0.9) — recommended",
    do_sample=True, top_p=0.92, temperature=0.9)

# Strategy 5: Beam search (for factual/precise output)
generate_text(prompt, "Beam Search (num_beams=4)",
    do_sample=False, num_beams=4, no_repeat_ngram_size=2)

# Generate multiple completions for the same prompt
print("\n── 3 Different Completions with top_p=0.9, T=0.8 ──")
inputs = tokenizer(prompt, return_tensors='pt')
with torch.no_grad():
    multi_output = model.generate(
        inputs['input_ids'],
        max_new_tokens=60,
        do_sample=True, top_p=0.9, temperature=0.8,
        num_return_sequences=3,
        pad_token_id=tokenizer.eos_token_id
    )
for i, seq in enumerate(multi_output):
    new_tokens = seq[inputs['input_ids'].shape[1]:]
    text = tokenizer.decode(new_tokens, skip_special_tokens=True)
    print(f"  [{i+1}] {text[:100]}...")

from transformers import TextStreamer

# Stream tokens as they are generated (ChatGPT-style output)
streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
inputs = tokenizer("The key to understanding deep learning is", return_tensors='pt')

print("Streaming output:")
with torch.no_grad():
    model.generate(
        inputs['input_ids'],
        max_new_tokens=100,
        do_sample=True, top_p=0.9, temperature=0.8,
        streamer=streamer
    )

import torch
import math
from transformers import AutoTokenizer, AutoModelForCausalLM

def compute_perplexity(text: str, model, tokenizer, max_length: int = 512) -> float:
    """Compute perplexity of a text sequence under the model."""
    model.eval()
    # Tokenize and truncate
    inputs = tokenizer(text, return_tensors='pt', truncation=True, max_length=max_length)
    input_ids = inputs['input_ids']

    with torch.no_grad():
        outputs = model(input_ids, labels=input_ids)
        # outputs.loss = mean cross-entropy loss per token (NLLoss)
        neg_log_likelihood = outputs.loss.item()

    # PPL = exp(mean NLL)
    ppl = math.exp(neg_log_likelihood)
    return ppl

tokenizer = AutoTokenizer.from_pretrained('gpt2')
model = AutoModelForCausalLM.from_pretrained('gpt2')

# Test texts — intuitively ordered by how "expected" they are
test_texts = {
    "Natural English prose": (
        "The Transformer architecture has revolutionized natural language processing "
        "since its introduction in 2017. Its key innovation is the self-attention mechanism, "
        "which allows every token to attend to every other token in the sequence."
    ),
    "Simple coherent text": (
        "The cat sat on the mat. The dog ran in the park. The bird flew over the trees."
    ),
    "Grammatically correct but unusual": (
        "Colourless green ideas sleep furiously. The prime minister decided to eat the moon."
    ),
    "Random English words": (
        "banana telephone gradient purple library suitcase thunder bicycle coffee quantum"
    ),
    "Random tokens / nonsense": (
        "xzq blfp wrtk mzp yrjq zkk pptm wl gfpb"
    ),
}

print(f"{'Text Type':<40} {'Perplexity':>12}")
print("-" * 55)
for name, text in test_texts.items():
    ppl = compute_perplexity(text, model, tokenizer)
    bar = "█" * min(int(ppl / 50), 30)
    print(f"{name:<40} {ppl:>10.1f}  {bar}")

print("\nReference points:")
print("  Humans writing Wikipedia:  ~7-10 PPL")
print("  GPT-2 on WikiText-103:     ~29 PPL")
print("  GPT-3 on WikiText-103:     ~12 PPL")

from transformers import (AutoModelForCausalLM, AutoTokenizer, TrainingArguments,
                           Trainer, DataCollatorForLanguageModeling)
from datasets import Dataset
import torch

# ── Prepare custom domain dataset ──
# Simulate product descriptions (in practice, load from your data source)
product_descriptions = [
    "Organic Bamboo Cutting Board: This premium cutting board is crafted from "
    "sustainably sourced bamboo. Ultra-hard surface resists knife marks. "
    "Naturally antibacterial. Dimensions: 18x12 inches.",

    "Noise-Canceling Wireless Earbuds: Crystal-clear 24kHz audio with "
    "active noise cancellation. 30-hour battery life with charging case. "
    "IPX5 waterproof rating. Compatible with iOS and Android.",

    "Stainless Steel Water Bottle: Double-wall vacuum insulation keeps drinks "
    "cold for 24 hours, hot for 12 hours. BPA-free. 32oz capacity. "
    "Available in 8 colors. Leak-proof lid.",

    # ... add 100+ real product descriptions for production use
] * 50  # duplicate to have more training data for demo

# Create a Hugging Face Dataset
dataset = Dataset.from_dict({'text': product_descriptions})
train_val = dataset.train_test_split(test_size=0.1, seed=42)

# ── Load model and tokenizer ──
model_name = 'gpt2'
tokenizer = AutoTokenizer.from_pretrained(model_name)
tokenizer.pad_token = tokenizer.eos_token  # GPT-2 has no pad token
model = AutoModelForCausalLM.from_pretrained(model_name)

# ── Tokenize ──
def tokenize(batch):
    return tokenizer(
        batch['text'],
        truncation=True,
        max_length=256,
        padding='max_length'
    )

tokenised_train = train_val['train'].map(tokenize, batched=True, remove_columns=['text'])
tokenised_val   = train_val['test'].map(tokenize, batched=True, remove_columns=['text'])
tokenised_train.set_format('torch')
tokenised_val.set_format('torch')

# ── DataCollator for Causal LM ──
# mlm=False means we do causal LM (next token prediction), not masked LM
data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)

# ── Training arguments ──
training_args = TrainingArguments(
    output_dir='./gpt2_products',
    num_train_epochs=5,
    per_device_train_batch_size=8,
    per_device_eval_batch_size=16,
    learning_rate=5e-5,
    warmup_ratio=0.05,
    weight_decay=0.01,
    evaluation_strategy='epoch',
    save_strategy='epoch',
    load_best_model_at_end=True,
    metric_for_best_model='eval_loss',
    greater_is_better=False,
    fp16=torch.cuda.is_available(),
    report_to='none',
    logging_steps=10,
)

# ── Train ──
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenised_train,
    eval_dataset=tokenised_val,
    data_collator=data_collator,
)
trainer.train()

# ── Generate product descriptions with fine-tuned model ──
model.eval()
prompt = "Ergonomic Office Chair:"
inputs = tokenizer(prompt, return_tensors='pt')
with torch.no_grad():
    output = model.generate(
        inputs['input_ids'],
        max_new_tokens=100,
        do_sample=True, top_p=0.9, temperature=0.8,
        pad_token_id=tokenizer.eos_token_id
    )
print(tokenizer.decode(output[0], skip_special_tokens=True))

from transformers import pipeline

# Use a larger, instruction-tuned model for better prompt following
# (GPT-2 raw is not great at following instructions; these examples
#  work better with GPT-3.5 or instruction-tuned models like Llama-2-chat)

generator = pipeline('text-generation', model='gpt2', max_new_tokens=100,
                     do_sample=True, top_p=0.9, temperature=0.7,
                     pad_token_id=50256)

# Zero-shot: describe the task
zero_shot_prompt = """Classify the sentiment of the following review as Positive, Negative, or Neutral.

Review: The food was absolutely amazing but the service was incredibly slow.
Sentiment:"""

result = generator(zero_shot_prompt, num_return_sequences=1)
print("Zero-shot result:")
print(result[0]['generated_text'][len(zero_shot_prompt):].strip()[:100])

# Few-shot: give 3 examples before asking for the real one
few_shot_prompt = """Classify the sentiment. Examples:

Review: Best pizza I've ever had!
Sentiment: Positive

Review: Waited an hour for cold food.
Sentiment: Negative

Review: Decent place, nothing special.
Sentiment: Neutral

Review: The staff was rude and the room was dirty.
Sentiment:"""

result = generator(few_shot_prompt, num_return_sequences=1)
print("\nFew-shot result:")
print(result[0]['generated_text'][len(few_shot_prompt):].strip()[:100])

# Chain-of-thought for a math word problem
cot_prompt = """Solve this step by step.

Problem: A store sells apples for $0.50 each and oranges for $0.75 each.
If Maria buys 4 apples and 3 oranges, how much does she spend?

Let's think step by step:"""

result = generator(cot_prompt, max_new_tokens=150, do_sample=False)
answer = result[0]['generated_text'][len(cot_prompt):].strip()
print("Chain-of-thought answer:")
print(answer[:300])

import torch
import torch.nn.functional as F

def dpo_loss(policy_chosen_logps, policy_rejected_logps,
             ref_chosen_logps, ref_rejected_logps, beta=0.1):
    """
    DPO loss -- no reward model, no RL loop, just log-probabilities from
    the model being trained (policy) and a frozen copy of it (reference).

    policy_*_logps:  log P(response | prompt) under the model being trained
    ref_*_logps:      log P(response | prompt) under a FROZEN reference model
                       (typically the instruction-tuned checkpoint before DPO)
    "chosen"/"rejected" = the human-preferred and human-rejected responses
    to the SAME prompt.
    """
    # How much MORE the policy favors chosen over rejected, relative to
    # how much the (frozen) reference model already favored it
    policy_logratios = policy_chosen_logps - policy_rejected_logps
    ref_logratios = ref_chosen_logps - ref_rejected_logps

    logits = beta * (policy_logratios - ref_logratios)
    loss = -F.logsigmoid(logits).mean()   # ordinary binary classification-style loss
    return loss

# In each training step:
#   1. Run the prompt+chosen and prompt+rejected through BOTH the policy
#      model (being trained) and the frozen reference model
#   2. Sum each response's per-token log-probabilities
#   3. Compute dpo_loss() above and backpropagate -- through the POLICY
#      model only; the reference model never updates
print("DPO turns preference alignment into a single supervised loss --")
print("no reward model, no PPO rollouts, no RL instability.")

# Demonstrate instruction-tuned model vs raw GPT-2
from transformers import pipeline

# Raw GPT-2 (document completion)
raw_gpt2 = pipeline('text-generation', model='gpt2', max_new_tokens=80,
                    do_sample=True, top_p=0.9, temperature=0.8,
                    pad_token_id=50256)

question = "What are the three main types of machine learning?"

print("Raw GPT-2 response (prompt continuation):")
result = raw_gpt2(question)
print(result[0]['generated_text'])
# GPT-2 will likely continue the question or generate unrelated text

print("\n---")
print("What an instruction-tuned model (e.g. Llama-3-8B-Instruct) would return:")
print("""1. Supervised Learning: training on labeled (input, output) pairs to learn
   a mapping from inputs to outputs.
2. Unsupervised Learning: finding patterns in unlabelled data — clustering,
   dimensionality reduction, generative modeling.
3. Reinforcement Learning: an agent learns to maximize a reward signal through
   trial-and-error interaction with an environment.""")

# Open-source instruction-tuned models you can use locally
print("\nOpen-source instruction-tuned models (via Hugging Face):")
models = [
    ("meta-llama/Llama-3.2-3B-Instruct", "3B params — good for local inference"),
    ("mistralai/Mistral-7B-Instruct-v0.3", "7B params — strong instruction following"),
    ("google/gemma-2-9b-it",              "9B params — Google's latest instruction model"),
    ("microsoft/Phi-3.5-mini-instruct",   "3.8B params — small but very capable"),
]
for name, desc in models:
    print(f"  {name:<45} ({desc})")