#!/usr/bin/env python
# coding: utf-8

# 
# # Exercise 3: Fine-Tuning Pretrained Transformer on Text Classification Task
# 
# In this lab, you will apply the concepts learned by fine-tuning a pre-trained Transformer model on a text classification task using the Hugging Face `transformers` library.
# 
# ## Objectives:
# - Learn to load a pre-trained model from Hugging Face.
# - Fine-tune the model on a text classification dataset.
# - Evaluate and save the fine-tuned model.
# 

# ## Installing Necessary Libraries
# 
# In this lab, we will use the following Python libraries:
# 
# - **Transformers**: Hugging Face's `transformers` library provides pre-trained models for various Natural Language Processing (NLP) tasks. We will use this to load and fine-tune a pre-trained transformer model.
#   
# - **Datasets**: Hugging Face's `datasets` library allows easy access to a wide range of datasets and provides tools for efficient data processing.
#   
# - **Scikit-learn**: A widely-used library for machine learning, which includes tools for metrics, evaluation, and preprocessing. In this lab, we'll use it to calculate evaluation metrics such as accuracy, precision, recall, and F1 score.
# 

# In[1]:


# Install necessary libraries


# # Setup and Import Libraries
# 
# In this section, we import all the necessary libraries for our IMDb sentiment analysis lab.
# 
# - **os, random, numpy, torch**: For system operations, reproducibility, and PyTorch support.
# - **datasets**: To load the IMDb dataset from Hugging Face.
# - **transformers**: Includes tokenizer, model, training utilities, data collator, and callbacks for early stopping.
# 

# In[2]:


# Set random seeds for reproducibility
import os
import random
import numpy as np
import torch

SEED = 42
random.seed(SEED)
np.random.seed(SEED)
torch.manual_seed(SEED)
if torch.cuda.is_available():
    torch.cuda.manual_seed_all(SEED)

# Import dataset and transformer utilities
from datasets import load_dataset
from transformers import (
    AutoTokenizer,
    AutoModelForSequenceClassification,
    DataCollatorWithPadding,
    TrainingArguments,
    Trainer,
    EarlyStoppingCallback
)


USE_MPS = torch.backends.mps.is_available()
device = torch.device("mps" if USE_MPS else "cpu")
print("Using device:", device)


# # Load IMDb Dataset
# 
# We load the IMDb movie review dataset from Hugging Face:
# 
# - The dataset contains **25,000 training examples** and **25,000 test examples**.
# - Each example has a `text` field (the movie review) and a `label` field (`0` = negative, `1` = positive).
# - We print the dataset to verify its structure.
# 

# In[3]:


# Load the ag_news dataset
raw = load_dataset("SetFit/ag_news")
print(raw)


# # Tokenization and Dataset Preparation
# 
# In this section, we:
# 
# 1. Load the pre-trained BERT tokenizer (`bert-base-uncased`).
# 2. Tokenize the IMDb text data with truncation to a maximum sequence length of 128 tokens.
# 3. Remove the original text column to ensure Trainer only works with tensor inputs.
# 4. Set the dataset format to PyTorch tensors for compatibility with the Trainer.
# 5. Split the training set to create a validation set (5,000 examples) for monitoring validation loss during fine-tuning.
# 

# In[4]:


# Load BERT tokenizer
MODEL_NAME = "bert-base-uncased"
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)

# Tokenization function
def tokenize_fn(examples):
    return tokenizer(examples["text"], truncation=True, max_length=128)

cols_to_remove = [c for c in raw["train"].column_names if c not in ("label",)]

# Apply tokenization to the dataset
tokenized = raw.map(tokenize_fn, batched=True, remove_columns=cols_to_remove)


# Remove original text column to avoid issues during batching
if "text" in tokenized["train"].column_names:
    tokenized = tokenized.remove_columns(["text"])

# Set dataset format to PyTorch tensors
tokenized.set_format("torch")

# Shuffle and split the training dataset to create a validation set
train_dataset = tokenized["train"].shuffle(seed=SEED)
val_split = train_dataset.train_test_split(test_size=5000, seed=SEED)
train_dataset = val_split["train"]
eval_dataset = val_split["test"]


