Labels are now added to model config under id2label and label2id (#2945)
This commit is contained in:
@@ -586,6 +586,8 @@ def main():
|
|||||||
config = config_class.from_pretrained(
|
config = config_class.from_pretrained(
|
||||||
args.config_name if args.config_name else args.model_name_or_path,
|
args.config_name if args.config_name else args.model_name_or_path,
|
||||||
num_labels=num_labels,
|
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,
|
cache_dir=args.cache_dir if args.cache_dir else None,
|
||||||
)
|
)
|
||||||
tokenizer = tokenizer_class.from_pretrained(
|
tokenizer = tokenizer_class.from_pretrained(
|
||||||
|
|||||||
Reference in New Issue
Block a user