Unverified Commit 4452b44b authored by Martin Malmsten's avatar Martin Malmsten Committed by GitHub
Browse files

Labels are now added to model config under id2label and label2id (#2945)

parent 53ce3854
......@@ -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(
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment