# Custom class for token classification

**URL:** <https://discuss.huggingface.co/t/custom-class-for-token-classification/104339>\
**Category:** 🤗Transformers\
**Created:** [August 28, 2024, 9:00am UTC](https://discuss.huggingface.co/t/custom-class-for-token-classification/104339 "2024-08-28T09:00:13Z")\
**Posts on this page:** 2\
**Page:** 1

<div class="post-metadata">

**Author:** ![rychlik](https://avatars.discourse-cdn.com/v4/letter/r/7c8e57/32.png) [@rychlik](https://discuss.huggingface.co/u/rychlik)\
**Post date:** [August 28, 2024, 9:00am UTC](https://discuss.huggingface.co/t/custom-class-for-token-classification/104339/1 "2024-08-28T09:00:13Z")

</div>

I am trying to define my own token classification model class in PyTorch  
so that it can be used in a similar way to RobertaForTokenClassification  
class from tranformers and that the trained model can be saved using save\_pretrained  
and reloaded using from\_pretrained.

I thought it would be easiest to slightly modify the code of RobertForTokenClassification from tranformers.  
The definition of MyModelForTokenClassification class is as follows:

class MyConfig(RobertaConfig):  
def **init** (self, \*\*kwargs):  
self.additional\_parameter = …  
super(). **init** (\*\*kwargs)

class MyModelForTokenClassification(RobertaPreTrainedModel):  
config\_class = MyConfig

```
def __init__ (self, config, additional_data = None):
    super(). __init__ (config)
    self.additional_data = additional_data
    self.additional_parameter = config.additional_parameter
    self.num_labels = config.num_labels
    self.model = RobertaModel(config, add_pooling_layer=False)
    classifier_dropout = (
        config.classifier_dropout if config.classifier_dropout is not None else config.hidden_dropout_prob
    )
    self.dropout = torch.nn.Dropout(classifier_dropout)
    hidden_size = config.hidden_size
    if self.additional_data is not None:
    	hidden_size = .....
    self.classifier = torch.nn.Linear(hidden_size, config.num_labels)
    self.post_init()

def forward(.....):
	......

```

In order to use AutoModelForTokenClassification, I added these 2 lines of code:

AutoConfig.register(“roberta”, MyConfig, exist\_ok=True)  
AutoModelForTokenClassification.register(MyConfig, MyModelForTokenClassification, exist\_ok=True)

When creating a model for training:

model = AutoModelForTokenClassification.from\_pretrained(“roberta-base”,  
num\_labels=3,  
id2label=id2label,  
label2id=label2id,  
additional\_parameter=…,  
additional\_data=…)

this message appears:  
Some weights of MyModelForTokenClassification were not initialized from the model checkpoint  
at roberta-base and are newly initialized:  
[‘roberta.model.encoder.layer.9.attention.self.query.weight’,  
‘roberta.model.encoder.layer.9.attention.output.LayerNorm.weight’,  
‘roberta.model.encoder.layer.11.output.LayerNorm.bias’,  
‘roberta.model.encoder.layer.3.output.LayerNorm.weight’,  
‘roberta.model.encoder.layer.4.attention.self.value.bias’,  
‘roberta.model.encoder.layer.5.intermediate.dense.bias’,  
‘roberta.model.encoder.layer.10.output.dense.weight’,  
‘roberta.model.encoder.layer.5.attention.output.dense.weight’,  
…]

It shows that the model has not been initialised with the weights from ‘robert-base’.  
What else should I do to initialise MyModel with the weights from the ‘robert-base’ model?

---

<div class="post-metadata">

**Author:** ![rychlik](https://avatars.discourse-cdn.com/v4/letter/r/7c8e57/32.png) [@rychlik](https://discuss.huggingface.co/u/rychlik)\
**Post date:** [August 30, 2024, 7:28am UTC](https://discuss.huggingface.co/t/custom-class-for-token-classification/104339/2 "2024-08-30T07:28:27Z")

</div>

The problem has been resolved. ‘additional\_data’ should be sent to the forward method.