2
$\begingroup$

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
$\endgroup$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Browse other questions tagged or ask your own question.