sdg-hub 0.1.0__py3-none-any.whl

sdg_hub/flows/generation/knowledge/synth_knowledge1.5.yaml

@@ -0,0 +1,136 @@
+- block_type: DuplicateColumns
+  block_config:
+    block_name: duplicate_document_col
+    columns_map:
+      document: base_document
+
+- block_type: LLMBlock
+  block_config:
+    block_name: gen_detailed_summary
+    config_path: configs/knowledge/detailed_summary.yaml
+    model_id: mistralai/Mixtral-8x7B-Instruct-v0.1
+    output_cols:
+      - summary_detailed
+  gen_kwargs:
+    max_tokens: 2048
+
+- block_type: LLMBlock
+  block_config:
+    block_name: gen_atomic_facts
+    config_path: configs/knowledge/atomic_facts.yaml
+    model_id: mistralai/Mixtral-8x7B-Instruct-v0.1
+    output_cols:
+      - summary_atomic_facts
+  gen_kwargs:
+    max_tokens: 2048
+
+- block_type: LLMBlock
+  block_config:
+    block_name: gen_extractive_summary
+    config_path: configs/knowledge/extractive_summary.yaml
+    model_id: mistralai/Mixtral-8x7B-Instruct-v0.1
+    output_cols:
+      - summary_extractive
+  gen_kwargs:
+    max_tokens: 2048
+
+- block_type: FlattenColumnsBlock
+  block_config:
+    block_name: flatten_summary_columns
+    var_cols:
+      - summary_detailed
+      - summary_extractive
+      - summary_atomic_facts
+      - base_document
+    value_name: summary
+    var_name: dataset_type
+
+- block_type: RenameColumns
+  block_config:
+    block_name: rename_to_document_column
+    columns_map:
+      document: raw_document
+      summary: document
+
+- block_type: LLMBlock
+  block_config:
+    block_name: knowledge generation
+    config_path: configs/knowledge/generate_questions_responses.yaml
+    model_id: mistralai/Mixtral-8x7B-Instruct-v0.1
+    output_cols:
+      - question
+      - response
+    parser_kwargs:
+      parser_name: custom
+      parsing_pattern: "\\[(?:Question|QUESTION)\\]\\s*(.*?)\\s*\\[(?:Answer|ANSWER)\\]\\s*(.*?)\\s*(?=\\[(?:Question|QUESTION)\\]|$)"
+      parser_cleanup_tags:
+        - "[END]"
+  gen_kwargs:
+    temperature: 0.0
+    max_tokens: 2048
+
+- block_type: LLMBlock
+  block_config:
+    block_name: eval_faithfulness_qa_pair
+    config_path: configs/knowledge/evaluate_faithfulness.yaml
+    model_id: mistralai/Mixtral-8x7B-Instruct-v0.1
+    output_cols:
+      - explanation
+      - judgment
+  gen_kwargs:
+    max_tokens: 2048
+
+- block_type: FilterByValueBlock
+  block_config:
+    block_name: filter_faithfulness
+    filter_column: judgment
+    filter_value: "YES"
+    operation: operator.eq
+  drop_columns:
+    - judgment
+    - explanation
+
+- block_type: LLMBlock
+  block_config:
+    block_name: eval_relevancy_qa_pair
+    config_path: configs/knowledge/evaluate_relevancy.yaml
+    model_id: mistralai/Mixtral-8x7B-Instruct-v0.1
+    output_cols:
+      - feedback
+      - score
+  gen_kwargs:
+    max_tokens: 2048
+
+- block_type: FilterByValueBlock
+  block_config:
+    block_name: filter_relevancy
+    filter_column: score
+    filter_value: 2.0
+    operation: operator.eq
+    convert_dtype: float
+  drop_columns:
+    - feedback
+    - score
+
+- block_type: LLMBlock
+  block_config:
+    block_name: eval_verify_question
+    config_path: configs/knowledge/evaluate_question.yaml
+    model_id: mistralai/Mixtral-8x7B-Instruct-v0.1
+    output_cols:
+      - explanation
+      - rating
+  gen_kwargs:
+    max_tokens: 2048
+
+- block_type: FilterByValueBlock
+  block_config:
+    block_name: filter_verify_question
+    filter_column: rating
+    filter_value: 1.0
+    operation: operator.eq
+    convert_dtype: float
+  drop_columns:
+    - explanation
+    - rating
+    - __index_level_0__

