Error with document-level model for NER task

Hello,
First of all, I would like to thank you for making this work available :)
I'm having an issue trying to replicate experiments documented here: https://github.com/Alibaba-NLP/ACE/blob/main/resources/docs/named_entity_recognition.md
I had success replicating the tutorial for Sentence-level Model for CoNLL 2003 English. However, when I try to do the same for Document-level Model, I'm having an error:
(ace) azeredo@ix-ws28:~/test_ace/test_recent_ace/ACE$ CUDA_VISIBLE_DEVICES=0 python train.py --config /home/azeredo/test_ace/test_recent_ace/ACE/config/xlnet-doc-en-ner-finetune.yaml
/home/azeredo/test_ace/test_recent_ace/ACE/flair/utils/params.py:104: YAMLLoadWarning: calling yaml.load() without Loader=... is deprecated, as the default Loader is unsafe. Please read https://msg.pyyaml.org/load for full details.
  dict_merge.dict_merge(params_dict, yaml.load(f))
2021-10-21 15:51:47,652 Reading data from /home/azeredo/.flair/datasets/conll_03_english
2021-10-21 15:51:47,652 Train: /home/azeredo/.flair/datasets/conll_03_english/eng.train
2021-10-21 15:51:47,652 Dev: /home/azeredo/.flair/datasets/conll_03_english/eng.testa
2021-10-21 15:51:47,652 Test: /home/azeredo/.flair/datasets/conll_03_english/eng.testb
2021-10-21 15:51:53,054 {b'<unk>': 0, b'O': 1, b'B-PER': 2, b'E-PER': 3, b'S-LOC': 4, b'B-MISC': 5, b'I-MISC': 6, b'E-MISC': 7, b'S-MISC': 8, b'S-PER': 9, b'B-ORG': 10, b'E-ORG': 11, b'S-ORG': 12, b'I-ORG': 13, b'B-LOC': 14, b'E-LOC': 15, b'I-PER': 16, b'I-LOC': 17, b'<START>': 18, b'<STOP>': 19}
2021-10-21 15:51:53,054 Corpus: 14987 train + 3466 dev + 3684 test sentences
[2021-10-21 15:51:53,562 INFO] https://s3.amazonaws.com/models.huggingface.co/bert/xlnet-large-cased-config.json not found in cache or force_download set to True, downloading to /home/azeredo/.cache/torch/transformers/tmpdwnsntaa
Downloading: 100%|███████████████████████████████████████| 761/761 [00:00<00:00, 391kB/s]
[2021-10-21 15:51:54,051 INFO] storing https://s3.amazonaws.com/models.huggingface.co/bert/xlnet-large-cased-config.json in cache at /home/azeredo/.cache/torch/transformers/df92a75c0ebbeb195065fe16fafa54ccd72e8362692cca884303a56788bd4bfc.0163e810fe4bdef52282bd9ddcbded8accbaa97a3ea7d89737ee7ce87511c587
[2021-10-21 15:51:54,051 INFO] creating metadata file for /home/azeredo/.cache/torch/transformers/df92a75c0ebbeb195065fe16fafa54ccd72e8362692cca884303a56788bd4bfc.0163e810fe4bdef52282bd9ddcbded8accbaa97a3ea7d89737ee7ce87511c587
[2021-10-21 15:51:54,052 INFO] loading configuration file https://s3.amazonaws.com/models.huggingface.co/bert/xlnet-large-cased-config.json from cache at /home/azeredo/.cache/torch/transformers/df92a75c0ebbeb195065fe16fafa54ccd72e8362692cca884303a56788bd4bfc.0163e810fe4bdef52282bd9ddcbded8accbaa97a3ea7d89737ee7ce87511c587
[2021-10-21 15:51:54,052 INFO] Model config XLNetConfig {
  "architectures": [
    "XLNetLMHeadModel"
  ],
  "attn_type": "bi",
  "bi_data": false,
  "bos_token_id": 1,
  "clamp_len": -1,
  "d_head": 64,
  "d_inner": 4096,
  "d_model": 1024,
  "dropout": 0.1,
  "end_n_top": 5,
  "eos_token_id": 2,
  "ff_activation": "gelu",
  "initializer_range": 0.02,
  "layer_norm_eps": 1e-12,
  "mem_len": null,
  "model_type": "xlnet",
  "n_head": 16,
  "n_layer": 24,
  "pad_token_id": 5,
  "reuse_len": null,
  "same_length": false,
  "start_n_top": 5,
  "summary_activation": "tanh",
  "summary_last_dropout": 0.1,
  "summary_type": "last",
  "summary_use_proj": true,
  "task_specific_params": {
    "text-generation": {
      "do_sample": true,
      "max_length": 250
    }
  },
  "untie_r": true,
  "vocab_size": 32000
}

