v1_chunk.py

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

Automatically generated by Colab.

Original file is located at
    https://colab.research.google.com/drive/1Ifg-kGn0HiUH-Uaadwu4yYNYDW7nxydG

# Starter Notebook

Install and import required libraries
"""

!pip install transformers datasets evaluate accelerate peft trl bitsandbytes
!pip install nvidia-ml-py3

import os
import re
import string
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.nn.functional import softmax

from datasets import load_dataset
from transformers import (
    RobertaTokenizer,
    RobertaForSequenceClassification,
    TrainingArguments,
    Trainer,
    DataCollatorWithPadding,
    AutoConfig
)

from sklearn.metrics import (
    accuracy_score,
    precision_score,
    recall_score,
    f1_score,
    confusion_matrix,
    ConfusionMatrixDisplay
)

from peft import PeftModel
from torch.utils.data import DataLoader
import evaluate
from tqdm import tqdm

"""## Load Tokenizer and Preprocess Data"""

base_model = 'roberta-base'
max_length = 384
dataset = load_dataset('ag_news', split='train')
tokenizer = RobertaTokenizer.from_pretrained(base_model)


# Apply tokenizer after cleaning
def preprocess(examples):
    tokenized = tokenizer(examples['text'], truncation=True, padding='max_length',max_length=max_length)
    return tokenized

tokenized_dataset = dataset.map(preprocess, batched=True,  remove_columns=["text"])
tokenized_dataset = tokenized_dataset.rename_column("label", "labels")
print(tokenizer.model_max_length)

"""## Token Count Histogram for AGNews"""

# Load AG News raw data
#dataset = load_dataset("ag_news", split="train")

# Tokenize a sample of the text (e.g., first 1000 rows)
texts = dataset["text"][:1000]
token_lengths = [len(tokenizer(text)["input_ids"]) for text in texts]

# Print a few examples
for i in range(5):
    print(f"Sample {i}: {token_lengths[i]} tokens → {texts[i][:60]}...")

plt.hist(token_lengths, bins=30, color="skyblue", edgecolor="black")
plt.axvline(384, color="orange", linestyle="--", label="max_length=384")
plt.axvline(512, color="red", linestyle="--", label="model_max_length=512")
plt.title("Token Length Distribution (First 1000 AG News Samples)")
plt.xlabel("Token Count")
plt.ylabel("Number of Samples")
plt.legend()
plt.grid(True)
plt.show()

# Extract the number of classess and their names
num_labels = dataset.features['label'].num_classes
class_names = dataset.features["label"].names
print(f"number of labels: {num_labels}")
print(f"the labels: {class_names}")

# Create an id2label mapping
# We will need this for our classifier.
id2label = {i: label for i, label in enumerate(class_names)}

data_collator = DataCollatorWithPadding(tokenizer=tokenizer, return_tensors="pt")

"""## Load Pre-trained Model
Set up config for pretrained model and download it from hugging face
"""

model = RobertaForSequenceClassification.from_pretrained(
    base_model,
    id2label=id2label)
model

"""### Applying inference chunks for unlabelled dataset"""

def chunked_predict(text, model, tokenizer, chunk_size=128, stride=64, device="cuda"):

    tokens = tokenizer(text, return_tensors="pt", truncation=False)
    input_ids = tokens["input_ids"][0]

    chunks = [input_ids[i : i + chunk_size] for i in range(0, len(input_ids), stride)]
    all_logits = []
    model.eval()

    for chunk in chunks:
        chunk = chunk.unsqueeze(0).to(device)
        attention_mask = (chunk != tokenizer.pad_token_id).long()

        with torch.no_grad():
            output = model(input_ids=chunk, attention_mask=attention_mask)
            probs = softmax(output.logits, dim=-1)
            all_logits.append(probs.squeeze(0))

    avg_probs = torch.stack(all_logits).mean(dim=0)
    return avg_probs.argmax().item(), avg_probs

"""## Anything from here on can be modified"""

# Split the original training set
split_datasets = tokenized_dataset.train_test_split(test_size=640, seed=42)
train_dataset = split_datasets['train']
eval_dataset = split_datasets['test']

"""## Setup LoRA Config
Setup PEFT config and get peft model for finetuning
"""

# PEFT Config
peft_config = LoraConfig(
    r=8,
    lora_alpha=16,
    lora_dropout=0.1,
    bias = 'none',
    target_modules = ['query', 'value'],
    task_type="SEQ_CLS",
)

peft_model = get_peft_model(model, peft_config)
peft_model

#print("Trainable parameters:")
#for name, param in peft_model.named_parameters():
#  if param.requires_grad:
#    print(name)

