ssl-pii-data-detection-deberta / code /training-actualorig-withotherdata.py

Upload DeBERTa-v3-large PII models, tokenizers, and code

9b08c94 verified 2 days ago

14.4 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'

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('/export/home/acs/stud/t/tudor.farcasanu/SSL_Assignment/pii-data-detection-3rd-place-solution/datasets/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])

# %%
all_labels = sorted(list(set(chain(*[x["labels"] for x in data]))))
label2id = {l: i for i,l in enumerate(all_labels)}
id2label = {v:k for k,v in label2id.items()}

target = [
    '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(len(all_labels), id2label)

# %% [markdown]
# # Dataset

# %%
def tokenize(example, tokenizer, label2id, max_length):

    # rebuild text from tokens
    text = []
    labels = []

    for t, l, ws in zip(
        example["tokens"], example["provided_labels"], example["trailing_whitespace"]
    ):
        text.append(t)
        labels.extend([l] * len(t))

        if ws:
            text.append(" ")
            labels.append("O")

    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:]

    labels = np.array(labels)

    text = "".join(text)
    token_labels = []

    for start_idx, end_idx in tokenized.offset_mapping:
        # CLS token
        if start_idx == 0 and end_idx == 0:
            token_labels.append(label2id["O"])
            continue

        # case when token starts with whitespace
        if text[start_idx].isspace():
            start_idx += 1

        token_labels.append(label2id[labels[start_idx]])

    length = len(tokenized.input_ids)

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

# %%
tokenizer = AutoTokenizer.from_pretrained(backbone)

class PIIDataset(Dataset):
    def __init__(self, tokenizer, data, label2id, max_length):
        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']
                self.data[i] = tokenize(
                    self.data[i],
                    self.tokenizer,
                    self.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=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,
    }

ds = PIIDataset(tokenizer, data[:20], label2id, TRAINING_MAX_LENGTH)
len(ds[9]['input_ids'])

# %% [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),
        id2label=id2label,
        label2id=label2id,
        ignore_mismatched_sizes=True
    )

# %% [markdown]
# # Loss

# %%
weights = torch.ones(13) * pos_weight
weights[12] = 1

CE = nn.CrossEntropyLoss(weight=weights.to(device), ignore_index=-100)
def criterion(logits, labels):
    return CE(logits.view(-1, len(all_labels)), labels.view(-1))

# %% [markdown]
# # Metric

