def get_args():
default_model_parameter = "pretrained_name"
default_model = "punctuation_en_bert"
parser = argparse.ArgumentParser(
formatter_class=argparse.ArgumentDefaultsHelpFormatter,
description=
"The script is for restoring punctuation and capitalization in text. Long strings are split into "
"segments of length `--max_seq_length`. `--max_seq_length` is the length which includes [CLS] and [SEP] "
"tokens. Parameter `--step` controls segments overlapping. `--step` is a distance between beginnings of "
"consequent segments. Model outputs for tokens near the borders of tensors are less accurate and can be "
"discarded before final predictions computation. Parameter `--margin` is number of discarded outputs near "
"segments borders. Probabilities of tokens in overlapping parts of segments multiplied before selecting the "
"best prediction. Default values of parameters `--max_seq_length`, `--step`, and `--margin` are optimal for "
"IWSLT 2019 test dataset.",
)
input_ = parser.add_mutually_exclusive_group(required=True)
input_.add_argument(
"--input_manifest",
"-m",
type=Path,
help=
"Path to the file with NeMo manifest which needs punctuation and capitalization. If the first element "
"of manifest contains key 'pred_text', 'pred_text' values are passed for tokenization. Otherwise 'text' "
"values are passed for punctuation and capitalization. Exactly one parameter of `--input_manifest` and "
"`--input_text` should be provided.",
)
input_.add_argument(
"--input_text",
"-t",
type=Path,
help=
"Path to file with text which needs punctuation and capitalization. Exactly one parameter of "
"`--input_manifest` and `--input_text` should be provided.",
)
output = parser.add_mutually_exclusive_group(required=True)
output.add_argument(
"--output_manifest",
"-M",
type=Path,
help=
"Path to output NeMo manifest. Text with restored punctuation and capitalization will be saved in "
"'pred_text' elements if 'pred_text' key is present in the input manifest. Otherwise text with restored "
"punctuation and capitalization will be saved in 'text' elements. Exactly one parameter of `--output_manifest` "
"and `--output_text` should be provided.",
)
output.add_argument(
"--output_text",
"-T",
type=Path,
help=
"Path to file with text with restored punctuation and capitalization. Exactly one parameter of "
"`--output_manifest` and `--output_text` should be provided.",
)
model = parser.add_mutually_exclusive_group(required=False)
model.add_argument(
"--pretrained_name",
"-p",
help=
f"The name of NGC pretrained model. No more than one of parameters `--pretrained_name`, `--model_path`"
f"should be provided. If neither of parameters `--pretrained_name` and `--model_path` are provided, then the "
f"script is run with `--{default_model_parameter}={default_model}`.",
choices=[
m.pretrained_model_name
for m in PunctuationCapitalizationModel.list_available_models()
],
)
model.add_argument(
"--model_path",
"-P",
type=Path,
help=
f"Path to .nemo checkpoint of punctuation and capitalization model. No more than one of parameters "
f"`--pretrained_name` and `--model_path` should be provided. If neither of parameters `--pretrained_name` and "
f"`--model_path` are provided, then the script is run with `--{default_model_parameter}={default_model}`.",
)
parser.add_argument(
"--max_seq_length",
"-L",
type=int,
default=64,
help=
"Length of segments into which queries are split. `--max_seq_length` includes [CLS] and [SEP] tokens.",
)
parser.add_argument(
"--step",
"-s",
type=int,
default=8,
help=
"Relative shift of consequent segments into which long queries are split. Long queries are split into "
"segments which can overlap. Parameter `step` controls such overlapping. Imagine that queries are "
"tokenized into characters, `max_seq_length=5`, and `step=2`. In such a case query 'hello' is tokenized "
"into segments `[['[CLS]', 'h', 'e', 'l', '[SEP]'], ['[CLS]', 'l', 'l', 'o', '[SEP]']]`.",
)
parser.add_argument(
"--margin",
"-g",
type=int,
default=16,
help=
"A number of subtokens in the beginning and the end of segments which output probabilities are not used "
"for prediction computation. The first segment does not have left margin and the last segment does not have "
"right margin. For example, if input sequence is tokenized into characters, `max_seq_length=5`, `step=1`, "
"and `margin=1`, then query 'hello' will be tokenized into segments `[['[CLS]', 'h', 'e', 'l', '[SEP]'], "
"['[CLS]', 'e', 'l', 'l', '[SEP]'], ['[CLS]', 'l', 'l', 'o', '[SEP]']]`. These segments are passed to the "
"model. Before final predictions computation, margins are removed. In the next list, subtokens which logits "
"are not used for final predictions computation are marked with asterisk: `[['[CLS]'*, 'h', 'e', 'l'*, "
"'[SEP]'*], ['[CLS]'*, 'e'*, 'l', 'l'*, '[SEP]'*], ['[CLS]'*, 'l'*, 'l', 'o', '[SEP]'*]]`.",
)
parser.add_argument(
"--batch_size",
"-b",
type=int,
default=128,
help="Number of segments which are processed simultaneously.",
)
parser.add_argument(
"--save_labels_instead_of_text",
"-B",
action="store_true",
help=
"If this option is set save punctuation and capitalization labels instead text with restored punctuation "
"and capitalization. Labels are saved in format described here https://docs.nvidia.com/deeplearning/nemo/"
"user-guide/docs/en/main/nlp/punctuation_and_capitalization.html#nemo-data-format",
)
parser.add_argument(
"--device",
"-d",
choices=['cpu', 'cuda'],
help=
"Which device to use. If device is not set and CUDA is available GPU will be used. If device is not set "
"and CUDA is not available CPU is used.",
)
args = parser.parse_args()
if args.input_manifest is None and args.output_manifest is not None:
parser.error("--output_manifest requires --input_manifest")
if args.pretrained_name is None and args.model_path is None:
setattr(args, default_model_parameter, default_model)
for name in [
"input_manifest", "input_text", "output_manifest", "output_text",
"model_path"
]:
if getattr(args, name) is not None:
setattr(args, name, getattr(args, name).expanduser())
return args
import time
from nemo.collections.nlp.models import PunctuationCapitalizationModel
import argparse
import sys
import torch
# to get the list of pre-trained models
#PunctuationCapitalizationModel.list_available_models()
# Read text sentences from input file
parser = argparse.ArgumentParser()
parser.add_argument(
"--model", "-m", type=str, required=False, default="punctuation_en_distilbert", help="Path to the model",
choices=[m.pretrained_model_name for m in PunctuationCapitalizationModel.list_available_models()]
)
parser.add_argument(
"--inpath", "-i", type=str, required=True, help="Input Path"
)
parser.add_argument(
"--outpath", "-o", type=str, required=True, help="Output Path"
)
parser.add_argument(
"--batch_size", "-b", type=int, default=128, help="Number of segments which are processed simultaneously.",
)
parser.add_argument(
"--device",
"-d",
choices=['cpu', 'cuda'],
help="Which device to use. If device is not set and CUDA is available, then GPU will be used. If device is "
"not set and CUDA is not available, then CPU is used.",
)