DSE-45114: Add Data Augmentation

app/client/src/pages/DataGenerator/Configure.tsx

@@ -8,9 +8,12 @@ import { File, WorkflowType } from './types';
 import { useFetchModels } from '../../api/api';
 import { MODEL_PROVIDER_LABELS } from './constants';
 import { ModelProviders, ModelProvidersDropdownOpts } from './types';
-import { useWizardCtx } from './utils';
+import { getWizardModel, getWizardModeType, useWizardCtx } from './utils';
 import FileSelectorButton from './FileSelectorButton';
 import UseCaseSelector from './UseCaseSelector';
+import { useLocation } from 'react-router-dom';
+import { WizardModeType } from '../../types';
+import { get } from 'lodash';
 
 
 const StepContainer = styled(Flex)`
@@ -39,7 +42,7 @@ export const USECASE_OPTIONS = [
 export const WORKFLOW_OPTIONS = [
     { label: 'Supervised Fine-Tuning', value: 'supervised-fine-tuning' },
     { label: 'Custom Data Generation', value: 'custom' },
-    { label: 'Freeform Data Generation', value: 'freeform' }
+    // { label: 'Freeform Data Generation', value: 'freeform' }
 ];
 
 export const MODEL_TYPE_OPTIONS: ModelProvidersDropdownOpts = [
@@ -48,6 +51,18 @@ export const MODEL_TYPE_OPTIONS: ModelProvidersDropdownOpts = [
 ];
 
 const Configure = () => {
+    const location = useLocation();
+    const [wizardModeType, setWizardModeType] = useState(getWizardModeType(location));
+
+    useEffect(() => {
+        if (wizardModeType === WizardModeType.DATA_AUGMENTATION) {
+            setWizardModeType(WizardModeType.DATA_AUGMENTATION);
+            form.setFieldValue('workflow_type', 'freeform');
+        } else {
+            setWizardModeType(WizardModeType.DATA_GENERATION);
+        }
+    }, [location, wizardModeType]);
+
     const form = Form.useFormInstance();
     const formData = Form.useWatch((values) => values, form);
     const { setIsStepValid } = useWizardCtx();
@@ -141,8 +156,10 @@ const Configure = () => {
                     label='Model Provider'
                     rules={[{ required: true }]}
                     labelCol={labelCol}
+                    shouldUpdate
                 >
                     <Select
+
                         onChange={() => form.setFieldValue('model_id', undefined)}
                         placeholder={'Select a model provider'}
                     >
@@ -210,6 +227,7 @@ const Configure = () => {
                     label='Workflow'
                     tooltip='A specialized workflow for your dataset'
                     labelCol={labelCol}
+                    hidden={wizardModeType === WizardModeType.DATA_AUGMENTATION}
                     shouldUpdate
                     rules={[
                             { required: true }

app/client/src/pages/DataGenerator/DataGenerator.tsx

@@ -1,6 +1,6 @@
 import isEmpty from 'lodash/isEmpty';
 import isString from 'lodash/isString';
-import { useEffect, useRef, useState } from 'react';
+import { FunctionComponent, useEffect, useRef, useState } from 'react';
 import { useLocation, useParams } from 'react-router-dom';
 
 import { Button, Flex, Form, Layout, Steps } from 'antd';
@@ -20,10 +20,15 @@ import { DataGenWizardSteps, WizardStepConfig, WorkflowType } from './types';
 import { WizardCtx } from './utils';
 import { fetchDatasetDetails, useGetDatasetDetails } from '../DatasetDetails/hooks';
 import { useMutation } from '@tanstack/react-query';
+import { WizardModeType } from '../../types';
 
 const { Content } = Layout;
 // const { Title } = Typography;
 
+interface Props {
+    mode?: WizardModeType;
+}
+
 const StyledTitle = styled.div`
     margin-top: 10px;
     font-family: Roboto, -apple-system, 'Segoe UI', sans-serif;