sdg_hub/flows/generation/skills/improve_responses.yaml

@@ -0,0 +1,103 @@
+- block_type: LLMBlock
+  block_config:
+    block_name: router
+    config_path: configs/skills/router.yaml
+    model_id: meta-llama/Llama-3.3-70B-Instruct
+    output_cols:
+      - route
+  gen_kwargs:
+    temperature: 0
+    max_tokens: 5
+    extra_body:
+      guided_choice:
+        - "coding"
+        - "extraction"
+        - "humanities"
+        - "math"
+        - "reasoning"
+        - "roleplay"
+        - "STEM"
+        - "writing"
+- block_type: SamplePopulatorBlock
+  block_config:
+    block_name: icl_populator
+    config_paths: 
+      - configs/skills/icl_examples/coding.yaml
+      - configs/skills/icl_examples/extraction.yaml
+      - configs/skills/icl_examples/humanities.yaml
+      - configs/skills/icl_examples/math.yaml
+      - configs/skills/icl_examples/reasoning.yaml
+      - configs/skills/icl_examples/roleplay.yaml
+      - configs/skills/icl_examples/STEM.yaml
+      - configs/skills/icl_examples/writing.yaml
+    column_name: route
+    batch_kwargs:
+      num_procs: 8
+- block_type: LLMBlock
+  block_config:
+    block_name: analyzer
+    config_path: configs/skills/analyzer.yaml
+    model_id: meta-llama/Llama-3.3-70B-Instruct
+    output_cols:
+      - analysis
+      - rubric
+- block_type: LLMBlock
+  block_config:
+    block_name: critic
+    config_path: configs/skills/critic.yaml
+    model_id: meta-llama/Llama-3.3-70B-Instruct
+    output_cols: 
+      - critique
+- block_type: LLMBlock
+  block_config:
+    block_name: planner
+    config_path: configs/skills/planner.yaml
+    model_id: meta-llama/Llama-3.3-70B-Instruct
+    output_cols: 
+      - plan
+- block_type: LLMBlock
+  block_config:
+    block_name: revised_responder
+    config_path: configs/skills/revised_responder.yaml
+    model_id: meta-llama/Llama-3.3-70B-Instruct
+    output_cols: 
+      - revised_response
+  drop_columns:
+    - icl_query
+    - icl_response
+    - icl_analysis
+    - icl_rubric
+    - icl_critique
+    - icl_plan
+    - icl_revised_response
+- block_type: LLMBlock
+  block_config:
+    block_name: judge
+    config_path: configs/skills/judge.yaml
+    model_id: meta-llama/Llama-3.3-70B-Instruct
+    output_cols:
+      - judgement
+      - verdict
+- block_type: FilterByValueBlock
+  block_config:
+    block_name: filter_judgement
+    filter_column: verdict
+    filter_value:
+      - Assistant A
+      - Assistant B
+    operation: operator.contains
+    batch_kwargs:
+      num_procs: 8
+- block_type: SelectorBlock
+  block_config:
+    block_name: response_selector
+    choice_map:
+       Assistant A: "response"
+       Assistant B: "revised_response"
+    choice_col: verdict
+    output_col: chosen_response
+    batch_kwargs:
+      num_procs: 8
+  drop_columns:
+    - judgemnent
+    - verdict

sdg_hub/flows/generation/skills/simple_freeform_skill.yaml

