# %% [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) # %%