Sklearn type transformer should be automatically loaded with import f…

…lytekit (#1423) * add flytekit.extras.sklearn to main __init__ import Signed-off-by: Niels Bantilan <niels.bantilan@gmail.com> * fix docs Signed-off-by: Niels Bantilan <niels.bantilan@gmail.com> * add temporary docs/requirements.txt for onnx plugins Signed-off-by: Niels Bantilan <niels.bantilan@gmail.com> --------- Signed-off-by: Niels Bantilan <niels.bantilan@gmail.com> Signed-off-by: Samhita Alla <aallasamhita@gmail.com>
flyteorg · Feb 2, 2023 · e7a0a65 · e7a0a65
1 parent 2b25665
commit e7a0a65
Show file tree

Hide file tree

Showing 6 changed files with 25 additions and 5 deletions.
diff --git a/.readthedocs.yml b/.readthedocs.yml
@@ -18,3 +18,4 @@ sphinx:
 python:
   install:
     - requirements: doc-requirements.txt
+    - requirements: docs/requirements.txt
diff --git a/docs/requirements.txt b/docs/requirements.txt
@@ -0,0 +1,7 @@
+# TODO: Remove after buf migration is done and packages updated, see doc-requirements.in
+# skl2onnx and tf2onnx added here so that the plugin API reference is rendered,
+# with the caveat that the docs build environment has the environment variable
+# PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=python set so that protobuf can be parsed
+# using Python, which is acceptable for docs building.
+skl2onnx
+tf2onnx
diff --git a/flytekit/__init__.py b/flytekit/__init__.py
@@ -184,7 +184,7 @@
 from flytekit.core.workflow import ImperativeWorkflow as Workflow
 from flytekit.core.workflow import WorkflowFailurePolicy, reference_workflow, workflow
 from flytekit.deck import Deck
-from flytekit.extras import pytorch, tensorflow
+from flytekit.extras import pytorch, sklearn, tensorflow
 from flytekit.extras.persistence import GCSPersistence, HttpPersistence, S3Persistence
 from flytekit.loggers import logger
 from flytekit.models.common import Annotations, AuthRole, Labels

diff --git a/flytekit/extras/pytorch/__init__.py b/flytekit/extras/pytorch/__init__.py
@@ -1,13 +1,14 @@
 """
-Flytekit PyTorch
-=========================================
 .. currentmodule:: flytekit.extras.pytorch
 
 .. autosummary::
    :template: custom.rst
    :toctree: generated/
 
     PyTorchCheckpoint
+    PyTorchCheckpointTransformer
+    PyTorchModuleTransformer
+    PyTorchTensorTransformer
 """
 from flytekit.loggers import logger
 

diff --git a/flytekit/extras/sklearn/__init__.py b/flytekit/extras/sklearn/__init__.py
@@ -1,11 +1,11 @@
 """
-Flytekit Sklearn
-=========================================
 .. currentmodule:: flytekit.extras.sklearn
 
 .. autosummary::
    :template: custom.rst
    :toctree: generated/
+
+   SklearnEstimatorTransformer
 """
 from flytekit.loggers import logger
 

diff --git a/flytekit/extras/tensorflow/__init__.py b/flytekit/extras/tensorflow/__init__.py
@@ -1,3 +1,14 @@
+"""
+.. currentmodule:: flytekit.extras.tensorflow
+
+.. autosummary::
+   :template: custom.rst
+   :toctree: generated/
+
+   TensorFlowRecordFileTransformer
+   TensorFlowRecordsDirTransformer
+"""
+
 from flytekit.loggers import logger
 
 # TODO: abstract this out so that there's an established pattern for registering plugins