How to use pipeline for Custom token-classification model?
Author: pratikchhapolika
Pull request: no

**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'`



Comments:
You should use subclasses of `PreTrainedModel` (and `PretrainedConfig` if necessary). Here is the [documentation](https://huggingface.co/docs/transformers/create_a_model) on that and [here](https://huggingface.co/docs/transformers/add_new_pipeline) is an example of custom pipeline if you need it.

@sgugger , I am kind of in a similar situation, I looked up the custom `Pipeline` class, I wanted to understand if there is additional documentation of the methods in the class, I do see `_sanitary_parameter`, is it mostly to add preprocessing inputs ?,  for context I am trying to train a custom NER model as well.

@sgugger , never mind I was able to write my own postprocess and preprocess function and also inherit a of the implementation from task specific classes, thank you.

> @sgugger , never mind I was able to write my own postprocess and preprocess function and also inherit a of the implementation from task specific classes, thank you.

Hi @rajathpatel23 could you please upload your code/notebook for reference?

Hi @pratikchhapolika , I will try to upload the the notebook soon over the weekend

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.