feat: Add torchaudio dependency and implement audiohandler for audio …

pyproject.toml

@@ -90,6 +90,7 @@ dependencies = [
     "open-clip-torch>=3.2.0",
     "mlflow>=3.5.0",
     "torch>=2.6.0",
+    "torchaudio",
 ]
 
 
@@ -110,6 +111,7 @@ no-build-isolation-package = [
 prerelease = "allow"
 override-dependencies = [
     "torch; sys_platform == 'never'",
+    "torchaudio; sys_platform == 'never'",
     "torchvision; sys_platform == 'never'",
     "triton; sys_platform == 'never'",
     "transformer-engine @ git+https://github.com/NVIDIA/TransformerEngine.git@5671fd3675906cda1ade26c24a65d3dedd88eb89",

src/megatron/bridge/recipes/qwen_vl/data/energon/task_encoder.py

@@ -23,6 +23,8 @@
 
 import numpy as np
 import torch
+
+import torchaudio
 from megatron.energon import Batch, DefaultTaskEncoder
 from megatron.energon.epathlib.epath import EPath
 from megatron.energon.flavors.base_dataset import Sample
@@ -206,6 +208,7 @@ class ChatMLSample(Sample):
     conversation: str  # JSON string of GPT-format conversations
     imgs: Optional[List[torch.Tensor]] = None
     videos: Optional[List[List[torch.Tensor]]] = None
+    audios: Optional[List[torch.Tensor]] = None
 
 
 class videohandler:
@@ -230,6 +233,26 @@ def __call__(self, key, data):
         return data
 
 
+class audiohandler:
+    """Create an audio handler."""
+
+    def __init__(self):
+        self.extensions = ["wavs", "mp3s"]
+
+    def __call__(self, key, data):
+        """Perform audio decoding."""
+
+        extension = re.sub(r".*[.]", "", key)
+        if extension not in self.extensions:
+            return None
+
+        data_list = pickle.loads(data)
+        audio_list = []
+        for data in data_list:
+            audio_list.append(torchaudio.load(io.BytesIO(data)))
+        return audio_list
+
+
 class ChatMLWebdataset(DefaultDecoderWebdatasetFactory[ChatMLSample]):
     """Webdataset factory for multi-turn ChatML samples with multimodal support.
 
@@ -246,6 +269,7 @@ def __init__(self, path: EPath, *, auto_decode: bool = True, **kwargs):
                 [
                     imagehandler(self.image_decode),
                     videohandler(self.image_decode),
+                    audiohandler(),
                 ]
             )