@@ -0,0 +1,12 @@
+- block_type: LLMBlock
+  block_config:
+    block_name: gen_skill_freeform
+    config_path: configs/skills/simple_generate_qa_freeform.yaml
+    model_id: mistralai/Mixtral-8x7B-Instruct-v0.1
+    output_cols:
+      - output
+  gen_kwargs:
+    temperature: 0.7
+    max_tokens: 2048
+  drop_duplicates:
+    - output

sdg_hub/flows/generation/skills/simple_grounded_skill.yaml

@@ -0,0 +1,12 @@
+- block_type: LLMBlock
+  block_config:
+    block_name: gen_skill_grounded
+    config_path: configs/skills/simple_generate_qa_grounded.yaml
+    model_id: mistralai/Mixtral-8x7B-Instruct-v0.1
+    output_cols:
+      - output
+  gen_kwargs:
+    temperature: 0.7
+    max_tokens: 2048
+  drop_duplicates:
+    - output

sdg_hub/flows/generation/skills/synth_grounded_skills.yaml

@@ -0,0 +1,80 @@
+- block_type: LLMBlock
+  block_config:
+    block_name: gen_contexts
+    config_path: configs/skills/contexts.yaml
+    model_id: mistralai/Mixtral-8x7B-Instruct-v0.1
+    output_cols:
+      - context
+  gen_kwargs:
+    temperature: 0.7
+    max_tokens: 2048
+    n: 10
+    seed: 42
+  drop_duplicates:
+    - context
+- block_type: LLMBlock
+  block_config:
+    block_name: gen_grounded_questions
+    config_path: configs/skills/grounded_questions.yaml
+    model_id: mistralai/Mixtral-8x7B-Instruct-v0.1
+    output_cols:
+      - question
+    batch_kwargs:
+      num_samples: 3
+  drop_duplicates:
+    - question
+- block_type: LLMBlock
+  block_config:
+    block_name: eval_grounded_questions
+    config_path: configs/skills/evaluate_grounded_questions.yaml
+    model_id: mistralai/Mixtral-8x7B-Instruct-v0.1
+    output_cols:
+      - evaluation
+      - score
+- block_type: FilterByValueBlock
+  block_config:
+    block_name: filter_grounded_questions
+    filter_column: score
+    filter_value: 1.0
+    operation: operator.eq
+    convert_dtype: float
+    batch_kwargs:
+      num_procs: 8
+  drop_columns:
+    - evaluation
+    - score
+    - num_samples
+- block_type: LLMBlock
+  block_config:
+    block_name: gen_grounded_responses
+    config_path: configs/skills/grounded_responses.yaml
+    model_id: mistralai/Mixtral-8x7B-Instruct-v0.1
+    output_cols:
+      - response
+- block_type: LLMBlock
+  block_config:
+    block_name: evaluate_grounded_qa_pair
+    config_path: configs/skills/evaluate_grounded_pair.yaml
+    model_id: mistralai/Mixtral-8x7B-Instruct-v0.1
+    output_cols:
+      - evaluation
+      - score
+- block_type: FilterByValueBlock
+  block_config:
+    block_name: filter_grounded_qa_pair
+    filter_column: score
+    filter_value: 2.0
+    operation: operator.ge
+    convert_dtype: float
+    batch_kwargs:
+      num_procs: 8
+- block_type: CombineColumnsBlock
+  block_config:
+    block_name: combine_question_and_context
+    columns:
+      - context
+      - question
+    output_col: question
+    batch_kwargs:
+      num_procs: 8
+      batched: True

sdg_hub/flows/generation/skills/synth_skills.yaml

