import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

# ===========================
# 1. Load GPT-2
# ===========================
model_name = "gpt2"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

# GPT-2 sizes:
# "gpt2"        → 117M params, ~500MB  (small, fits Colab easily)
# "gpt2-medium" → 345M params, ~1.4GB  (medium)
# "gpt2-large"  → 774M params, ~3.1GB  (large, tight on T4)
# "gpt2-xl"     → 1.5B params, ~6.2GB  (XL, needs >16GB VRAM)

# ===========================
# 2. Tokenize & compute loss
# ===========================
text = "The capital of France is Paris"
inputs = tokenizer(text, return_tensors="pt")

# For CLM: labels = input_ids (shifted internally by the model!)
# Model predicts token[i+1] from tokens[0:i]
outputs = model(**inputs, labels=inputs["input_ids"])

print(f"Loss: {outputs.loss.item():.4f}")
# Loss: 3.2145 (lower = better at predicting next tokens)
print(f"Perplexity: {torch.exp(outputs.loss).item():.2f}")
# Perplexity: 24.89 (lower = better, 1.0 = perfect prediction)

# ===========================
# 3. What the model sees internally:
# ===========================
tokens = tokenizer.tokenize(text)
print(f"Tokens: {tokens}")
# ['The', 'Ġcapital', 'Ġof', 'ĠFrance', 'Ġis', 'ĠParis']
# Ġ = space prefix (GPT-2 BPE tokenizer)

# Model internally shifts labels:
# Position 0: sees "The"         → should predict "capital"
# Position 1: sees "The capital"  → should predict "of"
# Position 2: sees "The capital of" → should predict "France"
# Position 3: sees "The capital of France" → should predict "is"
# Position 4: sees "The capital of France is" → should predict "Paris"
# Loss = average cross-entropy over all positions

# ===========================
# 4. Check what GPT-2 predicts at each position
# ===========================
with torch.no_grad():
    outputs = model(**inputs)
    logits = outputs.logits  # (1, seq_len, vocab_size=50257)

for i in range(len(tokens) - 1):
    predicted_id = logits[0, i].argmax().item()
    predicted_token = tokenizer.decode(predicted_id)
    actual_token = tokens[i + 1]
    match = "✅" if predicted_token.strip() == actual_token.replace("Ġ", "") else "❌"
    context = " ".join(tokens[:i+1]).replace("Ġ", "")
    print(f"  '{context}' → predicted: '{predicted_token}' | actual: '{actual_token}' {match}")
# 'The'                        → predicted: ' first'   | actual: 'Ġcapital' ❌
# 'The capital'                → predicted: ' of'      | actual: 'Ġof'      ✅
# 'The capital of'             → predicted: ' the'     | actual: 'ĠFrance'  ❌
# 'The capital of France'      → predicted: ' is'      | actual: 'Ġis'      ✅
# 'The capital of France is'   → predicted: ' Paris'   | actual: 'ĠParis'   ✅
# GPT-2 knows Paris is the capital of France! 🎉

# Cell 1: Verify GPU
!nvidia-smi
import torch
print(f"GPU: {torch.cuda.get_device_name(0)}, VRAM: {torch.cuda.get_device_properties(0).total_mem/1e9:.1f} GB")

# Cell 2: Install
!pip install -q transformers datasets accelerate

# Cell 3: Test GPT-2 inference (< 1 menit download)
from transformers import pipeline
generator = pipeline("text-generation", model="gpt2", device=0)
print(generator("Artificial intelligence will", max_new_tokens=30)[0]["generated_text"])
# ✅ Ready! GPT-2 small = ~500MB, fits easily on T4

# PENTING: GPT-2 tokenizer TIDAK punya pad token!
# Harus set manual sebelum fine-tuning:
from transformers import AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained("gpt2")
tokenizer.pad_token = tokenizer.eos_token  # ← WAJIB untuk GPT-2!
# Tanpa ini: error "Cannot handle padding" saat fine-tuning

from transformers import pipeline

# ===========================
# 1. Basic generation
# ===========================
generator = pipeline("text-generation", model="gpt2", device=0)

