From bcb60c943e211b1c3faaef598a6060a9960306cf Mon Sep 17 00:00:00 2001
From: Felix Altenberger <felix@zenml.io>
Date: Tue, 23 Aug 2022 11:41:32 +0200
Subject: [PATCH] Adjust ZenBytes to zenml version 0.13.0.

---
 1-1_Pipelines.ipynb | 25 ++++++++++++++++++-------
 steps/importer.py   |  9 +++++++--
 2 files changed, 25 insertions(+), 9 deletions(-)

diff --git a/1-1_Pipelines.ipynb b/1-1_Pipelines.ipynb
index 898a09d..3aec353 100644
--- a/1-1_Pipelines.ipynb
+++ b/1-1_Pipelines.ipynb
@@ -86,12 +86,17 @@
     "import numpy as np\n",
     "from sklearn.base import ClassifierMixin\n",
     "from sklearn.svm import SVC\n",
-    "from zenml.integrations.sklearn.helpers.digits import get_digits\n",
+    "from sklearn.datasets import load_digits\n",
+    "from sklearn.model_selection import train_test_split\n",
     "\n",
     "\n",
     "def train_test() -> None:\n",
     "    \"\"\"Train and test a Scikit-learn SVC classifier on digits\"\"\"\n",
-    "    X_train, X_test, y_train, y_test = get_digits()\n",
+    "    digits = load_digits()\n",
+    "    data = digits.images.reshape((len(digits.images), -1))\n",
+    "    X_train, X_test, y_train, y_test = train_test_split(\n",
+    "        data, digits.target, test_size=0.2, shuffle=False\n",
+    "    )\n",
     "    model = SVC(gamma=0.001)\n",
     "    model.fit(X_train, y_train)\n",
     "    test_acc = model.score(X_test, y_test)\n",
@@ -136,7 +141,11 @@
     "    y_test=np.ndarray,\n",
     "):\n",
     "    \"\"\"Load the digits dataset as numpy arrays.\"\"\"\n",
-    "    X_train, X_test, y_train, y_test = get_digits()\n",
+    "    digits = load_digits()\n",
+    "    data = digits.images.reshape((len(digits.images), -1))\n",
+    "    X_train, X_test, y_train, y_test = train_test_split(\n",
+    "        data, digits.target, test_size=0.2, shuffle=False\n",
+    "    )\n",
     "    return X_train, X_test, y_train, y_test\n",
     "\n",
     "\n",
@@ -220,11 +229,8 @@
   }
  ],
  "metadata": {
-  "interpreter": {
-   "hash": "a35bb4b4bceaf970a493ff7351e9d97180ab3fe9951c21e9e29c55a687242182"
-  },
   "kernelspec": {
-   "display_name": "Python 3.8.13 64-bit ('zenbytes-latest')",
+   "display_name": "Python 3.8.13 64-bit ('zenbytes')",
    "language": "python",
    "name": "python3"
   },
@@ -239,6 +245,11 @@
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
    "version": "3.8.13"
+  },
+  "vscode": {
+   "interpreter": {
+    "hash": "9f70ec6e6bd16014ded89c8222361cbe53cd9507d51ebdcdf3ab6e494d45cf74"
+   }
   }
  },
  "nbformat": 4,
diff --git a/steps/importer.py b/steps/importer.py
index b99832c..565bbcf 100644
--- a/steps/importer.py
+++ b/steps/importer.py
@@ -1,6 +1,7 @@
 import numpy as np
 import pandas as pd
-from zenml.integrations.sklearn.helpers.digits import get_digits
+from sklearn.datasets import load_digits
+from sklearn.model_selection import train_test_split
 from zenml.steps import Output, step
 
 
@@ -9,7 +10,11 @@ def importer() -> Output(
     X_train=np.ndarray, X_test=np.ndarray, y_train=np.ndarray, y_test=np.ndarray
 ):
     """Loads the digits array as normal numpy arrays."""
-    X_train, X_test, y_train, y_test = get_digits()
+    digits = load_digits()
+    data = digits.images.reshape((len(digits.images), -1))
+    X_train, X_test, y_train, y_test = train_test_split(
+        data, digits.target, test_size=0.2, shuffle=False
+    )
     return X_train, X_test, y_train, y_test