How to use pipeline for Custom token-classification model

🤗Transformers
1

Model description

I add simple custom pytorch-crf layer on top of TokenClassification model. It will make the model more robust.

I train the model successfully but when I test the mode. The folder doesn’t have config.json file inside it. So the pipeline function gives error as

Error:AttributeError: 'BERT_CRF' object has no attribute 'config'

CODE

class BERT_CRF(nn.Module):
    
  
   def __init__(self, bert_model, num_labels):
        
        super(BERT_CRF, self).__init__()
        self.bert = bert_model
        self.dropout = nn.Dropout(0.25)
        self.classifier = nn.Linear(768, num_labels)
        self.crf = CRF(num_labels, batch_first = True)
        
    
    def forward(self, input_ids, attention_mask,  labels=None, token_type_ids=None):
        
        outputs = self.bert(input_ids, attention_mask=attention_mask)
        sequence_output = torch.stack((outputs[1][-1], outputs[1][-2], outputs[1][-3], outputs[1][-4])).mean(dim=0)
        sequence_output = self.dropout(sequence_output)
        emission = self.classifier(sequence_output) # [32,256,17]
        if labels is not None:
            labels=labels.reshape(attention_mask.size()[0],attention_mask.size()[1])
            loss = -self.crf(log_soft(emission, 2), labels, mask=attention_mask.type(torch.uint8), reduction='mean')
            prediction = self.crf.decode(emission, mask=attention_mask.type(torch.uint8))
            return [loss, prediction]
                
        else:
            prediction = self.crf.decode(emission, mask=attention_mask.type(torch.uint8))
            return prediction

tokenizer = AutoTokenizer.from_pretrained("fine-tuned_model",model_max_length=256)

bert_model = BertForTokenClassification.from_pretrained('spanbert_base',id2label=id2label,label2id=label2id)
bert_model.config.output_hidden_states=True
model = BERT_CRF(bert_model, num_labels=21)
model.load_state_dict(torch.load("fine-tuned_model/pytorch_model.bin"))
model.eval()

token_classifier = pipeline("token-classification", model=model, aggregation_strategy="max",tokenizer=tokenizer,grouped_entities=True)

AttributeError: 'BERT_CRF' object has no attribute 'config'