Support flat native NeMo lists in idxpacks

Signed-off-by: Piotr Żelasko <pzelasko@nvidia.com>

Author: pzelasko

nemo/collections/common/data/lhotse/cutset.py

@@ -1662,23 +1662,45 @@ def read_nemo_manifest(config) -> tuple[CutSet, bool]:
     }
     tar_kwargs_extra = {"indexed": indexed, **indexed_extra, **pack_extra} if indexed else {}
     is_tarred = config.get("tarred_audio_filepaths") is not None
+    manifest_filepath = config.manifest_filepath
+    manifest_is_scalar = isinstance(manifest_filepath, (str, Path))
+    manifest_is_flat_list = (
+        isinstance(manifest_filepath, (list, tuple, ListConfig))
+        and bool(manifest_filepath)
+        and all(isinstance(item, (str, Path)) for item in manifest_filepath)
+    )
+    tarred_audio_filepaths = config.get("tarred_audio_filepaths")
+    tar_is_scalar = isinstance(tarred_audio_filepaths, (str, Path))
+    tar_is_flat_list = (
+        isinstance(tarred_audio_filepaths, (list, tuple, ListConfig))
+        and bool(tarred_audio_filepaths)
+        and all(isinstance(item, (str, Path)) for item in tarred_audio_filepaths)
+    )
     if index_pack is not None:
-        if not isinstance(config.manifest_filepath, (str, Path)):
+        if not (manifest_is_scalar or manifest_is_flat_list):
             raise ValueError(
                 "Packed native NeMo datasets require manifest_filepath to be "
-                "a string/Path (brace expansion is supported); list forms are not."
+                "a string/Path or a non-empty flat list of strings/Paths; nested "
+                "and weighted list forms are not supported."
             )
-        if is_tarred and not isinstance(config.tarred_audio_filepaths, (str, Path)):
+        if is_tarred and not (tar_is_scalar or tar_is_flat_list):
             raise ValueError(
                 "Packed native NeMo datasets require tarred_audio_filepaths to "
-                "be a string/Path (brace expansion is supported); list forms are not."
+                "be a string/Path or a non-empty flat list of strings/Paths; nested "
+                "list forms are not supported."
+            )
+        if is_tarred and manifest_is_flat_list != tar_is_flat_list:
+            raise ValueError(
+                "Packed native NeMo manifest_filepath and tarred_audio_filepaths "
+                "must both use scalar path specs or both use non-empty flat lists."
             )
-    if isinstance(config.manifest_filepath, (str, Path)):
+    packed_flat_list = index_pack is not None and (manifest_is_flat_list or (is_tarred and tar_is_flat_list))
+    if manifest_is_scalar or packed_flat_list:
         if is_tarred and not metadata_only:
             cuts = CutSet(
                 LazyNeMoTarredIterator(
-                    config.manifest_filepath,
-                    tar_paths=config.tarred_audio_filepaths,
+                    manifest_filepath,
+                    tar_paths=tarred_audio_filepaths,
                     skip_missing_manifest_entries=config.get("skip_missing_manifest_entries", False),
                     slice_length=config.get("slice_length", None),
                     **tar_kwargs_extra,
@@ -1688,9 +1710,7 @@ def read_nemo_manifest(config) -> tuple[CutSet, bool]:
             if not force_finite:
                 cuts = cuts.repeat(preserve_id=True)
         else:
-            cuts = CutSet(
-                LazyNeMoIterator(config.manifest_filepath, **notar_kwargs, **notar_kwargs_extra, **common_kwargs)
-            )
+            cuts = CutSet(LazyNeMoIterator(manifest_filepath, **notar_kwargs, **notar_kwargs_extra, **common_kwargs))
     else:
         # Format option 1:
         #   Assume it's [[path1], [path2], ...] (same for tarred_audio_filepaths).

scripts/dataloading/convert_indexes_to_idxpack.py

@@ -230,6 +230,25 @@ def _require_scalar_spec(value, field: str) -> None:
         )
 
 
+def _is_nonempty_flat_path_list(value) -> bool:
+    return (
+        isinstance(value, (list, tuple, ListConfig))
+        and bool(value)
+        and all(isinstance(item, (str, Path)) for item in value)
+    )
+
+
+def _require_scalar_or_flat_path_list(value, field: str) -> None:
+    if isinstance(value, (str, Path)):
+        return
+    if _is_nonempty_flat_path_list(value):
+        return
+    raise ValueError(
+        f"Packed native NeMo {field} must be a string/Path or a non-empty flat "
+        "list of strings/Paths; nested and weighted list forms are not supported."
+    )
+
+
 def _shard_number(path: str) -> int | None:
     matches = re.findall(r"\d+", Path(path).stem)
     return int(matches[-1]) if matches else None
@@ -256,6 +275,29 @@ def _validate_native_pair(manifests: list[str], tars: list[str]) -> None:
         )
 
 
