Update run_ner.py with id2label config (#12001)

2021-06-09 20:27:05 +09:00
parent 0e82f0cbc2
commit 82a2b76c95
1 changed files with 2 additions and 0 deletions
--- a/examples/pytorch/token-classification/run_ner.py
+++ b/examples/pytorch/token-classification/run_ner.py
@@ -281,6 +281,8 @@ def main():
    config = AutoConfig.from_pretrained(
        model_args.config_name if model_args.config_name else model_args.model_name_or_path,
        num_labels=num_labels,
+        label2id=label_to_id,
+        id2label={i: l for l, i in label_to_id.items()},
        finetuning_task=data_args.task_name,
        cache_dir=model_args.cache_dir,
        revision=model_args.model_revision,