Blame for code/training-single-model-v9.py

# %% [markdown]
# **Notice: You should run this notebook on a GPU with at least 24GB RAM for avoiding OOM**
# 
# # Introduction
# 
# Hi all,
# 
# In this notebook, I will share the comprehensive steps and strategies I employed to achieve `0.963` on the Private Leaderboard using a single model.
# 
# Throughout this competition, I explored various methodologies and refined my approaches to push the model to its limits. This notebook is structured to guide you through each phase of my process.
# 
# Whether you're a beginner eager to learn the ropes of competitive machine learning or an experienced practitioner looking for new insights, this notebook aims to provide valuable perspectives and techniques that can be adapted to your future projects.
# 
# Let's dive into the details and uncover the techniques that led to a successful submission!
# 
# [Submission Notebook PrivateLB 0.963](https://www.kaggle.com/code/haqishen/single-model-no-pp-privatelb-0-963)
# 
# # Summary
# 
# * With [2k external training samples](https://www.kaggle.com/code/valentinwerner/fix-punctuation-tokenization-external-dataset)
# * `TRAINING_MAX_LENGTH = 1024`, which should be afforded by many people
#     * When the text length exceeds 1024, only the first 512 and the last 512 characters of the text will be used for training.
#     * For inference, no truncation.
# * `downsample rate = 0.3`, randomly downsample negative samples (articles without any PII tokens) to 30%
# * `attention_probs_dropout_prob = 0.4`, default value was `0.1`
# * `weight_decay = 0.`, do not use weight decay for optimizer. default value was `0.01`
# * `pos_weight = 10`, means for positive samples, weighted 10x for loss calculation.
#     * As many people knows, for most public submission notebooks, they use predictions with `th > 0.9` to get higher score.
#     * By adapting `pos_weight`, I was able to achieve best performance with `th around 0.5`
# * cosine with warmup learning rate schedule -- almost became a standard setting.
#     * warmup 0.5 epo
#     * cosine 4.5 epo
# 
# # Acknowledge
# 
# I was provided Z8G5 Fury workstation with dual RTX6000 ada GPU.
# 
# For training a single model this notebook shows, it takes less than 40min!
# 

# %%
import os
os.environ['TOKENIZERS_PARALLELISM'] = 'false'
import warnings
warnings.filterwarnings("ignore")
DEBUG=False

# %%
# !pip install -q seqeval

# %%
import gc
import json
import copy
import time
import pickle
import numpy as np
import pandas as pd
from itertools import chain
from transformers import AutoTokenizer
from transformers import AutoModelForTokenClassification
from transformers import get_cosine_schedule_with_warmup
from tqdm import tqdm

from seqeval.metrics import recall_score, precision_score
from seqeval.metrics import f1_score

import torch
from torch import nn
from torch.utils.data import Dataset, DataLoader
from torch.nn.utils.rnn import pad_sequence
from torch.nn import CrossEntropyLoss
from torch.optim import AdamW

import matplotlib.pyplot as plt
from pylab import rcParams

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# %% [markdown]
# # Config

# %%
exp_name = 'deberta_v3_large_w10_ado4_wd0_bs2_lr1e5_5epo_simplified_labels' # Updated exp_name

backbone = 'microsoft/deberta-v3-large'

TRAINING_MAX_LENGTH = 4096
batch_size = 2
lr = 1e-5
ds_rate = 0.3
attention_probs_dropout_prob = 0.4
hidden_dropout_prob = 0.1
wd = 0.
pos_weight = 10
n_epochs = 5

use_amp = True

# %% [markdown]
# # Data

# %%
from sklearn.model_selection import train_test_split

data = json.load(open('/export/home/acs/stud/t/tudor.farcasanu/SSL_Assignment/pii-data-detection-3rd-place-solution/datasets/train.json'))
for sample in data:
    if len(set(sample['labels'])) > 1:
        sample['is_pos'] = 1
    else:
        sample['is_pos'] = 0