result = generator("The future of artificial intelligence is", max_new_tokens=50)
print(result[0]["generated_text"])
# "The future of artificial intelligence is not just about the technology,
#  but about how we use it. The question is whether we can build systems..."

# ===========================
# 2. Multiple completions
# ===========================
results = generator(
    "Once upon a time in Jakarta,",
    max_new_tokens=80,
    num_return_sequences=3,     # generate 3 different completions!
    do_sample=True,              # enable random sampling
    temperature=0.8,             # creativity level
)
for i, r in enumerate(results):
    print(f"\\n--- Completion {i+1} ---")
    print(r["generated_text"])
# Setiap completion berbeda! (karena sampling random)

# ===========================
# 3. Different generation strategies
# ===========================
# Deterministic (greedy — selalu sama)
result_greedy = generator("AI is", max_new_tokens=20, do_sample=False)

# Creative (high temperature sampling)
result_creative = generator("AI is", max_new_tokens=20,
    do_sample=True, temperature=1.2, top_p=0.9)

# Focused (low temperature)
result_focused = generator("AI is", max_new_tokens=20,
    do_sample=True, temperature=0.3)

print(f"Greedy:   {result_greedy[0]['generated_text']}")
print(f"Creative: {result_creative[0]['generated_text']}")
print(f"Focused:  {result_focused[0]['generated_text']}")
# Greedy:   "AI is a very important part of the future of the world."
# Creative: "AI is an existential rollercoaster of digital consciousness..."
# Focused:  "AI is a field of computer science that focuses on..."

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

model_name = "gpt2"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name).to("cuda")

# ===========================
# 1. Basic generate()
# ===========================
prompt = "Indonesia is a beautiful country with"
inputs = tokenizer(prompt, return_tensors="pt").to("cuda")

with torch.no_grad():
    output_ids = model.generate(
        **inputs,
        max_new_tokens=50,       # generate 50 NEW tokens
        do_sample=True,           # enable sampling
        temperature=0.7,          # creativity
        top_k=50,                 # consider top 50 tokens
        top_p=0.9,                # nucleus sampling
        repetition_penalty=1.2,   # penalize repetition
        pad_token_id=tokenizer.eos_token_id,
    )

# Decode — skip the prompt tokens
generated_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
print(generated_text)

# ===========================
# 2. Stream generation (token by token) — like ChatGPT!
# ===========================
from transformers import TextIteratorStreamer
from threading import Thread

streamer = TextIteratorStreamer(tokenizer, skip_special_tokens=True)

inputs = tokenizer("The meaning of life is", return_tensors="pt").to("cuda")
gen_kwargs = {**inputs, "max_new_tokens": 100, "streamer": streamer,
              "do_sample": True, "temperature": 0.7}

# Run in separate thread (non-blocking)
thread = Thread(target=model.generate, kwargs=gen_kwargs)
thread.start()

# Print tokens as they arrive!
for text in streamer:
    print(text, end="", flush=True)
# "The meaning of life is" ← appears token by token, like ChatGPT!

# ===========================
# 3. Access generation logits (for custom post-processing)
# ===========================
outputs = model.generate(
    **inputs,
    max_new_tokens=5,
    output_scores=True,          # return logits for each step!
    return_dict_in_generate=True,
)

# outputs.scores = tuple of (vocab_size,) tensors, one per generated token
for i, scores in enumerate(outputs.scores):
    probs = torch.softmax(scores[0], dim=-1)
    top5 = torch.topk(probs, 5)
    print(f"\\nStep {i+1} top-5 candidates:")
    for prob, idx in zip(top5.values, top5.indices):
        token = tokenizer.decode(idx)
        print(f"  '{token}': {prob:.1%}")

# ═══════════════════════════════════════════════════
# SETIAP PARAMETER GENERATION — EXPLAINED
# ═══════════════════════════════════════════════════

