# Imbalance classes in Named Entity Recognition

I am currently working on a NER problem which attempts to extract 2 entities - place-of-interest(POI) and street from an address string in the Indonesian language.

I used IndoBert (available here) and attached FC layers onto the BERT model, utilising cross-entropy loss to predict the class each word belongs to. However a typical sample sentence label in my dataset looks like this [1, 4, 5, 5, 1, 1, 2, 3, 3, 3, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]

Where 2-5 represents the POI and street class, 0 are the pad tokens and 1 codes for O tags(BIO tagging). However the model ends up constantly predicting majority class for all words.I then used ignore index and weights in my loss function which looks like this:

WEIGHTS = torch.Tensor([0.05, 0.2, 1, 1, 1, 1])
def loss(predicted, target):
predicted = torch.rot90(predicted, -1, (1,2))
criterion = nn.CrossEntropyLoss(weight = WEIGHTS,ignore_index=0, reduction='mean')
return criterion(predicted, target)


However while this solves the problem above. The model also does not learn and the losses simply fluctuates around the same level. I am therefore wondering if there is a way to adjust for this class imbalance or stop the model from predicting the classes for [PAD] tokens.

This is the code for my model:

class aem(nn.Module):
def __init__(self, no_class):
super().__init__()
self.bert = AutoModel.from_pretrained("sarahlintang/IndoBERT")
self.drop1 = nn.Dropout(p=0.1)
self.l1 = nn.Linear(self.bert.config.hidden_size, no_class)
self.out = nn.GELU()

def forward(self, inputs, attn):
hidden = self.bert(inputs, token_type_ids=None, attention_mask=attn, return_dict = True)
L1out = self.out(self.l1(self.drop1(hidden[0])))
return L1out