choices=["O0", "O1", "O2"])

args = parser.parse_args()
print(args)

if not os.path.exists(args.work_dir):
    raise ValueError(f'Work directory not found at {args.work_dir}')
if not os.path.exists(args.labels_dict):
    raise ValueError(
        f'Dictionary with ids to labels not found at {args.labels_dict}')

nf = nemo.core.NeuralModuleFactory(backend=nemo.core.Backend.PyTorch,
                                   optimization_level=args.amp_opt_level,
                                   log_dir=None)

labels_dict = get_vocab(args.labels_dict)
""" Load the pretrained BERT parameters
See the list of pretrained models, call:
nemo_nlp.huggingface.BERT.list_pretrained_models()
"""
pretrained_bert_model = nemo_nlp.huggingface.BERT(
    pretrained_model_name=args.pretrained_bert_model)
hidden_size = pretrained_bert_model.local_parameters["hidden_size"]
tokenizer = NemoBertTokenizer(args.pretrained_bert_model)

data_layer = nemo_nlp.BertTokenClassificationInferDataLayer(
    queries=args.queries,
    tokenizer=tokenizer,
    max_seq_length=args.max_seq_length,
    batch_size=1)
Exemple #2
0
args = parser.parse_args()

if not os.path.exists(args.checkpoints_dir):
    raise ValueError(f'Checkpoints folder not found at {args.checkpoints_dir}')
if not (os.path.exists(args.punct_labels_dict)
        and os.path.exists(args.capit_labels_dict)):
    raise ValueError(
        f'Dictionary with ids to labels not found at {args.punct_labels_dict} \
         or {args.punct_labels_dict}')

nf = nemo.core.NeuralModuleFactory(backend=nemo.core.Backend.PyTorch,
                                   optimization_level=args.amp_opt_level,
                                   log_dir=None)

punct_labels_dict = get_vocab(args.punct_labels_dict)

capit_labels_dict = get_vocab(args.capit_labels_dict)
""" Load the pretrained BERT parameters
See the list of pretrained models, call:
nemo_nlp.huggingface.BERT.list_pretrained_models()
"""
pretrained_bert_model = nemo_nlp.huggingface.BERT(
    pretrained_model_name=args.pretrained_bert_model)
hidden_size = pretrained_bert_model.local_parameters["hidden_size"]
tokenizer = NemoBertTokenizer(args.pretrained_bert_model)

data_layer = nemo_nlp.BertTokenClassificationInferDataLayer(
    queries=args.queries,
    tokenizer=tokenizer,
    max_seq_length=args.max_seq_length,