@@ -0,0 +1,59 @@
+- block_type: LLMBlock
+  block_config:
+    block_name: gen_questions
+    config_path: configs/skills/freeform_questions.yaml
+    model_id: meta-llama/Llama-3.3-70B-Instruct
+    output_cols:
+      - question
+    batch_kwargs:
+      num_samples: 30
+  drop_duplicates:
+    - question
+- block_type: LLMBlock
+  block_config:
+    block_name: eval_questions
+    config_path: configs/skills/evaluate_freeform_questions.yaml
+    model_id: meta-llama/Llama-3.3-70B-Instruct
+    output_cols:
+      - evaluation
+      - score
+- block_type: FilterByValueBlock
+  block_config:
+    block_name: filter_questions
+    filter_column: score
+    filter_value: 1.0
+    operation: operator.eq
+    convert_dtype: float
+    batch_kwargs:
+      num_procs: 8
+  drop_columns:
+    - evaluation
+    - score
+    - num_samples
+- block_type: LLMBlock
+  block_config:
+    block_name: gen_responses
+    config_path: configs/skills/freeform_responses.yaml
+    model_id: meta-llama/Llama-3.3-70B-Instruct
+    output_cols:
+      - response
+- block_type: LLMBlock
+  block_config:
+    block_name: evaluate_qa_pair
+    config_path: configs/skills/evaluate_freeform_pair.yaml
+    model_id: meta-llama/Llama-3.3-70B-Instruct
+    output_cols:
+      - evaluation
+      - score
+- block_type: FilterByValueBlock
+  block_config:
+    block_name: filter_qa_pair
+    filter_column: score
+    filter_value: 2.0
+    operation: operator.ge
+    convert_dtype: float
+    batch_kwargs:
+      num_procs: 8
+  drop_columns:
+    - evaluation
+    - score

sdg_hub/logger_config.py

@@ -0,0 +1,20 @@
+# SPDX-License-Identifier: Apache-2.0
+# Standard
+import os
+import logging
+
+# Third Party
+from rich.logging import RichHandler
+
+
+def setup_logger(name):
+    # Set up the logger
+    log_level = os.getenv("LOG_LEVEL", "INFO")
+    logging.basicConfig(
+        level=log_level,
+        format="%(message)s",
+        datefmt="[%X]",
+        handlers=[RichHandler()],
+    )
+    logger = logging.getLogger(name)
+    return logger

sdg_hub/pipeline.py