@@ -95,13 +100,15 @@ const steps: WizardStepConfig[] = [
 /**
  * Wizard component for Synthetic Data Generation workflow
  */
-const DataGenerator = () => {
+const DataGenerator: FunctionComponent<Props> = ({ mode }) => {
+    console.log('DataGenerator mode: ', mode);
     const [current, setCurrent] = useState(0);
     const [maxStep, setMaxStep] = useState(0);
     const [isStepValid, setIsStepValid] = useState<boolean>(false);
 
     // Data passed from listing table to prepopulate form
     const location = useLocation();
+    console.log('DataGenerator location: ', location);
     const { generate_file_name } = useParams();
     const initialData = location?.state?.data;
     const mutation = useMutation({

app/client/src/pages/DataGenerator/utils.ts

@@ -3,6 +3,7 @@ import { WizardCtxObj } from './types';
 import moment from 'moment';
 import toString from 'lodash/toString';
 import { File } from './types';
+import { WizardModeType } from '../../types';
 
 export const WizardCtx = createContext<WizardCtxObj | null>(null);
 export const useWizardCtx = (): WizardCtxObj => {
@@ -105,3 +106,15 @@ export const getHttpStatusCodeVerb = (statusCode: number) => {
       return null;
   }
 };
+
+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;
+  }
+}

app/client/src/pages/Home/HomePage.tsx

@@ -7,6 +7,7 @@ import EvaluationsTab from './EvaluationsTab';
 import DatasetIcon from '../../assets/ic-datasets.svg';
 import ArrowRightIcon from '../../assets/ic-arrow-right.svg';
 import EvaluateIcon from '../../assets/ic-evaluations.svg';
+import DataAugmentationIcon from '../../assets/ic-data-augmentation.svg';
 import EvaluateButton from './EvaluateButton';
 import ExportsTab from './ExportsTab';
 
@@ -116,6 +117,25 @@ const HomePage: React.FC = () => {
                             </div>
                         </div>
                     </HeaderSection>
+                    <HeaderSection style={{ marginLeft: '1rem' }}>
+                        <div className="left-section" style={{ padding: '5px' }}>
+                            <img src={DataAugmentationIcon} alt="Datasets" />
+                        </div>
+                        <div className="middle-section">
+                            <div className="section-title">Data Augmentation</div>
+                            <div className="section-description">
+                                <p>Generate multi-dimension datasets using LLM custom prompts</p>
+                            </div>
+                        </div>
+                        <div className="right-section">
+                            <div>
+                                <Button href="/data-augmentation">
+                                    Get Started
+                                    <img src={ArrowRightIcon} alt="Get Started" />
+                                </Button>
+                            </div>
+                        </div>
+                    </HeaderSection>
                     <HeaderSection style={{ marginLeft: '1rem' }}>
                         <div className="left-section evaluate-icon">
                             <img src={EvaluateIcon} alt="Datasets" />

app/client/src/routes.tsx

@@ -2,7 +2,7 @@ import { Navigate, createBrowserRouter } from "react-router-dom";
 import Layout from "./Container";
 import DataGenerator from "./pages/DataGenerator";
 import HomePage from "./pages/Home";
-import { Pages } from "./types";
+import { Pages, WizardModeType } from "./types";
 import EvaluatorPage from "./pages/Evaluator";
 import ReevaluatorPage from "./pages/Evaluator/ReevaluatorPage";
 import DatasetDetailsPage from "./pages/DatasetDetails/DatasetDetailsPage";
@@ -35,7 +35,13 @@ const router = createBrowserRouter([
       },
       { 
         path: Pages.GENERATOR, 
-        element: <DataGenerator key={Pages.GENERATOR}/>, 
+        element: <DataGenerator key={Pages.GENERATOR} mode={WizardModeType.DATA_GENERATION}/>, 
+        errorElement: <ErrorPage />,
+        loader: async () => null
+      },
+      { 
+        path: Pages.DATA_AUGMENTATION, 
+        element: <DataGenerator key={Pages.DATA_AUGMENTATION} mode={WizardModeType.DATA_AUGMENTATION}/>, 
         errorElement: <ErrorPage />,
         loader: async () => null
       },

app/client/src/types.ts

@@ -1,5 +1,6 @@
 export enum Pages {
     GENERATOR = 'data-generator',
+    DATA_AUGMENTATION = 'data-augmentation',
     REGENERATE = 're-generate',
     EVALUATOR = 'evaluator',
     HISTORY = 'history',
@@ -52,4 +53,9 @@ export interface UseCase {
     id: string;
     label: string;
     value: string;
+}
+
+export enum WizardModeType {
+    DATA_GENERATION = 'data-generation',
+    DATA_AUGMENTATION = 'data-augmention'
 }