# After tokenization, column removal, and setting the dataset format to PyTorch tensors, the training dataset looks like this:
# 
# - **Number of examples:** 20,000 (after splitting 5,000 examples for validation)
# - **Features:**
#   - `label` – the target sentiment (0 = negative, 1 = positive)
#   - `input_ids` – numerical token IDs representing the words/subwords of each sentence
#   - `token_type_ids` – segment IDs used by BERT to distinguish sentences (for single-sentence tasks, usually all zeros)
#   - `attention_mask` – indicates which tokens are real (1) and which are padding (0), allowing the model to ignore padding during self-attention
# 

# In[5]:


print(train_dataset)


# # Model Initialization and Data Collator
# 
# In this section, we:
# 
# 1. Load a pre-trained BERT model (`bert-base-uncased`) for sequence classification with 2 output labels (positive and negative sentiment).
# 2. Set up a `DataCollatorWithPadding` to dynamically pad batches during training and evaluation.
# 

# In[6]:


# Load pre-trained BERT model for sequence classification
model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME, num_labels=4)

# Create a data collator that dynamically pads input sequences in each batch
data_collator = DataCollatorWithPadding(tokenizer=tokenizer)


# # Metrics Calculation Using Scikit-Learn
# 
# In this section, we define a function to compute evaluation metrics for the model using **scikit-learn**:
# 
# - **Accuracy**: Fraction of correctly predicted examples.
# - **F1 Score**: Harmonic mean of precision and recall for binary classification.
# 
# Using scikit-learn simplifies metric computation and ensures correctness.
# 

# In[7]:


from sklearn.metrics import accuracy_score, f1_score
import numpy as np

# Define a metrics computation function using scikit-learn
def compute_metrics(eval_pred):
    logits, labels = eval_pred
    # Convert logits to predicted class indices
    preds = np.argmax(logits, axis=-1)

    # Compute accuracy and F1 score using scikit-learn
    acc = accuracy_score(labels, preds)
    f1 = f1_score(labels, preds, average='macro')

    return {"accuracy": acc, "f1_macro": f1}


# # Training Setup with Trainer and Early Stopping
# 
# In this section, we configure the training process:
# 
# 1. **TrainingArguments**:
#    - Sets output directory, batch sizes, number of epochs, learning rate, weight decay, warmup steps, and other training hyperparameters.
#    - Enables evaluation and checkpoint saving at the end of each epoch.
#    - Loads the best model at the end based on evaluation loss.
#    - Uses mixed-precision (fp16) if a GPU is available.
# 
# 2. **Trainer**:
#    - Combines the model, datasets, tokenizer, data collator, and metrics function.
#    - Includes `EarlyStoppingCallback` to stop training if validation loss does not improve for 2 evaluation steps (epochs in this case).
# 
# This setup makes fine-tuning BERT efficient and helps prevent overfitting.
# 

# In[8]:


from transformers import TrainingArguments, Trainer, EarlyStoppingCallback

# Define training arguments
training_args = TrainingArguments(
    output_dir="./results",
    eval_strategy="epoch",
    save_strategy="epoch",
    logging_strategy="epoch",
    #report_to=[],  # <- disable all integrations (no wandb, no tensorboard)
    per_device_train_batch_size=8,
    per_device_eval_batch_size=8,
    num_train_epochs=3,
    learning_rate=2e-5,
    weight_decay=0.1,
    warmup_steps=100,
    load_best_model_at_end=True,
    metric_for_best_model="eval_loss",
    greater_is_better=False,
    save_total_limit=3,
    fp16=torch.cuda.is_available(),
    dataloader_drop_last=False,
    gradient_accumulation_steps=1,
    seed=SEED,
)

# Create Trainer instance with early stopping
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
    tokenizer=tokenizer,
    data_collator=data_collator,
    compute_metrics=compute_metrics,
    callbacks=[EarlyStoppingCallback(early_stopping_patience=2)],
)


# # Start Training
# 
# Now we begin fine-tuning the BERT model on the IMDb dataset using the Hugging Face Trainer.
# 
# - The training process will:
#   1. Iterate over the training dataset for the specified number of epochs.
#   2. Evaluate on the validation set at the end of each epoch.
#   3. Save checkpoints for the best model based on validation loss.
#   4. Stop early if the validation loss does not improve for 2 consecutive evaluation steps (early stopping).
# - Training progress, loss, accuracy, and F1 score will be displayed in real time.
# 

# In[9]:


# Start model training
trainer.train()


# # Save the Fine-Tuned Model and Tokenizer
# 
# After training, we save the fine-tuned BERT model and its tokenizer so that they can be easily reloaded later for inference or further fine-tuning.
# 
# - The model weights and configuration will be saved in the folder `my-fine-tuned-bert`.
# - The tokenizer files are also saved in the same directory.
# 