output = model.generate(
    **inputs,

    # ── LENGTH CONTROL ──
    max_new_tokens=100,     # generate MAX 100 new tokens
    min_new_tokens=10,      # generate MIN 10 tokens (prevent empty output)
    # max_length=150,        # alternative: total length (prompt + generated)

    # ── SAMPLING vs GREEDY ──
    do_sample=True,         # True=random sampling, False=greedy (deterministic)
    # Greedy: always pick highest probability token → boring, repetitive
    # Sampling: randomly pick from probability distribution → creative, varied

    # ── TEMPERATURE ── (hanya berlaku jika do_sample=True)
    temperature=0.7,        # Controls randomness of sampling
    # temperature=0.1 → almost greedy (very focused, repetitive)
    # temperature=0.7 → balanced (creative but coherent) ← RECOMMENDED
    # temperature=1.0 → standard (model's natural distribution)
    # temperature=1.5 → very random (wild, often incoherent)
    # temperature=2.0 → chaos (mostly nonsense)
    #
    # HOW IT WORKS:
    # logits_adjusted = logits / temperature
    # probs = softmax(logits_adjusted)
    # Low temp → sharper distribution → top token dominates
    # High temp → flatter distribution → more variety

    # ── TOP-K SAMPLING ── (hanya jika do_sample=True)
    top_k=50,               # Only consider top K highest-probability tokens
    # top_k=1   → greedy (only top 1 token)
    # top_k=10  → conservative (limited vocabulary)
    # top_k=50  → balanced ← DEFAULT
    # top_k=0   → disabled (consider ALL tokens)
    #
    # PROBLEM: top_k=50 treats all distributions equally.
    # If model is very confident: top 5 tokens have 95% probability
    # → tokens 6-50 are almost random noise!
    # SOLUTION: use top_p instead (or together)

    # ── TOP-P (NUCLEUS) SAMPLING ── (hanya jika do_sample=True)
    top_p=0.9,              # Keep tokens until cumulative probability reaches P
    # top_p=0.9 → keep tokens that sum to 90% probability
    # If model confident: might keep only 3 tokens (they sum to 90%)
    # If model uncertain: might keep 50 tokens (all needed for 90%)
    # ADAPTS to model's confidence! Better than fixed top_k.
    #
    # top_p=1.0 → disabled (keep all tokens)
    # top_p=0.95 → slightly conservative
    # top_p=0.9  → balanced ← RECOMMENDED
    # top_p=0.5  → very focused

    # ── REPETITION CONTROL ──
    repetition_penalty=1.2, # Penalize tokens that already appeared
    # 1.0 = no penalty (can repeat freely)
    # 1.1 = mild (some repetition OK)
    # 1.2 = moderate ← RECOMMENDED for most use cases
    # 1.5 = strong (almost never repeats)
    # 2.0+ = too strong (forced to use rare words)

    no_repeat_ngram_size=3, # Never repeat the same 3-word phrase
    # 0 = disabled, 2 = no repeated bigrams, 3 = no repeated trigrams

    # ── BEAM SEARCH ── (alternative to sampling)
    # num_beams=5,           # explore 5 paths simultaneously
    # early_stopping=True,   # stop when all beams finish
    # length_penalty=1.0,    # >1 = prefer longer, <1 = prefer shorter
    # Beam search: more coherent but LESS creative than sampling
    # Good for: translation, summarization
    # Bad for: creative writing, chat (too boring)
    # NOTE: beam search is INCOMPATIBLE with do_sample=True!

    # ── STOP CONDITIONS ──
    pad_token_id=tokenizer.eos_token_id,
    eos_token_id=tokenizer.eos_token_id,
    # stop_strings=["Human:", "\n\n"],  # stop at these strings
)

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

# ═══════════════════════════════════════
# STEP 1: LOAD MODEL & TOKENIZER
# ═══════════════════════════════════════
model_name = "gpt2"  # 117M params, fits Colab T4
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

# CRITICAL: GPT-2 has NO pad token! Must set it!
tokenizer.pad_token = tokenizer.eos_token
model.config.pad_token_id = tokenizer.eos_token_id

print(f"Model params: {model.num_parameters():,}")  # 124,439,808

# ═══════════════════════════════════════
# STEP 2: LOAD & PREPARE DATASET
# ═══════════════════════════════════════
# Option A: From Hugging Face Hub
dataset = load_dataset("wikitext", "wikitext-2-raw-v1")
# Option B: From your own text file
# dataset = load_dataset("text", data_files="my_corpus.txt")
# Option C: From CSV with "text" column
# dataset = load_dataset("csv", data_files="poems.csv")

