Could the output label in turn with the text? - nlp

Question

I am new in NLP.
On Pytorch, I'm training BertForSequenceClassification for a multi-label task.
Is it possible for the output label to match the text?
nput is a text like : name, phon, address，output : label_name, label_phone, label_address.
In this case, input : phon, address, name，output : label_name, label_phone, label_address.
However, I'd like the output to look something like this: label_phone, label_address, label_name.
Is there anything I can change?
## Package
# PyTorch
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torch.utils.data.dataset import random_split
import torch.utils.data as data
# BERT Related Libraries
from transformers import BertTokenizer, BertForSequenceClassification
# Python
import pandas as pd
import numpy as np
import os
import time
import matplotlib.pyplot as plt
# ML Parameters
epoch = 8
device = 'cuda'
num_labels = 9
batch_size = 32
epsilon = 1e-8
learning_rate = 5e-5
## Define model
model = BertForSequenceClassification.from_pretrained('bert-base-uncased', num_labels=num_labels)
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
optimizer = optim.AdamW(model.parameters(), lr = learning_rate, eps = epsilon, weight_decay = 1e-2)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size = 1, gamma=0.1, last_epoch=-1)
criterion = nn.BCEWithLogitsLoss()
train_losses = []
train_acces = []
val_losses = []
val_acces = []
## Train with training set
def train(model, iterator, optimizer, criterion, device):
model.train()
train_loss = 0
correct = 0
total = 0
for batch_idx, (sentences, labels) in enumerate(iterator):
# tokenize the sentences
encoding = tokenizer(sentences, return_tensors='pt', padding=True, truncation=True)
input_ids = encoding['input_ids']
attention_mask = encoding['attention_mask']
# generate prediction
optimizer.zero_grad()
outputs = model(input_ids, attention_mask=attention_mask) # NOT USING INTERNAL CrossEntropyLoss
# compute gradients and update weights
loss = criterion(outputs.logits, labels) # BCEWithLogitsLoss has sigmoid
loss.backward()
optimizer.step()
# accumulate train loss
train_loss += loss.item()
# record processed data count
prob = outputs.logits.sigmoid()
total += (labels.size(0)*labels.size(1))
# take the index of the highest prob as prediction output
THRESHOLD = 0.7
prediction = prob.detach().clone()
prediction[prediction > THRESHOLD] = 1
prediction[prediction <= THRESHOLD] = 0
correct += prediction.eq(labels).sum().item()
train_acc = correct/total
print('train_loss: %f\t' % (train_loss))
print('train_acc: %f\t' % (train_acc))
return train_loss, train_acc
## Validate with testing set
def test(model, iterator, optimizer, criterion, device):
model.eval()
val_loss = 0
correct = 0
total = 0
with torch.no_grad():
for batch_idx, (sentences, labels) in enumerate(iterator):
# tokenize the sentences
encoding = tokenizer(sentences, return_tensors='pt', padding=True, truncation=True)
input_ids = encoding['input_ids']
attention_mask = encoding['attention_mask']
# generate prediction
outputs = model(input_ids, attention_mask=attention_mask)
loss = criterion(outputs.logits, labels)
val_loss += loss.item()
# record processed data count
prob = outputs.logits.sigmoid()
total += (labels.size(0)*labels.size(1))
# take the index of the highest prob as prediction output
THRESHOLD = 0.7
prediction = prob.detach().clone()
prediction[prediction > THRESHOLD] = 1
prediction[prediction <= THRESHOLD] = 0
correct += prediction.eq(labels).sum().item()
val_acc = correct/total
print('val_loss: %f\t' % (val_loss))
print('val_acc: %f\t' % (val_acc))
print('correct: %i , total: %i' % (correct, total))
return val_loss, val_acc
for e in range(epoch):
print("===== Epoch %i =====" % e)
print("Training started ...")
train_loss, train_acc = train(model, train_loader, optimizer, criterion, device)
train_losses.append(train_loss/len(train_loader))
train_acces.append(train_acc/len(train_loader))
# validation testing
print("Testing started ...")
val_loss, val_acc = test(model, test_loader, optimizer, criterion, device)
val_losses.append(val_loss/len(test_loader))
val_acces.append(val_acc/len(test_loader))
scheduler.step()
If you need more information,please let me konw.
Thanks for all.