NeMo Data Designer Integration

1
import data_designer.config as dd
2
3
config_builder = dd.DataDesignerConfigBuilder(model_configs=[model_config])
4
5
config_builder.add_column(
6
    dd.SamplerColumnConfig(
7
        name="patient_name",
8
        sampler_type=dd.SamplerType.PERSON_FROM_FAKER,
9
        params=dd.PersonFromFakerSamplerParams(),
10
    )
11
)
12
13
config_builder.add_column(
14
    dd.SamplerColumnConfig(
15
        name="patient_id",
16
        sampler_type=dd.SamplerType.UUID,
17
        params=dd.UUIDSamplerParams(prefix="PT-", short_form=True, uppercase=True),
18
    )
19
)

1
config_builder.add_column(
2
    dd.ExpressionColumnConfig(
3
        name="first_name",
4
        expr="{{ patient_name.first_name }}",
5
    )
6
)

1
config_builder.add_column(
2
    dd.LLMTextColumnConfig(
3
        name="physician_notes",
4
        prompt="""\
5
You are a primary-care physician who just had an appointment with {{ first_name }}.
6
{{ patient_summary }}
7
Write careful notes about your visit. Respond with only the notes.
8
""",
9
        model_alias="local-llm",
10
    )
11
)

1
import data_designer.config as dd
2
3
from nemo_curator.backends.ray_data import RayDataExecutor
4
from nemo_curator.core.client import RayClient
5
from nemo_curator.core.serve import InferenceModelConfig, InferenceServer
6
from nemo_curator.pipeline import Pipeline
7
from nemo_curator.stages.synthetic.nemo_data_designer.data_designer import DataDesignerStage
8
from nemo_curator.stages.text.io.reader.jsonl import JsonlReader
9
from nemo_curator.stages.text.io.writer.jsonl import JsonlWriter
10
11
# Start Ray cluster
12
client = RayClient(num_cpus=16, num_gpus=4)
13
client.start()
14
15
# Start local inference server
16
server_config = InferenceModelConfig(
17
    model_identifier="google/gemma-3-27b-it",
18
    deployment_config={"autoscaling_config": {"min_replicas": 1, "max_replicas": 1}},
19
    engine_kwargs={"tensor_parallel_size": 4},
20
)
21
inference_server = InferenceServer(models=[server_config])
22
inference_server.start()
23
24
# Configure NDD model
25
model_config = dd.ModelConfig(
26
    alias="local-llm",
27
    model="google/gemma-3-27b-it",
28
    provider="local",
29
    skip_health_check=True,
30
    inference_parameters=dd.ChatCompletionInferenceParams(
31
        temperature=1.0, top_p=1.0, max_tokens=2048,
32
    ),
33
)
34
35
model_provider = dd.ModelProvider(
36
    name="local",
37
    endpoint=inference_server.endpoint,
38
    api_key="unused",
39
)
40
41
# Build config with sampler and LLM columns
42
config_builder = dd.DataDesignerConfigBuilder(model_configs=[model_config])
43
44
config_builder.add_column(
45
    dd.SamplerColumnConfig(
46
        name="patient_name",
47
        sampler_type=dd.SamplerType.PERSON_FROM_FAKER,
48
        params=dd.PersonFromFakerSamplerParams(),
49
    )
50
)
51
52
config_builder.add_column(
53
    dd.LLMTextColumnConfig(
54
        name="physician_notes",
55
        prompt="You are a physician. Write notes for {{ patient_name.first_name }} "
56
               "who has {{ diagnosis }}. {{ patient_summary }}",
57
        model_alias="local-llm",
58
    )
59
)
60
61
# Build and run pipeline
62
pipeline = Pipeline(name="ndd_medical_notes")
63
pipeline.add_stage(JsonlReader(file_paths="seed_data/*.jsonl", fields=["diagnosis", "patient_summary"]))
64
pipeline.add_stage(DataDesignerStage(config_builder=config_builder, model_providers=[model_provider]))
65
pipeline.add_stage(JsonlWriter(path="./synthetic_output"))
66
67
pipeline.run(executor=RayDataExecutor())
68
69
inference_server.stop()
70
client.stop()

1
import os
2
3
import data_designer.config as dd
4
5
from nemo_curator.stages.synthetic.nemo_data_designer.data_designer import DataDesignerStage
6
7
model_config = dd.ModelConfig(
8
    alias="nim-llm",
9
    model="meta/llama-3.3-70b-instruct",
10
    provider="nvidia",
11
    inference_parameters=dd.ChatCompletionInferenceParams(
12
        temperature=0.5, top_p=0.9, max_tokens=1600,
13
    ),
14
)
15
16
model_provider = dd.ModelProvider(
17
    name="nvidia",
18
    endpoint="https://integrate.api.nvidia.com/v1",
19
    provider_type="openai",
20
    api_key=os.environ["NVIDIA_API_KEY"],
21
)
22
23
config_builder = dd.DataDesignerConfigBuilder(model_configs=[model_config])
24
# Add columns as needed...
25
26
stage = DataDesignerStage(
27
    config_builder=config_builder,
28
    model_providers=[model_provider],
29
)

1
import os
2
3
import data_designer.config as dd
4
5
from nemo_curator.stages.synthetic.nemotron_cc.nemo_data_designer.nemotron_cc import DiverseQAStage
6
7
model_config = dd.ModelConfig(
8
    alias="meta/llama-3.3-70b-instruct",
9
    model="meta/llama-3.3-70b-instruct",
10
    provider="nvidia",
11
    inference_parameters=dd.ChatCompletionInferenceParams(
12
        temperature=0.5, top_p=0.9, max_tokens=1600,
13
    ),
14
)
15
16
model_provider = dd.ModelProvider(
17
    name="nvidia",
18
    endpoint="https://integrate.api.nvidia.com/v1",
19
    provider_type="openai",
20
    api_key=os.environ["NVIDIA_API_KEY"],
21
)
22
23
stage = DiverseQAStage(
24
    input_field="text",
25
    output_field="diverse_qa",
26
    model_alias="meta/llama-3.3-70b-instruct",
27
    model_configs=[model_config],
28
    model_providers=[model_provider],
29
)

1
stage = DataDesignerStage(data_designer_config_file="config.yaml")