From f3b1769dd84a9d6c5159b61f7208c1e9a68c65f7 Mon Sep 17 00:00:00 2001
From: Nicolas Mowen <nickmowen213@gmail.com>
Date: Wed, 22 Oct 2025 14:34:33 -0600
Subject: [PATCH] Cleanup

---
 .../classification/ClassificationModelWizardDialog.tsx    | 8 +-------
 .../classification/wizard/Step3ChooseExamples.tsx         | 8 ++++----
 2 files changed, 5 insertions(+), 11 deletions(-)

diff --git a/web/src/components/classification/ClassificationModelWizardDialog.tsx b/web/src/components/classification/ClassificationModelWizardDialog.tsx
index 7efa57b66..9711fe246 100644
--- a/web/src/components/classification/ClassificationModelWizardDialog.tsx
+++ b/web/src/components/classification/ClassificationModelWizardDialog.tsx
@@ -19,14 +19,12 @@ import { isDesktop } from "react-device-detect";
 const OBJECT_STEPS = [
   "wizard.steps.nameAndDefine",
   "wizard.steps.chooseExamples",
-  "wizard.steps.train",
 ];
 
 const STATE_STEPS = [
   "wizard.steps.nameAndDefine",
   "wizard.steps.stateArea",
   "wizard.steps.chooseExamples",
-  "wizard.steps.train",
 ];
 
 type ClassificationModelWizardDialogProps = {
@@ -116,10 +114,6 @@ export default function ClassificationModelWizardDialog({
     dispatch({ type: "SET_STEP_2", payload: data });
   };
 
-  const handleStep3Next = (data: Step3FormData) => {
-    dispatch({ type: "SET_STEP_3", payload: data });
-  };
-
   const handleBack = () => {
     dispatch({ type: "PREVIOUS_STEP" });
   };
@@ -185,7 +179,7 @@ export default function ClassificationModelWizardDialog({
                 step1Data={wizardState.step1Data}
                 step2Data={wizardState.step2Data}
                 initialData={wizardState.step3Data}
-                onNext={handleStep3Next}
+                onClose={onClose}
                 onBack={handleBack}
               />
             )}
diff --git a/web/src/components/classification/wizard/Step3ChooseExamples.tsx b/web/src/components/classification/wizard/Step3ChooseExamples.tsx
index e373fde21..b26e07309 100644
--- a/web/src/components/classification/wizard/Step3ChooseExamples.tsx
+++ b/web/src/components/classification/wizard/Step3ChooseExamples.tsx
@@ -25,7 +25,7 @@ type Step3ChooseExamplesProps = {
   step1Data: Step1FormData;
   step2Data?: Step2FormData;
   initialData?: Partial<Step3FormData>;
-  onNext: (data: Step3FormData) => void;
+  onClose: () => void;
   onBack: () => void;
 };
 
@@ -33,7 +33,7 @@ export default function Step3ChooseExamples({
   step1Data,
   step2Data,
   initialData,
-  onNext,
+  onClose,
   onBack,
 }: Step3ChooseExamplesProps) {
   const { t } = useTranslation(["views/classificationModel"]);
@@ -202,7 +202,7 @@ export default function Step3ChooseExamples({
       await axios.post(`/classification/${step1Data.modelName}/train`);
 
       toast.success(t("wizard.step3.trainingStarted"));
-      onNext({ examplesGenerated: true, imageClassifications });
+      onClose();
     } catch (error) {
       const axiosError = error as {
         response?: { data?: { message?: string; detail?: string } };
@@ -218,7 +218,7 @@ export default function Step3ChooseExamples({
         t("wizard.step3.errors.classifyFailed", { error: errorMessage }),
       );
     }
-  }, [onNext, imageClassifications, step1Data, step2Data, t]);
+  }, [onClose, imageClassifications, step1Data, step2Data, t]);
 
   const allImagesClassified = useMemo(() => {
     if (!unknownImages || unknownImages.length === 0) return false;