ADLR/megatron-lm!2587 - Add safeguard to fail if video is empty

NVIDIA · Jan 28, 2025 · 0689058 · 0689058
1 parent 2167226
commit 0689058
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 1 deletion.
diff --git a/examples/multimodal/dataset_helpers.py b/examples/multimodal/dataset_helpers.py
@@ -409,6 +409,8 @@ def encode_llava_sft(self, sample: SimilarityInterleavedSample):
             # Grab the selected frames of the video as a tensor with shape
             # fhwc: (num_frames, num_channels, height, width).
             video_fchw = sample.images.frames
+            if video_fchw.shape[0] == 0:
+                raise ValueError(f"Video {sample.__key__} {sample.__restore_key__} {sample.texts} has no frames.")
             selected_frames = torch.linspace(
                 0, video_fchw.shape[0] - 1, self.args.num_frames).long()
             video_fchw = video_fchw[selected_frames]

diff --git a/examples/multimodal/run_text_generation.py b/examples/multimodal/run_text_generation.py
@@ -443,7 +443,7 @@ def get_conversation(task, question):
 
         conversation = [
             {"role": "system", "content": "Answer the questions."},
-            {"role": "user", "content": f"{IMAGE_TOKEN}\n{question}"},
+            {"role": "user", "content": f"{IMAGE_TOKEN}\n{q}"},
         ]
 
     return conversation