@@ -0,0 +1,121 @@
+"""
+Deprecated Pipeline class for data generation pipelines.
+
+Use the Flow class directly for new code.
+"""
+
+# SPDX-License-Identifier: Apache-2.0
+# Standard
+import warnings
+from typing import List, Dict, Any
+
+# Third Party
+from datasets import Dataset
+from datasets.data_files import EmptyDatasetError
+
+# Local
+from .logger_config import setup_logger
+
+logger = setup_logger(__name__)
+
+
+class Pipeline:
+    """A class representing a data generation pipeline.
+
+    This class is deprecated and will be removed in a future version.
+    Use the Flow class directly instead.
+
+    Parameters
+    ----------
+    chained_blocks : List[Dict[str, Any]]
+        List of block configurations to execute in sequence.
+
+    Attributes
+    ----------
+    chained_blocks : List[Dict[str, Any]]
+        List of block configurations to execute in sequence.
+    """
+
+    def __init__(self, chained_blocks: List[Dict[str, Any]]) -> None:
+        """
+        Initialize the Pipeline class with a configuration dictionary.
+        
+        DEPRECATED: This class is deprecated and will be removed in a future version.
+        Use the Flow class directly instead.
+        """
+        warnings.warn(
+            "Pipeline class is deprecated and will be removed in a future version. "
+            "Use Flow class directly instead of wrapping it with Pipeline.",
+            DeprecationWarning,
+            stacklevel=2
+        )
+        # pipeline config is the run configuration that consists of the pipeline steps
+        self.chained_blocks = chained_blocks
+
+    def _drop_duplicates(self, dataset: Dataset, cols: List[str]) -> Dataset:
+        """Drop duplicates from the dataset based on the columns provided.
+
+        Parameters
+        ----------
+        dataset : Dataset
+            The input dataset.
+        cols : List[str]
+            Columns to consider for duplicate detection.
+
+        Returns
+        -------
+        Dataset
+            Dataset with duplicates removed.
+        """
+        df = dataset.to_pandas()
+        df = df.drop_duplicates(subset=cols).reset_index(drop=True)
+        return Dataset.from_pandas(df)
+
+    def generate(self, dataset: Dataset) -> Dataset:
+        """Generate the dataset by running the pipeline steps.
+
+        Parameters
+        ----------
+        dataset : Dataset
+            The input dataset to process.
+
+        Returns
+        -------
+        Dataset
+            The processed dataset.
+
+        Raises
+        ------
+        EmptyDatasetError
+            If a block produces an empty dataset.
+        """
+        for block_prop in self.chained_blocks:
+            block_type = block_prop["block_type"]
+            block_config = block_prop["block_config"]
+            drop_columns = block_prop.get("drop_columns", [])
+            gen_kwargs = block_prop.get("gen_kwargs", {})
+            drop_duplicates_cols = block_prop.get("drop_duplicates", False)
+            block = block_type(**block_config)
+
+            logger.debug("------------------------------------\n")
+            logger.debug("Running block: %s", block_config["block_name"])
+            logger.debug("Input dataset: %s", dataset)
+
+            dataset = block.generate(dataset, **gen_kwargs)
+
+            if len(dataset) == 0:
+                raise EmptyDatasetError(
+                    f"Pipeline stopped: Empty dataset after running block: {block_config['block_name']}"
+                )
+
+            drop_columns_in_ds = [e for e in drop_columns if e in dataset.column_names]
+            if drop_columns:
+                dataset = dataset.remove_columns(drop_columns_in_ds)
+
+            if drop_duplicates_cols:
+                dataset = self._drop_duplicates(dataset, cols=drop_duplicates_cols)
+
+            logger.debug("Output dataset: %s", dataset)
+            logger.debug("------------------------------------\n\n")
+
+        return dataset

sdg_hub/prompts.py