print("original datapoints: ", len(data))


# moredata = json.load(open('/kaggle/input/fix-punctuation-tokenization-external-dataset/moredata_dataset_fixed.json'))
# print("moredata datapoints: ", len(moredata))


# %%
from sklearn.model_selection import GroupKFold
import random
random.seed(0)
kf = GroupKFold(n_splits=5)
folds = []
for i, (_, valid_index) in enumerate(kf.split(np.array(data), groups=[x['document'] for x in data])):
    folds.append([data[i]['document'] for i in valid_index])

# %%
# Define simplified labels
pii_types = [
    'EMAIL', 'ID_NUM', 'NAME_STUDENT', 'PHONE_NUM', 
    'STREET_ADDRESS', 'URL_PERSONAL', 'USERNAME'
]
# all_labels_simplified will be pii_types + ['O']
all_labels_simplified = sorted(pii_types + ['O'])
label2id_simplified = {l: i for i, l in enumerate(all_labels_simplified)}
id2label_simplified = {v: k for k, v in label2id_simplified.items()}

# Original target for reference or other logic if needed, but not for model head
original_target_bio = [
    'B-EMAIL', 'B-ID_NUM', 'B-NAME_STUDENT', 'B-PHONE_NUM', 
    'B-STREET_ADDRESS', 'B-URL_PERSONAL', 'B-USERNAME', 'I-ID_NUM', 
    'I-NAME_STUDENT', 'I-PHONE_NUM', 'I-STREET_ADDRESS', 'I-URL_PERSONAL'
]

print(f"Simplified labels ({len(all_labels_simplified)}): {id2label_simplified}")

# %% [markdown]
# # Dataset

# %%
def tokenize(example, tokenizer, label2id, max_length): # label2id here will be label2id_simplified

    # rebuild text from tokens
    text = []
    char_labels = [] # Renamed from 'labels' to avoid confusion with example['provided_labels']

    for t, l, ws in zip(
        example["tokens"], example["provided_labels"], example["trailing_whitespace"]
    ):
        text.append(t)
        # Simplify label: strip B-/I- prefix if it's a PII label
        current_label = l
        if l != "O":
            current_label = l[2:] # Strip "B-" or "I-"
        char_labels.extend([current_label] * len(t))

        if ws:
            text.append(" ")
            char_labels.append("O") # Whitespace is always "O"

    tokenized = tokenizer("".join(text), return_offsets_mapping=True, truncation=False)
    # Truncation strategy
    if len(tokenized['input_ids']) > max_length:
        # Implement the "first N/2 + last N/2" truncation strategy
        first_half_len = max_length // 2
        second_half_len = max_length - first_half_len # Ensures total length is max_length

        keys_to_truncate = ['input_ids', 'attention_mask', 'offset_mapping']
        for key in keys_to_truncate:
            if key in tokenized:
                original_list = tokenized[key]
                # This slicing takes the first 'first_half_len' elements and the last 'second_half_len' elements
                tokenized[key] = original_list[:first_half_len] + original_list[-second_half_len:]
            # Consider adding a warning if a key is expected but not found, though unlikely for these standard keys.

    char_labels_arr = np.array(char_labels) # This remains after truncation, using the original full char_labels

    text_str = "".join(text) # Renamed from 'text' to avoid conflict
    token_labels = []

    # Access offset_mapping using dictionary style access as tokenized might be a dict
    current_offset_mapping = tokenized["offset_mapping"] 

    for start_idx, end_idx in current_offset_mapping:
        # Handle special tokens like [CLS], [SEP] which have (0,0) offset_mapping
        # or any other token that doesn't map to actual characters in the original text.
        if start_idx == end_idx:
            token_labels.append(label2id["O"]) # Assign 'O' to special/empty tokens
            continue

        # Ensure indices are within the bounds of char_labels_arr and form a valid span
        # end_idx is exclusive for slicing, so it can be == len(char_labels_arr)
        if not (start_idx < end_idx and start_idx < len(char_labels_arr) and end_idx <= len(char_labels_arr)):
            token_labels.append(label2id["O"]) # Assign 'O' if token is outside char_labels range or invalid span
            continue
        
        char_span_labels = char_labels_arr[start_idx:end_idx]
        
        # If char_span_labels is somehow empty (e.g. if previous check was less strict)
        if not char_span_labels.size: 
            token_labels.append(label2id["O"])
            continue

        unique_labels, counts = np.unique(char_span_labels, return_counts=True)
        
        pii_labels_present_in_token = []
        pii_counts_present_in_token = []

        for label, count in zip(unique_labels, counts):
            if label != "O": # Assuming "O" is the string for non-PII
                pii_labels_present_in_token.append(label)
                pii_counts_present_in_token.append(count)
        
        if pii_labels_present_in_token:
            # If PII labels are present, choose the most frequent PII label
            most_frequent_pii_idx = np.argmax(pii_counts_present_in_token)
            most_frequent_pii_label = pii_labels_present_in_token[most_frequent_pii_idx]
            token_labels.append(label2id[most_frequent_pii_label])
        else:
            # If only 'O' is present or no PII labels found in the span
            token_labels.append(label2id["O"])

    length = len(tokenized["input_ids"]) # Use dictionary-style access

    return {**tokenized, "labels": token_labels, "length": length}

