From 379ee62f41e9a35ca6882df0425ae95c8f37bf02 Mon Sep 17 00:00:00 2001
From: Luca Tagliabue <32895300+lucataglia@users.noreply.github.com>
Date: Mon, 1 Jul 2024 15:25:37 +0200
Subject: [PATCH] feat: update wizard (#54)

* feat: enable all modelType

* fix: remove unused import

* fix: valid types

* refactor: add empty array for readability

---------

Co-authored-by: Luca Tagliabue <luca.tagliabue@radicalbit.io>
---
 .../add-new-model/modal-context-provider.jsx  |  4 +--
 .../add-new-model/step-four/form-fields.jsx   | 25 +++++++++++++------
 .../add-new-model/step-one/form-fields.jsx    | 15 ++++++++++-
 3 files changed, 34 insertions(+), 10 deletions(-)

diff --git a/ui/src/components/modals/add-new-model/modal-context-provider.jsx b/ui/src/components/modals/add-new-model/modal-context-provider.jsx
index 6762e0cf..6aba9bb4 100644
--- a/ui/src/components/modals/add-new-model/modal-context-provider.jsx
+++ b/ui/src/components/modals/add-new-model/modal-context-provider.jsx
@@ -1,4 +1,4 @@
-import { DataTypeEnum, ModelTypeEnum } from '@State/models/constants';
+import { DataTypeEnum } from '@State/models/constants';
 import useFormbit from '@radicalbit/formbit';
 import {
   createContext,
@@ -18,7 +18,7 @@ function ModalContextProvider({ children }) {
   const [isMaximize, setIsMaximize] = useState(false);
 
   const useFormbitStepOne = useFormbit({
-    initialValues: { modelType: ModelTypeEnum.BINARY_CLASSIFICATION, dataType: DataTypeEnum.TABULAR },
+    initialValues: { dataType: DataTypeEnum.TABULAR },
     yup: schemaStepOne,
   });
 
diff --git a/ui/src/components/modals/add-new-model/step-four/form-fields.jsx b/ui/src/components/modals/add-new-model/step-four/form-fields.jsx
index f2f65be7..7686c05e 100644
--- a/ui/src/components/modals/add-new-model/step-four/form-fields.jsx
+++ b/ui/src/components/modals/add-new-model/step-four/form-fields.jsx
@@ -1,9 +1,9 @@
+import { ModelTypeEnum } from '@Src/store/state/models/constants';
 import {
   FormField,
   Select,
   Tooltip,
 } from '@radicalbit/radicalbit-design-system';
-import { ModelTypeEnum } from '@Src/store/state/models/constants';
 import { useModalContext } from '../modal-context-provider';
 
 function Target() {
@@ -196,15 +196,18 @@ function Prediction() {
 }
 
 function Probability() {
-  const { useFormbit } = useModalContext();
+  const { useFormbit, useFormbitStepOne } = useModalContext();
+
   const {
     form, error, write, remove,
   } = useFormbit;
-
   const probabilities = useGetProbabilities();
   const predictionName = form?.prediction?.name;
   const value = form?.predictionProba?.name;
 
+  const { form: formStepOne } = useFormbitStepOne;
+  const { modelType } = formStepOne;
+
   const handleOnChange = (val) => {
     if (val === undefined) {
       remove('predictionProba');
@@ -218,6 +221,14 @@ function Probability() {
     }
   };
 
+  if (modelType === ModelTypeEnum.REGRESSION) {
+    return (
+      <FormField label="Probability" modifier="w-full">
+        <Select disabled readOnly value="Not available for Regression" />
+      </FormField>
+    );
+  }
+
   return (
     <FormField
       label="Probability"
@@ -267,7 +278,7 @@ function Probability() {
 
 const targetValidTypes = {
   [ModelTypeEnum.BINARY_CLASSIFICATION]: ['int', 'float', 'double'],
-  [ModelTypeEnum.MULTI_CLASSIFICATION]: ['int', 'float', 'double'],
+  [ModelTypeEnum.MULTI_CLASSIFICATION]: ['int', 'float', 'double', 'string'],
   [ModelTypeEnum.REGRESSION]: ['int', 'float', 'double'],
 };
 const useGetTargets = () => {
@@ -295,10 +306,10 @@ const useGetPredictions = () => {
   return form.outputs.filter(({ type }) => predictionValidTypes[modelType].includes(type));
 };
 
-const binaryClassificationProbabilityValidTypes = {
+const probabilityValidTypes = {
   [ModelTypeEnum.BINARY_CLASSIFICATION]: ['float', 'double'],
   [ModelTypeEnum.MULTI_CLASSIFICATION]: ['float', 'double'],
-  [ModelTypeEnum.REGRESSION]: ['float', 'double'],
+  [ModelTypeEnum.REGRESSION]: [],
 };
 const useGetProbabilities = () => {
   const { useFormbitStepOne, useFormbit } = useModalContext();
@@ -307,7 +318,7 @@ const useGetProbabilities = () => {
   const { form: formStepOne } = useFormbitStepOne;
   const { modelType } = formStepOne;
 
-  return form.outputs.filter(({ type }) => binaryClassificationProbabilityValidTypes[modelType].includes(type));
+  return form.outputs.filter(({ type }) => probabilityValidTypes[modelType].includes(type));
 };
 
 const timestampValidTypes = {
diff --git a/ui/src/components/modals/add-new-model/step-one/form-fields.jsx b/ui/src/components/modals/add-new-model/step-one/form-fields.jsx
index 5839be5d..e17dd8f7 100644
--- a/ui/src/components/modals/add-new-model/step-one/form-fields.jsx
+++ b/ui/src/components/modals/add-new-model/step-one/form-fields.jsx
@@ -39,9 +39,22 @@ function Name() {
 }
 
 function ModelType() {
+  const { useFormbit } = useModalContext();
+  const { form, write } = useFormbit;
+
+  const handleOnChange = (value) => {
+    write('modelType', value);
+  };
+
   return (
     <FormField label="Model type" modifier="w-full" required>
-      {ModelTypeEnumLabel[ModelTypeEnum.BINARY_CLASSIFICATION]}
+      <Select onChange={handleOnChange} value={form.modelType}>
+        {Object.values(ModelTypeEnum).map((value) => (
+          <Select.Option key={value}>
+            {ModelTypeEnumLabel[value]}
+          </Select.Option>
+        ))}
+      </Select>
     </FormField>
   );
 }