diff --git a/sdks/python/apache_beam/ml/inference/gemini_inference.py b/sdks/python/apache_beam/ml/inference/gemini_inference.py
index 8b6d0edfa996..2ba220d0162a 100644
--- a/sdks/python/apache_beam/ml/inference/gemini_inference.py
+++ b/sdks/python/apache_beam/ml/inference/gemini_inference.py
@@ -22,6 +22,7 @@
 from typing import Any
 from typing import Optional
 from typing import Union
+from typing import cast
 
 from google import genai
 from google.genai import errors
@@ -73,7 +74,7 @@ def generate_from_string(
       call.
   """
   return model.models.generate_content(
-      model=model_name, contents=batch, **inference_args)
+      model=model_name, contents=cast(Any, batch), **inference_args)
 
 
 def generate_image_from_strings_and_images(
@@ -96,7 +97,7 @@ def generate_image_from_strings_and_images(
       call.
   """
   return model.models.generate_content(
-      model=model_name, contents=batch, **inference_args)
+      model=model_name, contents=cast(Any, batch), **inference_args)
 
 
 class GeminiModelHandler(RemoteModelHandler[Any, PredictionResult,
diff --git a/sdks/python/apache_beam/ml/rag/embeddings/vertex_ai_test.py b/sdks/python/apache_beam/ml/rag/embeddings/vertex_ai_test.py
index aa294e98789d..e0da9499faf3 100644
--- a/sdks/python/apache_beam/ml/rag/embeddings/vertex_ai_test.py
+++ b/sdks/python/apache_beam/ml/rag/embeddings/vertex_ai_test.py
@@ -16,11 +16,12 @@
 
 """Tests for apache_beam.ml.rag.embeddings.vertex_ai."""
 
-import pytest
 import shutil
 import tempfile
 import unittest
 
+import pytest
+
 import apache_beam as beam
 from apache_beam.ml.rag.types import Chunk
 from apache_beam.ml.rag.types import Content
diff --git a/sdks/python/apache_beam/ml/rag/ingestion/milvus_search_it_test.py b/sdks/python/apache_beam/ml/rag/ingestion/milvus_search_it_test.py
index 38b497e8fa71..b6e5083ea728 100644
--- a/sdks/python/apache_beam/ml/rag/ingestion/milvus_search_it_test.py
+++ b/sdks/python/apache_beam/ml/rag/ingestion/milvus_search_it_test.py
@@ -22,12 +22,6 @@
 from typing import cast
 
 import pytest
-from pymilvus import CollectionSchema
-from pymilvus import DataType
-from pymilvus import FieldSchema
-from pymilvus import MilvusClient
-from pymilvus.exceptions import MilvusException
-from pymilvus.milvus_client import IndexParams
 
 import apache_beam as beam
 from apache_beam.ml.rag.ingestion.jdbc_common import WriteConfig
@@ -41,11 +35,21 @@
 from apache_beam.ml.rag.utils import unpack_dataclass_with_kwargs
 from apache_beam.testing.test_pipeline import TestPipeline
 
+# pylint: disable=wrong-import-order, wrong-import-position, ungrouped-imports
 try:
+  from pymilvus import CollectionSchema
+  from pymilvus import DataType
+  from pymilvus import FieldSchema
+  from pymilvus import MilvusClient
+  from pymilvus.exceptions import MilvusException
+  from pymilvus.milvus_client import IndexParams
+
   from apache_beam.ml.rag.ingestion.milvus_search import MilvusVectorWriterConfig
   from apache_beam.ml.rag.ingestion.milvus_search import MilvusWriteConfig
-except ImportError as e:
-  raise unittest.SkipTest(f'Milvus dependencies not installed: {str(e)}')
+  PYMILVUS_AVAILABLE = True
+except ImportError:
+  PYMILVUS_AVAILABLE = False
+# pylint: enable=wrong-import-order, wrong-import-position, ungrouped-imports
 
 
 def _construct_index_params():
@@ -158,6 +162,7 @@ def drop_collection(client: MilvusClient, collection_name: str):
 
 
 @pytest.mark.require_docker_in_docker
+@unittest.skipIf(not PYMILVUS_AVAILABLE, 'pymilvus is not installed.')
 @unittest.skipUnless(
     platform.system() == "Linux",
     "Test runs only on Linux due to lack of support, as yet, for nested "
diff --git a/sdks/python/apache_beam/ml/transforms/base_test.py b/sdks/python/apache_beam/ml/transforms/base_test.py
index 190381cc2f34..64d0afe955d0 100644
--- a/sdks/python/apache_beam/ml/transforms/base_test.py
+++ b/sdks/python/apache_beam/ml/transforms/base_test.py
@@ -51,8 +51,8 @@
   import PIL
   from PIL.Image import Image as PIL_Image
 except ImportError:
-  PIL = None
-  PIL_Image = Any
+  PIL = None  # type: ignore[assignment]
+  PIL_Image = Any  # type: ignore[misc, assignment]
 
 try:
 
diff --git a/sdks/python/apache_beam/ml/transforms/embeddings/huggingface_test.py b/sdks/python/apache_beam/ml/transforms/embeddings/huggingface_test.py
index a2358c544781..5e56dc710548 100644
--- a/sdks/python/apache_beam/ml/transforms/embeddings/huggingface_test.py
+++ b/sdks/python/apache_beam/ml/transforms/embeddings/huggingface_test.py
@@ -53,7 +53,7 @@
 try:
   from PIL import Image
 except ImportError:
-  Image = None
+  Image = None  # type: ignore[assignment]
 
 _HF_TOKEN = os.environ.get('HF_INFERENCE_TOKEN')
 test_query = "This is a test"
diff --git a/sdks/python/apache_beam/ml/transforms/embeddings/tensorflow_hub_test.py b/sdks/python/apache_beam/ml/transforms/embeddings/tensorflow_hub_test.py
index 0a4f8c8275c3..b9a14ddb8c53 100644
--- a/sdks/python/apache_beam/ml/transforms/embeddings/tensorflow_hub_test.py
+++ b/sdks/python/apache_beam/ml/transforms/embeddings/tensorflow_hub_test.py
@@ -52,7 +52,7 @@
   from apache_beam.ml.transforms.embeddings.tensorflow_hub import TensorflowHubImageEmbeddings
 except ImportError:
   TensorflowHubImageEmbeddings = None  # type: ignore
-  Image = None
+  Image = None  # type: ignore[assignment]
 
 
 @unittest.skipIf(