# %%
tokenizer = AutoTokenizer.from_pretrained(backbone)

class PIIDataset(Dataset):
    def __init__(self, tokenizer, data, label2id, max_length): # label2id will be simplified
        self.tokenizer = tokenizer
        self.data = copy.deepcopy(data)
        self.label2id = label2id
        self.max_length = max_length

        for i in tqdm(range(len(self.data))):
            if self.data[i].get('input_ids') is None:
                self.data[i]['provided_labels'] = self.data[i]['labels'] # Keep original BIO labels here for tokenize
                self.data[i] = tokenize(
                    self.data[i],
                    self.tokenizer,
                    self.label2id, # Pass simplified label2id
                    self.max_length
                )

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        return self.data[idx]


def collate_fn(batch):
    input_ids = [torch.tensor(x['input_ids']) for x in batch]
    attention_masks = [torch.tensor(x['attention_mask']) for x in batch]
    labels = [torch.tensor(x['labels']) for x in batch]

    input_ids_padded = pad_sequence(input_ids, batch_first=True, padding_value=tokenizer.pad_token_id if tokenizer.pad_token_id is not None else 0)
    attention_masks_padded = pad_sequence(attention_masks, batch_first=True, padding_value=0)
    labels_padded = pad_sequence(labels, batch_first=True, padding_value=-100)

    return {
        'input_ids': input_ids_padded,
        'attention_mask': attention_masks_padded,
        'labels': labels_padded,
    }

# Test with simplified labels
ds = PIIDataset(tokenizer, data[:20], label2id_simplified, TRAINING_MAX_LENGTH)
print(f"Example tokenized input IDs length: {len(ds[9]['input_ids'])}")
print(f"Example tokenized labels: {ds[9]['labels']}")


# %% [markdown]
# # Model

# %%
model = AutoModelForTokenClassification.from_pretrained(
        backbone,
        attention_probs_dropout_prob=attention_probs_dropout_prob,
        hidden_dropout_prob=hidden_dropout_prob,
        num_labels=len(all_labels_simplified), # Use length of simplified labels (8)
        id2label=id2label_simplified,         # Use simplified id2label
        label2id=label2id_simplified,         # Use simplified label2id
        ignore_mismatched_sizes=True      # Good practice if changing head
    )

# %% [markdown]
# # Loss