+def _expand_flat_native_pairs(manifest_specs, tar_specs) -> tuple[list[str], list[str]]:
+    if len(manifest_specs) != len(tar_specs):
+        raise ValueError(
+            "Packed native NeMo flat lists require one tar path spec per "
+            f"manifest path spec: manifests={len(manifest_specs)}, tars={len(tar_specs)}"
+        )
+
+    manifests: list[str] = []
+    tars: list[str] = []
+    for position, (manifest_spec, tar_spec) in enumerate(zip(manifest_specs, tar_specs)):
+        pair_manifests = _expand_jsonl(manifest_spec)
+        pair_tars = _expand_tars(tar_spec)
+        if len(pair_manifests) != len(pair_tars):
+            raise ValueError(
+                "Packed native NeMo flat lists require each positional manifest/tar "
+                f"pair to expand to the same number of shards; position={position}, "
+                f"manifests={len(pair_manifests)}, tars={len(pair_tars)}"
+            )
+        manifests.extend(pair_manifests)
+        tars.extend(pair_tars)
+    return manifests, tars
+
+
 def discover_pack_collections(
     entry,
     collections: Optional[list[IndexPackCollectionSpec]] = None,
@@ -305,25 +347,42 @@ def discover_pack_collections(
 
     if typ in {"nemo", "nemo_tarred", "share_gpt", *_TRANSFORM_TYPES} and entry.get("manifest_filepath") is not None:
         raw = entry.get("manifest_filepath")
-        _require_scalar_spec(raw, "manifest_filepath")
-        manifests = _expand_jsonl(raw)
+        if typ == "share_gpt":
+            _require_scalar_spec(raw, "manifest_filepath")
+        else:
+            _require_scalar_or_flat_path_list(raw, "manifest_filepath")
+        raw_tars = entry.get("tarred_audio_filepaths")
+        if raw_tars is None:
+            manifests = _expand_jsonl(raw)
+            tars = None
+        else:
+            if typ == "share_gpt":
+                raise NotImplementedError(
+                    "Packed ShareGPT supports JSONL manifests with direct/remote "
+                    "audio paths, not paired audio tar files."
+                )
+            _require_scalar_or_flat_path_list(raw_tars, "tarred_audio_filepaths")
+            manifest_is_flat_list = _is_nonempty_flat_path_list(raw)
+            tar_is_flat_list = _is_nonempty_flat_path_list(raw_tars)
+            if manifest_is_flat_list != tar_is_flat_list:
+                raise ValueError(
+                    "Packed native NeMo manifest_filepath and tarred_audio_filepaths "
+                    "must both use scalar path specs or both use non-empty flat lists."
+                )
+            if manifest_is_flat_list:
+                manifests, tars = _expand_flat_native_pairs(raw, raw_tars)
+            else:
+                manifests = _expand_jsonl(raw)
+                tars = _expand_tars(raw_tars)
+                _validate_native_pair(manifests, tars)
         _add_collection(
             collections,
             role="manifest",
             kind=JSONL,
             source_spec=raw,
             paths=manifests,
         )
-        if entry.get("tarred_audio_filepaths") is not None:
-            if typ == "share_gpt":
-                raise NotImplementedError(
-                    "Packed ShareGPT supports JSONL manifests with direct/remote "
-                    "audio paths, not paired audio tar files."
-                )
-            raw_tars = entry.get("tarred_audio_filepaths")
-            _require_scalar_spec(raw_tars, "tarred_audio_filepaths")
-            tars = _expand_tars(raw_tars)
-            _validate_native_pair(manifests, tars)
+        if tars is not None:
             _add_collection(
                 collections,
                 role="tar",

tests/collections/common/test_lhotse_index_pack.py

@@ -21,10 +21,12 @@
 from click.testing import CliRunner
 from lhotse.index_pack import IndexPack, IndexPackCollectionSpec, index_pack_collection_key, write_index_pack
 from lhotse.indexing import create_jsonl_index
+from omegaconf import OmegaConf
 from scripts.dataloading import convert_indexes_to_idxpack as converter
 from scripts.dataloading.convert_indexes_to_idxpack import main
 
 from nemo.collections.common.data.lhotse import nemo_adapters, text_adapters
+from nemo.collections.common.data.lhotse.cutset import read_nemo_manifest
 from nemo.collections.common.data.lhotse.indexed_adapters import create_tar_index as create_nemo_tar_index
 from nemo.collections.common.data.lhotse.nemo_adapters import LazyNeMoIterator, LazyNeMoTarredIterator
 
@@ -146,6 +148,117 @@ def test_convert_input_cfg_sidecars_to_one_index_pack(tmp_path):
         assert pack.num_segments == 2
 
 
+def test_flat_native_lists_use_aggregate_pack_and_preserve_positional_pairs(tmp_path, monkeypatch):
+    manifests = []
+    declared_tar_paths = []
+    expected_texts = []
+    for position, (manifest_id, tar_id) in enumerate(((3076, 2), (4100, 0), (5200, 1))):
+        manifest = tmp_path / f"manifest_{manifest_id}.jsonl"
+        member = f"sample-{position}.wav"
+        text = f"text-{position}"
+        manifest.write_text(
+            json.dumps(
+                {
+                    "audio_filepath": member,
+                    "duration": 1.0,
+                    "text": text,
+                    "lang": "en",
+                }
+            )
+            + "\n"
+        )
+        create_jsonl_index(manifest)
+        manifests.append(str(manifest))
+        declared_tar_paths.append(f"ais://bucket/audio_{tar_id}.tar")
+        expected_texts.append(text)
+
+    input_cfg = tmp_path / "flat-lists.yaml"
+    input_cfg.write_text(
+        yaml.safe_dump(
+            {
+                "type": "nemo_tarred",
+                "manifest_filepath": manifests,
+                "tarred_audio_filepaths": declared_tar_paths,
+            }
+        )
+    )
+    output = tmp_path / "flat-lists.idxpack"
+    result = CliRunner().invoke(
+        main,
+        [
+            "--output",
+            str(output),
+            "--native-tar-paths-only",
+            str(input_cfg),
+        ],
+    )
+    assert result.exit_code == 0, result.output
+
+    with IndexPack(output) as pack:
+        manifest_collection = pack.collection(index_pack_collection_key("manifest", "jsonl", manifests))
+        tar_collection = pack.collection(index_pack_collection_key("tar", "nemo_tar", declared_tar_paths))
+        assert manifest_collection.sequence_count == 3
+        assert tar_collection.sequence_count == 3
+        assert [tar_collection.path_for_shard(idx) for idx in range(3)] == declared_tar_paths
+
+    monkeypatch.setenv("USE_AIS_GET_BATCH", "true")
+    config = OmegaConf.create(
+        {
+            "manifest_filepath": manifests,
+            "tarred_audio_filepaths": declared_tar_paths,
+            "indexed": True,
+            "index_pack": str(output),
+            "force_finite": True,
+        }
+    )
+    cuts, is_tarred = read_nemo_manifest(config)
+    cuts = list(cuts)
+    assert is_tarred
+    assert [cut.supervisions[0].text for cut in cuts] == expected_texts
+    assert [cut.recording.sources[0].source for cut in cuts] == [
+        f"{tar_path}/sample-{position}.wav" for position, tar_path in enumerate(declared_tar_paths)
+    ]
+
+
+def test_converter_rejects_mixed_scalar_and_flat_native_pairs(tmp_path):
+    input_cfg = tmp_path / "mixed-path-forms.yaml"
+    input_cfg.write_text(
+        yaml.safe_dump(
+            {
+                "type": "nemo_tarred",
+                "manifest_filepath": [str(tmp_path / "manifest.jsonl")],
+                "tarred_audio_filepaths": str(tmp_path / "audio.tar"),
+            }
+        )
+    )
+
+    result = CliRunner().invoke(
+        main,
+        ["--dry-run", "--output", str(tmp_path / "mixed.idxpack"), str(input_cfg)],
+    )
+    assert result.exit_code != 0
+    assert "must both use scalar path specs or both use non-empty flat lists" in str(result.exception)
+
+
+def test_converter_rejects_nested_native_manifest_lists(tmp_path):
+    input_cfg = tmp_path / "nested-lists.yaml"
+    input_cfg.write_text(
+        yaml.safe_dump(
+            {
+                "type": "nemo",
+                "manifest_filepath": [[str(tmp_path / "manifest.jsonl"), 0.5]],
+            }
+        )
+    )
+
+    result = CliRunner().invoke(
+        main,
+        ["--dry-run", "--output", str(tmp_path / "nested.idxpack"), str(input_cfg)],
+    )
+    assert result.exit_code != 0
+    assert "nested and weighted list forms are not supported" in str(result.exception)
+
+
 def test_lazy_nemo_iterator_uses_pack_without_expanding_shards(tmp_path, monkeypatch):
     manifests = []
     source_spec = str(tmp_path / "manifest__OP_0..1_CL_.jsonl")