RASMUS/Finnish-ASR-Canary-v2

1# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.2#3# Licensed under the Apache License, Version 2.0 (the "License");4# you may not use this file except in compliance with the License.5# You may obtain a copy of the License at6#7#     http://www.apache.org/licenses/LICENSE-2.08#9# Unless required by applicable law or agreed to in writing, software10# distributed under the License is distributed on an "AS IS" BASIS,11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.12# See the License for the specific language governing permissions and13# limitations under the License.14 15import argparse16from itertools import islice17from pathlib import Path18 19from lhotse.cut import Cut20from lhotse.dataset.sampling.dynamic_bucketing import estimate_duration_buckets21from omegaconf import OmegaConf22 23from nemo.collections.common.data.lhotse.cutset import read_cutset_from_config24from nemo.collections.common.data.lhotse.dataloader import LhotseDataLoadingConfig25 26 27def parse_args():28    parser = argparse.ArgumentParser(29        description="Estimate duration bins for Lhotse dynamic bucketing using a sample of the input dataset. "30        "The dataset is read either from one or more manifest files and supports data weighting.",31        formatter_class=argparse.ArgumentDefaultsHelpFormatter,32    )33    parser.add_argument(34        "input",35        help='Data input. Options: '36        '1) "path.json" - any single NeMo manifest; '37        '2) "[[path1.json],[path2.json],...]" - any collection of NeMo manifests; '38        '3) "[[path1.json,weight1],[path2.json,weight2],...]" - any collection of weighted NeMo manifests; '39        '4) "input_cfg.yaml" - a new option supporting input configs, same as in model training \'input_cfg\' arg; '40        '5) "path/to/shar_data" - a path to Lhotse Shar data directory; '41        '6) "key=val" - in case none of the previous variants cover your case: "key" is the key you\'d use in NeMo training config with its corresponding value ',42    )43    parser.add_argument("-b", "--buckets", type=int, default=30, help="The desired number of buckets.")44    parser.add_argument(45        "-n",46        "--num_examples",47        type=int,48        default=-1,49        help="The number of examples (utterances) to estimate the bins. -1 means use all data "50        "(be careful: it could be iterated over infinitely).",51    )52    parser.add_argument(53        "-l",54        "--min_duration",55        type=float,56        default=-float("inf"),57        help="If specified, we'll filter out utterances shorter than this.",58    )59    parser.add_argument(60        "-u",61        "--max_duration",62        type=float,63        default=float("inf"),64        help="If specified, we'll filter out utterances longer than this.",65    )66    parser.add_argument(67        "-q", "--quiet", type=bool, default=False, help="When specified, only print the estimated duration bins."68    )69    return parser.parse_args()70 71 72def main():73    args = parse_args()74    if '=' in args.input:75        inp_arg = args.input76    elif args.input.endswith(".yaml"):77        inp_arg = f"input_cfg={args.input}"78    elif Path(args.input).is_dir():79        inp_arg = f"shar_path={args.input}"80    else:81        inp_arg = f"manifest_filepath={args.input}"82    config = OmegaConf.merge(83        OmegaConf.structured(LhotseDataLoadingConfig),84        OmegaConf.from_dotlist([inp_arg, "metadata_only=true"]),85    )86    cuts, _ = read_cutset_from_config(config)87    min_dur, max_dur = args.min_duration, args.max_duration88    nonaudio, discarded, tot = 0, 0, 089    observed_max_dur = 090 91    def duration_ok(cut) -> bool:92        nonlocal nonaudio, discarded, tot, observed_max_dur93        tot += 194        if not isinstance(cut, Cut):95            nonaudio += 196            return False97        if not (min_dur <= cut.duration <= max_dur):98            discarded += 199            return False100        observed_max_dur = max(cut.duration, observed_max_dur)101        return True102 103    cuts = cuts.filter(duration_ok)104    if (N := args.num_examples) > 0:105        cuts = islice(cuts, N)106    duration_bins = estimate_duration_buckets(cuts, num_buckets=args.buckets)107    duration_bins = f"[{','.join(str(round(b, ndigits=5)) for b in duration_bins)}]"108    if args.quiet:109        print(duration_bins)110        return111    if discarded:112        ratio = discarded / tot113        print(f"Note: we discarded {discarded}/{tot} ({ratio:.2%}) utterances due to min/max duration filtering.")114    if nonaudio:115        print(f"Note: we discarded {nonaudio} non-audio examples found during iteration.")116    print(f"Used {tot - nonaudio - discarded} examples for the estimation.")117    print("Use the following options in your config:")118    print(f"\tnum_buckets={args.buckets}")119    print(f"\tbucket_duration_bins={duration_bins}")120    print(f"\tmax_duration={observed_max_dur}")121 122 123if __name__ == "__main__":124    main()125