# %%
# Adjust weights for simplified labels
weights_simplified = torch.ones(len(all_labels_simplified)) * pos_weight
# Assuming 'O' is the label for non-PII and its ID is known via label2id_simplified
if 'O' in label2id_simplified:
    weights_simplified[label2id_simplified['O']] = 1
else:
    # Fallback if 'O' is not named exactly like that, though it should be.
    # This assumes 'O' is the last class if not found by name.
    print("Warning: 'O' label not found in label2id_simplified for weight assignment. Assuming last index for 'O'.")
    weights_simplified[-1] = 1


CE_simplified = nn.CrossEntropyLoss(weight=weights_simplified.to(device), ignore_index=-100)
def criterion(logits, labels):
    # logits shape: (batch_size, seq_len, num_simplified_labels)
    # labels shape: (batch_size, seq_len)
    return CE_simplified(logits.view(-1, len(all_labels_simplified)), labels.view(-1))

# %% [markdown]
# # Metric

# %%
def compute_metrics(p, current_all_labels, th=0.5): # Pass simplified all_labels
    predictions, labels = p # predictions are logits
    
    # Determine index of 'O' label in the current_all_labels (simplified)
    o_label_str = "O" # Assuming 'O' is always named this way
    o_label_idx = -1
    for i, lbl_str in enumerate(current_all_labels): # current_all_labels is a list of strings like ['EMAIL', ..., 'O']
        if lbl_str == o_label_str:
            o_label_idx = i
            break
    if o_label_idx == -1:
        print(f"Warning: Could not find 'O' label in {current_all_labels} for metric thresholding. Using last index.")
        o_label_idx = len(current_all_labels) - 1


    true_predictions = []
    true_labels = []
    for tmp, batch_l in zip(predictions, labels): # tmp contains logits for a batch item
        batch_p_probs = torch.softmax(torch.from_numpy(tmp), dim=-1).numpy() # Convert logits to probs
        
        # Get a view of the 'O' column probabilities
        # batch_p_probs shape is (batch_size, seq_len, num_simplified_labels)
        o_column_probs_slice = batch_p_probs[:, :, o_label_idx] # Shape (batch_size, seq_len)

        # Create the boolean mask based on this slice
        condition_mask = o_column_probs_slice < th # Shape (batch_size, seq_len)

        # Assign to the 'O' column slice using this mask. This modifies batch_p_probs in place.
        o_column_probs_slice[condition_mask] = -9999 

        batch_p_final_ids = np.argmax(batch_p_probs, axis=-1) # Get final predicted IDs

        true_predictions.extend([current_all_labels[p_id] for (p_id, l_id) in zip(batch_p_final_ids.reshape(-1), batch_l.reshape(-1)) if l_id != -100])
        true_labels.extend([current_all_labels[l_id] for l_id in batch_l.reshape(-1) if l_id != -100])

    # Using seqeval for P, R, F1 (F5)
    # seqeval expects lists of lists of strings.
    # Here true_labels and true_predictions are flat lists.
    # For proper seqeval, we'd need to group them back by sequence.
    # However, the original script's compute_metrics calculates overall P, R, F5 on flat lists.
    # Let's keep that structure for now, as it's what the script used for its internal F5.
    # Note: This F5 is on simplified labels, not BIO.
    
    recall = recall_score([true_labels], [true_predictions], average='micro') # Use micro for overall
    precision = precision_score([true_labels], [true_predictions], average='micro')
    # Calculate F5 score manually based on P and R
    beta = 5
    f5_score = (1 + beta**2) * precision * recall / ((beta**2 * precision) + recall + 1e-100)
    if precision + recall == 0: # Handle case where P and R are both zero
        f5_score = 0.0

    results = {
        'recall': recall,
        'precision': precision,
        'f1': f5_score # Reporting F5 as 'f1' to match original key, or add 'f5'
    }
    return results


# %% [markdown]
# # Train & Valid Func