[2021-10-21 15:51:54,532 INFO] https://s3.amazonaws.com/models.huggingface.co/bert/xlnet-large-cased-spiece.model not found in cache or force_download set to True, downloading to /home/azeredo/.cache/torch/transformers/tmp6pula_yo
Downloading: 100%|████████████████████████████████████| 798k/798k [00:00<00:00, 1.48MB/s]
[2021-10-21 15:51:55,568 INFO] storing https://s3.amazonaws.com/models.huggingface.co/bert/xlnet-large-cased-spiece.model in cache at /home/azeredo/.cache/torch/transformers/5b125ba222ff82664771f63cd8fac9696c24b403fc1ab720d537fe2ceaaf0576.8b10bd978b5d01c21303cc761fc9ecd464419b3bf921864a355ba807cfbfafa8
[2021-10-21 15:51:55,568 INFO] creating metadata file for /home/azeredo/.cache/torch/transformers/5b125ba222ff82664771f63cd8fac9696c24b403fc1ab720d537fe2ceaaf0576.8b10bd978b5d01c21303cc761fc9ecd464419b3bf921864a355ba807cfbfafa8
[2021-10-21 15:51:55,568 INFO] loading file https://s3.amazonaws.com/models.huggingface.co/bert/xlnet-large-cased-spiece.model from cache at /home/azeredo/.cache/torch/transformers/5b125ba222ff82664771f63cd8fac9696c24b403fc1ab720d537fe2ceaaf0576.8b10bd978b5d01c21303cc761fc9ecd464419b3bf921864a355ba807cfbfafa8
[2021-10-21 15:51:56,047 INFO] loading configuration file https://s3.amazonaws.com/models.huggingface.co/bert/xlnet-large-cased-config.json from cache at /home/azeredo/.cache/torch/transformers/df92a75c0ebbeb195065fe16fafa54ccd72e8362692cca884303a56788bd4bfc.0163e810fe4bdef52282bd9ddcbded8accbaa97a3ea7d89737ee7ce87511c587
[2021-10-21 15:51:56,048 INFO] Model config XLNetConfig {
  "architectures": [
    "XLNetLMHeadModel"
  ],
  "attn_type": "bi",
  "bi_data": false,
  "bos_token_id": 1,
  "clamp_len": -1,
  "d_head": 64,
  "d_inner": 4096,
  "d_model": 1024,
  "dropout": 0.1,
  "end_n_top": 5,
  "eos_token_id": 2,
  "ff_activation": "gelu",
  "initializer_range": 0.02,
  "layer_norm_eps": 1e-12,
  "mem_len": null,
  "model_type": "xlnet",
  "n_head": 16,
  "n_layer": 24,
  "output_hidden_states": true,
  "pad_token_id": 5,
  "reuse_len": null,
  "same_length": false,
  "start_n_top": 5,
  "summary_activation": "tanh",
  "summary_last_dropout": 0.1,
  "summary_type": "last",
  "summary_use_proj": true,
  "task_specific_params": {
    "text-generation": {
      "do_sample": true,
      "max_length": 250
    }
  },
  "untie_r": true,
  "vocab_size": 32000
}

[2021-10-21 15:51:56,274 INFO] https://cdn.huggingface.co/xlnet-large-cased-pytorch_model.bin not found in cache or force_download set to True, downloading to /home/azeredo/.cache/torch/transformers/tmpbnvrd8_z
Downloading: 100%|██████████████████████████████████| 1.44G/1.44G [00:31<00:00, 46.2MB/s]
[2021-10-21 15:52:27,536 INFO] storing https://cdn.huggingface.co/xlnet-large-cased-pytorch_model.bin in cache at /home/azeredo/.cache/torch/transformers/7fc554c19ef7bc74f1f74603c10156d751d2f99b09e8e38f91ed88e8c9ec6294.db8dc8babedbb75a56c36fca3e02b016e19fd682e79fb1a928e03c2df977cace
[2021-10-21 15:52:27,536 INFO] creating metadata file for /home/azeredo/.cache/torch/transformers/7fc554c19ef7bc74f1f74603c10156d751d2f99b09e8e38f91ed88e8c9ec6294.db8dc8babedbb75a56c36fca3e02b016e19fd682e79fb1a928e03c2df977cace
[2021-10-21 15:52:27,537 INFO] loading weights file https://cdn.huggingface.co/xlnet-large-cased-pytorch_model.bin from cache at /home/azeredo/.cache/torch/transformers/7fc554c19ef7bc74f1f74603c10156d751d2f99b09e8e38f91ed88e8c9ec6294.db8dc8babedbb75a56c36fca3e02b016e19fd682e79fb1a928e03c2df977cace
[2021-10-21 15:52:32,104 INFO] All model checkpoint weights were used when initializing XLNetModel.

[2021-10-21 15:52:32,104 INFO] All the weights of XLNetModel were initialized from the model checkpoint at xlnet-large-cased.
If your task is similar to the task the model of the ckeckpoint was trained on, you can already use XLNetModel for predictions without further training.
2021-10-21 15:52:33,874 Model Size: 360289300
Traceback (most recent call last):
  File "train.py", line 121, in <module>
    trainer: trainer_func = trainer_func(student, None, corpus, config=config.config, **config.config[trainer_name], is_test=args.test)
TypeError: __init__() got an unexpected keyword argument 'train_with_doc'
This happens if I try to use any of the available pre-defined configuration files: xlnet-doc-en-ner-finetune.yaml, xlmr-doc-en-ner-finetune.yaml or roberta-doc-en-ner-finetune.yaml
Thank you!
Alibaba-NLP / ACE

Error with document-level model for NER task #21
