From 4452b44b9062f71363d593e7ed506c702722ad63 Mon Sep 17 00:00:00 2001
From: Martin Malmsten <martin.malmsten@kb.se>
Date: Fri, 21 Feb 2020 14:53:05 +0100
Subject: [PATCH] Labels are now added to model config under id2label and
 label2id (#2945)

---
 examples/ner/run_ner.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/examples/ner/run_ner.py b/examples/ner/run_ner.py
index a1992d4df8..70a7e9f8a7 100644
--- a/examples/ner/run_ner.py
+++ b/examples/ner/run_ner.py
@@ -586,6 +586,8 @@ def main():
     config = config_class.from_pretrained(
         args.config_name if args.config_name else args.model_name_or_path,
         num_labels=num_labels,
+        id2label={str(i): label for i, label in enumerate(labels)},
+        label2id={label: i for i, label in enumerate(labels)},
         cache_dir=args.cache_dir if args.cache_dir else None,
     )
     tokenizer = tokenizer_class.from_pretrained(