|
import torch |
|
import numpy as np |
|
import pandas as pd |
|
|
|
from sklearn.model_selection import StratifiedKFold |
|
from sklearn.metrics import classification_report, confusion_matrix |
|
|
|
import torch.nn as nn |
|
from torch.optim import AdamW |
|
from torch.utils.data import Dataset, DataLoader |
|
|
|
from transformers import get_linear_schedule_with_warmup, AutoTokenizer, AutoModel, logging |
|
|
|
import warnings |
|
import time |
|
import pickle |
|
warnings.filterwarnings("ignore") |
|
|
|
logging.set_verbosity_error() |
|
|
|
|
|
def seed_everything(seed_value): |
|
np.random.seed(seed_value) |
|
torch.manual_seed(seed_value) |
|
|
|
if torch.cuda.is_available(): |
|
torch.cuda.manual_seed(seed_value) |
|
torch.cuda.manual_seed_all(seed_value) |
|
torch.backends.cudnn.deterministic = True |
|
torch.backends.cudnn.benchmark = True |
|
|
|
seed_everything(86) |
|
|
|
model_name = "bluenguyen/longformer-phobert-base-4096" |
|
max_len = 512 |
|
n_classes = 13 |
|
tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=False) |
|
|
|
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu') |
|
EPOCHS = 5 |
|
N_SPLITS = 5 |
|
|
|
TRAIN_PATH = "data/train_data_162k.json" |
|
TEST_PATH = "data/test_data_162k.json" |
|
VAL_PATH = "data/val_data_162k.json" |
|
|
|
|
|
def get_data(path): |
|
df = pd.read_json(path, lines=True) |
|
return df |
|
|
|
|
|
train_df = get_data(TRAIN_PATH) |
|
test_df = get_data(TEST_PATH) |
|
valid_df = get_data(VAL_PATH) |
|
|
|
|
|
train_df = pd.concat([train_df, valid_df], ignore_index=True) |
|
|
|
|
|
skf = StratifiedKFold(n_splits=N_SPLITS) |
|
for fold, (_, val_) in enumerate(skf.split(X=train_df, y=train_df.category)): |
|
train_df.loc[val_, "kfold"] = fold |
|
|
|
class NewsDataset(Dataset): |
|
def __init__(self, df, tokenizer, max_len): |
|
self.df = df |
|
self.max_len = max_len |
|
self.tokenizer = tokenizer |
|
|
|
def __len__(self): |
|
return len(self.df) |
|
|
|
def __getitem__(self, index): |
|
""" |
|
To customize dataset, inherit from Dataset class and implement |
|
__len__ & __getitem__ |
|
__getitem__ should return |
|
data: |
|
input_ids |
|
attention_masks |
|
text |
|
targets |
|
""" |
|
row = self.df.iloc[index] |
|
text, label = self.get_input_data(row) |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
encoding = self.tokenizer.encode_plus( |
|
text, |
|
truncation=True, |
|
add_special_tokens=True, |
|
max_length=self.max_len, |
|
padding='max_length', |
|
return_attention_mask=True, |
|
return_token_type_ids=False, |
|
return_tensors='pt', |
|
) |
|
|
|
return { |
|
'text': text, |
|
'input_ids': encoding['input_ids'].flatten(), |
|
'attention_masks': encoding['attention_mask'].flatten(), |
|
'targets': torch.tensor(label, dtype=torch.long), |
|
} |
|
|
|
|
|
def labelencoder(self, text): |
|
label_map = { |
|
'Cong nghe': 0, 'Doi song': 1, 'Giai tri': 2, 'Giao duc': 3, 'Khoa hoc': 4, |
|
'Kinh te': 5, 'Nha dat': 6, 'Phap luat': 7, 'The gioi': 8, 'The thao': 9, |
|
'Van hoa': 10, 'Xa hoi': 11, 'Xe co': 12 |
|
} |
|
return label_map.get(text, -1) |
|
|
|
def get_input_data(self, row): |
|
text = row['processed_content'] |
|
label = self.labelencoder(row['category']) |
|
return text, label |
|
|
|
class NewsClassifier(nn.Module): |
|
def __init__(self, n_classes, model_name): |
|
super(NewsClassifier, self).__init__() |
|
|
|
self.bert = AutoModel.from_pretrained(model_name) |
|
|
|
self.drop = nn.Dropout(p=0.3) |
|
|
|
self.fc = nn.Linear(self.bert.config.hidden_size, n_classes) |
|
|
|
nn.init.normal_(self.fc.weight, std=0.02) |
|
nn.init.normal_(self.fc.bias, 0) |
|
|
|
def forward(self, input_ids, attention_mask): |
|
|
|
last_hidden_state, output = self.bert( |
|
input_ids=input_ids, |
|
attention_mask=attention_mask, |
|
return_dict=False |
|
) |
|
|
|
x = self.drop(output) |
|
|
|
x = self.fc(x) |
|
return x |
|
|
|
def prepare_loaders(df, fold): |
|
df_train = df[df.kfold != fold].reset_index(drop=True) |
|
df_valid = df[df.kfold == fold].reset_index(drop=True) |
|
|
|
train_dataset = NewsDataset(df_train, tokenizer, max_len) |
|
valid_dataset = NewsDataset(df_valid, tokenizer, max_len) |
|
|
|
train_loader = DataLoader(train_dataset, batch_size=16, shuffle=True, num_workers=2) |
|
valid_loader = DataLoader(valid_dataset, batch_size=16, shuffle=True, num_workers=2) |
|
|
|
return train_loader, valid_loader |
|
|
|
|
|
def train(model, criterion, optimizer, train_loader, lr_scheduler): |
|
model.train() |
|
losses = [] |
|
correct = 0 |
|
|
|
|
|
for batch_idx, data in enumerate(train_loader): |
|
input_ids = data['input_ids'].to(device) |
|
attention_mask = data['attention_masks'].to(device) |
|
targets = data['targets'].to(device) |
|
|
|
optimizer.zero_grad() |
|
outputs = model( |
|
input_ids=input_ids, |
|
attention_mask=attention_mask |
|
) |
|
|
|
loss = criterion(outputs, targets) |
|
_, pred = torch.max(outputs, dim=1) |
|
|
|
correct += torch.sum(pred == targets) |
|
losses.append(loss.item()) |
|
loss.backward() |
|
nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0) |
|
optimizer.step() |
|
lr_scheduler.step() |
|
|
|
|
|
if batch_idx % 1000 == 0: |
|
print(f'Batch {batch_idx}/{len(train_loader)} - Loss: {loss.item():.4f}, Accuracy: {correct.double() / ((batch_idx + 1) * train_loader.batch_size):.4f}') |
|
|
|
train_accuracy = correct.double() / len(train_loader.dataset) |
|
avg_loss = np.mean(losses) |
|
print(f'Train Accuracy: {train_accuracy:.4f} Loss: {avg_loss:.4f}') |
|
|
|
|
|
def eval(model, criterion, valid_loader, test_loader=None): |
|
model.eval() |
|
losses = [] |
|
correct = 0 |
|
|
|
with torch.no_grad(): |
|
data_loader = test_loader if test_loader else valid_loader |
|
for batch_idx, data in enumerate(data_loader): |
|
input_ids = data['input_ids'].to(device) |
|
attention_mask = data['attention_masks'].to(device) |
|
targets = data['targets'].to(device) |
|
|
|
outputs = model( |
|
input_ids=input_ids, |
|
attention_mask=attention_mask |
|
) |
|
|
|
loss = criterion(outputs, targets) |
|
_, pred = torch.max(outputs, dim=1) |
|
|
|
correct += torch.sum(pred == targets) |
|
losses.append(loss.item()) |
|
|
|
dataset_size = len(test_loader.dataset) if test_loader else len(valid_loader.dataset) |
|
accuracy = correct.double() / dataset_size |
|
avg_loss = np.mean(losses) |
|
|
|
|
|
if test_loader: |
|
print(f'Test Accuracy: {accuracy:.4f} Loss: {avg_loss:.4f}') |
|
else: |
|
print(f'Valid Accuracy: {accuracy:.4f} Loss: {avg_loss:.4f}') |
|
|
|
return accuracy |
|
|
|
total_start_time = time.time() |
|
|
|
|
|
for fold in range(skf.n_splits): |
|
print(f'----------- Fold: {fold + 1} ------------------') |
|
train_loader, valid_loader = prepare_loaders(train_df, fold=fold) |
|
model = NewsClassifier(n_classes=13).to(device) |
|
criterion = nn.CrossEntropyLoss() |
|
optimizer = AdamW(model.parameters(), lr=2e-5) |
|
|
|
lr_scheduler = get_linear_schedule_with_warmup( |
|
optimizer, |
|
num_warmup_steps=0, |
|
num_training_steps=len(train_loader) * EPOCHS |
|
) |
|
best_acc = 0 |
|
|
|
for epoch in range(EPOCHS): |
|
print(f'Epoch {epoch + 1}/{EPOCHS}') |
|
print('-' * 30) |
|
|
|
train(model, criterion, optimizer, train_loader, lr_scheduler) |
|
val_acc = eval(model, criterion, valid_loader) |
|
|
|
if val_acc > best_acc: |
|
torch.save(model.state_dict(), f'phobert_fold{fold + 1}.pth') |
|
best_acc = val_acc |
|
print(f'Best Accuracy for Fold {fold + 1}: {best_acc:.4f}') |
|
print() |
|
print(f'Finished Fold {fold + 1} with Best Accuracy: {best_acc:.4f}') |
|
print('--------------------------------------') |
|
|
|
|
|
total_end_time = time.time() |
|
|
|
total_duration = total_end_time - total_start_time |
|
print(f'Total training time: {total_duration:.2f} seconds') |