🧾 NeMo Data Designer: W-2 Dataset Generator

📚 What you'll learn

The notebook demonstrates how you can combine numerical samplers, the person sampler and LLMs to create a synthetic
dataset of W-2 forms (US Wage & Tax Statements).

• We will use generate numerical fields using statistics published by the IRS for the year 2021:

• We will use the person sampler to generate realistic US taxpayers. When the US locale is chosen, statistics
  for generated persons reflect real-world census data.

👋 IMPORTANT – Environment Setup

• If you haven't already, follow the instructions in the README to install the necessary dependencies.

• You may need to restart your notebook's kernel after setting up the environment.

• In this notebook, we assume you have a self-hosted instance of Data Designer up and running.

• For deployment instructions, see the Installation Options section of the NeMo Data Designer documentation.

📦 Import the essentials

• The data_designer module of nemo_microservices exposes Data Designer's high-level SDK.

• The essentials module provides quick access to the most commonly used objects.

⚙️ Initialize the NeMo Data Designer Client

• NeMoDataDesignerClient is responsible for submitting generation requests to the microservice.

🎛️ Define model configurations

• Each ModelConfig defines a model that can be used during the generation process.

• The "model alias" is used to reference the model in the Data Designer config (as we will see below).

• The "model provider" is the external service that hosts the model (see the model config docs for more details).

• By default, the microservice uses build.nvidia.com as the model provider.

🏗️ Initialize the Data Designer Config Builder

• The Data Designer config defines the dataset schema and generation process.

• The config builder provides an intuitive interface for building this configuration.

• The list of model configs is provided to the builder at initialization.

🎲 Setting Up Taxpayer and Employer Sampling

• Sampler columns offer non-LLM based generation of synthetic data.

• They are particularly useful for steering the diversity of the generated data, as we demonstrate below.

• The persona samplers allow you to sample realistic details of individuals using a model trained on the US Census.
  If the locale of the persona you are generating is anything other than en_US, then the personas will be generated using Faker

⚡️ Defining the Fields

We will focus on the following:

• Box 1 (Wages, tips, and other compensation)
• Box 2 (Federal income tax withheld)
• Box 3 (Social security wages)
• Box 4 (Social security tax withheld)
• Box 5 (Medicare wages and tips)
• Box 6 (Medicare tax withheld)
• Box 7 (Social security tips)
• Box a (Employee's social security number)
• Box c (Employer's name, address and zip code)
• Box e (Employee's fist name, initial, and last name)
• Box f (Employee's address and zip code)

Numerical fields

Here, we'll define how to generate numerical samples for the currency fields of the W-2 (Boxes 1-7).
We'll use the W-2 statistics from the IRS linked above to generate realistic samples.

🦜 Non-numerical Fields

The remaining fields contain information about the employee (taxpayer) and the employer.
We'll use the person sampler in combination with an LLM to generate values here.

🔁 Iteration is key – preview the dataset!

1. Use the preview method to generate a sample of records quickly.

2. Inspect the results for quality and format issues.

3. Adjust column configurations, prompts, or parameters as needed.

4. Re-run the preview until satisfied.

📊 Analyze the generated data

• Data Designer automatically generates a basic statistical analysis of the generated data.

• This analysis is available via the analysis property of generation result objects.

🆙 Scale up!

• Happy with your preview data?

• Use the create method to submit larger Data Designer generation jobs.