Nemotron-CC Task Reference

1
from nemo_curator.stages.synthetic.nemotron_cc.nemotron_cc import WikipediaParaphrasingStage
2
3
stage = WikipediaParaphrasingStage(
4
    client=llm_client,
5
    model_name="meta/llama-3.3-70b-instruct",
6
    generation_config=generation_config,
7
    input_field="text",
8
    output_field="rephrased",
9
)

1
from nemo_curator.stages.synthetic.nemotron_cc.nemotron_cc import DiverseQAStage
2
3
stage = DiverseQAStage(
4
    client=llm_client,
5
    model_name="meta/llama-3.3-70b-instruct",
6
    generation_config=generation_config,
7
    input_field="text",
8
    output_field="diverse_qa",
9
)

1
from nemo_curator.stages.synthetic.nemotron_cc.nemotron_cc import DiverseQAPostProcessingStage
2
3
post_stage = DiverseQAPostProcessingStage(
4
    input_field="text",
5
    qa_field="diverse_qa",
6
    tokenizer=tokenizer,  # For length-based sampling
7
    prefix="Here are the questions and answers based on the provided text:",
8
    max_num_pairs=10,
9
)

1
num_pairs = random.randint(1, max(1, int(max_num_pairs * num_tokens / 150)))

1
from nemo_curator.stages.synthetic.nemotron_cc.nemotron_cc import DistillStage
2
3
stage = DistillStage(
4
    client=llm_client,
5
    model_name="meta/llama-3.3-70b-instruct",
6
    generation_config=generation_config,
7
    input_field="text",
8
    output_field="distill",
9
)

1
from nemo_curator.stages.synthetic.nemotron_cc.nemotron_cc import ExtractKnowledgeStage
2
3
stage = ExtractKnowledgeStage(
4
    client=llm_client,
5
    model_name="meta/llama-3.3-70b-instruct",
6
    generation_config=generation_config,
7
    input_field="text",
8
    output_field="extract_knowledge",
9
)

1
from nemo_curator.stages.synthetic.nemotron_cc.nemotron_cc import KnowledgeListStage
2
3
stage = KnowledgeListStage(
4
    client=llm_client,
5
    model_name="meta/llama-3.3-70b-instruct",
6
    generation_config=generation_config,
7
    input_field="text",
8
    output_field="knowledge_list",
9
)

1
from nemo_curator.stages.synthetic.nemotron_cc.nemotron_cc import KnowledgeListPostProcessingStage
2
3
post_stage = KnowledgeListPostProcessingStage(
4
    input_field="knowledge_list",
5
)

1
from dataclasses import dataclass
2
from nemo_curator.stages.synthetic.nemotron_cc.base import BaseSyntheticStage
3
4
@dataclass
5
class CustomSyntheticStage(BaseSyntheticStage):
6
    system_prompt: str = "You are a helpful assistant specialized in..."
7
    prompt: str = """Your custom prompt template here.
8
9
Text: {document}
10
11
Instructions: ..."""
12
    input_field: str = "text"
13
    output_field: str = "custom_output"
14
15
    @property
16
    def name(self) -> str:
17
        return "CustomSyntheticStage"

1
TASK_CONFIG = {
2
    "diverse_qa": {
3
        "system_prompt": None,  # DiverseQAStage uses no system prompt
4
        "prompt_template": DIVERSE_QA_PROMPT_TEMPLATE,
5
        "min_document_tokens": 30,
6
        "min_segment_tokens": 30,
7
        "max_input_tokens": 1000,
8
        "max_output_tokens": 598,
9
    },
10
    "distill": {
11
        "system_prompt": NEMOTRON_CC_DISTILL_SYSTEM_PROMPT,
12
        "prompt_template": DISTILL_PROMPT_TEMPLATE,
13
        "min_document_tokens": 30,
14
        "min_segment_tokens": 10,
15
        "max_input_tokens": 2000,
16
        "max_output_tokens": 1598,
17
    },
18
    "extract_knowledge": {
19
        "system_prompt": None,  # ExtractKnowledgeStage uses no system prompt
20
        "prompt_template": EXTRACT_KNOWLEDGE_PROMPT_TEMPLATE,
21
        "min_document_tokens": 30,
22
        "min_segment_tokens": 30,
23
        "max_input_tokens": 1400,
24
        "max_output_tokens": 1400,
25
    },
26
    "knowledge_list": {
27
        "system_prompt": None,  # KnowledgeListStage uses no system prompt
28
        "prompt_template": KNOWLEDGE_LIST_PROMPT_TEMPLATE,
29
        "min_document_tokens": 30,
30
        "min_segment_tokens": 30,
31
        "max_input_tokens": 1000,
32
        "max_output_tokens": 598,
33
    },
34
    "wikipedia_paraphrasing": {
35
        "system_prompt": NEMOTRON_CC_SYSTEM_PROMPT,
36
        "prompt_template": WIKIPEDIA_REPHRASING_PROMPT_TEMPLATE,
37
        "min_document_tokens": 5,
38
        "min_segment_tokens": 5,
39
        "max_input_tokens": 512,
40
        "max_output_tokens": 510,
41
    },
42
}
43
44
GENERATION_CONFIG = {
45
    "MAX_INPUT_TOKENS": 2000,
46
    "MAX_OUTPUT_TOKENS": 1600,
47
    "TOP_K": 0,
48
    "TOP_P": 0.9,
49
    "TEMPERATURE": 0.5,
50
}