# In[10]:


# Save the fine-tuned model
trainer.save_model('my-fine-tuned-bert')

# Save the tokenizer
tokenizer.save_pretrained('my-fine-tuned-bert')


# # Inference with the Fine-Tuned Model
# 
# In this section, we demonstrate how to load the fine-tuned BERT model and tokenizer, and perform sentiment prediction on new text.
# 
# - We use Hugging Face's `TextClassificationPipeline` for convenient text classification.
# - The model predicts either `LABEL_0` (negative) or `LABEL_1` (positive).
# - We map these labels to more meaningful sentiment labels for clarity.
# - Finally, we test the model on a sample sentence.
# 

# In[11]:


from transformers import AutoModelForSequenceClassification, AutoTokenizer, TextClassificationPipeline

# Load the fine-tuned model and tokenizer
new_model = AutoModelForSequenceClassification.from_pretrained('my-fine-tuned-bert')
new_tokenizer = AutoTokenizer.from_pretrained('my-fine-tuned-bert')

# Create a text classification pipeline
classifier = TextClassificationPipeline(
    model=new_model, 
    tokenizer=new_tokenizer,     )

# Define label mapping
label_mapping = {
    0: 'World',
    1: 'Sports',
    2: 'Business',
    3: 'Sci/Tech'
}

# Test the classifier on a sample sentence
sample_text = "This movie was good"
result = classifier(sample_text)

# Map the predicted label to a meaningful sentiment
mapped_result = {
    'label': label_mapping[int(result[0]['label'].split('_')[1])],
    'score': result[0]['score']
}

print(mapped_result)


# In[12]:


from datasets import load_dataset
from transformers import AutoTokenizer, AutoModelForSequenceClassification, Trainer
import numpy as np
from sklearn.metrics import accuracy_score, f1_score, classification_report, confusion_matrix

MODEL_DIR = "my-fine-tuned-bert"
MODEL_NAME = "bert-base-uncased"

raw = load_dataset("SetFit/ag_news")
tok = AutoTokenizer.from_pretrained(MODEL_NAME)

def tokenize_fn(ex):
    return tok(ex["text"], truncation=True, max_length=128)

test_tok = raw["test"].map(tokenize_fn, batched=True)
test_tok = test_tok.rename_column("label", "labels")
cols = ["input_ids", "attention_mask", "labels"]
if "token_type_ids" in test_tok.column_names:
    cols.append("token_type_ids")
test_tok.set_format(type="torch", columns=cols)

model = AutoModelForSequenceClassification.from_pretrained(MODEL_DIR)

trainer = Trainer(model=model, tokenizer=tok)
pred_out = trainer.predict(test_tok)

y_prob = pred_out.predictions
y_pred = np.argmax(y_prob, axis=1)
y_true = pred_out.label_ids

acc = accuracy_score(y_true, y_pred)
f1m = f1_score(y_true, y_pred, average="macro")
print({"test_accuracy": acc, "test_f1_macro": f1m})
print(classification_report(y_true, y_pred, digits=4))

from sklearn.metrics import ConfusionMatrixDisplay
import matplotlib.pyplot as plt
ConfusionMatrixDisplay.from_predictions(y_true, y_pred)
plt.title("AG News – Confusion Matrix (Test)")
plt.show()


# In[13]:


import torch, gradio as gr
from transformers import AutoTokenizer, AutoModelForSequenceClassification, pipeline

MODEL_ID = "my-fine-tuned-bert"

tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForSequenceClassification.from_pretrained(MODEL_ID)

label_names = {0: "World", 1: "Sports", 2: "Business", 3: "Sci/Tech"}

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
clf = pipeline("text-classification", model=model, tokenizer=tokenizer, top_k=None, device=device)

def predict(text):
    text = text.strip()
    out = clf(text, truncation=True)
    out = out[0] if isinstance(out[0], list) else out
    results = {}
    for o in sorted(out, key=lambda x: -x["score"]):
        idx = int(o["label"].split("_")[1])
        results[label_names[idx]] = o["score"]
    return results

demo = gr.Interface(
    fn=predict,
    inputs=gr.Textbox(lines=3, label="Enter news headline"),
    outputs=gr.Label(num_top_classes=4, label="Predicted topic"),
    title="AG News Topic Classifier (BERT-base)"
)

demo.launch(share=True)