print(dataset)
print(f"Sample: {dataset['train'][0]['text'][:100]}...")

# ═══════════════════════════════════════
# STEP 3: TOKENIZE
# ═══════════════════════════════════════
def tokenize_function(examples):
    return tokenizer(
        examples["text"],
        truncation=True,
        max_length=512,          # GPT-2 max context = 1024
        return_overflowing_tokens=True,  # split long texts into chunks!
        return_length=True,
    )

tokenized = dataset.map(tokenize_function, batched=True,
                         remove_columns=dataset["train"].column_names)

# Filter out very short sequences
tokenized = tokenized.filter(lambda x: len(x["input_ids"]) > 10)
print(f"Training examples: {len(tokenized['train'])}")

# ═══════════════════════════════════════
# STEP 4: DATA COLLATOR (special for CLM!)
# ═══════════════════════════════════════
data_collator = DataCollatorForLanguageModeling(
    tokenizer=tokenizer,
    mlm=False,                  # False = Causal LM (GPT)
    # mlm=True → Masked LM (BERT) — NOT for GPT!
)
# DataCollatorForLanguageModeling automatically:
# 1. Pads sequences in each batch
# 2. Creates labels = input_ids (shifted by 1 internally)
# 3. Sets label=-100 for padding tokens (ignored in loss)

# ═══════════════════════════════════════
# STEP 5: TRAINING
# ═══════════════════════════════════════
args = TrainingArguments(
    output_dir="./gpt2-finetuned",
    num_train_epochs=3,
    per_device_train_batch_size=4,       # smaller batch for generation models
    gradient_accumulation_steps=8,       # effective batch = 4 × 8 = 32
    learning_rate=5e-5,                  # slightly higher than BERT (5e-5 vs 2e-5)
    weight_decay=0.01,
    warmup_ratio=0.1,
    fp16=True,
    logging_steps=50,
    save_strategy="epoch",
    save_total_limit=2,
    prediction_loss_only=True,           # don't compute metrics (CLM only needs loss)
    report_to="none",
)

trainer = Trainer(
    model=model,
    args=args,
    train_dataset=tokenized["train"],
    eval_dataset=tokenized["validation"],
    data_collator=data_collator,
    tokenizer=tokenizer,
)

print("🏋️ Training GPT-2...")
trainer.train()

# ═══════════════════════════════════════
# STEP 6: EVALUATE (Perplexity)
# ═══════════════════════════════════════
import math
eval_results = trainer.evaluate()
perplexity = math.exp(eval_results["eval_loss"])
print(f"Perplexity: {perplexity:.2f}")
# Lower = better. GPT-2 base on WikiText: ~30. Fine-tuned: ~20-25.
# Human text: ~20-50 depending on domain.

# ═══════════════════════════════════════
# STEP 7: SAVE & GENERATE!
# ═══════════════════════════════════════
trainer.save_model("./gpt2-finetuned-final")
tokenizer.save_pretrained("./gpt2-finetuned-final")

# Test generation!
from transformers import pipeline
gen = pipeline("text-generation", model="./gpt2-finetuned-final", device=0)

prompts = ["In the field of machine learning,", "The history of Indonesia"]
for p in prompts:
    result = gen(p, max_new_tokens=80, do_sample=True, temperature=0.7,
                 top_p=0.9, repetition_penalty=1.2)
    print(f"\\nPrompt: {p}")
    print(f"Output: {result[0]['generated_text']}")

print("\\n🏆 GPT-2 fine-tuning complete!")

# ===========================
# 1. Format data instruction tuning
# ===========================
# Setiap training example = instruksi + respons dalam satu string

