def from_pretrained(model_name_or_path,
checkpoint_file='model.pt',
data_name_or_path='.',
archive_map=None,
**kwargs):
from fairseq import checkpoint_utils, file_utils
if archive_map is not None:
if model_name_or_path in archive_map:
model_name_or_path = archive_map[model_name_or_path]
if data_name_or_path is not None and data_name_or_path in archive_map:
data_name_or_path = archive_map[data_name_or_path]
# allow archive_map to set default arg_overrides (e.g., tokenizer, bpe)
# for each model
if isinstance(model_name_or_path, dict):
for k, v in model_name_or_path.items():
if k == 'checkpoint_file':
checkpoint_file = v
elif (k != 'path'
# only set kwargs that don't already have overrides
and k not in kwargs):
kwargs[k] = v
model_name_or_path = model_name_or_path['path']
model_path = file_utils.load_archive_file(model_name_or_path)
# convenience hack for loading data and BPE codes from model archive
if data_name_or_path.startswith('.'):
kwargs['data'] = os.path.abspath(
os.path.join(model_path, data_name_or_path))
else:
kwargs['data'] = file_utils.load_archive_file(data_name_or_path)
for file, arg in {
'code': 'bpe_codes',
'bpecodes': 'bpe_codes',
'sentencepiece.bpe.model': 'sentencepiece_vocab',
}.items():
print(model_path, file)
path = os.path.join(model_path, file)
if os.path.exists(path):
kwargs[arg] = path
if 'user_dir' in kwargs:
utils.import_user_module(
argparse.Namespace(user_dir=kwargs['user_dir']))
models, args, task = checkpoint_utils.load_model_ensemble_and_task(
[
os.path.join(model_path, cpt)
for cpt in checkpoint_file.split(os.pathsep)
],
arg_overrides=kwargs,
)
return {
'args': args,
'task': task,
'models': models,
}
def setUp(self, args):
if args.task == "speech_to_speech":
args.normalize_waveform = False
self.task = tasks.setup_task(args)
self.pre_tokenizer = self.task.build_tokenizer(args)
self.bpe_tokenizer = self.task.build_bpe(args)
try:
self.src_dict = self.task.src_dict
except Exception:
self.src_dict = None
ensemble, saved_cfg, task = checkpoint_utils.load_model_ensemble_and_task(
[args.path],
arg_overrides=ast.literal_eval(args.model_overrides),
task=self.task,
strict=False,
)
self.model = ensemble[0]
if self.use_cuda:
self.model.cuda()
# criterion.cuda()
self.model.eval()
self.generator = self.task.build_generator(
[self.model],
args,
)
def from_pretrained(
model_name_or_path,
checkpoint_file="model.pt",
data_name_or_path=".",
archive_map=None,
**kwargs
):
from fairseq import checkpoint_utils, file_utils
if archive_map is not None:
if model_name_or_path in archive_map:
model_name_or_path = archive_map[model_name_or_path]
if data_name_or_path is not None and data_name_or_path in archive_map:
data_name_or_path = archive_map[data_name_or_path]
# allow archive_map to set default arg_overrides (e.g., tokenizer, bpe)
# for each model
if isinstance(model_name_or_path, dict):
for k, v in model_name_or_path.items():
if k == "checkpoint_file":
checkpoint_file = v
elif (
k != "path"
# only set kwargs that don't already have overrides
and k not in kwargs
):
kwargs[k] = v
model_name_or_path = model_name_or_path["path"]
model_path = file_utils.load_archive_file(model_name_or_path)
# convenience hack for loading data and BPE codes from model archive
if data_name_or_path.startswith("."):
kwargs["data"] = os.path.abspath(os.path.join(model_path, data_name_or_path))
else:
kwargs["data"] = file_utils.load_archive_file(data_name_or_path)
for file, arg in {
"code": "bpe_codes",
"bpecodes": "bpe_codes",
"sentencepiece.bpe.model": "sentencepiece_model",
"merges.txt": "bpe_merges",
"vocab.json": "bpe_vocab",
}.items():
path = os.path.join(model_path, file)
if os.path.exists(path):
kwargs[arg] = path
if "user_dir" in kwargs:
utils.import_user_module(argparse.Namespace(user_dir=kwargs["user_dir"]))
models, args, task = checkpoint_utils.load_model_ensemble_and_task(
[os.path.join(model_path, cpt) for cpt in checkpoint_file.split(os.pathsep)],
arg_overrides=kwargs,
)
return {
"args": args,
"task": task,
"models": models,
}
def main(args):
logging.basicConfig(
format='%(asctime)s | %(levelname)s | %(name)s | %(message)s',
datefmt='%Y-%m-%d %H:%M:%S',
level=logging.INFO,
stream=sys.stdout,
)
logger = logging.getLogger('fairnr_cli.extract')
logger.info(args)
use_cuda = torch.cuda.is_available() and not args.cpu
models, model_args, task = checkpoint_utils.load_model_ensemble_and_task(
[args.path], suffix=getattr(args, "checkpoint_suffix", ""))
model = models[0]
if use_cuda:
model.cuda()
if args.format == 'mc_mesh':
plydata = model.encoder.export_surfaces(
model.field, th=args.mc_threshold,
bits=2 * args.mc_num_samples_per_halfvoxel)
elif args.format == 'voxel_center':
plydata = model.encoder.export_voxels(False)
elif args.format == 'voxel_mesh':
plydata = model.encoder.export_voxels(True)
else:
raise NotImplementedError
# write to ply file.
if not os.path.exists(args.output):
os.makedirs(args.output)
plydata.text = args.savetext
plydata.write(open(os.path.join(args.output, args.name + '.ply'), 'wb'))
def cli_main():
parser = argparse.ArgumentParser(
description=
'Extract geometry from a trained model (only for learnable embeddings).'
)
parser.add_argument('--path', type=str, required=True)
parser.add_argument('--output', type=str, required=True)
parser.add_argument('--name', type=str, default='sparsevoxel')
parser.add_argument('--user-dir', default='fairnr')
args = options.parse_args_and_arch(parser)
models, model_args, task = checkpoint_utils.load_model_ensemble_and_task(
[args.path],
suffix=getattr(args, "checkpoint_suffix", ""),
)
model = models[0]
voxel_idx, voxel_pts = model.encoder.extract_voxels()
# write to ply file.
points = [(voxel_pts[k, 0], voxel_pts[k, 1], voxel_pts[k, 2], voxel_idx[k])
for k in range(voxel_idx.size(0))]
vertex = np.array(points,
dtype=[('x', 'f4'), ('y', 'f4'), ('z', 'f4'),
('quality', 'f4')])
PlyData([PlyElement.describe(vertex, 'vertex')],
text=True).write(os.path.join(args.output, args.name + '.ply'))
def __init__(
self,
language,
temperature=1.0,
model_dir=telescope_cache_folder() + "m39v1",
**kwargs,
):
"""
model_dir should contain:
1) checkpoint.pt: the fairseq model
2) spm.model: the sentencepiece model
3) dict.src.txt: the fairseq source dictionary
4) dict.tgt.txt: the fairseq target dictionary (likely a copy of the source)
lang: ISO 639-1 Code (e.g. "en"). Must be a language compatable with the model.
"""
if not os.path.isdir(model_dir):
download_file_maybe_extract(
url="http://data.statmt.org/prism/m39v1.tar",
directory=telescope_cache_folder(),
)
self.sp = spm.SentencePieceProcessor()
self.sp.Load(model_dir + "/spm.model")
self.lang = language
self.temperature = temperature
# this prints things and I can't figure out how to disable it
with open(os.devnull, "w") as sys.stdout:
(
self.models,
self.args,
self.task,
) = checkpoint_utils.load_model_ensemble_and_task(
[
model_dir + "/checkpoint.pt",
],
arg_overrides=dict(data=model_dir + "/"),
)
sys.stdout = sys.__stdout__
self.use_cuda = torch.cuda.is_available()
self.generator = SequenceScorer(self.task.target_dictionary,
temperature=temperature)
for model in self.models:
if self.use_cuda:
model.cuda()
model.make_generation_fast_(
beamable_mm_beam_size=None,
need_attn=False,
)
# hash model
self.model_hash = hash_model(model_dir)
if not self.language_support(language):
raise Exception(f"{language} is not supported by {self.name}.")
def main(args):
assert (args.dump_features or args.dump_waveforms or args.dump_attentions
or args.dump_eos_probs or args.dump_plots)
if args.max_tokens is None and args.batch_size is None:
args.max_tokens = 8000
logger.info(args)
use_cuda = torch.cuda.is_available() and not args.cpu
task = tasks.setup_task(args)
models, saved_cfg, task = checkpoint_utils.load_model_ensemble_and_task(
[args.path],
task=task,
arg_overrides=ast.literal_eval(args.model_overrides),
)
model = models[0].cuda() if use_cuda else models[0]
# use the original n_frames_per_step
task.args.n_frames_per_step = saved_cfg.task.n_frames_per_step
task.load_dataset(args.gen_subset, task_cfg=saved_cfg.task)
data_cfg = task.data_cfg
sample_rate = data_cfg.config.get("features", {}).get("sample_rate", 22050)
resample_fn = {
False:
lambda x: x,
True:
lambda x: torchaudio.sox_effects.apply_effects_tensor(
x.detach().cpu().unsqueeze(0), sample_rate,
[['rate', str(args.output_sample_rate)]])[0].squeeze(0)
}.get(args.output_sample_rate != sample_rate)
if args.output_sample_rate != sample_rate:
logger.info(f"resampling to {args.output_sample_rate}Hz")
generator = task.build_generator([model], args)
itr = task.get_batch_iterator(
dataset=task.dataset(args.gen_subset),
max_tokens=args.max_tokens,
max_sentences=args.batch_size,
max_positions=(sys.maxsize, sys.maxsize),
ignore_invalid_inputs=args.skip_invalid_size_inputs_valid_test,
required_batch_size_multiple=args.required_batch_size_multiple,
num_shards=args.num_shards,
shard_id=args.shard_id,
num_workers=args.num_workers,
data_buffer_size=args.data_buffer_size,
).next_epoch_itr(shuffle=False)
Path(args.results_path).mkdir(exist_ok=True, parents=True)
is_na_model = getattr(model, "NON_AUTOREGRESSIVE", False)
dataset = task.dataset(args.gen_subset)
vocoder = task.args.vocoder
with progress_bar.build_progress_bar(args, itr) as t:
for sample in t:
sample = utils.move_to_cuda(sample) if use_cuda else sample
hypos = generator.generate(model,
sample,
has_targ=args.dump_target)
for result in postprocess_results(dataset, sample, hypos,
resample_fn, args.dump_target):
dump_result(is_na_model, args, vocoder, *result)
def __init__(self, model_dir, lang, temperature):
'''
model_dir should contain:
1) checkpoint.pt: the fairseq model
2) spm.model: the sentencepiece model
3) dict.src.txt: the fairseq source dictionary
4) dict.tgt.txt: the fairseq target dictionary (likely a copy of the source)
lang: ISO 639-1 Code (e.g. "en"). Must be a language compatable with the model.
'''
self.sp = spm.SentencePieceProcessor()
self.sp.Load(model_dir + '/spm.model')
self.lang = lang
self.temperature = temperature
# this prints things and I can't figure out how to disable it
sys.stdout = open(os.devnull, 'w')
self.models, self.args, self.task = checkpoint_utils.load_model_ensemble_and_task(
[
model_dir + '/checkpoint.pt',
],
arg_overrides=dict(data=model_dir + '/'),
)
sys.stdout = sys.__stdout__
self.use_cuda = torch.cuda.is_available()
self.generator = SequenceScorer(self.task.target_dictionary,
temperature=temperature)
for model in self.models:
if self.use_cuda:
model.cuda()
model.make_generation_fast_(
beamable_mm_beam_size=None,
need_attn=False,
)
# if model.args.fp16:
# model.half()
# hash model
self.model_hash = hash_model(model_dir)
if self.model_hash in MODELS:
model_langs = MODELS[self.model_hash]['langs']
if lang not in model_langs:
model_name = MODELS[self.model_hash]['name']
logger.warning(
f'Language "{lang}" is unsupported for model "{model_name}"'
)
logger.warning(
f'Supported languages for {model_name}: {", ".join(model_langs)}'
)
else:
logger.warning('unrecognized model, so cannot check language')
def from_pretrained(cls, model_name_or_path, checkpoint_file='model.pt', data_name_or_path=None, **kwargs):
"""
Load a :class:`~fairseq.models.FairseqModel` from a pre-trained model
file. Downloads and caches the pre-trained model file if needed.
The base implementation returns a :class:`fairseq.hub_utils.Generator`,
which can be used to generate translations or sample from language
models. The underlying :class:`~fairseq.models.FairseqModel` can be
accessed via the *generator.models* attribute.
Other models may override this to implement custom PyTorch Hub APIs.
Args:
model_name_or_path (str): either the name of a pre-trained model to
load or a path/URL to a pre-trained model state dict
checkpoint_file (str, optional): colon-separated list of checkpoint
files in the model archive to ensemble (default: 'model.pt')
data_name_or_path (str, optional): point args.data to the archive
at the given path/URL. Can start with '.' or './' to reuse the
model archive path.
"""
from fairseq import checkpoint_utils, file_utils, hub_utils
if hasattr(cls, 'hub_models'):
archive_map = cls.hub_models()
if model_name_or_path in archive_map:
model_name_or_path = archive_map[model_name_or_path]
if data_name_or_path is not None and data_name_or_path in archive_map:
data_name_or_path = archive_map[data_name_or_path]
model_path = file_utils.load_archive_file(model_name_or_path)
# convenience hack for loading data and BPE codes from model archive
if data_name_or_path is not None:
if data_name_or_path.startswith('.'):
kwargs['data'] = os.path.abspath(os.path.join(model_path, data_name_or_path))
else:
kwargs['data'] = file_utils.load_archive_file(data_name_or_path)
for file, arg in {
'code': 'bpe_codes',
'bpecodes': 'bpe_codes',
'sentencepiece.bpe.model': 'sentencepiece_vocab',
}.items():
path = os.path.join(model_path, file)
if os.path.exists(path):
kwargs[arg] = path
models, args, task = checkpoint_utils.load_model_ensemble_and_task(
[os.path.join(model_path, cpt) for cpt in checkpoint_file.split(':')],
arg_overrides=kwargs,
)
print(args)
return hub_utils.Generator(args, task, models)
def _load_model(self, model_path, target_dict):
class BasicTask(object):
# Class to get around creating a full task object.
def __init__(self, target_dict):
self.target_dictionary = target_dict
data_folder = self.model_path.replace(
self.model_path.split('/')[-1], '')
args_overrides = {
'data': data_folder,
}
try:
model, cfg, task = checkpoint_utils.load_model_ensemble_and_task(
[model_path], arg_overrides=args_overrides)
except MissingMandatoryValue as e:
args_overrides['w2v_path'] = model_path
model, cfg, task = checkpoint_utils.load_model_ensemble_and_task(
[model_path], arg_overrides=args_overrides)
return model
def download_and_load_checkpoint(
self,
checkpoint_filename: str,
arg_overrides: Optional[Dict[str, str]] = None):
path = self.download(self.base_url, self.root, checkpoint_filename)
_arg_overrides = arg_overrides or {}
_arg_overrides["data"] = self.root.as_posix()
models, cfg, task = load_model_ensemble_and_task(
[path.as_posix()], arg_overrides=_arg_overrides)
if self.use_cuda:
for model in models:
model.cuda()
generator = task.build_generator(models, cfg)
return models, cfg, task, generator
def test_load_model_ensemble_and_task(self):
with contextlib.redirect_stdout(StringIO()):
with self._train_transformer(seed=123) as model1:
with self._train_transformer(seed=456) as model2:
ensemble, cfg, task = checkpoint_utils.load_model_ensemble_and_task(
filenames=[model1, model2])
self.assertEqual(len(ensemble), 2)
# after Transformer has been migrated to Hydra, this will probably
# become cfg.common.seed
self.assertEqual(ensemble[0].args.seed, 123)
self.assertEqual(ensemble[1].args.seed, 456)
# the task from the first model should be returned
self.assertEqual(task.args.seed, 123)
def test_prune_state_dict(self):
with contextlib.redirect_stdout(StringIO()):
extra_args = [
"--encoder-layerdrop", "0.01", "--decoder-layerdrop", "0.01"
]
with self._train_transformer(seed=1,
extra_args=extra_args) as model:
ensemble, cfg, task = checkpoint_utils.load_model_ensemble_and_task(
filenames=[model],
arg_overrides={
"encoder_layers_to_keep": "0,2",
"decoder_layers_to_keep": "1",
},
)
self.assertEqual(len(ensemble), 1)
self.assertEqual(len(ensemble[0].encoder.layers), 2)
self.assertEqual(len(ensemble[0].decoder.layers), 1)
def from_pretrained(
model_name_or_path,
checkpoint_file='model.pt',
data_name_or_path='.',
archive_map=None,
**kwargs
):
from fairseq import checkpoint_utils, file_utils
if archive_map is not None:
if model_name_or_path in archive_map:
model_name_or_path = archive_map[model_name_or_path]
if data_name_or_path is not None and data_name_or_path in archive_map:
data_name_or_path = archive_map[data_name_or_path]
model_path = file_utils.load_archive_file(model_name_or_path)
# convenience hack for loading data and BPE codes from model archive
if data_name_or_path.startswith('.'):
kwargs['data'] = os.path.abspath(os.path.join(model_path, data_name_or_path))
else:
kwargs['data'] = file_utils.load_archive_file(data_name_or_path)
for file, arg in {
'code': 'bpe_codes',
'bpecodes': 'bpe_codes',
'sentencepiece.bpe.model': 'sentencepiece_vocab',
}.items():
path = os.path.join(model_path, file)
if os.path.exists(path):
kwargs[arg] = path
if 'user_dir' in kwargs:
utils.import_user_module(argparse.Namespace(user_dir=kwargs['user_dir']))
models, args, task = checkpoint_utils.load_model_ensemble_and_task(
[os.path.join(model_path, cpt) for cpt in checkpoint_file.split(':')],
arg_overrides=kwargs,
)
return {
'args': args,
'task': task,
'models': models,
}
def main():
parser = argparse.ArgumentParser()
# rescoring_task.RescoringTask.add_args(parser, add_def_opts=False, add_data_opts=False)
WordTokenizer.add_args(parser)
#LatsDataSet.add_args(parser)
LatsOracleAlignDataSet.add_args(parser)
parser.add_argument('--no_progress_bar',
action='store_true',
help="disable progress bar")
parser.add_argument('--device',
type=torch.device,
default=torch.device('cpu'),
help="PyTorch device")
parser.add_argument('--format',
type=str,
default='pdf',
help="Render file type. Default .pdf")
parser.add_argument('--norm',
action='store_true',
default=False,
help="Normalize lt probs")
parser.add_argument("model", type=str, help="Saved model")
parser.add_argument("out_dir", type=str, help='Out dir')
parser.add_argument("utts", nargs="+", help="Utt_ids for draw")
args = parser.parse_args()
setup_logger()
tokenizer = WordTokenizer.build_from_args(args)
#dataset = LatsDataSet.build_from_args(args, tokenizer)
dataset = LatsOracleAlignDataSet.build_from_args(args, tokenizer)
models, model_args, task = checkpoint_utils.load_model_ensemble_and_task(
[args.model], arg_overrides=vars(args))
model = models[-1].to(args.device)
logger.info(f"Drawing {args.utts}")
draw_rescoring_lats(model,
dataset,
args.out_dir,
args.utts,
norm=args.norm,
device=args.device,
format=args.format,
progress_bar=(not args.no_progress_bar))
def setup_ar(checkpoint, batch_size, tokens_per_target=None, fp16=False):
models, args, task = checkpoint_utils.load_model_ensemble_and_task(
[checkpoint])
if tokens_per_target is None:
tokens_per_target = args.tokens_per_sample
args_overrides = {
'cpu': False,
'no_beamable_mm': False,
'max_sentences': batch_size,
'fp16': fp16,
'tokens_per_target': tokens_per_target,
'sampling': True,
'beam': 1
}
for arg_name, arg_val in args_overrides.items():
setattr(args, arg_name, arg_val)
use_cuda = torch.cuda.is_available() and not args.cpu
for model in models:
model.make_generation_fast_(
beamable_mm_beam_size=None
if args.no_beamable_mm else args.max_sentences,
need_attn=False,
)
if args.fp16:
model.half()
if use_cuda:
model.cuda()
generator = task.build_generator(args)
def generate_fn(n):
bs = task.args.max_sentences
results = []
for i in range(0, n, bs):
real_bs = min(bs, n - i)
sample = build_sample(task, bs=real_bs)
tokens, _ = task.inference_step(generator, models, sample)
results.append(tokens)
return np.concatenate(results, axis=0)
return generate_fn
def set_up_model(args, path, override_args=None):
use_fp16 = args.fp16
use_cuda = torch.cuda.is_available() and not args.cpu
if override_args is not None:
overrides = vars(override_args)
overrides.update(eval(getattr(override_args, 'model_overrides', '{}')))
else:
overrides = None
# Load ensemble
print('| loading model(s) from {}'.format(path))
models, model_args, task = checkpoint_utils.load_model_ensemble_and_task(
[path],
arg_overrides=overrides,
)
model = models[0]
# Move models to GPU
for model in models:
if use_fp16:
model.half()
if use_cuda:
model.cuda()
# Print args
print(model_args)
# Build criterion
criterion = task.build_criterion(model_args)
criterion.eval()
if model_args.task == 'language_modeling':
model_args.sampling = True
model_args.sampling_topk = 20
model_args.max_len_b = 80
generator = task.build_generator(model_args)
return model, task, criterion, generator, model_args
def load_models_and_task(args, use_cuda, task=None):
# if `task` is not provided, it will be from the saved model args
models, model_args, task = checkpoint_utils.load_model_ensemble_and_task(
args.path.split(':'),
arg_overrides=eval(args.model_overrides),
task=task,
)
# Optimize ensemble for generation
for model in models:
model.make_generation_fast_(
beamable_mm_beam_size=None if args.no_beamable_mm else args.beam,
need_attn=args.print_alignment,
)
if args.fp16:
print('using fp16 for models')
model.half()
if use_cuda:
print('using GPU for models')
model.cuda()
else:
print('using CPU for models')
# model = Model(models, task.target_dictionary)
return models, model_args, task
type=torch.device,
default=torch.device('cpu'),
help="PyTorch device")
parser.add_argument('--keep_tmp', action='store_true', help="keep tmp hyp")
parser.add_argument('--acwt',
type=float,
default=1,
help="Acoustic weight")
parser.add_argument("model", type=str, help="Saved model")
parser.add_argument("ref_fname", type=str, help="Reference text")
args = parser.parse_args()
setup_logger()
tokenizer = WordTokenizer.build_from_args(args)
dataset = LatsDataSet.build_from_args(args, tokenizer)
models, model_args, task = checkpoint_utils.load_model_ensemble_and_task(
[args.model], arg_overrides=vars(args))
model = models[-1].to(args.device)
ref_fname = args.ref_fname
wer = compute_model_wer(model=model,
dataset=dataset,
ref_fname=ref_fname,
device=args.device,
btz=args.infer_btz,
model_weight=args.model_weight,
progress_bar=(not args.no_progress_bar),
keep_tmp=args.keep_tmp,
dataloader_nj=args.dataloader_nj,
hyp_filter=args.hyp_filter,
acwt=args.acwt,
lmwt=args.lmwt,
strategy=args.strategy)
def decode_rerank_scores(args):
if args.max_tokens is None and args.batch_size is None:
args.batch_size = 1
logger.info(args)
use_cuda = torch.cuda.is_available() and not args.cpu
# Load ensemble
logger.info("loading model(s) from {}".format(args.path))
models, _model_args, task = checkpoint_utils.load_model_ensemble_and_task(
[args.path],
arg_overrides=eval(args.model_overrides),
)
for model in models:
if args.fp16:
model.half()
if use_cuda:
model.cuda()
# Initialize generator
generator = task.build_generator(args)
# Handle tokenization and BPE
tokenizer = task.build_tokenizer(args)
bpe = task.build_bpe(args)
def encode_fn(x):
if tokenizer is not None:
x = tokenizer.encode(x)
if bpe is not None:
x = bpe.encode(x)
return x
max_positions = utils.resolve_max_positions(
task.max_positions(), *[model.max_positions() for model in models])
src, hyp, mt_scores = parse_fairseq_gen(args.in_text, task)
model_scores = {}
logger.info("decode reranker score")
for batch in make_batches(args, src, hyp, task, max_positions, encode_fn):
src_tokens = batch.src_tokens
src_lengths = batch.src_lengths
if use_cuda:
src_tokens = src_tokens.cuda()
src_lengths = src_lengths.cuda()
sample = {
"net_input": {
"src_tokens": src_tokens,
"src_lengths": src_lengths
},
}
scores = task.inference_step(generator, models, sample)
for id, sc in zip(batch.ids.tolist(), scores.tolist()):
model_scores[id] = sc[0]
model_scores = [model_scores[i] for i in range(len(model_scores))]
return src, hyp, mt_scores, model_scores
def main(cfg: DictConfig, override_args=None):
if isinstance(cfg, Namespace):
cfg = convert_namespace_to_omegaconf(cfg)
utils.import_user_module(cfg.common)
assert (
cfg.dataset.max_tokens is not None
or cfg.dataset.batch_size is not None
), "Must specify batch size either with --max-tokens or --batch-size"
use_fp16 = cfg.common.fp16
use_cuda = torch.cuda.is_available() and not cfg.common.cpu
if use_cuda:
torch.cuda.set_device(cfg.distributed_training.device_id)
if cfg.distributed_training.distributed_world_size > 1:
data_parallel_world_size = distributed_utils.get_data_parallel_world_size(
)
data_parallel_rank = distributed_utils.get_data_parallel_rank()
else:
data_parallel_world_size = 1
data_parallel_rank = 0
if override_args is not None:
overrides = vars(override_args)
overrides.update(eval(getattr(override_args, "model_overrides", "{}")))
else:
overrides = None
# Load ensemble
logger.info("loading model(s) from {}".format(cfg.common_eval.path))
models, saved_cfg, task = checkpoint_utils.load_model_ensemble_and_task(
[cfg.common_eval.path],
arg_overrides=overrides,
suffix=cfg.checkpoint.checkpoint_suffix,
)
model = models[0]
# Move models to GPU
for model in models:
if use_fp16:
model.half()
if use_cuda:
model.cuda()
# Print args
logger.info(saved_cfg)
# Build criterion
criterion = task.build_criterion(saved_cfg.criterion)
criterion.eval()
for subset in cfg.dataset.valid_subset.split(","):
try:
task.load_dataset(subset,
combine=False,
epoch=1,
task_cfg=saved_cfg.task)
dataset = task.dataset(subset)
except KeyError:
raise Exception("Cannot find dataset: " + subset)
# Initialize data iterator
itr = task.get_batch_iterator(
dataset=dataset,
max_tokens=cfg.dataset.max_tokens,
max_sentences=cfg.dataset.batch_size,
max_positions=utils.resolve_max_positions(
task.max_positions(),
*[m.max_positions() for m in models],
),
ignore_invalid_inputs=cfg.dataset.
skip_invalid_size_inputs_valid_test,
required_batch_size_multiple=cfg.dataset.
required_batch_size_multiple,
seed=cfg.common.seed,
num_shards=data_parallel_world_size,
shard_id=data_parallel_rank,
num_workers=cfg.dataset.num_workers,
data_buffer_size=cfg.dataset.data_buffer_size,
).next_epoch_itr(shuffle=False)
progress = progress_bar.progress_bar(
itr,
log_format=cfg.common.log_format,
log_interval=cfg.common.log_interval,
prefix=f"valid on '{subset}' subset",
default_log_format=("tqdm" if not cfg.common.no_progress_bar else
"simple"),
)
log_outputs = []
for i, sample in enumerate(progress):
sample = utils.move_to_cuda(sample) if use_cuda else sample
_loss, _sample_size, log_output = task.valid_step(
sample, model, criterion)
progress.log(log_output, step=i)
log_outputs.append(log_output)
if data_parallel_world_size > 1:
log_outputs = distributed_utils.all_gather_list(
log_outputs,
max_size=cfg.common.all_gather_list_size,
group=distributed_utils.get_data_parallel_group(),
)
log_outputs = list(chain.from_iterable(log_outputs))
with metrics.aggregate() as agg:
task.reduce_metrics(log_outputs, criterion)
log_output = agg.get_smoothed_values()
progress.print(log_output, tag=subset, step=i)
def main(args, override_args=None):
utils.import_user_module(args)
use_fp16 = args.fp16
use_cuda = torch.cuda.is_available() and not args.cpu
if override_args is not None:
overrides = vars(override_args)
overrides.update(eval(getattr(override_args, 'model_overrides', '{}')))
else:
overrides = None
# Load ensemble
print('| loading model(s) from {}'.format(args.path))
models, model_args, task = checkpoint_utils.load_model_ensemble_and_task(
[args.path],
arg_overrides=overrides,
)
model = models[0]
# Move models to GPU
for model in models:
if use_fp16:
model.half()
if use_cuda:
model.cuda()
# Print args
print(model_args)
# Build criterion
criterion = task.build_criterion(model_args)
criterion.eval()
# Load valid dataset (we load training data below, based on the latest checkpoint)
for subset in args.valid_subset.split(','):
try:
task.load_dataset(subset, combine=False, epoch=0)
dataset = task.dataset(subset)
except KeyError:
raise Exception('Cannot find dataset: ' + subset)
# Initialize data iterator
itr = task.get_batch_iterator(
dataset=dataset,
max_tokens=args.max_tokens,
max_sentences=args.max_sentences,
max_positions=utils.resolve_max_positions(
task.max_positions(),
*[m.max_positions() for m in models],
),
ignore_invalid_inputs=args.skip_invalid_size_inputs_valid_test,
required_batch_size_multiple=args.required_batch_size_multiple,
seed=args.seed,
num_workers=args.num_workers,
).next_epoch_itr(shuffle=False)
progress = progress_bar.build_progress_bar(
args,
itr,
prefix='valid on \'{}\' subset'.format(subset),
no_progress_bar='simple')
log_outputs = []
for i, sample in enumerate(progress):
sample = utils.move_to_cuda(sample) if use_cuda else sample
_loss, _sample_size, log_output = task.valid_step(
sample, model, criterion)
progress.log(log_output, step=i)
log_outputs.append(log_output)
log_output = task.aggregate_logging_outputs(log_outputs, criterion)
progress.print(log_output, tag=subset, step=i)
def setup_ar(checkpoint, batch_size, tokens_per_target=None, fp16=False):
models, args, task = checkpoint_utils.load_model_ensemble_and_task([checkpoint])
if tokens_per_target is None:
def main(rank, world_size, args):
if world_size > 1:
torch.distributed.init_process_group(backend="gloo",
init_method="env://",
world_size=world_size,
rank=rank)
torch.cuda.set_device(rank)
raw_args = args
args = convert_namespace_to_omegaconf(args)
if args.common.seed is not None:
random.seed(args.common.seed)
np.random.seed(args.common.seed)
utils.set_torch_seed(args.common.seed)
models, model_args, task = checkpoint_utils.load_model_ensemble_and_task(
[raw_args.path], arg_overrides={"data": args.task.data})
tgt_dict = task.target_dictionary
for model in models:
model.prepare_for_inference_(args)
model.cuda().eval()
if raw_args.fp16:
model = model.half()
model = models[0]
config = ExpressiveCodeDataConfig(args.task.data)
dataset = CodeDataset(
manifest=config.manifests[raw_args.subset],
dictionary=task.source_dictionary,
dur_dictionary=task.source_duration_dictionary,
f0_dictionary=task.source_f0_dictionary,
config=config,
discrete_dur=task.cfg.discrete_duration,
discrete_f0=task.cfg.discrete_f0,
log_f0=task.cfg.log_f0,
normalize_f0_mean=task.cfg.normalize_f0_mean,
normalize_f0_std=task.cfg.normalize_f0_std,
interpolate_f0=task.cfg.interpolate_f0,
shifts=task.cfg.stream_shifts,
return_filename=True,
strip_filename=False,
)
tgt_dict = task.target_dictionary
shifts = dataset.shifts.dur, dataset.shifts.f0
max_shift = max(shifts)
fname = raw_args.output
if world_size > 1:
fname += f"_{rank}"
output_file = open(fname, "w")
if raw_args.filter_names:
dataset = FilterNamesDataset(dataset, raw_args.filter_names)
dataset = InferenceDataset(dataset,
raw_args.prefix_length,
filter_short=True)
print(f"Dataset size {len(dataset)}")
sampler = (None if world_size == 1 else DistributedSampler(
dataset, num_replicas=world_size, rank=rank, shuffle=False))
dataloader = DataLoader(
dataset,
batch_size=1,
shuffle=False,
collate_fn=dataset.collater,
sampler=sampler,
)
Ts = raw_args.T_token, raw_args.T_duration, raw_args.T_f0
decoder = TemperatureDecoder(Ts,
discrete_dur=task.cfg.discrete_duration,
discrete_f0=task.cfg.discrete_f0)
dataset_size = len(dataset)
f0_decoder = None
if raw_args.f0_discretization_bounds:
assert task.cfg.discrete_f0
f0_decoder = Naive_F0_Decoder(raw_args.f0_discretization_bounds).cuda()
pbar = (tqdm.tqdm(
total=dataset_size if raw_args.max_samples is None else min(
raw_args.max_samples, dataset_size)) if world_size == 1 else None)
samples_produced = 0
for batch in dataloader:
if (raw_args.max_samples is not None
and samples_produced >= raw_args.max_samples):
break
prefix = batch["prefix"][0]
batch = explode_batch(batch, raw_args.batch_explosion_rate)
batch = move_to_cuda(batch)
if not raw_args.short_curcuit:
produced_tokens, produced_durations, produced_f0, _ = do_sampling(
models[0],
batch,
tgt_dict.eos(),
decoder,
autoregressive_steps=raw_args.max_length - prefix + max_shift,
teacher_force_tokens=raw_args.teacher_force_tokens,
match_duration=raw_args.match_duration,
teacher_force_duration=raw_args.teacher_force_duration,
teacher_force_f0=raw_args.teacher_force_f0,
)
# stip entries corresponding to <s>
produced_tokens = produced_tokens[:, 1:]
produced_durations = produced_durations[:, 1:]
produced_f0 = produced_f0[:, 1:]
else:
max_length = raw_args.max_length + max_shift
produced_tokens, produced_durations, produced_f0 = (
batch["target"][:, :max_length],
batch["dur_target"][:, :max_length],
batch["f0_target"][:, :max_length],
)
if f0_decoder is not None:
produced_f0 = f0_decoder(produced_f0)
produced_tokens, produced_durations, produced_f0 = (
produced_tokens.cpu().tolist(),
produced_durations.cpu().tolist(),
produced_f0.cpu().tolist(),
)
bsz = batch["target"].size(0)
assert bsz == raw_args.batch_explosion_rate
for i in range(bsz):
if (raw_args.max_samples is not None
and samples_produced >= raw_args.max_samples):
break
produced_tokens_i = produced_tokens[i]
produced_durations_i = produced_durations[i]
produced_f0_i = produced_f0[i]
(
produced_tokens_i,
produced_durations_i,
produced_f0_i,
) = realign_shifted_streams(produced_tokens_i,
produced_durations_i, produced_f0_i,
shifts)
produced_tokens_i, produced_durations_i, produced_f0_i = maybe_cut_eos(
produced_tokens_i, produced_durations_i, produced_f0_i,
tgt_dict.eos())
produced_tokens_i, produced_durations_i, produced_f0_i = maybe_filter_pad(
produced_tokens_i, produced_durations_i, produced_f0_i,
tgt_dict.pad())
if raw_args.match_duration:
# NB: here we cheat a bit and use that padding has duration 0
# so no need to re-align and remove padding
dur_target_i = batch["dur_target"][i, :].sum().item()
produced_tokens_i, produced_durations_i, produced_f0_i = match_duration(
produced_tokens_i, produced_durations_i, produced_f0_i,
dur_target_i)
if raw_args.cut_prompt:
produced_tokens_i, produced_durations_i, produced_f0_i = (
produced_tokens_i[prefix:],
produced_durations_i[prefix:],
produced_f0_i[prefix:],
)
prompt_fname = batch["filename"][0]
fname = str(
pathlib.Path(prompt_fname).with_suffix("")) + f"__{i}.wav"
token_stream = unroll_duration(produced_tokens_i,
produced_durations_i)
f0_stream = unroll_duration(produced_f0_i, produced_durations_i)
output_line = json.dumps({
"audio":
fname,
"prompt":
prompt_fname,
raw_args.code_type:
" ".join(map(str, token_stream)),
"duration":
round(
sum(produced_durations_i) *
CODETYPE_TO_FRAMETIME[raw_args.code_type],
3,
),
"raw_duration":
produced_durations_i,
"raw_f0":
produced_f0_i,
"f0": [round(f0, 3) for f0 in f0_stream],
})
print(output_line, file=output_file)
if pbar:
pbar.update(1)
samples_produced += 1
if raw_args.debug:
break
output_file.close()
if world_size > 1:
# important that everything is flushed before aggregating
torch.distributed.barrier()
if world_size > 1 and rank == 0:
with open(raw_args.output, "w") as fout:
for i in range(world_size):
f = raw_args.output + f"_{i}"
with open(f, "r") as fin:
fout.write(fin.read())
os.remove(f)
def main(args, task=None, model_state=None):
check_args(args)
if args.max_tokens is None and args.batch_size is None:
args.max_tokens = 4000000
logger.info(args)
use_cuda = torch.cuda.is_available() and not args.cpu
logger.info("| decoding with criterion {}".format(args.criterion))
task = tasks.setup_task(args)
# Load ensemble
if args.load_emissions:
models, criterions = [], []
task.load_dataset(args.gen_subset)
else:
logger.info("| loading model(s) from {}".format(args.path))
models, saved_cfg, task = checkpoint_utils.load_model_ensemble_and_task(
utils.split_paths(args.path),
arg_overrides=ast.literal_eval(args.model_overrides),
task=task,
suffix=args.checkpoint_suffix,
strict=(args.checkpoint_shard_count == 1),
num_shards=args.checkpoint_shard_count,
state=model_state,
)
optimize_models(args, use_cuda, models)
task.load_dataset(args.gen_subset, task_cfg=saved_cfg.task)
# Set dictionary
tgt_dict = task.target_dictionary
logger.info(
"| {} {} {} examples".format(
args.data, args.gen_subset, len(task.dataset(args.gen_subset))
)
)
# hack to pass transitions to W2lDecoder
if args.criterion == "asg_loss":
raise NotImplementedError("asg_loss is currently not supported")
# trans = criterions[0].asg.trans.data
# args.asg_transitions = torch.flatten(trans).tolist()
# Load dataset (possibly sharded)
itr = get_dataset_itr(args, task, models)
# Initialize generator
gen_timer = StopwatchMeter()
def build_generator(args):
w2l_decoder = getattr(args, "w2l_decoder", None)
if w2l_decoder == "viterbi":
from examples.speech_recognition.w2l_decoder import W2lViterbiDecoder
return W2lViterbiDecoder(args, task.target_dictionary)
elif w2l_decoder == "kenlm":
from examples.speech_recognition.w2l_decoder import W2lKenLMDecoder
return W2lKenLMDecoder(args, task.target_dictionary)
elif w2l_decoder == "fairseqlm":
from examples.speech_recognition.w2l_decoder import W2lFairseqLMDecoder
return W2lFairseqLMDecoder(args, task.target_dictionary)
else:
print(
"only flashlight decoders with (viterbi, kenlm, fairseqlm) options are supported at the moment"
)
# please do not touch this unless you test both generate.py and infer.py with audio_pretraining task
generator = build_generator(args)
if args.load_emissions:
generator = ExistingEmissionsDecoder(
generator, np.load(args.load_emissions, allow_pickle=True)
)
logger.info("loaded emissions from " + args.load_emissions)
num_sentences = 0
if args.results_path is not None and not os.path.exists(args.results_path):
os.makedirs(args.results_path)
max_source_pos = (
utils.resolve_max_positions(
task.max_positions(), *[model.max_positions() for model in models]
),
)
if max_source_pos is not None:
max_source_pos = max_source_pos[0]
if max_source_pos is not None:
max_source_pos = max_source_pos[0] - 1
if args.dump_emissions:
emissions = {}
if args.dump_features:
features = {}
models[0].bert.proj = None
else:
res_files = prepare_result_files(args)
errs_t = 0
lengths_t = 0
with progress_bar.build_progress_bar(args, itr) as t:
wps_meter = TimeMeter()
for sample in t:
sample = utils.move_to_cuda(sample) if use_cuda else sample
if "net_input" not in sample:
continue
prefix_tokens = None
if args.prefix_size > 0:
prefix_tokens = sample["target"][:, : args.prefix_size]
gen_timer.start()
if args.dump_emissions:
with torch.no_grad():
encoder_out = models[0](**sample["net_input"])
emm = models[0].get_normalized_probs(encoder_out, log_probs=True)
emm = emm.transpose(0, 1).cpu().numpy()
for i, id in enumerate(sample["id"]):
emissions[id.item()] = emm[i]
continue
elif args.dump_features:
with torch.no_grad():
encoder_out = models[0](**sample["net_input"])
feat = encoder_out["encoder_out"].transpose(0, 1).cpu().numpy()
for i, id in enumerate(sample["id"]):
padding = (
encoder_out["encoder_padding_mask"][i].cpu().numpy()
if encoder_out["encoder_padding_mask"] is not None
else None
)
features[id.item()] = (feat[i], padding)
continue
hypos = task.inference_step(generator, models, sample, prefix_tokens)
num_generated_tokens = sum(len(h[0]["tokens"]) for h in hypos)
gen_timer.stop(num_generated_tokens)
for i, sample_id in enumerate(sample["id"].tolist()):
speaker = None
# id = task.dataset(args.gen_subset).ids[int(sample_id)]
id = sample_id
toks = (
sample["target"][i, :]
if "target_label" not in sample
else sample["target_label"][i, :]
)
target_tokens = utils.strip_pad(toks, tgt_dict.pad()).int().cpu()
# Process top predictions
errs, length = process_predictions(
args,
hypos[i],
None,
tgt_dict,
target_tokens,
res_files,
speaker,
id,
)
errs_t += errs
lengths_t += length
wps_meter.update(num_generated_tokens)
t.log({"wps": round(wps_meter.avg)})
num_sentences += (
sample["nsentences"] if "nsentences" in sample else sample["id"].numel()
)
wer = None
if args.dump_emissions:
emm_arr = []
for i in range(len(emissions)):
emm_arr.append(emissions[i])
np.save(args.dump_emissions, emm_arr)
logger.info(f"saved {len(emissions)} emissions to {args.dump_emissions}")
elif args.dump_features:
feat_arr = []
for i in range(len(features)):
feat_arr.append(features[i])
np.save(args.dump_features, feat_arr)
logger.info(f"saved {len(features)} emissions to {args.dump_features}")
else:
if lengths_t > 0:
wer = errs_t * 100.0 / lengths_t
logger.info(f"WER: {wer}")
logger.info(
"| Processed {} sentences ({} tokens) in {:.1f}s ({:.2f}"
"sentences/s, {:.2f} tokens/s)".format(
num_sentences,
gen_timer.n,
gen_timer.sum,
num_sentences / gen_timer.sum,
1.0 / gen_timer.avg,
)
)
logger.info("| Generate {} with beam={}".format(args.gen_subset, args.beam))
return task, wer
def main(args, override_args=None):
utils.import_user_module(args)
assert args.max_tokens is not None or args.max_sentences is not None, \
'Must specify batch size either with --max-tokens or --max-sentences'
use_fp16 = args.fp16
use_cuda = torch.cuda.is_available() and not args.cpu
if use_cuda:
torch.cuda.set_device(args.device_id)
if override_args is not None:
overrides = vars(override_args)
overrides.update(eval(getattr(override_args, 'model_overrides', '{}')))
else:
overrides = None
# Load ensemble
logger.info('loading model(s) from {}'.format(args.path))
models, model_args, task = checkpoint_utils.load_model_ensemble_and_task(
[args.path],
arg_overrides=overrides,
suffix=getattr(args, "checkpoint_suffix", ""),
)
model = models[0]
# Move models to GPU
for model in models:
if use_fp16:
model.half()
if use_cuda:
model.cuda()
# Print args
logger.info(model_args)
# Build criterion
criterion = task.build_criterion(model_args)
criterion.eval()
for subset in args.valid_subset.split(','):
try:
task.load_dataset(subset, combine=False, epoch=1)
dataset = task.dataset(subset)
except KeyError:
raise Exception('Cannot find dataset: ' + subset)
# Initialize data iterator
itr = task.get_batch_iterator(
dataset=dataset,
max_tokens=args.max_tokens,
max_sentences=args.max_sentences,
max_positions=utils.resolve_max_positions(
task.max_positions(),
*[m.max_positions() for m in models],
),
ignore_invalid_inputs=args.skip_invalid_size_inputs_valid_test,
required_batch_size_multiple=args.required_batch_size_multiple,
seed=args.seed,
num_shards=args.distributed_world_size,
shard_id=args.distributed_rank,
num_workers=args.num_workers,
data_buffer_size=args.data_buffer_size,
).next_epoch_itr(shuffle=False)
progress = progress_bar.progress_bar(
itr,
log_format=args.log_format,
log_interval=args.log_interval,
prefix=f"valid on '{subset}' subset",
default_log_format=('tqdm'
if not args.no_progress_bar else 'simple'),
)
log_outputs = []
for i, sample in enumerate(progress):
sample = utils.move_to_cuda(sample) if use_cuda else sample
_loss, _sample_size, log_output = task.valid_step(
sample, model, criterion)
progress.log(log_output, step=i)
log_outputs.append(log_output)
if args.distributed_world_size > 1:
log_outputs = distributed_utils.all_gather_list(
log_outputs,
max_size=getattr(args, 'all_gather_list_size', 16384),
)
log_outputs = list(chain.from_iterable(log_outputs))
with metrics.aggregate() as agg:
task.reduce_metrics(log_outputs, criterion)
log_output = agg.get_smoothed_values()
progress.print(log_output, tag=subset, step=i)
def main(cfg: DictConfig, override_args=None, **unused_kwargs):
if isinstance(cfg, Namespace):
cfg = convert_namespace_to_omegaconf(cfg)
utils.import_user_module(cfg.common)
use_fp16 = cfg.common.fp16
use_cuda = torch.cuda.is_available() and not cfg.common.cpu
if use_cuda:
torch.cuda.set_device(cfg.distributed_training.device_id)
if override_args is not None:
overrides = vars(override_args)
overrides.update(eval(getattr(override_args, "model_overrides", "{}")))
else:
overrides = None
logger.info(cfg)
# Load ensemble
logger.info("loading model(s) from {}".format(cfg.common_eval.path))
# reduce tokens per sample by the required context window size
cfg.task.tokens_per_sample -= cfg.eval_lm.context_window
models, model_args, task = checkpoint_utils.load_model_ensemble_and_task(
[cfg.common_eval.path],
arg_overrides=overrides,
suffix=cfg.checkpoint.checkpoint_suffix,
strict=(cfg.checkpoint.checkpoint_shard_count == 1),
num_shards=cfg.checkpoint.checkpoint_shard_count,
)
# Load dataset splits
gen_subset = cfg.dataset.gen_subset
task.load_dataset(gen_subset)
dataset = task.dataset(gen_subset)
if cfg.eval_lm.context_window > 0:
dataset = LMContextWindowDataset(
dataset=dataset,
tokens_per_sample=cfg.task.tokens_per_sample,
context_window=cfg.eval_lm.context_window,
pad_idx=task.source_dictionary.pad(),
)
logger.info("{} {} {} examples".format(cfg.task.data, gen_subset, len(dataset)))
# Optimize ensemble for generation and set the source and dest dicts on the model (required by scorer)
for model in models:
if use_fp16:
model.half()
if use_cuda and not cfg.distributed_training.pipeline_model_parallel:
model.cuda()
model.prepare_for_inference_(cfg)
assert len(models) > 0
logger.info(
"num. model params: {}".format(sum(p.numel() for p in models[0].parameters()))
)
itr = task.get_batch_iterator(
dataset=dataset,
max_tokens=cfg.dataset.max_tokens or 36000,
max_sentences=cfg.dataset.batch_size,
max_positions=utils.resolve_max_positions(
*[model.max_positions() for model in models]
),
ignore_invalid_inputs=True,
num_shards=max(
cfg.dataset.num_shards,
cfg.distributed_training.distributed_world_size,
),
shard_id=max(
cfg.dataset.shard_id,
cfg.distributed_training.distributed_rank,
),
num_workers=cfg.dataset.num_workers,
data_buffer_size=cfg.dataset.data_buffer_size,
).next_epoch_itr(shuffle=False)
progress = progress_bar.progress_bar(
itr,
log_format=cfg.common.log_format,
log_interval=cfg.common.log_interval,
default_log_format=("tqdm" if not cfg.common.no_progress_bar else "simple"),
)
gen_timer = StopwatchMeter()
scorer = SequenceScorer(task.target_dictionary, cfg.eval_lm.softmax_batch)
score_sum = 0.0
count = 0
if cfg.common_eval.remove_bpe is not None:
if cfg.common_eval.remove_bpe == "sentencepiece":
raise NotImplementedError
else:
bpe_cont = cfg.common_eval.remove_bpe.rstrip()
bpe_toks = {
i
for i in range(len(task.source_dictionary))
if task.source_dictionary[i].endswith(bpe_cont)
}
bpe_len = len(bpe_cont)
else:
bpe_toks = None
bpe_len = 0
word_stats = dict()
wps_meter = TimeMeter()
for sample in progress:
if "net_input" not in sample:
continue
sample = utils.move_to_cuda(sample) if use_cuda else sample
gen_timer.start()
hypos = scorer.generate(models, sample)
gen_timer.stop(sample["ntokens"])
for i, hypos_i in enumerate(hypos):
hypo = hypos_i[0]
sample_id = sample["id"][i]
tokens = hypo["tokens"]
tgt_len = tokens.numel()
pos_scores = hypo["positional_scores"].float()
if cfg.task.add_bos_token:
assert hypo["tokens"][0].item() == task.target_dictionary.bos()
tokens = tokens[1:]
pos_scores = pos_scores[1:]
skipped_toks = 0
if bpe_toks is not None:
for i in range(tgt_len - 1):
if tokens[i].item() in bpe_toks:
skipped_toks += 1
pos_scores[i + 1] += pos_scores[i]
pos_scores[i] = 0
inf_scores = pos_scores.eq(float("inf")) | pos_scores.eq(float("-inf"))
if inf_scores.any():
logger.info(
"skipping tokens with inf scores:",
task.target_dictionary.string(tokens[inf_scores.nonzero()]),
)
pos_scores = pos_scores[(~inf_scores).nonzero()]
score_sum += pos_scores.sum().cpu()
count += pos_scores.numel() - skipped_toks
if cfg.eval_lm.output_word_probs or cfg.eval_lm.output_word_stats:
w = ""
word_prob = []
is_bpe = False
for i in range(len(tokens)):
w_ind = tokens[i].item()
w += task.source_dictionary[w_ind]
if bpe_toks is not None and w_ind in bpe_toks:
w = w[:-bpe_len]
is_bpe = True
else:
word_prob.append((w, pos_scores[i].item()))
next_prob = None
ind = i + 1
while ind < len(tokens):
if pos_scores[ind].item() != 0:
next_prob = pos_scores[ind]
break
ind += 1
word_stats.setdefault(w, WordStat(w, is_bpe)).add(
pos_scores[i].item(), next_prob
)
is_bpe = False
w = ""
if cfg.eval_lm.output_word_probs:
logger.info(
str(int(sample_id))
+ " "
+ (
"\t".join(
"{} [{:2f}]".format(x[0], x[1]) for x in word_prob
)
)
)
wps_meter.update(sample["ntokens"])
progress.log({"wps": round(wps_meter.avg)})
avg_nll_loss = -score_sum / count / math.log(2) # convert to base 2
logger.info(
"Evaluated {} tokens in {:.1f}s ({:.2f} tokens/s)".format(
gen_timer.n, gen_timer.sum, 1.0 / gen_timer.avg
)
)
logger.info(
"Loss (base 2): {:.4f}, Perplexity: {:.2f}".format(
avg_nll_loss, 2 ** avg_nll_loss
)
)
if cfg.eval_lm.output_word_stats:
for ws in sorted(word_stats.values(), key=lambda x: x.count, reverse=True):
logger.info(ws)
def main(rank, world_size, args):
start = time.time()
if world_size > 1:
torch.distributed.init_process_group(backend="gloo",
init_method="env://",
world_size=world_size,
rank=rank)
torch.cuda.set_device(rank % torch.cuda.device_count())
raw_args = args
args = convert_namespace_to_omegaconf(args)
if args.common.seed is not None:
np.random.seed(args.common.seed)
utils.set_torch_seed(args.common.seed)
models, model_args, task = checkpoint_utils.load_model_ensemble_and_task(
[raw_args.path], arg_overrides={"data": args.task.data})
tgt_dict = task.target_dictionary
for model in models:
model.prepare_for_inference_(args)
model.cuda().eval()
if raw_args.fp16:
model = model.half()
model = models[0]
config = ExpressiveCodeDataConfig(args.task.data)
dataset = CodeDataset(
manifest=config.manifests[raw_args.eval_subset],
dictionary=task.source_dictionary,
dur_dictionary=task.source_duration_dictionary,
f0_dictionary=task.source_f0_dictionary,
config=config,
discrete_dur=task.cfg.discrete_duration,
discrete_f0=task.cfg.discrete_f0,
log_f0=task.cfg.log_f0,
normalize_f0_mean=task.cfg.normalize_f0_mean,
normalize_f0_std=task.cfg.normalize_f0_std,
interpolate_f0=task.cfg.interpolate_f0,
shifts=task.cfg.stream_shifts,
return_filename=True,
strip_filename=False,
return_continuous_f0=raw_args.dequantize_prosody,
)
if raw_args.filter_names:
dataset = FilterNamesDataset(dataset, raw_args.filter_names)
criterion = task.build_criterion(model_args.criterion)
name2metric = {
"continuation": continuation,
"teacher_force_everything": teacher_force_everything,
"correlation": correlation,
}
name2keys = {
"continuation": (
"Token BLEU3",
"Duration NLL",
"Duration MAE",
"F0 NLL",
"F0 MAE",
"F0 sum",
"F0 sum_sq",
"Dur sum",
"Dur sum_sq",
),
"teacher_force_everything": ("token_loss", "duration_loss", "f0_loss"),
"correlation": ("Duration corr", "F0 corr"),
}
metric_name = raw_args.metric
metric = name2metric[metric_name]
results = metric(raw_args, dataset, model, criterion, tgt_dict, rank,
world_size)
values = None
if metric_name not in [
"correlation",
]:
values, normalizers = results
values = maybe_aggregate_normalize(values, normalizers, world_size)
elif metric_name == "correlation":
values = maybe_aggregate_correlations(results, world_size)
else:
assert False
assert values is not None
summary = dict(zip(name2keys[raw_args.metric], values.tolist()))
if metric_name == "continuation":
summary["F0 Std"] = np.sqrt(-summary["F0 sum"]**2 +
summary["F0 sum_sq"])
summary["Dur Std"] = np.sqrt(-summary["Dur sum"]**2 +
summary["Dur sum_sq"])
del summary["F0 sum"]
del summary["F0 sum_sq"]
del summary["Dur sum"]
del summary["Dur sum_sq"]
summary["metric"] = metric_name
if rank == 0:
print(summary)
if raw_args.wandb:
wandb_results(summary, raw_args)
print("# finished in ", time.time() - start, "seconds")
def main(cfg: DictConfig, **unused_kwargs):
if isinstance(cfg, Namespace):
cfg = convert_namespace_to_omegaconf(cfg)
utils.import_user_module(cfg.common)
logger.info(cfg)
if cfg.eval_lm.context_window > 0:
# reduce tokens per sample by the required context window size
cfg.task.tokens_per_sample -= cfg.eval_lm.context_window
# Initialize the task using the current *cfg*
task = tasks.setup_task(cfg.task)
# Load ensemble
logger.info("loading model(s) from {}".format(cfg.common_eval.path))
models, model_args, task = checkpoint_utils.load_model_ensemble_and_task(
[cfg.common_eval.path],
arg_overrides=eval(cfg.common_eval.model_overrides),
suffix=cfg.checkpoint.checkpoint_suffix,
strict=(cfg.checkpoint.checkpoint_shard_count == 1),
num_shards=cfg.checkpoint.checkpoint_shard_count,
task=task,
)
use_fp16 = cfg.common.fp16
use_cuda = torch.cuda.is_available() and not cfg.common.cpu
if use_cuda:
torch.cuda.set_device(cfg.distributed_training.device_id)
# Optimize ensemble for generation and set the source and dest dicts on the model
# (required by scorer)
for model in models:
if use_fp16:
model.half()
if use_cuda and not cfg.distributed_training.pipeline_model_parallel:
model.cuda()
model.prepare_for_inference_(cfg)
assert len(models) > 0
logger.info(
"num. model params: {:,}".format(sum(p.numel() for p in models[0].parameters()))
)
# Load dataset splits
task.load_dataset(cfg.dataset.gen_subset)
dataset = task.dataset(cfg.dataset.gen_subset)
logger.info(
"{} {} {:,} examples".format(
cfg.task.data, cfg.dataset.gen_subset, len(dataset)
)
)
itr = task.eval_lm_dataloader(
dataset=dataset,
max_tokens=cfg.dataset.max_tokens or 36000,
batch_size=cfg.dataset.batch_size,
max_positions=utils.resolve_max_positions(
*[model.max_positions() for model in models]
),
num_shards=max(
cfg.dataset.num_shards,
cfg.distributed_training.distributed_world_size,
),
shard_id=max(
cfg.dataset.shard_id,
cfg.distributed_training.distributed_rank,
),
num_workers=cfg.dataset.num_workers,
data_buffer_size=cfg.dataset.data_buffer_size,
context_window=cfg.eval_lm.context_window,
)
itr = progress_bar.progress_bar(
itr,
log_format=cfg.common.log_format,
log_interval=cfg.common.log_interval,
default_log_format=("tqdm" if not cfg.common.no_progress_bar else "simple"),
)
results = eval_lm(
models=models,
source_dictionary=task.source_dictionary,
batch_iterator=itr,
post_process=cfg.common_eval.post_process,
output_word_probs=cfg.eval_lm.output_word_probs,
output_word_stats=cfg.eval_lm.output_word_stats,
target_dictionary=task.target_dictionary,
softmax_batch=cfg.eval_lm.softmax_batch,
remove_bos_token=getattr(cfg.task, "add_bos_token", False),
)
logger.info(
"Loss (base 2): {:.4f}, Perplexity: {:.2f}".format(
results["loss"], results["perplexity"]
)
)
return results
def main(args, override_args, init_distributed=False):
utils.import_user_module(args)
assert args.max_tokens is not None or args.max_sentences is not None, \
'Must specify batch size either with --max-tokens or --max-sentences'
# Initialize CUDA and distributed training
if torch.cuda.is_available() and not args.cpu:
torch.cuda.set_device(args.device_id)
np.random.seed(args.seed)
torch.manual_seed(args.seed)
if init_distributed:
args.distributed_rank = distributed_utils.distributed_init(args)
# Print args
print(args)
use_fp16 = args.fp16
use_cuda = torch.cuda.is_available() and not args.cpu
if override_args is not None:
overrides = vars(override_args)
overrides.update(eval(getattr(override_args, 'model_overrides', '{}')))
else:
overrides = None
# Load ensemble
print('| loading model(s) from {}'.format(args.path))
models, model_args, task = checkpoint_utils.load_model_ensemble_and_task(
[args.path],
arg_overrides=overrides,
)
model = models[0]
# Move models to GPU
for model in models:
if use_fp16:
model.half()
if use_cuda:
model.cuda()
# Print args
print(model_args)
# Build criterion
criterion = task.build_criterion(model_args)
criterion.eval()
# Load valid dataset (we load training data below, based on the latest checkpoint)
for valid_sub_split in args.valid_subset.split(','):
task.load_dataset(valid_sub_split, combine=False, epoch=0)
# Build trainer
trainer = validator(args, task, model, criterion)
print('| training on {} GPUs'.format(args.distributed_world_size))
print('| max tokens per GPU = {} and max sentences per GPU = {}'.format(
args.max_tokens,
args.max_sentences,
))
# Load the latest checkpoint if one is available and restore the
# corresponding train iterator
valid_subsets = args.valid_subset.split(',')
valid_losses = validate(args, trainer, task, valid_subsets)
