Hi everyone,
i’m using the script run_ner 8 from huggingface transformers to perform PoS tagging task with conll-u dataset.
Now I would like to do two tasks together: predict both the PoS tag and the head of each word, always in the conll-u format.
To do this my idea is to use the pre-trained bert model as it is made available by the run_ner script, remove its last layer and add two dense layers to do the two classification tasks at the same time.
The model is as follows:
config = AutoConfig.from_pretrained(
        model_args.config_name if model_args.config_name else model_args.model_name_or_path,
        num_labels=num_labels,
        id2label=label_map,
        label2id={label: i for i, label in enumerate(labels)},
        cache_dir=model_args.cache_dir,
    )
tokenizer = AutoTokenizer.from_pretrained(
        model_args.tokenizer_name if model_args.tokenizer_name else model_args.model_name_or_path,
        cache_dir=model_args.cache_dir,
        use_fast=model_args.use_fast,
    )
model = AutoModelForTokenClassification.from_pretrained(
        model_args.model_name_or_path,
        from_tf=bool(".ckpt" in model_args.model_name_or_path),
        config=config,
        cache_dir=model_args.cache_dir,
    )

to remove the last layer I did:
model = nn.Sequential(*list(model.children())[:-1])

but now I have two problems:

I don’t know how to add the two nn.Linear() layers to the model
I don’t know how to pass the label to the model for both the Pos-tag and the head

Thanks in advance!
Edit:
I added a new model in the following way:
model = nn.Sequential(*list(model.children())[:-1], NewModule(num_labels))

where NewModule() is:
class NewModule(nn.Module):
    def __init__(self, num_labels):
        super(NewModule, self).__init__()
        self.classifierDeprel = nn.Linear(768, num_labels)
        self.classifierRelPos = nn.Linear(768, num_labels)

    def forward(self, x):
        x1 = torch.softmax(self.fc1(x))
        y1 = torch.softmax(self.fc2(x))
        return x1, y1

Now the problem is that after instantiating the trainer in this way
trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset,
        eval_dataset=eval_dataset,
        compute_metrics=compute_metrics,
    )

I can’t train the model
trainer.train(
            model_path=model_args.model_name_or_path if os.path.isdir(model_args.model_name_or_path) else None
        )

I think because the model_path is no longer Bert’s one, but it’s my new model, but i don’t know how to specify it.