# Experiment P: Practice: DistilGPT2 fine-tuning (Hugging Face)
# Requires: torch, transformers, datasets, accelerate | downloads distilgpt2 (internet)
# Exact code of Prac.ipynb. Notebook-only lines (!pip / %pip) are kept as comments so the file runs as normal Python.
# ===========================================================================

import numpy as np
import torch
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."
]

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

model_name = "distilgpt2"

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

tokenizer.pad_token = tokenizer.eos_token
model.config.pad_token_id = tokenizer.pad_token_id

def tokenize(example):
    return tokenizer(
        example['text'],
        truncation=True,
        max_length=128
    )

tokenized_dataset = data.map(
    tokenize,
    batched=True,
    remove_columns=['text']
)

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

train_dataset = splited_dataset["train"]
test_dataset = splited_dataset["test"]

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

training_args = TrainingArguments(
    num_train_epochs=5,
    learning_rate=5e-5,
    weight_decay=0.01,
    logging_strategy="epoch",
)

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

trainer.train()

results = trainer.evaluate()

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