[Prototype] Multimodal Audio

fast_llm/data/data/gpt/data.py

@@ -34,6 +34,8 @@ class GPTBatch:
     sequence_lengths: list[torch.Tensor] | None = None
     images: list[torch.Tensor] | None = None
     image_positions: list[torch.Tensor] | None = None
+    audio: list[torch.Tensor] | None = None
+    audio_positions: list[torch.Tensor] | None = None
 
 
 def gpt_data_collate_fn(batch: list[GPTSample], sampling_parameters: GPTSamplingParameters) -> GPTBatch:
@@ -54,16 +56,34 @@ def gpt_data_collate_fn(batch: list[GPTSample], sampling_parameters: GPTSampling
             batch_images.append([])
     batch_image_positions = []
     for sample in batch:
-        if sample.image_positions is not None:
+        if sample.image_positions is not None and len(sample.image_positions) > 0:
             batch_image_positions.append(torch.from_numpy(sample.image_positions))
         else:
             batch_image_positions.append([])
+
+    has_audio = False
+    batch_audio = []
+    for sample in batch:
+        if sample.audio is not None and sample.audio_positions is not None:
+            batch_audio.append([torch.from_numpy(audio) for audio in sample.audio])
+            has_audio = True
+        else:
+            batch_audio.append(None)
+    batch_audio_positions = []
+    for sample in batch:
+        if sample.audio_positions is not None:
+            batch_audio_positions.append(torch.from_numpy(sample.audio_positions))
+        else:
+            batch_audio_positions.append([])
+
     return GPTBatch(
         token_ids=torch.from_numpy(stacked_ids),
         loss_masking_spans=stacked_spans,
         sequence_lengths=sequence_lengths,
         images=batch_images if has_images else None,
         image_positions=batch_image_positions if has_images else None,
+        audio=batch_audio if has_audio else None,
+        audio_positions=batch_audio_positions if has_audio else None,
     )
 
 

fast_llm/data/dataset/gpt/config.py

@@ -76,8 +76,13 @@ class GPTSamplingParameters(SamplingParameters):
     cross_document_attention: bool = True
     patch_size: int | None = None
     image_size: int | None = None
+    aud_downsampling_k: int | None = None
+    aud_padding_duration: int | None = None
+    aud_sampling_rate: int | None = None
     image_break_token: int | None = None
     image_end_token: int | None = None
+    audio_start_token: int | None = None
+    audio_end_token: int | None = None
     # How many extra tokens to add to the sequence length.
     # This is used to provide labels even for the last tokens in the sequence.
     extra_tokens: int = 1
@@ -204,6 +209,11 @@ class GPTMemmapDatasetConfig(GPTIndexedDatasetConfig):
         desc="Expected number of pixels in the dataset.",
         hint=FieldHint.optional,
     )
+    num_audio: int | None = Field(
+        default=None,
+        desc="Expected number of audio in the dataset.",
+        hint=FieldHint.optional,
+    )
 
     def build(self) -> "GPTMemmapDataset":
         from fast_llm.data.dataset.gpt.memmap import GPTMemmapDataset

fast_llm/data/dataset/gpt/indexed.py

@@ -44,8 +44,12 @@ class GPTDatasetSlice[IndexedDatasetType: GPTIndexedDataset](DatasetSlice[Indexe
 
     def get_document_sizes(self) -> np.ndarray:
         # TODO: This can be really big.
-        doc_sizes, im_sizes = self._dataset.get_document_sizes()
-        return doc_sizes[self._begin : self._end], im_sizes[self._begin : self._end] if im_sizes else []
+        doc_sizes, im_sizes, aud_sizes = self._dataset.get_document_sizes()
+        return (
+            doc_sizes[self._begin : self._end],
+            im_sizes[self._begin : self._end] if im_sizes else [],
+            aud_sizes[self._begin : self._end] if aud_sizes else [],
+        )
 
     def get_document_size(self, index: int) -> int:
         return self._dataset.get_document_size(self._begin + index)