# Experiment 3: Fine-tuning DistilGPT2 on AI text + sampling strategies (Hugging Face)
# Requires: torch, transformers, datasets, accelerate, matplotlib | downloads the pretrained distilgpt2 (internet)
# Exact code of Exp3.ipynb. Notebook-only lines (!pip / %pip) are kept as comments so the file runs as normal Python.
# ===========================================================================

# !pip install -q transformers datasets matplotlib

# %pip install -U "accelerate>=1.1.0"

import torch
import numpy as np
import matplotlib.pyplot as plt

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

texts = [
    "Artificial intelligence enables machines to perform tasks that normally require human intelligence.",
    "Machine learning is a branch of artificial intelligence that learns patterns from data.",
    "Deep learning uses neural networks with multiple layers to learn complex representations.",
    "Natural language processing allows computers to understand and generate human language.",
    "Generative AI creates new content such as text, images, music, and computer code.",
    "Large language models are trained on large amounts of text to predict the next token.",
    "Transformers use attention mechanisms to learn relationships between words and tokens.",
    "GPT models generate text by predicting one token at a time based on previous tokens.",
    "Fine-tuning adapts a pretrained language model to a specific task or domain.",
    "Tokenization converts text into tokens and maps them to numerical token IDs.",
    "Embedding layers convert token IDs into dense numerical vectors.",
    "Training a language model involves minimizing the difference between predicted and actual next tokens.",
    "Artificial intelligence can support healthcare, education, robotics, and scientific research.",
    "Responsible AI requires attention to fairness, privacy, safety, and reliability.",
    "Data quality plays an important role in the performance of machine learning models.",
    "A neural network learns useful patterns by updating its parameters during training.",
    "Transfer learning reuses knowledge learned by a pretrained model.",
    "Text generation quality depends on training data, model parameters, and decoding strategies."
]

dataset = Dataset.from_dict({"text": texts})

print("Number of text samples:", len(dataset))

model_name = "distilgpt2"

tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

# GPT-2 has no separate padding token by default
tokenizer.pad_token = tokenizer.eos_token
model.config.pad_token_id = tokenizer.pad_token_id

print("Model loaded successfully!")

def tokenize_function(examples):
    return tokenizer(
        examples["text"],
        truncation=True,
        max_length=128
    )

tokenized_dataset = dataset.map(
    tokenize_function,
    batched=True,
    remove_columns=["text"]
)

# Display tokenization example
sample_text = texts[0]
tokens = tokenizer.tokenize(sample_text)
token_ids = tokenizer.convert_tokens_to_ids(tokens)

print("Original text:", sample_text)
print("Tokens:", tokens)
print("Token IDs:", token_ids)

# Embedding vectors are produced inside the model
input_ids = torch.tensor([token_ids])
with torch.no_grad():
    embeddings = model.get_input_embeddings()(input_ids)

print("Embedding shape:", embeddings.shape)

split_dataset = tokenized_dataset.train_test_split(
    test_size=0.2,
    seed=42
)

train_dataset = split_dataset["train"]
eval_dataset = split_dataset["test"]

data_collator = DataCollatorForLanguageModeling(
    tokenizer=tokenizer,
    mlm=False
)

training_args = TrainingArguments(
    output_dir="./ai_gpt2",
    num_train_epochs=5,
    per_device_train_batch_size=2,
    per_device_eval_batch_size=2,
    learning_rate=5e-5,
    weight_decay=0.01,
    logging_strategy="epoch",
    eval_strategy="epoch",
    save_strategy="no",
    report_to="none",
    seed=42
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
    data_collator=data_collator
)

trainer.train()

results = trainer.evaluate()

eval_loss = results["eval_loss"]
perplexity = np.exp(eval_loss)

print("Evaluation Loss:", round(eval_loss, 4))
print("Perplexity:", round(perplexity, 4))

history = trainer.state.log_history

train_losses = [
    item["loss"] for item in history
    if "loss" in item and "eval_loss" not in item
]

eval_losses = [
    item["eval_loss"] for item in history
    if "eval_loss" in item
]

epochs = range(1, len(eval_losses) + 1)

eval_perplexities = [np.exp(loss) for loss in eval_losses]

plt.figure(figsize=(7, 4))
plt.plot(epochs, eval_losses, marker="o", label="Evaluation Loss")
plt.plot(epochs, eval_perplexities, marker="s", label="Evaluation Perplexity")
plt.xlabel("Epoch")
plt.ylabel("Value")
plt.title("Loss and Perplexity vs Epoch")
plt.legend()
plt.grid(True)
plt.show()

def generate_text(prompt, temperature=0.8, top_k=50, top_p=0.95):
    model.eval()

    inputs = tokenizer(prompt, return_tensors="pt")

    inputs = {k: v.to(model.device) for k, v in inputs.items()}

    with torch.no_grad():
        output = model.generate(
            **inputs,
            max_new_tokens=60,
            do_sample=True,
            temperature=temperature,
            top_k=top_k,
            top_p=top_p,
            pad_token_id=tokenizer.eos_token_id
        )

    return tokenizer.decode(output[0], skip_special_tokens=True)


prompt = "Artificial intelligence"

print("Temperature 0.3:")
print(generate_text(prompt, temperature=0.3))

print("\nTemperature 0.8:")
print(generate_text(prompt, temperature=0.8))

print("\nTemperature 1.2:")
print(generate_text(prompt, temperature=1.2))

print("\nTop-k sampling:")
print(generate_text(prompt, temperature=0.8, top_k=20, top_p=1.0))

print("\nTop-p sampling:")
print(generate_text(prompt, temperature=0.8, top_k=0, top_p=0.9))
