import os
import pickle
import numpy as np
import pandas as pd
from gensim.models import KeyedVectors
from collections import Counter
import torch
import torch.nn as nn
import torch.nn.functional as F
from sklearn.metrics import roc_auc_score, precision_recall_curve
import tqdm
from copy import deepcopy
import matplotlib.pyplot as plt
from transformers import DistilBertTokenizer, DistilBertModel
def get_roc_aucs(y, probas):
y_onehot = pd.get_dummies(y)
roc_auc_scores = []
if y_onehot.shape[1] > 2:
for i in range(y_onehot.shape[1]):
roc_auc_scores.append(roc_auc_score(y_onehot[i], probas[:, i]))
roc_auc_scores.append(roc_auc_score(y, probas, multi_class='ovo', average='macro'))
else:
roc_auc_scores.append(roc_auc_score(y, probas[:, 1]))
return roc_auc_scores
def get_max_f1_score(y, probas):
if probas.shape[1] != 2:
raise ValueError('Expected probabilities for 2 classes would be given')
y_onehot = pd.get_dummies(y)
f1_score = []
threshold = []
p, r, t = precision_recall_curve(y, probas[:, 1])
f1_scores = 2 * p * r / (p + r + 0.001)
threshold.append(t[np.argmax(f1_scores)])
f1_score.append(np.max(f1_scores))
return f1_score, threshold
class RNN(nn.Module):
def __init__(self, vectors, n_of_words, n_of_classes, num_layers, bidirectional):
dim = vectors.shape[1]
d = 2 if bidirectional else 1
super().__init__()
self.emb = nn.Embedding(n_of_words, dim)
self.emb.load_state_dict({'weight': torch.tensor(vectors)})
self.emb.weight.requires_grad = False
self.gru = nn.GRU(input_size=dim, hidden_size=dim, batch_first=True,
num_layers=num_layers, bidirectional=bidirectional)
self.linear = nn.Linear(dim * num_layers * d, n_of_classes)
def forward(self, batch):
emb = self.emb(batch)
_, last_state = self.gru(emb)
last_state = torch.permute(last_state, (1, 0, 2)).reshape(1, batch.shape[0], -1).squeeze()
out = self.linear(last_state.squeeze())
if len(out.size()) == 1:
out = out.unsqueeze(0)
return out
class DistilBERTClass(torch.nn.Module):
def __init__(self, n_classes):
super().__init__()
self.l1 = DistilBertModel.from_pretrained('DeepPavlov/distilrubert-small-cased-conversational')
self.linear = torch.nn.Linear(768, n_classes)
def forward(self, input_ids, attention_mask, token_type_ids):
output_1 = self.l1(input_ids=input_ids, attention_mask=attention_mask)
hidden_state = output_1[0]
pooler = hidden_state[:, 0]
output = self.linear(pooler)
return output
class BaseClassifier:
def __init__(self, batch_size=16, epochs=100):
self.batch_size = batch_size
self.epochs = epochs
self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
def preprocess_with_random_initialization(self, train_tokens):
self.pad_idx = 0
self.unk_idx = 1
set_of_words = set()
for tokens_string in train_tokens:
set_of_words.update(tokens_string)
self.idx_to_word = ['PADDING', 'UNK'] + list(set_of_words)
self.word_to_idx = {key: i for i, key in enumerate(self.idx_to_word)}
self.amount_of_words = len(self.idx_to_word)
self.vectors = np.zeros((len(self.idx_to_word), 300))
self.vectors[0, :] = np.zeros(300)
self.vectors[1:len(self.idx_to_word), :] = (np.random.rand(len(self.idx_to_word) - 1, 300) - 0.5) / 300
def preprocess(self, vectors_file_path):
self.emb = KeyedVectors.load_word2vec_format(vectors_file_path)
self.pad_idx = 0
self.unk_idx = 1
self.idx_to_word = ['PADDING', 'UNK'] + list(self.emb.index_to_key)
self.word_to_idx = {key: i for i, key in enumerate(self.idx_to_word)}
self.amount_of_words = len(self.idx_to_word)
self.vectors = np.zeros((len(self.idx_to_word), 300))
self.vectors[0, :] = np.zeros(300)
self.vectors[1, :] = (np.random.rand(300) - 0.5) / 300
for i in range(2, len(self.idx_to_word)):
self.vectors[i, :] = self.emb.get_vector(self.idx_to_word[i])
def fit(self, train_tokens, y_train, test_tokens=None, y_test=None,
reinitialize=True, stop_epochs=None, show_logs=False):
if reinitialize:
self.n_of_classes = y_train.nunique()
self.initialize_nnet()
self.print_test = test_tokens and y_test
self.stop_epochs = stop_epochs
train_scores = []
self.train_scores_mean = []
self.test_scores = []
self.test_aucs = []
self.test_f1 = []
criterion = nn.CrossEntropyLoss()
for epoch in tqdm.tqdm(range(self.epochs)):
self.epoch = epoch
self.nnet.train()
train_batches = self.batch_generator(train_tokens, y_train)
test_batches = self.batch_generator(test_tokens, y_test)
for i, batch in tqdm.tqdm(
enumerate(train_batches),
total=len(train_tokens) // self.batch_size
):
pred = self.nnet(batch['tokens'])
loss = criterion(pred, batch['labels'])
self.optimizer.zero_grad()
loss.backward()
self.optimizer.step()
if show_logs and i % 400 == 0:
train_score = criterion(self.nnet(batch['tokens']), batch['labels'])
print(train_score.item())
train_scores.append(train_score.item())
if show_logs:
self.train_scores_mean.append(sum(train_scores) / len(train_scores))
train_scores = []
if self.print_test:
test_pred_prob = torch.tensor([], device='cpu')
with torch.no_grad():
self.nnet.eval()
for batch in test_batches:
test_batch_pred_prob = self.nnet(batch['tokens'])
test_batch_pred_prob_cpu = test_batch_pred_prob.to('cpu')
test_pred_prob = torch.cat((test_pred_prob, test_batch_pred_prob_cpu), 0)
test_score = criterion(test_pred_prob, torch.tensor(y_test.values, device='cpu'))
self.test_scores.append(test_score.item())
test_pred_probas = F.softmax(test_pred_prob).detach().cpu().numpy()
self.test_aucs.append(get_roc_aucs(y_test, test_pred_probas))
self.test_f1.append(get_max_f1_score(y_test, test_pred_probas)[0])
self.print_metrics()
if self.early_stopping_check():
break
def count_tokens(self, tokens):
self.words_counter = Counter()
self.amount_of_tokens = 0
for s in tokens:
self.words_counter.update(s)
self.amount_of_tokens += len(s)
def index_tokens(self, tokens_string):
return [self.word_to_idx.get(token, self.unk_idx) for token in tokens_string]
def fill_with_pads(self, tokens):
tokens = deepcopy(tokens)
max_len = 0
for tokens_string in tokens:
max_len = max(max_len, len(tokens_string))
for tokens_string in tokens:
for i in range(len(tokens_string), max_len):
tokens_string.append(self.pad_idx)
return tokens
def as_matrix(self, tokens):
tokens = deepcopy(tokens)
for j, s in enumerate(tokens):
tokens[j] = self.index_tokens(s)
tokens = self.fill_with_pads(tokens)
return tokens
def batch_generator(self, tokens, labels=None):
for i in range(0, len(tokens), self.batch_size):
batch_tokens = tokens[i: i + self.batch_size]
if labels:
batch_labels = torch.tensor(labels.values[i: i + self.batch_size],
dtype=torch.long,
device=self.device)
else:
batch_labels = None
batch_tokens_idx = torch.tensor(self.as_matrix(batch_tokens),
dtype=torch.int,
device=self.device)
if len(batch_tokens_idx.size()) == 1:
batch_tokens_idx = torch.unsqueeze(batch_tokens_idx, 0)
batch = {
'tokens': batch_tokens_idx,
'labels': batch_labels
}
yield batch
def print_metrics(self, print_test=True):
if self.print_test:
print(f'epoch {self.epoch}/{self.epochs}')
print('auc', self.test_aucs[-1])
print('score', self.test_scores[-1])
print('f1 score', self.test_f1[-1])
legend_labels = []
if self.n_of_classes > 2:
for i in range(self.n_of_classes):
legend_labels.append(f'Class {i}')
legend_labels.append('General')
plt.figure(figsize=(5, 15))
plt.clf()
plt.subplot(3, 1, 1)
plt.plot(np.arange(1, self.epoch + 2), self.test_aucs)
plt.grid()
plt.title('Test ROC AUC')
plt.xlabel('Num. of epochs')
plt.ylabel('ROC AUC')
plt.legend(legend_labels)
plt.subplot(3, 1, 2)
plt.plot(np.arange(1, self.epoch + 2), self.test_f1)
plt.grid()
plt.title('Test F1-score')
plt.xlabel('Num. of epochs')
plt.ylabel('F1-score')
plt.legend(legend_labels)
plt.subplot(3, 1, 3)
plt.plot(np.arange(1, self.epoch + 2), self.train_scores_mean, label='Train loss')
plt.plot(np.arange(1, self.epoch + 2), self.test_scores, label='Test loss')
plt.title('Loss')
plt.xlabel('Num. of epochs')
plt.ylabel('Loss')
plt.legend()
plt.grid()
plt.draw()
else:
plt.figure(figsize=(5, 15))
plt.plot(np.arange(1, self.epoch + 2), self.train_scores_mean, label='Train loss')
plt.title('Loss')
plt.xlabel('Num. of epochs')
plt.ylabel('Loss')
plt.legend()
plt.grid()
plt.show()
def early_stopping_check(self):
if self.stop_epochs is None or self.stop_epochs >= len(self.test_scores):
return False
else:
print(self.test_scores)
first_score = np.array(self.test_scores)[-self.stop_epochs - 1]
last_scores = np.array(self.test_scores)[-self.stop_epochs:]
return np.all(last_scores >= first_score)
def predict_proba(self, tokens, labels):
batches = self.batch_generator(tokens, labels)
pred_probas = torch.tensor([], device=self.device)
with torch.no_grad():
self.nnet.eval()
for batch in batches:
batch_prob = self.nnet(batch['tokens'])
pred_probas = torch.cat((pred_probas, batch_prob))
return F.softmax(pred_probas).detach().cpu().numpy()
class RNNClassifier(BaseClassifier):
def __init__(self, batch_size=16, epochs=100,
num_layers=1, bidirectional=False):
self.batch_size = batch_size
self.epochs = epochs
self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
self.num_layers = num_layers
self.bidirectional = bidirectional
def initialize_nnet(self):
self.nnet = RNN(self.vectors, self.amount_of_words,
n_of_classes=self.n_of_classes,
num_layers=self.num_layers,
bidirectional=self.bidirectional).to(self.device)
self.optimizer = torch.optim.Adam(self.nnet.parameters())
def save_model(self, filepath):
with open(filepath, 'wb') as file:
torch.save(self.nnet.state_dict(), file)
def load_model(self, filepath, amount_of_words):
self.amount_of_words = amount_of_words
self.vectors = np.zeros((amount_of_words, 300))
self.n_of_classes = 2
self.nnet = RNN(self.vectors, self.amount_of_words,
n_of_classes=self.n_of_classes,
num_layers=self.num_layers,
bidirectional=self.bidirectional).to(self.device)
self.nnet.load_state_dict(torch.load(filepath, map_location=self.device))
class DBERTClassifier(BaseClassifier):
def __init__(self, batch_size=16, epochs=100):
self.batch_size = batch_size
self.epochs = epochs
self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
def initialize_nnet(self):
self.nnet = DistilBERTClass(self.n_of_classes).to(self.device)
self.optimizer = torch.optim.Adam(self.nnet.parameters(), lr=2e-6)
# 'DeepPavlov/rubert-base-cased' 'DeepPavlov/distilrubert-small-cased-conversational',
self.tokenizer = DistilBertTokenizer.from_pretrained('DeepPavlov/distilrubert-small-cased-conversational',
do_lower_case=True)
def batch_generator(self, tokens, labels=None):
for i in range(0, len(tokens), self.batch_size):
batch_tokens = tokens[i: i + self.batch_size]
batch_tokens = [' '.join(s) for s in batch_tokens]
if labels:
batch_labels = torch.tensor(labels.values[i: i + self.batch_size],
dtype=torch.long,
device=self.device)
else:
batch_labels = None
if len(batch_tokens) == 1:
inputs = self.tokenizer.encode_plus(
batch_tokens,
None,
add_special_tokens=True,
max_length=512,
truncation=True,
pad_to_max_length=True,
return_token_type_ids=True
)
else:
inputs = self.tokenizer.batch_encode_plus(
batch_tokens,
add_special_tokens=True,
max_length=512,
truncation=True,
pad_to_max_length=True,
return_token_type_ids=True
)
batch_token_ids = torch.tensor(inputs['input_ids'], device=self.device, dtype=torch.long)
batch_mask = torch.tensor(inputs['attention_mask'], device=self.device, dtype=torch.long)
batch_token_type_ids = torch.tensor(inputs["token_type_ids"], device=self.device, dtype=torch.long)
if len(batch_tokens) == 1:
batch_token_ids = batch_token_ids.unsqueeze(0)
batch_mask = batch_mask.unsqueeze(0)
batch_token_type_ids = batch_token_type_ids.unsqueeze(0)
batch = {
'tokens': batch_token_ids,
'mask': batch_mask,
'token_type_ids': batch_token_type_ids,
'labels': batch_labels
}
yield batch
def fit(self, train_tokens, y_train, test_tokens=None, y_test=None,
reinitialize=True, stop_epochs=None, show_logs=False):
if reinitialize:
self.n_of_classes = y_train.nunique()
self.initialize_nnet()
self.stop_epochs = stop_epochs
self.print_test = test_tokens and y_test
train_scores = []
self.train_scores_mean = []
self.test_scores = []
self.test_aucs = []
self.test_f1 = []
criterion = nn.CrossEntropyLoss()
for epoch in tqdm.tqdm(range(self.epochs)):
self.epoch = epoch
self.nnet.train()
train_batches = self.batch_generator(train_tokens, y_train)
test_batches = self.batch_generator(test_tokens, y_test)
for i, batch in tqdm.tqdm(
enumerate(train_batches),
total=len(train_tokens) // self.batch_size
):
pred = self.nnet(batch['tokens'], batch['mask'], batch['token_type_ids'])
loss = criterion(pred, batch['labels'])
self.optimizer.zero_grad()
loss.backward()
self.optimizer.step()
if show_logs and i % 400 == 0:
train_score = criterion(self.nnet(batch['tokens'], batch['mask'], batch['token_type_ids']),
batch['labels'])
print(train_score.item())
train_scores.append(train_score.item())
if show_logs:
self.train_scores_mean.append(sum(train_scores) / len(train_scores))
train_scores = []
if self.print_test:
test_pred_prob = torch.tensor([], device='cpu')
with torch.no_grad():
self.nnet.eval()
for batch in test_batches:
test_batch_pred_prob = self.nnet(batch['tokens'], batch['mask'], batch['token_type_ids'])
test_batch_pred_prob_cpu = test_batch_pred_prob.to('cpu')
test_pred_prob = torch.cat((test_pred_prob, test_batch_pred_prob_cpu), 0)
test_score = criterion(test_pred_prob, torch.tensor(y_test.values, device='cpu'))
self.test_scores.append(test_score.item())
test_pred_probas = F.softmax(test_pred_prob).detach().cpu().numpy()
self.test_aucs.append(get_roc_aucs(y_test, test_pred_probas))
self.test_f1.append(get_max_f1_score(y_test, test_pred_probas)[0])
self.print_metrics()
if self.early_stopping_check():
break
def predict_proba(self, tokens, labels):
batches = self.batch_generator(tokens, labels)
pred_probas = torch.tensor([], device=self.device)
with torch.no_grad():
self.nnet.eval()
for batch in batches:
batch_prob = self.nnet(batch['tokens'], batch['mask'],
batch['token_type_ids'])
pred_probas = torch.cat((pred_probas, batch_prob))
return F.softmax(pred_probas).detach().cpu().numpy()
def predict(self, tokens, labels):
return np.argmax(self.predict_proba(tokens, labels), axis=1)
def save_model(self, filepath):
with open(filepath, 'wb') as file:
torch.save(self.nnet.state_dict(), file)
def load_model(self, filepath):
self.n_of_classes = 2
self.nnet = DistilBERTClass(self.n_of_classes).to(self.device)
self.optimizer = torch.optim.Adam(self.nnet.parameters(), lr=2e-6)
self.tokenizer = DistilBertTokenizer.from_pretrained(
'DeepPavlov/distilrubert-small-cased-conversational',
do_lower_case=True
)
self.nnet.load_state_dict(torch.load(filepath, map_location=self.device))
class AdClassifier:
def __init__(self, weights_folder, dictionary_path):
self.batch_size = 16
self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
self.pad_idx = 0
self.unk_idx = 1
with open(dictionary_path, 'rb') as file:
self.word_to_idx = pickle.load(file)
self.tokenizer = DistilBertTokenizer.from_pretrained(
'DeepPavlov/distilrubert-small-cased-conversational',
do_lower_case=True
)
nationality_nn_path = os.path.join(weights_folder, 'model_nationality.pt')
families_nn_path = os.path.join(weights_folder, 'model_families.pt')
sex_nn_path = os.path.join(weights_folder, 'model_sex.pt')
limit_nn_path = os.path.join(weights_folder, 'model_limit.pt')
self.nationality_clf = DBERTClassifier()
self.nationality_clf.load_model(nationality_nn_path)
self.families_clf = DBERTClassifier()
self.families_clf.load_model(families_nn_path)
self.sex_clf = DBERTClassifier()
self.sex_clf.load_model(sex_nn_path)
self.limit_clf = RNNClassifier(bidirectional=True)
self.limit_clf.load_model(limit_nn_path, amount_of_words=len(self.word_to_idx))
def index_tokens(self, tokens_string):
return [self.word_to_idx.get(token, self.unk_idx) for token in tokens_string]
def fill_with_pads(self, tokens):
tokens = deepcopy(tokens)
max_len = 0
for tokens_string in tokens:
max_len = max(max_len, len(tokens_string))
for tokens_string in tokens:
for i in range(len(tokens_string), max_len):
tokens_string.append(self.pad_idx)
return tokens
def as_matrix(self, tokens):
tokens = deepcopy(tokens)
for j, s in enumerate(tokens):
tokens[j] = self.index_tokens(s)
tokens = self.fill_with_pads(tokens)
return tokens
def batch_generator(self, tokens):
for i in range(0, len(tokens), self.batch_size):
batch_tokens = tokens[i: i + self.batch_size]
batch_tokens = [' '.join(s) for s in batch_tokens]
inputs = self.tokenizer.batch_encode_plus(
batch_tokens,
add_special_tokens=True,
max_length=512,
truncation=True,
pad_to_max_length=True,
return_token_type_ids=True
)
batch_token_ids = torch.tensor(inputs['input_ids'], device=self.device, dtype=torch.long)
batch_mask = torch.tensor(inputs['attention_mask'], device=self.device, dtype=torch.long)
batch_token_type_ids = torch.tensor(inputs['token_type_ids'], device=self.device, dtype=torch.long)
batch_tokens_rnn = tokens[i: i + self.batch_size]
batch_tokens_rnn_ids = torch.tensor(self.as_matrix(batch_tokens_rnn),
dtype=torch.int,
device=self.device)
batch = {
'tokens': batch_token_ids,
'mask': batch_mask,
'token_type_ids': batch_token_type_ids,
'tokens_rnn': batch_tokens_rnn_ids
}
yield batch
def predict_probas(self, tokens):
batches = self.batch_generator(tokens)
pred_probas = {'nationality': torch.tensor([], device=self.device),
'families': torch.tensor([], device=self.device),
'sex': torch.tensor([], device=self.device),
'limit': torch.tensor([], device=self.device)}
batch_probas = dict()
with torch.no_grad():
self.nationality_clf.nnet.eval()
self.families_clf.nnet.eval()
self.sex_clf.nnet.eval()
self.limit_clf.nnet.eval()
for batch in batches:
batch_probas['nationality'] = self.nationality_clf.nnet(batch['tokens'], batch['mask'],
batch['token_type_ids'])
batch_probas['families'] = self.families_clf.nnet(batch['tokens'], batch['mask'],
batch['token_type_ids'])
batch_probas['sex'] = self.sex_clf.nnet(batch['tokens'], batch['mask'],
batch['token_type_ids'])
batch_probas['limit'] = self.limit_clf.nnet(batch['tokens_rnn'])
for batch_prob_label in batch_probas:
pred_probas[batch_prob_label] = torch.cat((pred_probas[batch_prob_label],
batch_probas[batch_prob_label]))
for pred_prob_label in pred_probas:
pred_probas[pred_prob_label] = F.softmax(pred_probas[pred_prob_label]).\
detach().cpu().numpy()
return pred_probas
def predict_labels(self, tokens):
predicted_probas = self.predict_probas(tokens)
predicted_labels = dict()
thresholds = {
'nationality': 0.75,
'families': 0.7,
'sex': 0.25,
'limit': 0.42
}
for label in predicted_probas:
predicted_labels[label] = predicted_probas[label][:, 1] >= thresholds[label]
return predicted_labels
def save_model(self, filepath):
with open(filepath, 'wb') as file:
torch.save(self, file)