Skip IV2 Unit Test if package not installed (#1111)

* Skip IV2 Unit Test if package not installed

Signed-off-by: Ao Tang <aot@nvidia.com>

* syntax

Signed-off-by: Ao Tang <aot@nvidia.com>

---------

Signed-off-by: Ao Tang <aot@nvidia.com>
Co-authored-by: Dong Hyuk Chang <thomaschang26@tutanota.com>

Author: suiyoubi

tests/models/test_internvideo2.py

@@ -18,12 +18,18 @@
 from unittest.mock import MagicMock, Mock, patch
 
 import numpy as np
+import pytest
 import torch
 
-from nemo_curator.stages.video.embedding.internvideo2 import (
-    InternVideo2EmbeddingStage,
-    InternVideo2FrameCreationStage,
-)
+try:
+    from nemo_curator.stages.video.embedding.internvideo2 import (
+        InternVideo2EmbeddingStage,
+        InternVideo2FrameCreationStage,
+    )
+except ImportError:
+    pytest.skip("InternVideo2 package is not available")
+
+
 from nemo_curator.tasks.video import Clip, Video, VideoTask
 
 # Create a random generator for consistent testing

tests/models/test_internvideo2_mm.py

@@ -21,14 +21,17 @@
 import torch
 from easydict import EasyDict
 
-from nemo_curator.models.internvideo2_mm import (
-    BERT_MODEL_ID,
-    INTERNVIDEO2_MODEL_FILE,
-    INTERNVIDEO2_MODEL_ID,
-    InternVideo2MultiModality,
-    _InternVideo2Stage2Wrapper,
-    _setup_internvideo2,
-)
+try:
+    from nemo_curator.models.internvideo2_mm import (
+        BERT_MODEL_ID,
+        INTERNVIDEO2_MODEL_FILE,
+        INTERNVIDEO2_MODEL_ID,
+        InternVideo2MultiModality,
+        _InternVideo2Stage2Wrapper,
+        _setup_internvideo2,
+    )
+except ImportError:
+    pytest.skip("InternVideo2 package is not available")
 
 # Create a random generator for consistent testing
 rng = np.random.default_rng(42)