diff --git a/app/client/src/pages/DataGenerator/Configure.tsx b/app/client/src/pages/DataGenerator/Configure.tsx
index 1326a1e..91aeb6d 100644
--- a/app/client/src/pages/DataGenerator/Configure.tsx
+++ b/app/client/src/pages/DataGenerator/Configure.tsx
@@ -55,20 +55,28 @@ const Configure: FunctionComponent = () => {
     const location = useLocation();
     const { template_name, generate_file_name } = useParams();
     const [wizardModeType, setWizardModeType] = useState(getWizardModeType(location));
+    console.log('wizardModeType', wizardModeType);
 
     useEffect(() => {
         if (wizardModeType === WizardModeType.DATA_AUGMENTATION) {
             setWizardModeType(WizardModeType.DATA_AUGMENTATION);
-            form.setFieldValue('workflow_type', 'freeform');
+            form.setFieldValue('workflow_type', 'custom');
         } else {
             setWizardModeType(WizardModeType.DATA_GENERATION);
-            form.setFieldValue('workflow_type', 'custom');
+            form.setFieldValue('workflow_type', 'freeform');
         }
     }, [location, wizardModeType]);
 
+    console.log('wizardModeType', wizardModeType);
+
     useEffect(() => {
+        console.log('useEffect 2');
         if (template_name) {
-            form.setFieldValue('use_case', template_name);
+            setTimeout(() => {
+                console.log('setting template name');
+                form.setFieldValue('use_case', template_name);
+            }, 1000);
+            
         }
     }, [template_name]);
 
@@ -85,6 +93,7 @@ const Configure: FunctionComponent = () => {
         delete values.doc_paths;
         delete values.output_key;
         delete values.output_value;
+        console.log('validateForm', values);
         
         const allFieldsFilled = Object.values(values).every(value => Boolean(value));
         if (allFieldsFilled && isFunction(setIsStepValid)) {
@@ -99,10 +108,11 @@ const Configure: FunctionComponent = () => {
 
     
     useEffect(() => {
+        console.log('useEffect 1');
         if (formData && formData?.inference_type === undefined && isEmpty(generate_file_name)) {
             form.setFieldValue('inference_type', ModelProviders.CAII);
             setTimeout(() => {
-                form.setFieldValue('use_case','custom');
+                form.setFieldValue('use_case', template_name ? template_name : 'custom');
             }, 1000);
         }
     }, [formData]);
@@ -253,18 +263,31 @@ const Configure: FunctionComponent = () => {
                         )}
                     </Select>
                 </Form.Item>
-                {(formData?.workflow_type === WorkflowType.SUPERVISED_FINE_TUNING || 
-                 formData?.workflow_type === WorkflowType.FREE_FORM_DATA_GENERATION) && 
-                 <UseCaseSelector form={form} />}
+                 <UseCaseSelector form={form} hidden={formData?.workflow_type === WorkflowType.CUSTOM_DATA_GENERATION} />
 
                 {(
-                    formData?.workflow_type === WorkflowType.SUPERVISED_FINE_TUNING || 
-                    formData?.workflow_type === WorkflowType.CUSTOM_DATA_GENERATION) && 
+                    formData?.workflow_type === WorkflowType.FREE_FORM_DATA_GENERATION || 
+                    (formData?.workflow_type === WorkflowType.CUSTOM_DATA_GENERATION &&
+                    formData?.use_case === 'custom'))  && 
+                    <Form.Item
+                    noStyle
+                    shouldUpdate={(prevValues, currentValues) =>
+                      prevValues.doc_paths !== currentValues.doc_paths ||
+                      prevValues.use_case !== currentValues.use_case
+                    }
+                  >
+                    {({}) => {
+                           const useCase = form.getFieldValue('use_case');
+                           if (useCase === 'custom') {
+   
+                           }
+                      return (
+                     
                 <Form.Item
                     name='doc_paths'
-                    label='Input File'
+                    label={(useCase === 'custom' && formData?.workflow_type === WorkflowType.FREE_FORM_DATA_GENERATION)  ? 'Context' : 'Input File'}
                     labelCol={labelCol}
-                    dependencies={['workflow_type']}
+                    dependencies={['workflow_type', 'use_case]']}
                     shouldUpdate
                     validateTrigger="['onBlur','onChange']"
                     tooltip='Select a file from your project that contains the initial data to be augmented.'
@@ -302,9 +325,9 @@ const Configure: FunctionComponent = () => {
                 >
                     <Flex>
                         <Select placeholder={'Select project files'} mode="multiple" value={selectedFiles || []} onChange={onFilesChange} allowClear/>    
-                        <FileSelectorButton onAddFiles={onAddFiles} workflowType={form.getFieldValue('workflow_type')} allowFileTypes={['pdf', 'docx']}/>
+                        <FileSelectorButton onAddFiles={onAddFiles} workflowType={form.getFieldValue('workflow_type')} allowFileTypes={['pdf', 'docx', 'json']}/>
                     </Flex>
-                </Form.Item>}
+                </Form.Item>)}}</Form.Item>}
                 {formData?.workflow_type === WorkflowType.CUSTOM_DATA_GENERATION && 
                 <>
                     <Form.Item
diff --git a/app/client/src/pages/DataGenerator/Examples.tsx b/app/client/src/pages/DataGenerator/Examples.tsx
index 31e61a6..9b46f38 100644
--- a/app/client/src/pages/DataGenerator/Examples.tsx
+++ b/app/client/src/pages/DataGenerator/Examples.tsx
@@ -149,7 +149,7 @@ const Examples: FunctionComponent = () => {
                     </Space>
                 </StyledTitle>
                 <Flex align='center' gap={15}>       
-                    {workflowType === WorkflowType.FREE_FORM_DATA_GENERATION && 
+                    {(workflowType === WorkflowType.CUSTOM_DATA_GENERATION || workflowType === WorkflowType.FREE_FORM_DATA_GENERATION) && 
                       <>
                         <Form.Item
                             name="example_path"
@@ -158,7 +158,7 @@ const Examples: FunctionComponent = () => {
                             style={{ display: 'none' }}
                             shouldUpdate
                             rules={[
-                               { required: true }
+                               { required: false }
                             ]}
                         >
                             <Input disabled />
diff --git a/app/client/src/pages/DataGenerator/Finish.tsx b/app/client/src/pages/DataGenerator/Finish.tsx
index 5a8197b..b092b2c 100644
--- a/app/client/src/pages/DataGenerator/Finish.tsx
+++ b/app/client/src/pages/DataGenerator/Finish.tsx
@@ -202,7 +202,8 @@ const Finish = () => {
     }
 
     let topicTabs = [];
-    if (!hasDocSeeds && formValues.workflow_type !== WorkflowType.CUSTOM_DATA_GENERATION && 
+    console.log('Preparing topics...', formValues.workflow_type === WorkflowType.CUSTOM_DATA_GENERATION)
+    if (!hasDocSeeds && formValues.workflow_type === WorkflowType.CUSTOM_DATA_GENERATION && formValues.use_case !== 'custom' && 
         hasTopics(genDatasetResp) && !isEmpty(genDatasetResp?.results)) {
             const values = Object.values(genDatasetResp?.results);
             
@@ -297,6 +298,10 @@ const Finish = () => {
             </>
         )
     }
+    console.log('Finish >> ');
+    console.log('hasTopics', hasTopics(genDatasetResp));
+    console.log('formValues', formValues);
+    console.log('isDemo', isDemo, topicTabs);
 
     return (
         <div>
@@ -322,12 +327,12 @@ const Finish = () => {
                     </StyledButton>
                 </Flex>
             )}
-            {isDemo && formValues.workflow_type !== WorkflowType.CUSTOM_DATA_GENERATION && hasTopics(genDatasetResp && !hasDocSeeds) && (
+            {(isDemo && formValues.workflow_type === WorkflowType.CUSTOM_DATA_GENERATION && formValues.use_case !== 'custom' && hasTopics(genDatasetResp) && !hasDocSeeds) && (
                 <TabsContainer title={'Generated Dataset'}>
                     <Tabs tabPosition='left' items={topicTabs}/>
                 </TabsContainer>
             )}
-            {formValues.workflow_type === WorkflowType.CUSTOM_DATA_GENERATION && 
+            {formValues.workflow_type === WorkflowType.CUSTOM_DATA_GENERATION && formValues.use_case === 'custom' && 
                 <CustomResultTable results={genDatasetResp?.results || []} />
             }
             {hasDocSeeds && <SeedResultTable results={genDatasetResp?.results || {}} />}
diff --git a/app/client/src/pages/DataGenerator/FreeFormTable.tsx b/app/client/src/pages/DataGenerator/FreeFormTable.tsx
index 24fb9f8..d09bdb9 100644
--- a/app/client/src/pages/DataGenerator/FreeFormTable.tsx
+++ b/app/client/src/pages/DataGenerator/FreeFormTable.tsx
@@ -45,11 +45,13 @@ interface Props {
 }
 
 const FreeFormTable: FunctionComponent<Props> = ({ data }) => {
+    console.log('FreeFormTable', data);
     const [colDefs, setColDefs] = useState([]);
     const [rowData, setRowData] = useState([]);
     
     useEffect(() => {
         if (!isEmpty(data)) {
+            console.log('data', data);
             const columnNames = Object.keys(first(data));
             const columnDefs = columnNames.map((colName) => ({
                 field: colName,
diff --git a/app/client/src/pages/DataGenerator/Prompt.tsx b/app/client/src/pages/DataGenerator/Prompt.tsx
index f89be15..082b15c 100644
--- a/app/client/src/pages/DataGenerator/Prompt.tsx
+++ b/app/client/src/pages/DataGenerator/Prompt.tsx
@@ -76,6 +76,7 @@ const SeedsFormItem = styled(StyledFormItem)`
 
 const Prompt = () => {
     const form = Form.useFormInstance();
+    console.log('form', form);
     const selectedTopics = Form.useWatch('topics');
     const numQuestions = Form.useWatch('num_questions');
     const datasetSize = Form.useWatch('num_questions');
@@ -108,6 +109,12 @@ const Prompt = () => {
         input_value,
         output_key
     );
+
+    console.log('values', form.getFieldsValue(true));
+    console.log('values', form);
+    console.log('workflow_type', workflow_type);
+    console.log('useCase', useCase);
+    console.log('dataset_size', dataset_size);
     const mutation = useMutation({
         mutationFn: fetchSeedList
     });
@@ -176,7 +183,7 @@ const Prompt = () => {
                 isStepValid = true;
             }
             setIsStepValid(isStepValid)
-        } else if(!isEmpty(doc_paths) && workflow_type === WorkflowType.SUPERVISED_FINE_TUNING) {
+        } else if(!isEmpty(doc_paths) && workflow_type === WorkflowType.FREE_FORM_DATA_GENERATION) {
             setIsStepValid(isNumber(datasetSize) && datasetSize > 0);
         } else {
             setIsStepValid(true);
diff --git a/app/client/src/pages/DataGenerator/Success.tsx b/app/client/src/pages/DataGenerator/Success.tsx
index 639f200..20f9fd0 100644
--- a/app/client/src/pages/DataGenerator/Success.tsx
+++ b/app/client/src/pages/DataGenerator/Success.tsx
@@ -52,7 +52,9 @@ const TopicsTable: FC<TopicsTableProps> = ({ formData, topic }) => {
             render: (_text: QuestionSolution, record: QuestionSolution) => <>{record.solution}</>
         },
     ]
-    const dataSource = formData.results[topic]
+    console.log('topic', topic);
+    const dataSource = formData.results[topic];
+    console.log('dataSource', dataSource);
     return (
         <StyledTable
             columns={cols}
@@ -87,12 +89,14 @@ interface SuccessProps {
     isDemo?: boolean;
 }
 const Success: FC<SuccessProps> = ({ formData, isDemo = true }) => {
+    console.log('Success', formData);
     const topicTabs = formData?.results && Object.keys(formData.results).map((topic, i) => ({
         key: `tab-${i}`,
         label: <Popover content={topic}><Text>{topic}</Text></Popover>,
         value: topic.replace(/\n/g, ' '),
         children: <TopicsTable formData={formData} topic={topic} />
     }));
+    console.log('topicTabs', topicTabs);
     const nextStepsList = [
         {
             avatar: '',
diff --git a/app/client/src/pages/DataGenerator/UseCaseSelector.tsx b/app/client/src/pages/DataGenerator/UseCaseSelector.tsx
index 16ae1e0..c159bc5 100644
--- a/app/client/src/pages/DataGenerator/UseCaseSelector.tsx
+++ b/app/client/src/pages/DataGenerator/UseCaseSelector.tsx
@@ -6,10 +6,11 @@ import get from "lodash/get";
 
 interface Props {
     form: FormInstance<any>;
+    hidden?: boolean;
 }
 
 
-const UseCaseSelector: FunctionComponent<Props> = ({ form }) => {
+const UseCaseSelector: FunctionComponent<Props> = ({ form, hidden }) => {
   const [useCases, setUseCases] = useState<UseCase[]>([]);
   const useCasesReq = useGetUseCases();  
 
@@ -45,6 +46,7 @@ const UseCaseSelector: FunctionComponent<Props> = ({ form }) => {
         labelCol={{
             span: 8
         }}
+        hidden={hidden}
         shouldUpdate
         >
             <Select placeholder={'Select a template'} onChange={onChange}>
diff --git a/app/client/src/pages/DataGenerator/hooks.ts b/app/client/src/pages/DataGenerator/hooks.ts
index 2980205..2f672f6 100644
--- a/app/client/src/pages/DataGenerator/hooks.ts
+++ b/app/client/src/pages/DataGenerator/hooks.ts
@@ -213,13 +213,13 @@ export const useDatasetSize = (
     input_value: string,
     output_key: string
  ) => {
-    if (workflow_type !== WorkflowType.CUSTOM_DATA_GENERATION || isEmpty(doc_paths)) {
+    if (workflow_type !== WorkflowType.CUSTOM_DATA_GENERATION && isEmpty(doc_paths)) {
         return {
           data: 0
         }
     }
     const params = {
-        input_path: doc_paths.map(item => item.value),
+        input_path: Array.isArray(doc_paths) ? doc_paths.map(item => item.value) : [],
         input_key,
         input_value,
         output_key
diff --git a/app/client/src/pages/DataGenerator/utils.ts b/app/client/src/pages/DataGenerator/utils.ts
index 699ce95..9b6cda4 100644
--- a/app/client/src/pages/DataGenerator/utils.ts
+++ b/app/client/src/pages/DataGenerator/utils.ts
@@ -109,12 +109,11 @@ export const getHttpStatusCodeVerb = (statusCode: number) => {
 
 export const getWizardModeType = (location: any) => {
   const pathname = location?.pathname || '';
-  switch (pathname) {
-    case '/data-augmentation':
-      return WizardModeType.DATA_AUGMENTATION;
-    case '/data-generator':
-      return WizardModeType.DATA_GENERATION;
-    default:
-      return null;
+  console.log('pathname', pathname);
+  if (pathname.includes('/data-augmentation')) {
+    return WizardModeType.DATA_AUGMENTATION;
+  } else if (pathname.includes('/data-generator')) {
+    return WizardModeType.DATA_GENERATION;
   }
+  return null;
 }