# %%
def train_one_epoch(model, loader, optimizer, scheduler, scaler=None):
    model.train()
    running_loss = []
    bar = tqdm(loader, total=len(loader))
    for batch in bar:
        optimizer.zero_grad()
        
        input_ids = batch['input_ids'].to(device)
        attention_mask = batch['attention_mask'].to(device)
        labels = batch['labels'].to(device)

        if scaler is not None:
            with torch.cuda.amp.autocast():
                outputs = model(input_ids, attention_mask=attention_mask, labels=labels)
                loss = criterion(outputs['logits'], labels) # criterion uses simplified label count
            scaler.scale(loss).backward()
            scaler.step(optimizer)
            scaler.update()
        else:
            outputs = model(input_ids, attention_mask=attention_mask, labels=labels)
            loss = criterion(outputs['logits'], labels) # criterion uses simplified label count
            loss.backward()
            optimizer.step()
        
        scheduler.step() # Moved scheduler step here, common practice per batch
        running_loss.append(loss.item())
        bar.set_description(f'train loss: {np.mean(running_loss[-20:]):.4f}')

    return np.mean(running_loss)


def valid_one_epoch(model, loader):
    model.eval()
    running_loss = []
    bar = tqdm(loader, total=len(loader))
    predictions_logits = [] # Store logits
    true_labels_list = [] # Store true labels

    with torch.no_grad():
        for batch in bar:
            input_ids = batch['input_ids'].to(device)
            attention_mask = batch['attention_mask'].to(device)
            labels_ = batch['labels'].to(device)
            
            outputs = model(input_ids, attention_mask=attention_mask, labels=labels_)
            loss = criterion(outputs['logits'], labels_) # criterion uses simplified label count
            
            running_loss.append(loss.item())
            predictions_logits.append(outputs['logits'].detach().cpu().numpy())
            true_labels_list.append(labels_.detach().cpu().numpy())

            bar.set_description(f'valid loss: {np.mean(running_loss[-20:]):.4f}')
    
    # Pass simplified all_labels_simplified to compute_metrics
    metrics_results = compute_metrics((predictions_logits, true_labels_list), all_labels_simplified)
    return np.mean(running_loss), metrics_results

# %% [markdown]
# # Scheduler

# %%
rcParams['figure.figsize'] = 20, 2

def get_scheduler(optimizer, num_training_steps, warmup_steps_ratio=0.1):
    return get_cosine_schedule_with_warmup(
        optimizer,
        int(num_training_steps * warmup_steps_ratio),
        num_training_steps
    )

opt = AdamW(model.parameters(), lr=lr)
sch = get_scheduler(opt, 100, 0.1)

lrs = []
for epoch in range(1, 100+1):
    sch.step()
    lrs.append(opt.param_groups[0]["lr"])

plt.plot(range(len(lrs)), lrs)

# %% [markdown]
# # Run