# %%
def compute_metrics(p, all_labels, th=0.5):
    predictions, labels = p
    true_predictions = []
    true_labels = []
    for tmp, batch_l in zip(predictions, labels):
        batch_p = copy.deepcopy(tmp)
        batch_p[:, :, -1][batch_p[:, :, -1] < th] = -9999
        batch_p = np.argmax(batch_p, axis=2)
        true_predictions.extend([all_labels[p] for (p, l) in zip(batch_p.reshape(-1), batch_l.reshape(-1)) if l != -100])
        true_labels.extend([all_labels[l] for l in batch_l.reshape(-1) if l != -100])

    recall = recall_score([true_labels], [true_predictions])
    precision = precision_score([true_labels], [true_predictions])
    f1_score = (1 + 5*5) * recall * precision / (5*5*precision + recall)

    results = {
        'recall': recall,
        'precision': precision,
        'f1': f1_score
    }
    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()
        scheduler.step()

        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)
            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)
            loss.backward()
            optimizer.step()
        
        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 = []
    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_)
            
            running_loss.append(loss.item())
            predictions.append(outputs['logits'].detach().cpu().numpy())
            labels.append(labels_.detach().cpu().numpy())

            bar.set_description(f'valid loss: {np.mean(running_loss[-20:]):.4f}')

    return np.mean(running_loss), compute_metrics((predictions, labels), all_labels)

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

    log_file = os.path.join('/export/home/acs/stud/t/tudor.farcasanu/SSL_Assignment/single_model/model/logs', f'{exp_name}.txt')
    model_file = os.path.join('/export/home/acs/stud/t/tudor.farcasanu/SSL_Assignment/single_model/model', f'{exp_name}_fold{fold}_best.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
    
    valid_data = [x for x in data if x['document'] in valid_doc_ids]

    dataset_train = PIIDataset(tokenizer, train_data, label2id, TRAINING_MAX_LENGTH)
    dataset_valid = PIIDataset(tokenizer, valid_data, label2id, TRAINING_MAX_LENGTH)
    print(len(dataset_train), len(dataset_valid))

    loader_train = DataLoader(dataset_train, batch_size=batch_size, shuffle=True, num_workers=8, collate_fn=collate_fn)
    loader_valid = DataLoader(dataset_valid, batch_size=batch_size, shuffle=False, num_workers=8, collate_fn=collate_fn)


    # model = get_model()
    model = model.to(device)

    optimizer = AdamW(model.parameters(), lr=lr, weight_decay=wd)
    scaler = torch.cuda.amp.GradScaler() if use_amp else None
    from_epoch = 0
    metric_best = 0

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

    for epoch in range(1, n_epochs+1):
        if epoch < from_epoch + 1:
            print(logs[epoch-1])
            continue

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

        train_loss = train_one_epoch(model, loader_train, optimizer, scheduler, scaler)
        valid_loss, result = valid_one_epoch(model, loader_valid)
        metric = result["f1"]

        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: {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 ...')
#             if not DEBUG:
            torch.save(model.state_dict(), model_file)
            metric_best = metric

        # Save Last
        if not DEBUG:
            torch.save(
                {
                    'epoch': epoch,
                    'model_state_dict': model.state_dict(),
                    'optimizer_state_dict': optimizer.state_dict(),
                    'scaler_state_dict': scaler.state_dict() if scaler else None,
                    'metric_best': metric_best,
                },
                model_file.replace('_best', '_last')
            )

    del model
    torch.cuda.empty_cache()
    gc.collect()

# %%
run(0,model)

# %%

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

	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('/export/home/acs/stud/t/tudor.farcasanu/SSL_Assignment/pii-data-detection-3rd-place-solution/datasets/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])

	# %%
	all_labels = sorted(list(set(chain(*[x["labels"] for x in data]))))
	label2id = {l: i for i,l in enumerate(all_labels)}
	id2label = {v:k for k,v in label2id.items()}

	target = [
	'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(len(all_labels), id2label)

	# %% [markdown]
	# # Dataset

	# %%
	def tokenize(example, tokenizer, label2id, max_length):

	# rebuild text from tokens
	text = []
	labels = []

	for t, l, ws in zip(
	example["tokens"], example["provided_labels"], example["trailing_whitespace"]
	):
	text.append(t)
	labels.extend([l] * len(t))

	if ws:
	text.append(" ")
	labels.append("O")

	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:]

	labels = np.array(labels)

	text = "".join(text)
	token_labels = []

	for start_idx, end_idx in tokenized.offset_mapping:
	# CLS token
	if start_idx == 0 and end_idx == 0:
	token_labels.append(label2id["O"])
	continue

	# case when token starts with whitespace
	if text[start_idx].isspace():
	start_idx += 1

	token_labels.append(label2id[labels[start_idx]])

	length = len(tokenized.input_ids)

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

	# %%
	tokenizer = AutoTokenizer.from_pretrained(backbone)

	class PIIDataset(Dataset):
	def __init__(self, tokenizer, data, label2id, max_length):
	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']
	self.data[i] = tokenize(
	self.data[i],
	self.tokenizer,
	self.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=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,
	}

	ds = PIIDataset(tokenizer, data[:20], label2id, TRAINING_MAX_LENGTH)
	len(ds[9]['input_ids'])

	# %% [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),
	id2label=id2label,
	label2id=label2id,
	ignore_mismatched_sizes=True
	)

	# %% [markdown]
	# # Loss

	# %%
	weights = torch.ones(13) * pos_weight
	weights[12] = 1

	CE = nn.CrossEntropyLoss(weight=weights.to(device), ignore_index=-100)
	def criterion(logits, labels):
	return CE(logits.view(-1, len(all_labels)), labels.view(-1))

	# %% [markdown]
	# # Metric

	# %%
	def compute_metrics(p, all_labels, th=0.5):
	predictions, labels = p
	true_predictions = []
	true_labels = []
	for tmp, batch_l in zip(predictions, labels):
	batch_p = copy.deepcopy(tmp)
	batch_p[:, :, -1][batch_p[:, :, -1] < th] = -9999
	batch_p = np.argmax(batch_p, axis=2)
	true_predictions.extend([all_labels[p] for (p, l) in zip(batch_p.reshape(-1), batch_l.reshape(-1)) if l != -100])
	true_labels.extend([all_labels[l] for l in batch_l.reshape(-1) if l != -100])

	recall = recall_score([true_labels], [true_predictions])
	precision = precision_score([true_labels], [true_predictions])
	f1_score = (1 + 55) recall * precision / (55precision + recall)

	results = {
	'recall': recall,
	'precision': precision,
	'f1': f1_score
	}
	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()
	scheduler.step()

	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)
	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)
	loss.backward()
	optimizer.step()

	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 = []
	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_)

	running_loss.append(loss.item())
	predictions.append(outputs['logits'].detach().cpu().numpy())
	labels.append(labels_.detach().cpu().numpy())

	bar.set_description(f'valid loss: {np.mean(running_loss[-20:]):.4f}')

	return np.mean(running_loss), compute_metrics((predictions, labels), all_labels)

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

	log_file = os.path.join('/export/home/acs/stud/t/tudor.farcasanu/SSL_Assignment/single_model/model/logs', f'{exp_name}.txt')
	model_file = os.path.join('/export/home/acs/stud/t/tudor.farcasanu/SSL_Assignment/single_model/model', f'{exp_name}_fold{fold}_best.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

	valid_data = [x for x in data if x['document'] in valid_doc_ids]

	dataset_train = PIIDataset(tokenizer, train_data, label2id, TRAINING_MAX_LENGTH)
	dataset_valid = PIIDataset(tokenizer, valid_data, label2id, TRAINING_MAX_LENGTH)
	print(len(dataset_train), len(dataset_valid))

	loader_train = DataLoader(dataset_train, batch_size=batch_size, shuffle=True, num_workers=8, collate_fn=collate_fn)
	loader_valid = DataLoader(dataset_valid, batch_size=batch_size, shuffle=False, num_workers=8, collate_fn=collate_fn)


	# model = get_model()
	model = model.to(device)

	optimizer = AdamW(model.parameters(), lr=lr, weight_decay=wd)
	scaler = torch.cuda.amp.GradScaler() if use_amp else None
	from_epoch = 0
	metric_best = 0

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

	for epoch in range(1, n_epochs+1):
	if epoch < from_epoch + 1:
	print(logs[epoch-1])
	continue

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

	train_loss = train_one_epoch(model, loader_train, optimizer, scheduler, scaler)
	valid_loss, result = valid_one_epoch(model, loader_valid)
	metric = result["f1"]

	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: {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 ...')
	# if not DEBUG:
	torch.save(model.state_dict(), model_file)
	metric_best = metric

	# Save Last
	if not DEBUG:
	torch.save(
	{
	'epoch': epoch,
	'model_state_dict': model.state_dict(),
	'optimizer_state_dict': optimizer.state_dict(),
	'scaler_state_dict': scaler.state_dict() if scaler else None,
	'metric_best': metric_best,
	},
	model_file.replace('_best', '_last')
	)

	del model
	torch.cuda.empty_cache()
	gc.collect()

	# %%
	run(0,model)

	# %%