Update __init__.py

koaning · May 14, 2024 · 5ef97f5 · 5ef97f5
1 parent b0e1c65
commit 5ef97f5
Showing 1 changed file with 5 additions and 5 deletions.
diff --git a/drawdata/__init__.py b/drawdata/__init__.py
@@ -44,11 +44,11 @@ def data_as_polars(self):
     def data_as_X_y(self, kind='classification'):
         import numpy as np
 
-        colors = [_['label'] for _ in self.data]
+        colors = [_['color'] for _ in self.data]
+        X = np.array([[_['x']] for _ in self.data])
+
+        # Assume that we're dealing with regression in this case
         if np.unique(colors).shape[0] == 1:
-            X = np.array([[_['x']] for _ in self.data])
             y = np.array([_['y'] for _ in self.data])
             return X, y
-        X = np.array([[_['x'], _['y']] for _ in self.data])
-        y = np.array([_['label'] for _ in self.data])
-        return X, y
+        return X, colors