# %%
def run(fold,model_to_train): # Renamed model to model_to_train to avoid conflict

    log_file = os.path.join('/export/home/acs/stud/t/tudor.farcasanu/SSL_Assignment/single_model/model/logs', f'{exp_name}.txt')
    # Define a directory for saving Hugging Face model (config, weights, tokenizer_config etc.)
    model_save_dir_best = os.path.join('/export/home/acs/stud/t/tudor.farcasanu/SSL_Assignment/single_model/model', f'{exp_name}_fold{fold}_best_hf')
    model_save_dir_last = os.path.join('/export/home/acs/stud/t/tudor.farcasanu/SSL_Assignment/single_model/model', f'{exp_name}_fold{fold}_last_hf')
    
    # Create directories if they don't exist
    os.makedirs(model_save_dir_best, exist_ok=True)
    os.makedirs(model_save_dir_last, exist_ok=True)
    # Path for optimizer and other states (still using torch.save for these with custom loop)
    optimizer_etc_file_last = os.path.join(model_save_dir_last, 'optimizer_etc_states.pth')


    valid_doc_ids = folds[fold]
    # train_data = [x for x in data    if (x['document'] not in valid_doc_ids) and ((x['is_pos'] == 1) or (x['is_pos'] == 0 and random.random() < ds_rate))] + moredata
    train_data = [x for x in data    if (x['document'] not in valid_doc_ids) and ((x['is_pos'] == 1) or (x['is_pos'] == 0 and random.random() < ds_rate))]
    valid_data = [x for x in data if x['document'] in valid_doc_ids]

    # Use simplified label2id
    dataset_train = PIIDataset(tokenizer, train_data, label2id_simplified, TRAINING_MAX_LENGTH)
    dataset_valid = PIIDataset(tokenizer, valid_data, label2id_simplified, TRAINING_MAX_LENGTH)
    print(len(dataset_train), len(dataset_valid))

    loader_train = DataLoader(dataset_train, batch_size=batch_size, shuffle=True, num_workers=2, collate_fn=collate_fn) # Reduced num_workers for potential local issues
    loader_valid = DataLoader(dataset_valid, batch_size=batch_size, shuffle=False, num_workers=2, collate_fn=collate_fn)


    # model = get_model() # Original script had this, but model is passed now
    current_model = model_to_train.to(device) # Use the passed model

    optimizer = AdamW(current_model.parameters(), lr=lr, weight_decay=wd)
    scaler = torch.cuda.amp.GradScaler() if use_amp else None
    from_epoch = 0 # Assuming fresh train, not resuming
    metric_best = 0

    scheduler = get_scheduler(optimizer, n_epochs*len(loader_train), 0.1) # warmup_steps_ratio is 0.1

    for epoch in range(1, n_epochs+1):
        # if epoch < from_epoch + 1:
        #     print(logs[epoch-1]) # Assuming logs is not used for fresh train
        #     continue

        print(time.ctime(), 'Epoch:', epoch)

        train_loss = train_one_epoch(current_model, loader_train, optimizer, scheduler, scaler)
        valid_loss, result = valid_one_epoch(current_model, loader_valid)
        metric = result["f1"] # This is F5 score on simplified labels

        content = time.ctime() + ' ' + \
                f'Fold {fold}, Epoch {epoch}, lr: {optimizer.param_groups[0]["lr"]:.7f}, ' + \
                f'train loss: {train_loss:.5f}, valid loss: {valid_loss:.5f}, ' + \
                f'recall: {result["recall"]:.4f} precision: {result["precision"]:.4f} metric (F5 on simplified): {metric:.4f}.'
        print(content)
        with open(log_file, 'a') as appender:
            appender.write(content + '\n')

        if metric > metric_best:
            print(f'metric_best ({metric_best:.6f} --> {metric:.6f}). Saving model to {model_save_dir_best} ...')
            if not DEBUG:
                current_model.save_pretrained(model_save_dir_best, safe_serialization=True)
                tokenizer.save_pretrained(model_save_dir_best) # Save tokenizer too
            metric_best = metric

        # Save Last states (model, optimizer, scheduler)
        if not DEBUG:
            current_model.save_pretrained(model_save_dir_last, safe_serialization=True)
            tokenizer.save_pretrained(model_save_dir_last)
            torch.save(
                {
                    'epoch': epoch,
                    # model_state_dict is not needed if using save_pretrained for the model itself
                    'optimizer_state_dict': optimizer.state_dict(),
                    'scheduler_state_dict': scheduler.state_dict(), # Save scheduler state
                    'scaler_state_dict': scaler.state_dict() if scaler else None,
                    'metric_best': metric_best,
                },
                optimizer_etc_file_last # Save optimizer and other states separately
            )
    
    del current_model # Use current_model
    torch.cuda.empty_cache()
    gc.collect()

# %%
# model is defined globally from "Model" cell. Pass it to run.
run(0, model)

# %%