1
import data_designer.config as dd
2
from data_designer.interface import DataDesigner
3
 

1
data_designer = DataDesigner()
2
 

1
# This name is set in the model provider configuration.
2
MODEL_PROVIDER = "nvidia"
3
 
4
# The model ID is from build.nvidia.com.
5
MODEL_ID = "nvidia/nemotron-3-nano-30b-a3b"
6
 
7
# We choose this alias to be descriptive for our use case.
8
MODEL_ALIAS = "nemotron-nano-v3"
9
 
10
model_configs = [
11
    dd.ModelConfig(
12
        alias=MODEL_ALIAS,
13
        model=MODEL_ID,
14
        provider=MODEL_PROVIDER,
15
        inference_parameters=dd.ChatCompletionInferenceParams(
16
            temperature=1.0,
17
            top_p=1.0,
18
            max_tokens=2048,
19
            extra_body={"chat_template_kwargs": {"enable_thinking": False}},
20
        ),
21
    )
22
]
23
 

1
config_builder = dd.DataDesignerConfigBuilder(model_configs=model_configs)
2
 

1
# Download sample dataset from Github
2
import urllib.request
3
 
4
url = "https://raw.githubusercontent.com/NVIDIA/GenerativeAIExamples/refs/heads/main/nemo/NeMo-Data-Designer/data/gretelai_symptom_to_diagnosis.csv"
5
local_filename, _ = urllib.request.urlretrieve(url, "gretelai_symptom_to_diagnosis.csv")
6
 
7
# Seed datasets are passed as reference objects to the config builder.
8
seed_source = dd.LocalFileSeedSource(path=local_filename)
9
 
10
config_builder.with_seed_dataset(seed_source)
11
 

1
config_builder.add_column(
2
    dd.SamplerColumnConfig(
3
        name="patient_sampler",
4
        sampler_type=dd.SamplerType.PERSON_FROM_FAKER,
5
        params=dd.PersonFromFakerSamplerParams(),
6
    )
7
)
8
 
9
config_builder.add_column(
10
    dd.SamplerColumnConfig(
11
        name="doctor_sampler",
12
        sampler_type=dd.SamplerType.PERSON_FROM_FAKER,
13
        params=dd.PersonFromFakerSamplerParams(),
14
    )
15
)
16
 
17
config_builder.add_column(
18
    dd.SamplerColumnConfig(
19
        name="patient_id",
20
        sampler_type=dd.SamplerType.UUID,
21
        params=dd.UUIDSamplerParams(
22
            prefix="PT-",
23
            short_form=True,
24
            uppercase=True,
25
        ),
26
    )
27
)
28
 
29
config_builder.add_column(dd.ExpressionColumnConfig(name="first_name", expr="{{ patient_sampler.first_name }}"))
30
 
31
config_builder.add_column(dd.ExpressionColumnConfig(name="last_name", expr="{{ patient_sampler.last_name }}"))
32
 
33
config_builder.add_column(dd.ExpressionColumnConfig(name="dob", expr="{{ patient_sampler.birth_date }}"))
34
 
35
config_builder.add_column(
36
    dd.SamplerColumnConfig(
37
        name="symptom_onset_date",
38
        sampler_type=dd.SamplerType.DATETIME,
39
        params=dd.DatetimeSamplerParams(start="2024-01-01", end="2024-12-31"),
40
    )
41
)
42
 
43
config_builder.add_column(
44
    dd.SamplerColumnConfig(
45
        name="date_of_visit",
46
        sampler_type=dd.SamplerType.TIMEDELTA,
47
        params=dd.TimeDeltaSamplerParams(dt_min=1, dt_max=30, reference_column_name="symptom_onset_date"),
48
    )
49
)
50
 
51
config_builder.add_column(dd.ExpressionColumnConfig(name="physician", expr="Dr. {{ doctor_sampler.last_name }}"))
52
 
53
config_builder.add_column(
54
    dd.LLMTextColumnConfig(
55
        name="physician_notes",
56
        prompt="""\
57
You are a primary-care physician who just had an appointment with {{ first_name }} {{ last_name }},
58
who has been struggling with symptoms from {{ diagnosis }} since {{ symptom_onset_date }}.
59
The date of today's visit is {{ date_of_visit }}.
60
 
61
{{ patient_summary }}
62
 
63
Write careful notes about your visit with {{ first_name }},
64
as Dr. {{ doctor_sampler.first_name }} {{ doctor_sampler.last_name }}.
65
 
66
Format the notes as a busy doctor might.
67
Respond with only the notes, no other text.
68
""",
69
        model_alias=MODEL_ALIAS,
70
    )
71
)
72
 
73
data_designer.validate(config_builder)
74
 

1
preview = data_designer.preview(config_builder, num_records=2)
2
 

1
# Run this cell multiple times to cycle through the 2 preview records.
2
preview.display_sample_record()
3
 

1
# The preview dataset is available as a pandas DataFrame.
2
preview.dataset
3
 

1
# Print the analysis as a table.
2
preview.analysis.to_report()
3
 

1
results = data_designer.create(config_builder, num_records=10, dataset_name="tutorial-3")
2
 

1
# Load the generated dataset as a pandas DataFrame.
2
dataset = results.load_dataset()
3
 
4
dataset.head()
5
 

1
# Load the analysis results into memory.
2
analysis = results.load_analysis()
3
 
4
analysis.to_report()
5