print('PEFT Model')
peft_model.print_trainable_parameters()

"""## Training Setup"""

# To track evaluation accuracy during training
def compute_metrics(pred):
    labels = pred.label_ids
    preds = pred.predictions.argmax(-1)

    return {
        'accuracy': accuracy_score(labels, preds),
        'precision': precision_score(labels, preds, average='weighted'),
        'recall': recall_score(labels, preds, average='weighted'),
        'f1': f1_score(labels, preds, average='weighted'),
    }

# Setup Training args
output_dir = "results"
training_args = TrainingArguments(
    output_dir=output_dir,
    report_to=None,
    eval_strategy='epoch',
    logging_steps=100,
    learning_rate=2e-4,
    num_train_epochs=5,
    #max_steps=1200,
    use_cpu=False,
    dataloader_num_workers=4,
    per_device_train_batch_size=16,
    per_device_eval_batch_size=64,
    optim="adamw_torch",
    gradient_checkpointing=False,
    gradient_checkpointing_kwargs={'use_reentrant':True}
)

def get_trainer(model):
      return  Trainer(
          model=model,
          args=training_args,
          compute_metrics=compute_metrics,
          train_dataset=train_dataset,
          eval_dataset=eval_dataset,
          data_collator=data_collator,
      )

"""### Start Training"""

peft_lora_finetuning_trainer = get_trainer(peft_model)

result = peft_lora_finetuning_trainer.train()

"""## Evaluate Finetuned Model

### Performing Inference on Custom Input
Uncomment following functions for running inference on custom inputs
"""

def classify(model, tokenizer, text):
  device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
  inputs = tokenizer(text, truncation=True, padding=True, return_tensors="pt").to(device)
  output = model(**inputs)

  prediction = output.logits.argmax(dim=-1).item()

  print(f'\n Class: {prediction}, Label: {id2label[prediction]}, Text: {text}')
  return id2label[prediction]

# classify( peft_model, tokenizer, "Kederis proclaims innocence Olympic champion Kostas Kederis today left hospital ahead of his date with IOC inquisitors claiming his ...")
# classify( peft_model, tokenizer, "Wall St. Bears Claw Back Into the Black (Reuters) Reuters - Short-sellers, Wall Street's dwindling\band of ultra-cynics, are seeing green again.")

#Load your unlabelled data
unlabelled_dataset = pd.read_pickle("test_unlabelled.pkl")
test_dataset = unlabelled_dataset.map(preprocess, batched=True, remove_columns=["text"])
unlabelled_dataset

"""## Evaluate model through inference chunks"""

device = "cuda" if torch.cuda.is_available() else "cpu"
peft_model.to(device)
peft_model.eval()

texts = unlabelled_dataset["text"]
preds = []

for i, text in enumerate(texts):
    label, _ = chunked_predict(text, model=peft_model, tokenizer=tokenizer, chunk_size=384, stride=384, device=device)
    preds.append(label)

print("Chunked inference complete.")

df_output = pd.DataFrame({
    'ID': range(len(preds)),
    'Label': preds
})
df_output.to_csv(os.path.join(output_dir, "v1_chunk_output.csv"), index=False)
print("Predictions saved to v1_chunk_output.csv")

"""### Evaluate Model (OG Function)"""

def evaluate_model(inference_model, dataset, labelled=True, batch_size=8, data_collator=None):
    """
    Evaluate a PEFT model on a dataset.

    Args:
        inference_model: The model to evaluate.
        dataset: The dataset (Hugging Face Dataset) to run inference on.
        labelled (bool): If True, the dataset includes labels and metrics will be computed.
                         If False, only predictions will be returned.
        batch_size (int): Batch size for inference.
        data_collator: Function to collate batches. If None, the default collate_fn is used.

    Returns:
        If labelled is True, returns a tuple (metrics, predictions)
        If labelled is False, returns the predictions.
    """
    # Create the DataLoader
    eval_dataloader = DataLoader(dataset, batch_size=batch_size, collate_fn=data_collator)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    inference_model.to(device)
    inference_model.eval()

    all_predictions = []
    if labelled:
        metric = evaluate.load('accuracy')

    # Loop over the DataLoader
    for batch in tqdm(eval_dataloader):
        # Move each tensor in the batch to the device
        batch = {k: v.to(device) for k, v in batch.items()}
        with torch.no_grad():
            outputs = inference_model(**batch)
        predictions = outputs.logits.argmax(dim=-1)
        all_predictions.append(predictions.cpu())

        if labelled:
            # Expecting that labels are provided under the "labels" key.
            references = batch["labels"]
            metric.add_batch(
                predictions=predictions.cpu().numpy(),
                references=references.cpu().numpy()
            )

    # Concatenate predictions from all batches
    all_predictions = torch.cat(all_predictions, dim=0)

    if labelled:
        eval_metric = metric.compute()
        print("Evaluation Metric:", eval_metric)
        return eval_metric, all_predictions
    else:
        return all_predictions

