Linter fixes.

sdks/python/apache_beam/ml/rag/chunking/base_test.py

@@ -17,15 +17,20 @@
 """Tests for apache_beam.ml.rag.chunking.base."""
 
 import unittest
+from typing import Any
+from typing import Dict
+from typing import Optional
+
 import pytest
 
 import apache_beam as beam
+from apache_beam.ml.rag.chunking.base import ChunkIdFn
+from apache_beam.ml.rag.chunking.base import ChunkingTransformProvider
+from apache_beam.ml.rag.types import Chunk
+from apache_beam.ml.rag.types import Content
 from apache_beam.testing.test_pipeline import TestPipeline
 from apache_beam.testing.util import assert_that
 from apache_beam.testing.util import equal_to
-from apache_beam.ml.rag.chunking.base import ChunkingTransformProvider, ChunkIdFn
-from apache_beam.ml.rag.types import Chunk, Content
-from typing import Optional, Dict, Any
 
 
 class WordSplitter(beam.DoFn):

sdks/python/apache_beam/ml/rag/chunking/langchain_test.py

@@ -19,18 +19,16 @@
 import unittest
 
 import apache_beam as beam
-
+from apache_beam.ml.rag.types import Chunk
 from apache_beam.testing.test_pipeline import TestPipeline
 from apache_beam.testing.util import assert_that
 from apache_beam.testing.util import equal_to
-from apache_beam.ml.rag.types import Chunk
 
 try:
   from apache_beam.ml.rag.chunking.langchain import LangChainChunker
+
   from langchain.text_splitter import (
-      RecursiveCharacterTextSplitter,
-      CharacterTextSplitter,
-  )
+      CharacterTextSplitter, RecursiveCharacterTextSplitter)
   LANGCHAIN_AVAILABLE = True
 except ImportError:
   LANGCHAIN_AVAILABLE = False

sdks/python/apache_beam/ml/rag/embeddings/base.py

@@ -14,10 +14,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from apache_beam.ml.transforms.base import EmbeddingTypeAdapter
-from apache_beam.ml.rag.types import Embedding, Chunk
-from typing import List
 from collections.abc import Sequence
+from typing import List
+
+from apache_beam.ml.rag.types import Chunk
+from apache_beam.ml.rag.types import Embedding
+from apache_beam.ml.transforms.base import EmbeddingTypeAdapter
 
 
 def create_rag_adapter() -> EmbeddingTypeAdapter[Chunk, Chunk]:

sdks/python/apache_beam/ml/rag/embeddings/base_test.py

@@ -15,8 +15,11 @@
 # limitations under the License.
 
 import unittest
-from apache_beam.ml.rag.types import Chunk, Content, Embedding
-from apache_beam.ml.rag.embeddings.base import (create_rag_adapter)
+
+from apache_beam.ml.rag.embeddings.base import create_rag_adapter
+from apache_beam.ml.rag.types import Chunk
+from apache_beam.ml.rag.types import Content
+from apache_beam.ml.rag.types import Embedding
 
 
 class RAGBaseEmbeddingsTest(unittest.TestCase):

sdks/python/apache_beam/ml/rag/embeddings/huggingface.py

@@ -22,10 +22,10 @@
 from apache_beam.ml.inference.base import RunInference
 from apache_beam.ml.rag.embeddings.base import create_rag_adapter
 from apache_beam.ml.rag.types import Chunk
-from apache_beam.ml.transforms.base import (
-    EmbeddingsManager, _TextEmbeddingHandler)
-from apache_beam.ml.transforms.embeddings.huggingface import (
-    SentenceTransformer, _SentenceTransformerModelHandler)
+from apache_beam.ml.transforms.base import EmbeddingsManager
+from apache_beam.ml.transforms.base import _TextEmbeddingHandler
+from apache_beam.ml.transforms.embeddings.huggingface import SentenceTransformer
+from apache_beam.ml.transforms.embeddings.huggingface import _SentenceTransformerModelHandler
 
 
 class HuggingfaceTextEmbeddings(EmbeddingsManager):

sdks/python/apache_beam/ml/rag/embeddings/huggingface_test.py

@@ -16,16 +16,20 @@
 
 """Tests for apache_beam.ml.rag.embeddings.huggingface."""
 
-import pytest
 import tempfile
 import unittest
 
+import pytest
+
 import apache_beam as beam
 from apache_beam.ml.rag.embeddings.huggingface import HuggingfaceTextEmbeddings
-from apache_beam.ml.rag.types import Chunk, Content, Embedding
+from apache_beam.ml.rag.types import Chunk
+from apache_beam.ml.rag.types import Content
+from apache_beam.ml.rag.types import Embedding
 from apache_beam.ml.transforms.base import MLTransform
 from apache_beam.testing.test_pipeline import TestPipeline
-from apache_beam.testing.util import assert_that, equal_to
+from apache_beam.testing.util import assert_that
+from apache_beam.testing.util import equal_to
 
 # pylint: disable=unused-import
 try:

sdks/python/apache_beam/ml/rag/types.py

@@ -21,9 +21,14 @@
 contracts between different stages of the pipeline.
 """
 
-from dataclasses import dataclass, field
-from typing import Dict, List, Optional, Tuple, Any
 import uuid
+from dataclasses import dataclass
+from dataclasses import field
+from typing import Any
+from typing import Dict
+from typing import List
+from typing import Optional
+from typing import Tuple
 
 
 @dataclass

sdks/python/apache_beam/ml/transforms/base.py

@@ -15,29 +15,28 @@
 # limitations under the License.
 
 import abc
+import functools
 import logging
 import os
 import tempfile
 import uuid
 from collections.abc import Callable
 from collections.abc import Mapping
 from collections.abc import Sequence
+from dataclasses import dataclass
 from typing import Any
-from typing import cast
 from typing import Dict
 from typing import Generic
 from typing import Iterable
 from typing import List
 from typing import Optional
 from typing import TypeVar
 from typing import Union
+from typing import cast
 
-import functools
 import jsonpickle
 import numpy as np
 
-from dataclasses import dataclass
-
 import apache_beam as beam
 from apache_beam.io.filesystems import FileSystems
 from apache_beam.metrics.metric import Metrics
@@ -258,11 +257,6 @@ def __init__(
       max_batch_size: Optional[int] = None,
       large_model: bool = False,
       **kwargs):
-    if columns is not None and type_adapter is not None:
-      raise ValueError(
-          "Cannot specify both 'columns' and 'type_adapter'. "
-          "Use either columns for dict processing or type_adapter "
-          "for custom types.")
     self.load_model_args = load_model_args or {}
     self.min_batch_size = min_batch_size
     self.max_batch_size = max_batch_size