# -*- coding: utf-8 -*-
"""text_summarization_finetune.ipynb

Automatically generated by Colab.

Original file is located at
    https://colab.research.google.com/drive/1DC3LFNnBCIfmnKp8DvUFF8Q2FIhwc7RP

# Text Summarization — Fine-tuning T5-small on CNN/DailyMail

**Dataset:** `cnn_dailymail` (v3.0.0) — news articles with human-written highlights (summaries)

**Model:** `t5-small` — lightweight encoder-decoder model, good fit for Colab's free GPU

**Steps:**
1. Install libraries
2. Load & explore dataset
3. Load tokenizer & model
4. Preprocess (tokenize) data
5. Set up training (Seq2SeqTrainer)
6. Train
7. Evaluate with ROUGE
8. Run inference on a custom example
9. Save & (optionally) push the model

> Tip: In Colab go to **Runtime > Change runtime type > T4 GPU** before running.

## 1. Install libraries
"""

!pip install -q transformers datasets evaluate rouge_score accelerate sentencepiece

!pip install -q -U datasets huggingface_hub transformers

"""## 2. Load & explore the dataset"""

from datasets import load_dataset
raw_datasets = load_dataset("abisee/cnn_dailymail", "3.0.0")

train_dataset = raw_datasets["train"].shuffle(seed=42).select(range(3000))
val_dataset   = raw_datasets["validation"].shuffle(seed=42).select(range(300))
test_dataset  = raw_datasets["test"].shuffle(seed=42).select(range(300))

print(train_dataset)
print(train_dataset[0]["article"][:500])
print("\n--- Summary ---")
print(train_dataset[0]["highlights"])

"""## 3. Load tokenizer & model"""

from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

model_checkpoint = "t5-small"

tokenizer = AutoTokenizer.from_pretrained(model_checkpoint)
model = AutoModelForSeq2SeqLM.from_pretrained(model_checkpoint)

prefix = "summarize: "

"""## 4. Preprocess (tokenize) the data"""

max_input_length = 512
max_target_length = 128

def preprocess_function(examples):
    inputs = [prefix + doc for doc in examples["article"]]
    model_inputs = tokenizer(inputs, max_length=max_input_length, truncation=True)

    labels = tokenizer(text_target=examples["highlights"], max_length=max_target_length, truncation=True)

    model_inputs["labels"] = labels["input_ids"]
    return model_inputs

tokenized_train = train_dataset.map(preprocess_function, batched=True, remove_columns=train_dataset.column_names)
tokenized_val   = val_dataset.map(preprocess_function, batched=True, remove_columns=val_dataset.column_names)
tokenized_test  = test_dataset.map(preprocess_function, batched=True, remove_columns=test_dataset.column_names)

"""## 5. Set up training"""

import numpy as np
import evaluate
from transformers import DataCollatorForSeq2Seq, Seq2SeqTrainingArguments, Seq2SeqTrainer

data_collator = DataCollatorForSeq2Seq(tokenizer=tokenizer, model=model)

rouge = evaluate.load("rouge")

def compute_metrics(eval_pred):
    predictions, labels = eval_pred
    decoded_preds = tokenizer.batch_decode(predictions, skip_special_tokens=True)

    labels = np.where(labels != -100, labels, tokenizer.pad_token_id)
    decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True)

    result = rouge.compute(predictions=decoded_preds, references=decoded_labels, use_stemmer=True)
    result = {k: round(v * 100, 2) for k, v in result.items()}

    prediction_lens = [np.count_nonzero(pred != tokenizer.pad_token_id) for pred in predictions]
    result["gen_len"] = round(np.mean(prediction_lens), 2)
    return result

training_args = Seq2SeqTrainingArguments(
    output_dir="./t5-summarization-cnn",
    eval_strategy="epoch",
    save_strategy="epoch",
    learning_rate=3e-4,
    per_device_train_batch_size=8,
    per_device_eval_batch_size=8,
    weight_decay=0.01,
    save_total_limit=2,
    num_train_epochs=3,
    predict_with_generate=True,
    fp16=True,
    logging_steps=50,
    report_to="none",
)

trainer = Seq2SeqTrainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_train,
    eval_dataset=tokenized_val,
    data_collator=data_collator,
    compute_metrics=compute_metrics,
)

"""## 6. Train"""

trainer.train()

"""## 7. Evaluate on the test set"""

test_results = trainer.predict(tokenized_test)
print(test_results.metrics)

"""## 8. Try it on a custom example"""

def summarize(text, max_length=128):
    inputs = tokenizer(prefix + text, return_tensors="pt", truncation=True, max_length=max_input_length).to(model.device)
    summary_ids = model.generate(
        **inputs,
        max_length=max_length,
        num_beams=4,
        length_penalty=2.0,
        early_stopping=True,
    )
    return tokenizer.decode(summary_ids[0], skip_special_tokens=True)

sample_article = test_dataset[0]["article"]
print("Original article:\n", sample_article[:800])
print("\nReference summary:\n", test_dataset[0]["highlights"])
print("\nModel summary:\n", summarize(sample_article))

"""## 9. Save the model (and optionally push to Hugging Face Hub)"""

save_dir = "./t5-summarization-cnn-final"
trainer.save_model(save_dir)
tokenizer.save_pretrained(save_dir)
print("Model saved to", save_dir)

"""## Notes & next steps
- **Scaling up:** increase `train_dataset`/`val_dataset` sizes and `num_train_epochs` for better ROUGE scores (full dataset training takes hours even on T4 — good for a final run, not quick iteration).
- **Bigger model:** swap `t5-small` for `t5-base`, `facebook/bart-base`, or `sshleifer/distilbart-cnn-12-6` if you have more GPU memory/time.
- **Different domain:** swap the dataset for `samsum` (dialogue summarization) or `xsum` (very short summaries) by changing the `load_dataset(...)` call and the column names (`dialogue`/`summary` for samsum).
"""