# Format Alpaca-style (paling populer):
training_examples = [
    """### Instruction:
Summarize the following text in one sentence.

### Input:
Hugging Face is a company that provides tools and platforms for machine learning. They are best known for their Transformers library, which provides thousands of pre-trained models for natural language processing, computer vision, and audio tasks.

### Response:
Hugging Face is an ML company known for their Transformers library offering thousands of pre-trained models for NLP, vision, and audio.""",

    """### Instruction:
Translate the following English text to Indonesian.

### Input:
I love learning about artificial intelligence.

### Response:
Saya suka belajar tentang kecerdasan buatan.""",

    """### Instruction:
What is the capital of Japan?

### Response:
The capital of Japan is Tokyo.""",
]

# Format ChatML (used by many chat models):
chat_examples = [
    """<|im_start|>system
You are a helpful assistant.<|im_end|>
<|im_start|>user
What is Python?<|im_end|>
<|im_start|>assistant
Python is a high-level programming language known for its readability and versatility.<|im_end|>""",
]

# ===========================
# 2. Prepare dataset
# ===========================
from datasets import Dataset

# Dari list of strings:
dataset = Dataset.from_dict({"text": training_examples})

# Atau dari JSONL file:
# {"instruction": "...", "input": "...", "output": "..."}
# dataset = load_dataset("json", data_files="instructions.jsonl")

# Format setiap row menjadi satu string
def format_instruction(example):
    if example.get("input"):
        text = f"""### Instruction:\n{example['instruction']}\n\n### Input:\n{example['input']}\n\n### Response:\n{example['output']}{tokenizer.eos_token}"""
    else:
        text = f"""### Instruction:\n{example['instruction']}\n\n### Response:\n{example['output']}{tokenizer.eos_token}"""
    return {"text": text}

# dataset = dataset.map(format_instruction)
# Then tokenize & train exactly like Section 8!

# ===========================
# 3. Inference with instruction format
# ===========================
def ask(instruction, input_text=""):
    if input_text:
        prompt = f"### Instruction:\n{instruction}\n\n### Input:\n{input_text}\n\n### Response:\n"
    else:
        prompt = f"### Instruction:\n{instruction}\n\n### Response:\n"

    result = gen(prompt, max_new_tokens=200, do_sample=True,
                 temperature=0.7, top_p=0.9,
                 repetition_penalty=1.2)
    response = result[0]["generated_text"][len(prompt):]
    # Stop at next "###" (prevent generating another instruction)
    if "###" in response:
        response = response[:response.index("###")]
    return response.strip()

print(ask("What is the largest planet in our solar system?"))
# "Jupiter is the largest planet in our solar system."

from transformers import pipeline

# Load fine-tuned model (atau GPT-2 biasa untuk demo)
gen = pipeline("text-generation",
    model="./gpt2-finetuned-final",  # atau "gpt2" untuk demo
    device=0)

def chat(user_input, history=""):
    """Generate chatbot response."""
    prompt = history + f"### Human: {user_input}\n### Assistant:"

    result = gen(
        prompt,
        max_new_tokens=150,
        do_sample=True,
        temperature=0.7,
        top_p=0.9,
        repetition_penalty=1.3,
        pad_token_id=gen.tokenizer.eos_token_id,
    )

    full_text = result[0]["generated_text"]
    response = full_text[len(prompt):].strip()

    # Stop at next "### Human:" or newline
    for stop in ["### Human:", "###", "\n\n"]:
        if stop in response:
            response = response[:response.index(stop)]

    # Update history for context
    new_history = prompt + " " + response + "\n"
    return response.strip(), new_history

# ═══════════════════════════════════════
# Interactive loop
# ═══════════════════════════════════════
print("🤖 GPT-2 Chatbot (type 'quit' to exit)")
print("=" * 50)

history = ""
while True:
    user_input = input("\\n👤 You: ")
    if user_input.lower() in ["quit", "exit", "q"]:
        print("👋 Bye!")
        break

    response, history = chat(user_input, history)
    print(f"🤖 Bot: {response}")

# ═══════════════════════════════════════
# Sample conversation:
# 👤 You: What is machine learning?
# 🤖 Bot: Machine learning is a subset of AI that enables systems
#         to learn from data without being explicitly programmed.
# 👤 You: Give me an example.
# 🤖 Bot: A spam filter that learns to identify spam emails by
#         analyzing thousands of examples is a common ML application.
# ═══════════════════════════════════════