# Check evaluation accuracy
_, _ = evaluate_model(peft_model, eval_dataset, True, 8, data_collator)

"""#### Run Inference on unlabelled dataset"""

# Run inference and save predictions
#preds = evaluate_model(peft_model, test_dataset, False, 8, data_collator)
df_output = pd.DataFrame({
    'ID': range(len(preds)),
    'Label': preds.numpy()  # or preds.tolist()
})
df_output.to_csv(os.path.join(output_dir,"v1_chunk_output.csv"), index=False)
print("Inference complete. Predictions saved to inference_output.csv")

"""Train & Validation Accuracy over time Plot"""

# Access training log history
history = peft_lora_finetuning_trainer.state.log_history

# Extract steps and accuracy
eval_steps = [entry['step'] for entry in history if 'eval_accuracy' in entry]
eval_acc = [entry['eval_accuracy'] for entry in history if 'eval_accuracy' in entry]

# Extract eval loss (optional)
eval_loss = [entry['eval_loss'] for entry in history if 'eval_loss' in entry]

# Plot validation accuracy
plt.plot(eval_steps, eval_acc, label='Validation Accuracy', marker='o')
plt.xlabel('Step')
plt.ylabel('Accuracy')
plt.title('Validation Accuracy Over Time')
plt.legend()
plt.grid(True)
plt.show()

# Optional: Plot validation loss too
plt.plot(eval_steps, eval_loss, label='Validation Loss', marker='o', color='orange')
plt.xlabel('Step')
plt.ylabel('Loss')
plt.title('Validation Loss Over Time')
plt.legend()
plt.grid(True)
plt.show()

"""## Confusion Matrix"""

# predict on validation/test set
preds = peft_lora_finetuning_trainer.predict(eval_dataset)
y_true = preds.label_ids
y_pred = preds.predictions.argmax(axis=1)

# plot
cm = confusion_matrix(y_true, y_pred)
disp = ConfusionMatrixDisplay(confusion_matrix=cm)
disp.plot(cmap='Blues')
plt.title("Confusion Matrix")
plt.show()

"""## Token Count Histogram for unlabelled dataset"""

# Tokenize a sample of the Kaggle text (first 1000 rows or full if small)
texts = unlabelled_dataset["text"][:1000]  # adjust slice if needed
token_lengths = [len(tokenizer(text)["input_ids"]) for text in texts]

# Print a few examples for verification
for i in range(5):
    print(f"Sample {i}: {token_lengths[i]} tokens → {texts[i][:60]}...")

# Plot the token length distribution
import matplotlib.pyplot as plt

plt.hist(token_lengths, bins=30, color="skyblue", edgecolor="black")
plt.axvline(384, color="orange", linestyle="--", label="max_length=384")
plt.axvline(512, color="red", linestyle="--", label="model_max_length=512")
plt.title("Token Length Distribution (First 1000 Kaggle Test Samples)")
plt.xlabel("Token Count")
plt.ylabel("Number of Samples")
plt.legend()
plt.grid(True)
plt.show()

"""## Length Distribution Error"""

raw_dataset = load_dataset("ag_news", split="train")

# Step 2: Reproduce your original split (must match the one used before training)
raw_split = raw_dataset.train_test_split(test_size=640, seed=42)
raw_eval_dataset = raw_split["test"]

# Step 3: Extract raw texts
val_texts = raw_eval_dataset["text"]

# Step 4: Use your existing predictions
# Make sure y_true and y_pred are already defined based on eval_dataset
# Example:
# preds = trainer.predict(eval_dataset)
# y_true = preds.label_ids
# y_pred = preds.predictions.argmax(axis=1)

# Step 5: Error analysis
df = pd.DataFrame({
    "text": val_texts,
    "true": y_true,
    "pred": y_pred
})

df["length"] = df["text"].apply(len)
errors = df[df["true"] != df["pred"]]

plt.figure(figsize=(10, 5))
plt.hist(errors["length"], bins=30, color='salmon', edgecolor='black')
plt.title("Length Distribution of Misclassified Examples")
plt.xlabel("Text Length (characters)")
plt.ylabel("Number of Errors")
plt.grid(True)
plt.show()

# View top misclassified examples
error_df = df[df["true"] != df["pred"]]
error_df[["text", "true", "pred"]].sample(10)

