diff --git a/app/client/src/pages/DataGenerator/Configure.tsx b/app/client/src/pages/DataGenerator/Configure.tsx
index 0bb50ce..1326a1e 100644
--- a/app/client/src/pages/DataGenerator/Configure.tsx
+++ b/app/client/src/pages/DataGenerator/Configure.tsx
@@ -1,8 +1,8 @@
 import endsWith from 'lodash/endsWith';
 import isEmpty from 'lodash/isEmpty';
 import isFunction from 'lodash/isFunction';
-import { useEffect, useState } from 'react';
-import { Flex, Form, Input, Select, Typography } from 'antd';
+import { FunctionComponent, useEffect, useState } from 'react';
+import { Flex, Form, FormInstance, Input, Select, Typography } from 'antd';
 import styled from 'styled-components';
 import { File, WorkflowType } from './types';
 import { useFetchModels } from '../../api/api';
@@ -49,7 +49,9 @@ export const MODEL_TYPE_OPTIONS: ModelProvidersDropdownOpts = [
     { label: MODEL_PROVIDER_LABELS[ModelProviders.CAII], value: ModelProviders.CAII },
 ];
 
-const Configure = () => {
+const Configure: FunctionComponent = () => {
+    const form = Form.useFormInstance();
+    const formData = Form.useWatch((values) => values, form);
     const location = useLocation();
     const { template_name, generate_file_name } = useParams();
     const [wizardModeType, setWizardModeType] = useState(getWizardModeType(location));
@@ -57,10 +59,10 @@ const Configure = () => {
     useEffect(() => {
         if (wizardModeType === WizardModeType.DATA_AUGMENTATION) {
             setWizardModeType(WizardModeType.DATA_AUGMENTATION);
-            form.setFieldValue('workflow_type', 'custom');
+            form.setFieldValue('workflow_type', 'freeform');
         } else {
             setWizardModeType(WizardModeType.DATA_GENERATION);
-            form.setFieldValue('workflow_type', 'freeform');
+            form.setFieldValue('workflow_type', 'custom');
         }
     }, [location, wizardModeType]);
 
@@ -70,8 +72,8 @@ const Configure = () => {
         }
     }, [template_name]);
 
-    const form = Form.useFormInstance();
-    const formData = Form.useWatch((values) => values, form);
+    
+    // let formData = Form.useWatch((values) => values, form);
     const { setIsStepValid } = useWizardCtx();
     const { data } = useFetchModels();
     const [selectedFiles, setSelectedFiles] = useState(
diff --git a/app/client/src/pages/DataGenerator/DataGenerator.tsx b/app/client/src/pages/DataGenerator/DataGenerator.tsx
index e4c6196..36f640d 100644
--- a/app/client/src/pages/DataGenerator/DataGenerator.tsx
+++ b/app/client/src/pages/DataGenerator/DataGenerator.tsx
@@ -67,35 +67,7 @@ const WizardFooter = styled(Flex)`
 
 `;
 
-const steps: WizardStepConfig[] = [
-    {
-        title: 'Configure',
-        key: DataGenWizardSteps.CONFIGURE,
-        content: <Configure/>,
-        required: true,
-    },
-    {
-        title: 'Examples',
-        key: DataGenWizardSteps.EXAMPLES,
-        content: <Examples/>
-    },
-    {
-        title: 'Prompt',
-        key: DataGenWizardSteps.PROMPT,
-        content: <Prompt/>,
-    },
-    {
-        title: 'Summary',
-        key: DataGenWizardSteps.SUMMARY,
-        content: <Summary/>
-    },
-    {
-        title: 'Finish',
-        key: DataGenWizardSteps.FINISH,
-        content: <Finish/>
-    },
-
-];
+
 
 /**
  * Wizard component for Synthetic Data Generation workflow
@@ -105,10 +77,12 @@ const DataGenerator: FunctionComponent<Props> = () => {
     const [maxStep, setMaxStep] = useState(0);
     const [isStepValid, setIsStepValid] = useState<boolean>(false);
     
+    
     // Data passed from listing table to prepopulate form
     const location = useLocation();
     const { generate_file_name } = useParams();
     const initialData = location?.state?.data;
+    
     const mutation = useMutation({
         mutationFn: fetchDatasetDetails
     });
@@ -118,14 +92,21 @@ const DataGenerator: FunctionComponent<Props> = () => {
         if (generate_file_name && !mutation.data) {
             mutation.mutate(generate_file_name);
         }
+    }, [generate_file_name]);
+
+    useEffect(() => {
         if (mutation.data && mutation?.data?.dataset) {
-            form.setFieldsValue({
+            const dataset = mutation?.data?.dataset as any;
+            const values = {
                 ...initialData,
-                ...(mutation?.data?.dataset as any)
-            });
+                ...dataset,
+                workflow_type: dataset.technique === 'freeform' ? 
+                    WorkflowType.FREE_FORM_DATA_GENERATION : WorkflowType.CUSTOM_DATA_GENERATION
+            }
+            form.setFieldsValue(values);
+            formData.current = values;
         }
-
-    }, [generate_file_name]);
+    }, [mutation.data]);
     
 
     if (initialData?.technique) {
@@ -157,11 +138,43 @@ const DataGenerator: FunctionComponent<Props> = () => {
         initialData.doc_paths = [];
     }
 
-
     const formData = useRef(initialData || { num_questions: 20, topics: [] });
 
     const [form] = Form.useForm<FormInstance>();
 
+
+
+
+    const steps: WizardStepConfig[] = [
+        {
+            title: 'Configure',
+            key: DataGenWizardSteps.CONFIGURE,
+            content: <Configure />,
+            required: true,
+        },
+        {
+            title: 'Examples',
+            key: DataGenWizardSteps.EXAMPLES,
+            content: <Examples />
+        },
+        {
+            title: 'Prompt',
+            key: DataGenWizardSteps.PROMPT,
+            content: <Prompt />,
+        },
+        {
+            title: 'Summary',
+            key: DataGenWizardSteps.SUMMARY,
+            content: <Summary/>
+        },
+        {
+            title: 'Finish',
+            key: DataGenWizardSteps.FINISH,
+            content: <Finish/>
+        },
+    
+    ];
+
     const onStepChange = (value: number) => {
         setCurrent(value);
     };
@@ -173,7 +186,7 @@ const DataGenerator: FunctionComponent<Props> = () => {
         }
     };
 
-    const prev = () => setCurrent(Math.max(0, current - 1))
+    const prev = () => setCurrent(Math.max(0, current - 1));
 
     return (
         <WizardCtx.Provider value={{ setIsStepValid }}>
diff --git a/app/client/src/pages/DataGenerator/Examples.tsx b/app/client/src/pages/DataGenerator/Examples.tsx
index d0ffe6c..31e61a6 100644
--- a/app/client/src/pages/DataGenerator/Examples.tsx
+++ b/app/client/src/pages/DataGenerator/Examples.tsx
@@ -18,6 +18,7 @@ import FreeFormExampleTable from './FreeFormExampleTable';
 import { L } from 'vitest/dist/chunks/reporters.DTtkbAtP.js';
 import Loading from '../Evaluator/Loading';
 import { isEqual, set } from 'lodash';
+import { useParams } from 'react-router-dom';
 
 const { Title, Text } = Typography;
 const Container = styled.div`
@@ -52,6 +53,7 @@ const StyledContainer = styled.div`
 
 const Examples: FunctionComponent = () => {
     const form = Form.useFormInstance();
+    const { generate_file_name } = useParams();
     const [records, setRecords] = useState<Record<string, string>[]>([]);
     const workflowType = form.getFieldValue('workflow_type');
 
@@ -64,9 +66,13 @@ const Examples: FunctionComponent = () => {
     });
 
     useEffect(() => {
-        const useCase = form.getFieldValue('use_case');
-        restore_mutation.mutate(useCase);
-    }, [form.getFieldValue('use_case')]);
+        if (isEmpty(generate_file_name)) {
+            const useCase = form.getFieldValue('use_case');
+            restore_mutation.mutate(useCase);
+        } else {
+            setRecords(form.getFieldValue('examples'));
+        }
+    }, [form.getFieldValue('use_case'), generate_file_name]);
 
 
     useEffect(() => {
@@ -95,6 +101,12 @@ const Examples: FunctionComponent = () => {
             setRecords(examples);
         }
     }, [restore_mutation.data]);
+
+    useEffect(() => {
+        if (generate_file_name) {
+            setRecords(form.getFieldValue('examples'));
+        }
+    }, [generate_file_name]);
     
     const onRestoreDefaults = async() => {
         const useCase = form.getFieldValue('use_case');
diff --git a/app/client/src/pages/DataGenerator/Prompt.tsx b/app/client/src/pages/DataGenerator/Prompt.tsx
index 45f8b79..f89be15 100644
--- a/app/client/src/pages/DataGenerator/Prompt.tsx
+++ b/app/client/src/pages/DataGenerator/Prompt.tsx
@@ -226,7 +226,6 @@ const Prompt = () => {
             }
         };
     }
-    console.log('mutation data:', mutation);
 
     return (
         <Row gutter={[50,0]}>
diff --git a/app/client/src/pages/DataGenerator/UseCaseSelector.tsx b/app/client/src/pages/DataGenerator/UseCaseSelector.tsx
index 1107bc9..16ae1e0 100644
--- a/app/client/src/pages/DataGenerator/UseCaseSelector.tsx
+++ b/app/client/src/pages/DataGenerator/UseCaseSelector.tsx
@@ -26,7 +26,6 @@ const UseCaseSelector: FunctionComponent<Props> = ({ form }) => {
   }, [useCasesReq.data]);
 
   const onChange = (value: string) => {
-    console.log('value', value);
     form.setFieldValue('use_case', value);
     if (value !== 'custom') {
         form.setFieldValue('example_path', null);
diff --git a/app/client/src/pages/DataGenerator/hooks.ts b/app/client/src/pages/DataGenerator/hooks.ts
index c64d0c3..2980205 100644
--- a/app/client/src/pages/DataGenerator/hooks.ts
+++ b/app/client/src/pages/DataGenerator/hooks.ts
@@ -280,7 +280,6 @@ export const fetchExamplesByUseCase = async (use_case: string) => {
 }
 
 export const useGetExamplesByUseCase = (use_case: string) => {
-    console.log('------------------> useGetExamplesByUseCase', use_case);
     const { data, isLoading, isError, error, isFetching, refetch } = useQuery(
         {
             queryKey: ['fetchUseCaseTopics', fetchExamplesByUseCase],
diff --git a/app/client/src/pages/Home/TemplatesSection.tsx b/app/client/src/pages/Home/TemplatesSection.tsx
index 6937951..a27b8cd 100644
--- a/app/client/src/pages/Home/TemplatesSection.tsx
+++ b/app/client/src/pages/Home/TemplatesSection.tsx
@@ -28,13 +28,11 @@ const StyledContainer = styled.div`
 
 const TemplatesSection: React.FC = () => {
     const useCasesReq = useGetUseCases();
-    console.log("useCasesReq", useCasesReq);
     if (useCasesReq.isLoading) {
         return <Container><Loading /></Container>;
     }
     
     const useCases: Template[] = get(useCasesReq, 'data.usecases', []);
-    console.log("useCases", useCases);
     
 
     return (