aryn-ai · bsowell · Dec 3, 2024 · Dec 3, 2024
diff --git a/lib/sycamore/sycamore/tests/unit/transforms/test_embed.py b/lib/sycamore/sycamore/tests/unit/transforms/test_embed.py
@@ -1,10 +1,11 @@
+import pickle
 import pytest
 import ray.data
 
 from sycamore.data import Document
 from sycamore.plan_nodes import Node
 from sycamore.transforms import Embed
-from sycamore.transforms.embed import SentenceTransformerEmbedder
+from sycamore.transforms.embed import OpenAIEmbedder, SentenceTransformerEmbedder
 
 
 class TestEmbedding:
@@ -77,3 +78,10 @@ def test_sentence_transformer_embedding(self, mocker):
         input_dataset.show()
         output_dataset = embedding.execute()
         output_dataset.show()
+
+    def test_openai_embedder_pickle(self):
+        obj = OpenAIEmbedder()
+        obj._client = obj.client_wrapper.get_client()
+
+        pickle.dumps(obj)
+        assert True
diff --git a/lib/sycamore/sycamore/transforms/embed.py b/lib/sycamore/sycamore/transforms/embed.py
@@ -177,6 +177,14 @@ def __init__(
         self._client: Optional[OpenAIClient] = None
         self.model_name = model_name
 
+    def __getstate__(self):
+        state = self.__dict__.copy()
+        state["_client"] = None
+        return state
+
+    def __setstate__(self, state):
+        self.__dict__.update(state)
+
     def generate_embeddings(self, doc_batch: list[Document]) -> list[Document]:
         # TODO: Add some input validation here.
         # The OpenAI docs are quite vague on acceptable values for model_batch_size.