ssl-pii-data-detection-deberta / code /training-single-model-v8.py
Upload DeBERTa-v3-large PII models, tokenizers, and code
9b08c94 verified 2 days ago
24 kB
# %% [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:
        # Truncate by taking first max_length//2 and last max_length//2 tokens
        # This strategy is mentioned in the original notebook's markdown for inputs longer than max_length
        # For 'offset_mapping', we need to be careful. If we truncate input_ids, we must truncate offset_mapping too.
        # The original code did this:
        # for key in tokenized.keys():
        #     tokenized[key] = tokenized[key][:max_length//2] + tokenized[key][-max_length//2:]
        # This is complex for offset_mapping as the second half's offsets are relative to the original string.
        # A simpler truncation for now, or ensure max_length is sufficient.
        # For simplicity, let's use Hugging Face's truncation, but it needs to be enabled in tokenizer call if not already.
        # The original code had truncation=False then a manual truncation.
        # Let's stick to the original manual truncation for input_ids and attention_mask,
        # and apply it carefully to offset_mapping.
        # However, the provided code snippet for tokenize in the original script does:
        # tokenized = tokenizer("".join(text), return_offsets_mapping=True, truncation=False)
        # if len(tokenized['input_ids']) > max_length:
        #     for key in tokenized.keys():
        #         tokenized[key] = tokenized[key][:max_length//2] + tokenized[key][-max_length//2:]
        # This manual truncation is applied *after* tokenization.
        # Let's re-enable truncation in the tokenizer call for safety, though the manual one might override.
        # For this modification, we'll assume the existing truncation logic is what the user wants to keep.
        # The key change is how labels are derived.
        
        # Re-applying the original truncation logic if input_ids are too long
        if len(tokenized['input_ids']) > max_length:
            first_half_len = max_length // 2
            second_half_start = len(tokenized['input_ids']) - (max_length - first_half_len) # Ensure total is max_length

            new_tokenized = {}
            for key in tokenized.keys():
                if key == 'offset_mapping': # Special handling for offset_mapping if needed, or just truncate
                     new_tokenized[key] = tokenized[key][:first_half_len] + tokenized[key][second_half_start:]
                else:
                    new_tokenized[key] = tokenized[key][:first_half_len] + tokenized[key][second_half_start:]
            tokenized = new_tokenized


    char_labels_arr = np.array(char_labels) # Use the simplified 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:
        # CLS token
        if start_idx == 0 and end_idx == 0:
            token_labels.append(label2id["O"]) # Use the passed simplified label2id
            continue

        # case when token starts with whitespace
        # Ensure start_idx is within bounds of text_str
        if start_idx < len(text_str) and text_str[start_idx].isspace():
            start_idx += 1
        
        # Ensure start_idx is within bounds of char_labels_arr after potential increment
        # And handle cases where start_idx might become >= len(char_labels_arr) due to offsets near end
        if start_idx >= len(char_labels_arr):
             # This can happen if offset_mapping points beyond the original text due to truncation or special tokens
             # A safe fallback is 'O' or to ensure char_labels_arr is long enough / handle offsets carefully
             # Given the original code's `while start_idx >= len(labels): start_idx -= 1`, let's adapt that.
             current_char_label_idx = min(start_idx, len(char_labels_arr) - 1)
        else:
            current_char_label_idx = start_idx

        token_labels.append(label2id[char_labels_arr[current_char_label_idx]])

    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)

# %%