@@ -0,0 +1,43 @@
+# Local
+from .registry import PromptRegistry
+
+
+@PromptRegistry.register("blank")
+def blank_chat_template():
+    return """{{ messages }}"""
+
+
+@PromptRegistry.register("instructlab")
+def instructlab_chat_template():
+    return """{% for message in messages %}{% if message['role'] == 'pretraining' %}{{ '<|pretrain|>' + message['content'] + '<|endoftext|>' + '<|/pretrain|>' }}{% elif message['role'] == 'system' %}{{ '<|system|>' + '\n' + message['content'] + '\n' }}{% elif message['role'] == 'user' %}{{ '<|user|>' + '\n' + message['content'] + '\n' }}{% elif message['role'] == 'assistant' %}{{ '<|assistant|>' + '\n' + message['content'] + '<|endoftext|>' + ('' if loop.last else '\n') }}{% endif %}{% if loop.last and add_generation_prompt %}{{ '<|assistant|>' + '\n' }}{% endif %}{% endfor %}"""
+
+
+@PromptRegistry.register("mistralai")
+def mistral_chat_template():
+    return """{%- if messages[0]['role'] == 'system' %}\n    {%- set system_message = messages[0]['content'] %}\n    {%- set loop_messages = messages[1:] %}\n{%- else %}\n    {%- set loop_messages = messages %}\n{%- endif %}\n\n<s>\n{%- for message in loop_messages %}\n    {%- if (message['role'] == 'user') != (loop.index0 % 2 == 0) %}\n        {{- raise_exception('After the optional system message, conversation roles must alternate user/assistant/user/assistant/...') }}\n    {%- endif %}\n    {%- if message['role'] == 'user' %}\n        {%- if loop.first and system_message is defined %}\n            {{- ' [INST] ' + system_message + '\\n\\n' + message['content'] + ' [/INST]' }}\n        {%- else %}\n            {{- ' [INST] ' + message['content'] + ' [/INST]' }}\n        {%- endif %}\n    {%- elif message['role'] == 'assistant' %}\n        {{- ' ' + message['content'] + '</s>'}}\n    {%- else %}\n        {{- raise_exception('Only user and assistant roles are supported, with the exception of an initial optional system message!') }}\n    {%- endif %}\n{%- endfor %}\n"""
+
+
+@PromptRegistry.register("meta-llama/Llama-3.3")
+def meta_llama_chat_template():
+    return """{{- bos_token }}\n{%- if custom_tools is defined %}\n    {%- set tools = custom_tools %}\n{%- endif %}\n{%- if not tools_in_user_message is defined %}\n    {%- set tools_in_user_message = true %}\n{%- endif %}\n{%- if not date_string is defined %}\n    {%- set date_string = \"26 Jul 2024\" %}\n{%- endif %}\n{%- if not tools is defined %}\n    {%- set tools = none %}\n{%- endif %}\n\n{#- This block extracts the system message, so we can slot it into the right place. #}\n{%- if messages[0]['role'] == 'system' %}\n    {%- set system_message = messages[0]['content']|trim %}\n    {%- set messages = messages[1:] %}\n{%- else %}\n    {%- set system_message = \"\" %}\n{%- endif %}\n\n{#- System message + builtin tools #}\n{{- \"<|start_header_id|>system<|end_header_id|>\\n\\n\" }}\n{%- if builtin_tools is defined or tools is not none %}\n    {{- \"Environment: ipython\\n\" }}\n{%- endif %}\n{%- if builtin_tools is defined %}\n    {{- \"Tools: \" + builtin_tools | reject('equalto', 'code_interpreter') | join(\", \") + \"\\n\\n\"}}\n{%- endif %}\n{{- \"Cutting Knowledge Date: December 2023\\n\" }}\n{{- \"Today Date: \" + date_string + \"\\n\\n\" }}\n{%- if tools is not none and not tools_in_user_message %}\n    {{- \"You have access to the following functions. To call a function, please respond with JSON for a function call.\" }}\n    {{- 'Respond in the format {\"name\": function name, \"parameters\": dictionary of argument name and its value}.' }}\n    {{- \"Do not use variables.\\n\\n\" }}\n    {%- for t in tools %}\n        {{- t | tojson(indent=4) }}\n        {{- \"\\n\\n\" }}\n    {%- endfor %}\n{%- endif %}\n{{- system_message }}\n{{- \"<|eot_id|>\" }}\n\n{#- Custom tools are passed in a user message with some extra guidance #}\n{%- if tools_in_user_message and not tools is none %}\n    {#- Extract the first user message so we can plug it in here #}\n    {%- if messages | length != 0 %}\n        {%- set first_user_message = messages[0]['content']|trim %}\n        {%- set messages = messages[1:] %}\n    {%- else %}\n        {{- raise_exception(\"Cannot put tools in the first user message when there's no first user message!\") }}\n{%- endif %}\n    {{- '<|start_header_id|>user<|end_header_id|>\\n\\n' -}}\n    {{- \"Given the following functions, please respond with a JSON for a function call \" }}\n    {{- \"with its proper arguments that best answers the given prompt.\\n\\n\" }}\n    {{- 'Respond in the format {\"name\": function name, \"parameters\": dictionary of argument name and its value}.' }}\n    {{- \"Do not use variables.\\n\\n\" }}\n    {%- for t in tools %}\n        {{- t | tojson(indent=4) }}\n        {{- \"\\n\\n\" }}\n    {%- endfor %}\n    {{- first_user_message + \"<|eot_id|>\"}}\n{%- endif %}\n\n{%- for message in messages %}\n    {%- if not (message.role == 'ipython' or message.role == 'tool' or 'tool_calls' in message) %}\n        {{- '<|start_header_id|>' + message['role'] + '<|end_header_id|>\\n\\n'+ message['content'] | trim + '<|eot_id|>' }}\n    {%- elif 'tool_calls' in message %}\n        {%- if not message.tool_calls|length == 1 %}\n            {{- raise_exception(\"This model only supports single tool-calls at once!\") }}\n        {%- endif %}\n        {%- set tool_call = message.tool_calls[0].function %}\n        {%- if builtin_tools is defined and tool_call.name in builtin_tools %}\n            {{- '<|start_header_id|>assistant<|end_header_id|>\\n\\n' -}}\n            {{- \"<|python_tag|>\" + tool_call.name + \".call(\" }}\n            {%- for arg_name, arg_val in tool_call.arguments | items %}\n                {{- arg_name + '=\"' + arg_val + '\"' }}\n                {%- if not loop.last %}\n                    {{- \", \" }}\n                {%- endif %}\n                {%- endfor %}\n            {{- \")\" }}\n        {%- else  %}\n            {{- '<|start_header_id|>assistant<|end_header_id|>\\n\\n' -}}\n            {{- '{\"name\": \"' + tool_call.name + '\", ' }}\n            {{- '\"parameters\": ' }}\n            {{- tool_call.arguments | tojson }}\n            {{- \"}\" }}\n        {%- endif %}\n        {%- if builtin_tools is defined %}\n            {#- This means we're in ipython mode #}\n            {{- \"<|eom_id|>\" }}\n        {%- else %}\n            {{- \"<|eot_id|>\" }}\n        {%- endif %}\n    {%- elif message.role == \"tool\" or message.role == \"ipython\" %}\n        {{- \"<|start_header_id|>ipython<|end_header_id|>\\n\\n\" }}\n        {%- if message.content is mapping or message.content is iterable %}\n            {{- message.content | tojson }}\n        {%- else %}\n            {{- message.content }}\n        {%- endif %}\n        {{- \"<|eot_id|>\" }}\n    {%- endif %}\n{%- endfor %}\n{%- if add_generation_prompt %}\n    {{- '<|start_header_id|>assistant<|end_header_id|>\\n\\n' }}\n{%- endif %}\n"""
+
+
+@PromptRegistry.register("microsoft/phi-4")
+def microsoft_phi_chat_template():
+    return """{% for message in messages %}{% if (message['role'] == 'system') %}{{'<|im_start|>system<|im_sep|>' + message['content'] + '<|im_end|>'}}{% elif (message['role'] == 'user') %}{{'<|im_start|>user<|im_sep|>' + message['content'] + '<|im_end|>'}}{% elif (message['role'] == 'assistant') %}{{'<|im_start|>assistant<|im_sep|>' + message['content'] + '<|im_end|>'}}{% endif %}{% endfor %}{% if add_generation_prompt %}{{ '<|im_start|>assistant<|im_sep|>' }}{% endif %}"""
+
+@PromptRegistry.register("nvidia/Llama-3_3-Nemotron-Super-49B-v1")
+def nemotron_chat_template():
+    return """{{- bos_token }}
+{{- "<|start_header_id|>system<|end_header_id|>\n\n" }}detailed thinking on{{- "<|eot_id|>" }}
+{%- for message in messages %}
+  {%- if message['role'] == 'assistant' and '</think>' in message['content'] %}
+    {%- set content = message['content'].split('</think>')[-1].lstrip() %}
+  {%- else %}
+    {%- set content = message['content'] %}
+  {%- endif %}
+  {{- '<|start_header_id|>' + message['role'] + '<|end_header_id|>\n\n' + content | trim + '<|eot_id|>' }}
+{%- endfor %}
+{%- if add_generation_prompt %}
+  {{- '<|start_header_id|>assistant<|end_header_id|>\n\n' }}
+{%- endif %}"""