"""## Validation F1 Score & Loss over time"""

# Get training log history
log_history = peft_lora_finetuning_trainer.state.log_history

# Extract values from log
eval_steps = [entry['step'] for entry in log_history if 'eval_loss' in entry]
eval_loss = [entry['eval_loss'] for entry in log_history if 'eval_loss' in entry]
eval_f1 = [entry['eval_f1'] for entry in log_history if 'eval_f1' in entry]

# Plotting
import matplotlib.pyplot as plt

plt.figure(figsize=(12, 5))

# F1 Plot
plt.subplot(1, 2, 1)
plt.plot(eval_steps, eval_f1, marker='o')
plt.title("Validation F1 Score Over Time")
plt.xlabel("Step")
plt.ylabel("F1 Score")
plt.grid(True)

# Loss Plot
plt.subplot(1, 2, 2)
plt.plot(eval_steps, eval_loss, marker='o', color='orange')
plt.title("Validation Loss Over Time")
plt.xlabel("Step")
plt.ylabel("Loss")
plt.grid(True)

plt.tight_layout()
plt.show()

"""Student Model (training off the LoRa model)"""

base_model = RobertaForSequenceClassification.from_pretrained("roberta-base", num_labels=4)
teacher_model = PeftModel.from_pretrained(base_model, "path/to/lora_checkpoint")
teacher_model.eval()


soft_labels = []

with torch.no_grad():
    for batch in train_dataloader:
        outputs = teacher_model(input_ids=batch["input_ids"], attention_mask=batch["attention_mask"])
        soft_labels.append(torch.nn.functional.softmax(outputs.logits / temperature, dim=-1))  # soft targets

class TinyTransformer(nn.Module):
    def __init__(self, vocab_size=30522, hidden_dim=128, n_heads=2, n_layers=2, num_labels=4, max_len=128):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, hidden_dim)
        self.pos_embedding = nn.Embedding(max_len, hidden_dim)
        encoder_layer = nn.TransformerEncoderLayer(d_model=hidden_dim, nhead=n_heads)
        self.encoder = nn.TransformerEncoder(encoder_layer, num_layers=n_layers)
        self.classifier = nn.Linear(hidden_dim, num_labels)

    def forward(self, input_ids, attention_mask=None):
        pos_ids = torch.arange(0, input_ids.size(1)).unsqueeze(0).to(input_ids.device)
        x = self.embedding(input_ids) + self.pos_embedding(pos_ids)
        x = self.encoder(x)
        x = x.mean(dim=1)
        return self.classifier(x)

"""Train Student with distillation"""

def distillation_loss(student_logits, soft_labels, hard_labels, alpha=0.5, temperature=2.0):
    ce_loss = nn.CrossEntropyLoss()(student_logits, hard_labels)
    kl_loss = nn.KLDivLoss(reduction="batchmean")(
        torch.nn.functional.log_softmax(student_logits / temperature, dim=1),
        soft_labels
    )
    return alpha * ce_loss + (1 - alpha) * (kl_loss * temperature * temperature)


# train the student
student_logits = student_model(input_ids=batch["input_ids"])
loss = distillation_loss(student_logits, soft_labels[idx], hard_labels=batch["labels"])

"""Save the Student predictions"""

student_preds = evaluate_model(student_model, test_dataset, False, 8, data_collator)

# Save predictions
df_student = pd.DataFrame({
    'ID': range(len(student_preds)),
    'Label': student_preds.numpy()
})
df_student.to_csv("results/student_predictions.csv", index=False)

# Save teacher and student model weights
torch.save(teacher_model.state_dict(), "results/teacher_model.pt")
torch.save(student_model.state_dict(), "results/student_model.pt")

df = pd.read_pickle("/content/test_unlabelled.pkl")

# Load the tokenizer
tokenizer = RobertaTokenizer.from_pretrained("roberta-base")

# Tokenize and measure input lengths
token_lengths = [len(tokenizer(text)["input_ids"]) for text in df["text"]]

# Plot the histogram
plt.figure(figsize=(10, 5))
plt.hist(token_lengths, bins=30, color="skyblue", edgecolor="black")
plt.axvline(128, color="green", linestyle="--", label="chunk_size=128")
plt.axvline(256, color="orange", linestyle="--", label="chunk_size=256")
plt.axvline(384, color="red", linestyle="--", label="chunk_size=384")
plt.axvline(512, color="purple", linestyle="--", label="model_max_length=512")
plt.title("Token Length Distribution of Kaggle Test Set")
plt.xlabel("Token Count")
plt.ylabel("Number of Samples")
plt.legend()
plt.grid(True)
plt.tight_layout()
plt.show()