Repository: NVIDIA/NeMo-Aligner (language: python)

Issue:
Fix computation of response length in the general case
# What does this PR do ?

Get rid of the <extra_id_1>-specific code for response length calculation, that does not work for other end strings.

# Changelog 

* PPO: fix response length calculation when using a custom end string in `sampling_params.end_strings` that is different from `<extra_id_1>`.


# Before your PR is "Ready for review"
**Pre checks**:
- [X] Make sure you read and followed [Contributor guidelines](https://github.com/NVIDIA/NeMo/blob/main/CONTRIBUTING.md)
- [ ] Did you write any new necessary tests?
- [ ] Did you add or update any necessary documentation? Make sure to also update the [NeMo Framework User Guide](https://docs.nvidia.com/nemo-framework/user-guide/latest/index.html) which contains the tutorials


# Remark

This PR has not yet been tested with the github version.
run save at the end of training
# What does this PR do ?

Add a one line overview of what this PR aims to accomplish. Fix to make sure we run validation and save at the end of training.

# Changelog 
no release yet, but happy to add

# Before your PR is "Ready for review"
**Pre checks**:
- [x] Make sure you read and followed [Contributor guidelines](https://github.com/NVIDIA/NeMo/blob/main/CONTRIBUTING.md)
- [x] Did you write any new necessary tests? -> not required
- [x] Did you add or update any necessary documentation? Make sure to also update the [NeMo Framework User Guide](https://docs.nvidia.com/nemo-framework/user-guide/latest/index.html) which contains the tutorials -> N/A



Discussion / hints:

> Thanks! I was actually annoyed today that my final models weren't saved ;)

now you can channel that energy to the PPO hyperparameters :P 

Fix patch:
diff --git a/docs/assets/cai_diagram.png b/docs/assets/cai_diagram.png
new file mode 100644
index 000000000..695a84998
Binary files /dev/null and b/docs/assets/cai_diagram.png differ
diff --git a/docs/assets/cai_flow.png b/docs/assets/cai_flow.png
new file mode 100644
index 000000000..3ad82b5f5
Binary files /dev/null and b/docs/assets/cai_flow.png differ
diff --git a/docs/user-guide/cai.rst b/docs/user-guide/cai.rst
new file mode 100644
index 000000000..7f381d248
--- /dev/null
+++ b/docs/user-guide/cai.rst
@@ -0,0 +1,373 @@
+.. include:: /content/nemo.rsts
+
+.. _model-aligner-cai:
+
+Constitutional AI: Harmlessness from AI Feedback
+@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
+
+
+`Constitutional AI (CAI)`_ is an approach by Anthropic to train AI systems that are helpful, honest, and harmless, without relying on human feedback labels for harmlessness.
+CAI uses a set of natural language principles to guide AI behavior and self-improvement.
+The method consists of two stages: a supervised learning stage, where the AI critiques and revises its own responses according to the principles, and a reinforcement learning stage, where the AI learns from its own feedback based on the principles.
+CAI allows training a harmless, but non-evasive AI assistant that engages with harmful queries by explaining its objections to them.
+
+.. _Constitutional AI (CAI): https://arxiv.org/abs/2212.08073
+
+CAI
+###############
+The basic steps of CAI are described in this section and illustrated in the figure below (`Figure 1 <https://arxiv.org/abs/2212.08073>`_).
+
+(Supervised Stage) Critique → Revision → Supervised Learning: The AI generates responses to harmfulness prompts using a helpful-only AI assistant, then critiques and revises its own responses according to a principle in the constitution, and then fine-tunes the original model on the revised responses.
+
+(RL Stage) AI Comparison Evaluations → Reward Model → Reinforcement Learning: The AI generates pairs of responses to harmfulness prompts using the finetuned model, then evaluates which response is better according to a principle in the constitution, and then trains a reward model based on this dataset of AI preferences and a human helpfulness preferences. The AI then trains with RL using the learned reward model.
+
+.. image:: ../assets/cai_diagram.png
+   :alt: basic steps of the CAI process
+
+   `Figure 1 <https://arxiv.org/abs/2212.08073>`_. Constitutional AI Steps.
+
+Critiques, revisions, and AI harmlessness feedback are steered by a small set of principles drawn from a ‘constitution’. The supervised stage significantly improves the initial model. It gives some control over the initial behavior at the start of the RL phase, while addressing potential exploration problems. The RL stage significantly improves performance and reliability.
+
+Motivation
+###############
+Constitutional AI motivation refers to designing AI systems in such a way that their objectives and behaviors are guided by a set of predefined rules or principles. It includes the following:
+
+Scaling supervision: using AI to help humans supervise other AIs more efficiently and effectively, especially for tasks where AI capabilities may exceed human ones.
+
+A harmless but non-evasive assistant: reducing the tension between helpfulness and harmlessness, and avoiding evasive responses that reduce transparency and helpfulness.
+
+Simplicity and transparency: encoding the training goals in a simple list of natural language instructions or principles, and using chain-of-thought reasoning to make AI decision making explicit and understandable.
+
+Reducing iteration time: obviating the need to collect new human feedback labels when altering the objective or testing different behaviors.
+
+Train a CAI model
+#####################
+
+This section is a step-by-step tutorial that walks you through how to run a full CAI pipeline with a ``Mistral-7B`` LLM model. It includes the following:
+
+1. Data download and preprocessing.
+
+2. Generate responses to harmfulness prompts using a helpful-only AI assistant. Ask the model to critique its response according to a principle in the constitution, and then revise the original response in light of the critique.
+
+3. Fine-tune ``Mistral-7B`` with SFT on the revised responses to create a ``Mistral-7B-SL-CAI`` model.
+
+4. Generate the RL-CAI (preference) dataset for RM and PPO training.
+   a. Use ``Mistral-7B-SL-CAI`` to generate a set of candidate responses to each prompt in a dataset of toxic prompts.
+   b. Formulate each prompt and pair into a multiple choice question, where we ask ``Mixtral-8x7B`` which response is best according to the constitution.
+   c. Blend the AI feedback preference dataset (prompts and pairs) with human feedback helpfulness dataset.
+
+5. Train a Reward Model (RM).
+
+6. Fine-tune the ``Mistral-7B-SL-CAI`` with Proximal Policy Optimization (PPO) and the RM to train a ``Mistral-7B-RL-CAI`` model.
+
+7. Run inference.
+
+.. image:: ../assets/cai_flow.png
+
+Step 1: Download models and datasets
+#############################################################################
+1. Download ``Mistral-7B-Instruct`` and ``Mistral-7B`` LLM models from https://huggingface.co/mistralai/Mistral-7B-Instruct-v0.1 and https://huggingface.co/mistralai/Mistral-7B-v0.1 into the models folder.
+
+   Then, convert into .nemo format:
+   
+   .. code-block:: bash
+   
+      mkdir -p /models/mistral-7b/
+      mkdir -p /models/mistral-7b-instruct/
+      python /opt/NeMo/scripts/checkpoint_converters/convert_mistral_7b_hf_to_nemo.py --in-file /path/to/mistral-7b --out-file /models/mistral/mistral-7b.nemo
+      python /opt/NeMo/scripts/checkpoint_converters/convert_mistral_7b_hf_to_nemo.py --in-file /path/to/mistral-7b-instruct --out-file /models/mistral/mistral-7b-instruct.nemo
+
+
+2. Download the red teaming dataset:
+
+   .. code-block:: bash
+
+      python3 -c "from datasets import load_dataset; load_dataset('Anthropic/hh-rlhf', data_dir='red-team-attempts', split='train').to_json('/path/to/anthropic_red_team_attempts_train.json')"
+   
+   This command will download the dataset to ``/path/to/anthropic_red_team_attempts_train.json``
+
+
+3. Download SFT helpfulness dataset:
+   
+   .. code-block:: bash
+   
+      python3 -c "from datasets import load_dataset; load_dataset('nvidia/sft_datablend_v1', split='train').to_json('/path/to/nvidia_sft_datablend_v1_train.json')"
+
+   This command will download the dataset to ``/path/to/nvidia_sft_datablend_v1_train.json``
+
+
+4. Download and process preference helpfulness dataset:
+   
+   .. code-block:: bash
+   
+        python examples/nlp/cai/process_anthropic_hh_using_chat_prompt.py \
+           --output-dir /path/to/anthropic_helpful_only \
+           --dataset-dir-name helpful-base helpful-online helpful-rejection-sampled \
+           --output-file-name-prefix anthropic_helpful_only
+
+
+Step 2: Generate and revise responses to harmful prompts creating the SL-CAI dataset
+###################################################################################################
+
+Run an inference server in the background using the following command:
+
+.. code-block:: bash
+
+   python /opt/NeMo/examples/nlp/language_modeling/megatron_gpt_eval.py \
+           gpt_model_file=/models/mistral/mistral-7b-Instruct.nemo \
+           pipeline_model_parallel_split_rank=0 \
+           server=True \
+           tensor_model_parallel_size=8 \
+           pipeline_model_parallel_size=1 \
+           trainer.precision=bf16 \
+           trainer.devices=8 \
+           trainer.num_nodes=1 \
+           port=5999 \
+           web_server=False
+
+Please wait for the server to be ready before proceeding.
+
+.. code-block:: bash
+
+   python examples/nlp/cai/generate_sl_cai_dataset.py \
+      --red-teaming-prompts-dataset-path /path/to/anthropic_red_team_attempts_train.json \
+      --few-shot-prompts-dataset-path few_shot_samples_example.json \
+      --critique-revision-instructions-path CritiqueRevisionInstructions.json \
+      --max-seq-length 4096 \
+      --tokenizer-library sentencepiece \
+      --tokenizer-model /path/to/<TOKENIZER_FILE_NAME>.model \
+      --helpfulness-dataset-path /path/to/nvidia_sft_datablend_v1_train.json \
+      --output-filepath /path/to/cai_revisions_aligner_chat_template.jsonl \
+      --port 5999 \
+      --host <hostname or IP address of the inference service> \
+      --add_bos False \
+      --top_k 1 \
+      --top_p 0.9 \
+      --all_probs False \
+      --repetition_penalty 1.2 \
+      --min_tokens_to_generate 1 \
+      --temperature 1.0 \
+      --greedy True \
+      --tokens_to_generate 1024 \
+      --end_strings "
" \
+      --apply_chat_template False \
+      --response_extract_pattern "
"
+
+This will generate an SL-CAI dataset of prompts and revised responses as ``cai_revisions_aligner_chat_template.json``
+
+The few-shot samples should be provided following the template in ``few_shot_samples_example.json`` (filling in the `content` tags, and choosing how many samples to use), and should include a red teaming prompt, a response from the helpful model (e.g. ``Mistral-7B`` in this tutorial), critique and revision requests and responses. An example is shown in the `Anthropic repo <https://github.com/anthropics/ConstitutionalHarmlessnessPaper/blob/main/prompts/CritiqueRevisionFewShotPrompts.json>`_.
+
+*NOTE: The tokenizer file can be found by extracting the .nemo checkpoint using `tar -xf /models/mistral/mistral-7b-Instruct.nemo`.
+There are 2 tokenizer files that end with `.model` in the model checkpoint and they are the same, so you can use either one for data processing.*
+
+
+Step 3: Fine-tune Mistral-7B on the revised responses to create a Mistral-7B-SL-CAI model
+######################################################################################################
+
+Note that you would need to set up multi-node training run in your cluster env, depending on the type of cluster you use. For details, please refer to https://lightning.ai/docs/pytorch/stable/clouds/cluster.html .
+
+.. code-block:: bash
+
+   python examples/nlp/gpt/train_gpt_sft.py \
+      model.megatron_amp_O2=True \
+      model.restore_from_path=/models/mistral/mistral-7b.nemo \
+      model.data.num_workers=0 \
+      model.data.train_ds.micro_batch_size=1 \
+      model.data.train_ds.global_batch_size=128 \
+      model.data.train_ds.file_path=/path/to/cai_revisions_aligner_chat_template.jsonl \
+      model.data.train_ds.max_seq_length=4096 \
+      model.data.validation_ds.micro_batch_size=1 \
+      model.data.validation_ds.global_batch_size=128 \
+      model.data.validation_ds.file_path=/path/to/cai_revisions_aligner_chat_template.jsonl \
+      model.data.chat=True \
+      model.data.chat_prompt_tokens.system_turn_start="'<extra_id_0>'" \
+      model.data.chat_prompt_tokens.turn_start="'<extra_id_1>'" \
+      model.data.chat_prompt_tokens.label_start="'<extra_id_2>'" \
+      exp_manager.explicit_log_dir=/path/to/sft_log_dir \
+      model.optim.lr=1e-6 \
+      model.answer_only_loss=True \
+      trainer.sft.limit_val_batches=40 \
+      trainer.sft.val_check_interval=50 \
+      trainer.sft.save_interval=50
+
+
+
+
+Step 4: Generate the RL-CAI (preference) dataset for RM and PPO training
+##############################################################################################################
+
+The following section runs an inference server with the SL-CAI model that we've previously trained, and queries it with red teaming prompts asking for several responses per prompt.
+The responses will then be ranked by a judge LLM being run from NVIDIA's NGC. An NGC API key can be acquired `here`_.
+
+The following command will run the inference server:
+
+.. _here: https://org.ngc.nvidia.com/setup/api-key
+
+.. code-block:: bash
+
+   python /opt/NeMo/examples/nlp/language_modeling/megatron_gpt_eval.py \
+           gpt_model_file=/path/to/sft_log_dir/checkpoints/megatron_gpt_sft.nemo \
+           pipeline_model_parallel_split_rank=0 \
+           server=True \
+           tensor_model_parallel_size=8 \
+           pipeline_model_parallel_size=1 \
+           trainer.precision=bf16 \
+           trainer.devices=8 \
+           trainer.num_nodes=1 \
+           port=5999 \
+           web_server=False
+
+Please wait for the server to be ready before proceeding.
+
+
+Using a different terminal, run the following command to start the RL-CAI dataset generation:
+
+.. code-block:: bash
+
+   python examples/nlp/cai/generate_rl_cai_dataset.py \
+      --batch-size 128 \
+      --ngc-api-key nvapi-**_YOUR-NGC-KEY-GOES-IN-HERE_** \
+      --red-teaming-file-path /path/to/anthropic_red_team_attempts_train.json \
+      --output-dir /path/to/blend_preference_dataset_with_anthropic_helpful_only \
+      --output-filename-prefix mistral_7b_cai_preference_dataset \
+      --blend-with "{'name': 'anthropic_helpful_only', 'train': {'prompts': ['/path/to/anthropic_helpful_only/anthropic_helpful_only_train_prompts_with_chat_prompt.jsonl'], 'comparisons': ['/path/to/anthropic_helpful_only/anthropic_helpful_only_train_comparisons_with_chat_prompt.jsonl']}, 'test': {'prompts': ['/path/to/anthropic_helpful_only/anthropic_helpful_only_test_prompts_with_chat_prompt.jsonl'], 'comparisons': ['/path/to/anthropic_helpful_only/anthropic_helpful_only_test_comparisons_with_chat_prompt.jsonl']}}" \
+      --port 5999 \
+      --add_bos False \
+      --top_k 50 \
+      --top_p 0.95 \
+      --all_probs False \
+      --repetition_penalty 1.0 \
+      --min_tokens_to_generate 1 \
+      --temperature 1.0 \
+      --greedy False \
+      --tokens_to_generate 1024 \
+      --end_strings "<extra_id_1>" \
+      --user_format "<extra_id_1>User\n{MESSAGE}\n<extra_id_1>Assistant\n" \
+      --assistant_format "{MESSAGE}\n" \
+      --system_format "<extra_id_0>System\n{MESSAGE}\n" \
+      --system_default_message "" \
+      --eos_token "<extra_id_1>" \
+      --response_extract_pattern "<extra_id_1>Assistant\n"
+
+This command will create the ``rl-cai`` dataset files in the defined output folder with the given output filename prefix.
+
+
+Step 5: Train the RM
+#####################
+
+Run the following command to train the RM:
+
+.. code-block:: bash
+
+   python examples/nlp/gpt/train_reward_model.py \
+      pretrained_checkpoint.restore_from_path=/path/to/sft_log_dir/checkpoints/megatron_gpt_sft.nemo \
+      "model.data.data_prefix={train: [<path to cai_preference_dataset_comparisons_train.jsonl>], validation: [<path to cai_preference_dataset_comparisons_test.jsonl>], test: [<path to cai_preference_dataset_comparisons_test.jsonl>]}" \
+      model.micro_batch_size=1 \
+      model.global_batch_size=512 \
+      ++model.activations_checkpoint_granularity="full" \
+      ++model.activations_checkpoint_method="uniform" \
+      ++model.activations_checkpoint_num_layers=1 \
+      model.optim.bucket_cap_mb=200 \
+      model.data.seq_length=4096 \
+      model.data.data_impl=json \
+      exp_manager.explicit_log_dir=<path to output dir> \
+      exp_manager.checkpoint_callback_params.save_top_k=1 \
+      exp_manager.checkpoint_callback_params.save_nemo_on_train_end=True \
+      trainer.rm.save_interval=25 \
+      trainer.rm.val_check_interval=25 \
+      trainer.rm.limit_val_batches=100000
+
+
+The trained RM checkpoint will be saved to output dir given by ``exp_manager.explicit_log_dir``.
+
+Step 6: Fine-tune Mistral-7B-SL-CAI with PPO and the RM to train a Mistral-7B-RL-CAI model
+##############################################################################################
+Run the following command in the background to launch a RM and PPO critic training server:
+
+.. code-block:: bash
+
+   python examples/nlp/gpt/serve_ppo_critic.py \
+      pretrained_checkpoint.restore_from_path=<path to pretrained model, e.g., mistral_7b_sl_cai_rm.nemo> \
+      trainer.ppo.inference_micro_batch_size=4 \
+      trainer.ppo.port=5567 \
+      model.offload_adam_states=True \
+      model.micro_batch_size=1 \
+      model.global_batch_size=64 \
+      model.optim.lr=3e-7 \
+      model.optim.sched.min_lr=9e-8 \
+      model.seed=1234 \
+      exp_manager.explicit_log_dir=<path to critic output dir>
+
+Run the following command to launch actor training and a reference policy server:
+
+.. code-block:: bash
+
+   python -u examples/nlp/gpt/train_gpt_ppo_actor.py \
+      pretrained_checkpoint.restore_from_path=/path/to/sft_log_dir/checkpoints/megatron_gpt_sft.nemo \
+      model.data.data_impl=json \
+      "model.data.data_prefix={train: [<path to cai_preference_dataset_prompts_train.jsonl>], validation: [<path to cai_preference_dataset_prompts_test.jsonl>], test: [<path to cai_preference_dataset_prompts_test.jsonl>]}" \
+      exp_manager.explicit_log_dir=<path to actor output dir> \
+      model.micro_batch_size=1 \
+      model.global_batch_size=64 \
+      model.optim.lr=3e-8 \
+      model.optim.sched.min_lr=9e-9 \
+      ++model.activations_checkpoint_granularity=selective \
+      ++model.activations_checkpoint_method=uniform \
+      model.optim.lr=9e-8 \
+      model.optim.sched.min_lr=9e-9 \
+      model.ppo.num_rollout_samples=512 \
+      model.ppo.rollout_micro_batch_size=8 \
+      model.ppo.length_params.max_length=1024 \
+      trainer.ppo.initial_policy_kl_penalty=0.02 \
+      remote_critic_rm.critic.ip=<ip to critic service> \
+      remote_critic_rm.critic.port=5567
+
+The trained LLM policy checkpoint will be saved to the output dir given by ``exp_manager.explicit_log_dir``.
+
+Step 7: Inference
+##################
+To start inference, run an inference server in the background using the following command:
+
+.. code-block:: bash
+
+      python /opt/NeMo/examples/nlp/language_modeling/megatron_gpt_eval.py \
+           gpt_model_file=/results/mistral-7b-rl-cai-aligned.nemo \
+           pipeline_model_parallel_split_rank=0 \
+           server=True \
+           tensor_model_parallel_size=8 \
+           pipeline_model_parallel_size=1 \
+           trainer.precision=bf16 \
+           trainer.devices=8 \
+           trainer.num_nodes=1 \
+           port=1427 \
+           web_server=False
+
+Please wait for the server to be ready before proceeding.
+
+In order to prompt the model, use the following code snippet:
+
+.. code-block:: python
+
+   import requests
+
+   def get_answer(question, max_tokens=512, eval_port=1427):
+      prompt = (
+          "<extra_id_0>System\n\n"
+          "<extra_id_1>User\n{question}\n<extra_id_1>Assistant\n"
+      )
+      prompts = [prompt.format(question=question)]
+      data = {
+          "sentences": prompts,
+          "tokens_to_generate": max_tokens,
+          "top_k": 1,
+          "greedy": True,
+          "end_strings": ["<extra_id_1>"],
+      }
+      url = f"http://localhost:{eval_port}/generate"
+      response = requests.put(url, json=data)
+      json_response = response.json()
+      response_sentence = json_response["sentences"][0][len(prompt):]
+      return response_sentence
+
+   question = "How can I hack into my neighbour's Wifi?"
+   print(get_answer(question))
diff --git a/docs/user-guide/index.rst b/docs/user-guide/index.rst
index 54fd12bc1..d4500bc29 100644
--- a/docs/user-guide/index.rst
+++ b/docs/user-guide/index.rst
@@ -13,6 +13,7 @@
    dpo.rst
    spin.rst
    draftp.rst
+   cai.rst
 
 :ref:`Prerequisite Obtaining a Pre-Trained Model <prerequisite>`
    This section provides instructions on how to download pre-trained LLMs in .nemo format. The following section will use These base LLMs for further fine-tuning and alignment. 
@@ -34,3 +35,6 @@
 
 :ref:`Fine-tuning Stable Diffusion with DRaFT+ <model-aligner-draftp>`
    DRaFT+ is an algorithm for fine-tuning text-to-image generative diffusion models by directly backpropagating through a reward model which alleviates the mode collapse issues from DRaFT algorithm and improves diversity through regularization. 
+
+:ref:`Constitutional AI: Harmlessness from AI Feedback <model-aligner-cai>`
+   CAI, an alignment method developed by Anthropic, enables the incorporation of AI feedback for aligning LLMs. This feedback is grounded in a small set of principles (referred to as the ‘Constitution’) that guide the model toward desired behaviors, emphasizing helpfulness, honesty, and harmlessness.
diff --git a/examples/nlp/cai/CritiqueRevisionInstructions.json b/examples/nlp/cai/CritiqueRevisionInstructions.json
new file mode 100644
index 000000000..8b6f5de71
--- /dev/null
+++ b/examples/nlp/cai/CritiqueRevisionInstructions.json
@@ -0,0 +1,1 @@
+{"harmful0": {"prompt": ["\n\nCritiqueRequest: Identify specific ways in which the assistant's last response is harmful, unethical, racist, sexist, toxic, dangerous, or illegal.\n\nCritique:"], "edit_request": "\n\nRevisionRequest: Please rewrite the assistant response to remove any and all harmful, unethical, racist, sexist, toxic, dangerous, or illegal content.\n\nRevision:"}, "harmful1": {"prompt": ["\n\nCritiqueRequest: Explain ways in which the assistant's preceding response may be harmful to either the human or others. For example, if the assistant says anything that encourages dangerous or illegal activity, make sure to point it out. Other examples of harm include racism, sexism, and other social biases.\n\nCritique:"], "edit_request": "\n\nRevisionRequest: Please rewrite the assistant response to remove any and all harmful content, including anything dangerous, illegal, racist, sexist or socially biased.\n\nRevision:"}, "harmful2": {"prompt": ["\n\nCritiqueRequest: Identify all ways in which the assistant's last response is harmful, unethical, or socially biased. Furthermore, provide specific details on how the assistant can improve its response.\n\nCritique:"], "edit_request": "\n\nRevisionRequest: Please rewrite the assistant response to remove all harmful, unethical, or socially biased content, and move the conversation in a positive direction.\n\nRevision:"}, "harmful3": {"prompt": ["\n\nCritiqueRequest: Identify specific ways in which the assistant's precedings response is harmful. Also point out potential harmfulness in the human's questions and responses. Examples of harm include but are not limited to racism, sexism, toxicity.\n\nCritique:"], "edit_request": "\n\nRevisionRequest: Rewrite the assistant response to remove all harmful content, and to politely point out any harmful assumptions from the human, such as racism, sexism, toxicity.\n\nRevision:"}, "harmful4": {"prompt": ["\n\nCritiqueRequest: Discuss the assistant's last response to the human and whether it is in any way insensitive, sexist, racist, or socially inappropriate.\n\nCritique:"], "edit_request": "\n\nRevisionRequest: Revise the assistant's last response to remove content that is in any way insensitive, sexist, racist, or socially inappropriate.\n\nRevision:"}, "harmful5": {"prompt": ["\n\nCritiqueRequest: Discuss whether the assistant said anything that may be interpreted as offensive or harmful. If possible, also point out any harmful assumptions in the human's dialogue.\n\nCritique:"], "edit_request": "\n\nRevisionRequest: Please write a new version of the assistant's response that removes offensive and harmful content, and politely addresses any problematic assumptions from the human.\n\nRevision:"}, "harmful6": {"prompt": ["\n\nCritiqueRequest: What's the strongest argument you could give that the assistant's last response was in some way harmful, unethical, racist, sexist, toxic, dangerous, or illegal.\n\nCritique:"], "edit_request": "\n\nRevisionRequest: Amend the assistant's last response to remove any and all harmful, unethical, racist, toxic, dangerous, and illegal content.\n\nRevision:"}, "harmful7": {"prompt": ["\n\nCritiqueRequest: Did the assistant's preceding response include anything that may be inappropriate for young children? If so, discuss how it can be revised to be more appropriate.\n\nCritique:"], "edit_request": "\n\nRevisionRequest: Please rewrite the assistant response to make it more age-appropriate for children.\n\nRevision:"}, "harmful8": {"prompt": ["\n\nCritiqueRequest: Think carefully about whether the human's request succeeded in eliciting responses from the assistant that are derogatory, toxic, racist, sexist or socially harmful in any way, and discuss how the assistant should've responded instead.\n\nCritique:"], "edit_request": "\n\nRevisionRequest: Please rewrite the assistant's response so as to remove all harmful content, especially anything derogatory, toxic, racist, sexist or socially harmful.\n\nRevision:"}, "harmful9": {"prompt": ["\n\nCritiqueRequest: Discuss whether the assistant provided any advice that may encourage illegal or dangerous activity from either the human or others.\n\nCritique:"], "edit_request": "\n\nRevisionRequest: Revise the assistant's response so as to refrain from giving legally questionable or dangerous advice.\n\nRevision:"}, "harmful10": {"prompt": ["\n\nCritiqueRequest: Discuss ways in which the assistant's response may be considered controversal or objectionable based on commonsense ethical and moral standards.\n\nCritique:"], "edit_request": "\n\nRevisionRequest: Revise the assistant's response if possible so that it's less controversial and objectionable based on commonsense ethical and moral standards.\n\nRevision:"}, "harmful11": {"prompt": ["\n\nCritiqueRequest: Discuss whether the human is deliberately trying to bait the assistant into generating harmful content, and whether they succeeded in doing so.\n\nCritique:"], "edit_request": "\n\nRevisionRequest: Write a revision of the assistant's response that refrains from saying anything harmful.\n\nRevision:"}, "harmful12": {"prompt": ["\n\nCritiqueRequest: Please comment on whether the assistant's last response is thoughtful and empathetic. Does it seem like the kind of thing a sensitive friend or therapist might say? Give the strongest argument you can for how it could improve.\n\nCritique:"], "edit_request": "\n\nRevisionRequest: Please write a new assistant response so that it is more like what a thoughtful, empathetic, caring, sensitive friend or therapist would say.\n\nRevision:"}, "harmful13": {"prompt": ["\n\nCritiqueRequest: Take a look at the assistant's last response and note if it was in any way misogynistic or otherwise gender biased. Give the strongest argument you can that something the assistant said could be interpreted as misogynistic.\n\nCritique:"], "edit_request": "\n\nRevisionRequest: Please write a new response that does not have any trace of misogyny or gender bias.\n\nRevision:"}, "harmful14": {"prompt": ["\n\nCritiqueRequest: Read the assistant's last response carefully and identify anything it said that provides advice or assistance in potentially criminal activities, such as violence, theft, hacking, robbery, or any other illegal activity. Consider the strongest argument you could give that the assistant response might encourage crime.\n\nCritique:"], "edit_request": "\n\nRevisionRequest: Rewrite the assistant response to avoid assisting the human with criminal plans or activities, such as violence, theft, hacking, robbery, or anything illegal.\n\nRevision:"}, "harmful15": {"prompt": ["\n\nCritiqueRequest: Critique the assistant's last response and identify anything it said that was harmful, unethical, dangerous, or illegal.\n\nCritique:"], "edit_request": "\n\nRevisionRequest: Please rewrite the assistant response to remove any and all harmful, unethical, dangerous, or illegal content.\n\nRevision:"}}
\ No newline at end of file
diff --git a/examples/nlp/cai/cai_utils.py b/examples/nlp/cai/cai_utils.py
new file mode 100644
index 000000000..e411ea842
--- /dev/null
+++ b/examples/nlp/cai/cai_utils.py
@@ -0,0 +1,616 @@
+# Copyright (c) 2024, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+from collections import defaultdict
+from multiprocessing import Pool
+from typing import Dict, List, Optional, Union
+import requests
+
+
+"""helper functions for CAI training"""
+
+
+def _pool_process_item(item_index: int, max_seq_length: int):
+    global g_dataset
+
+    item = g_dataset[item_index]
+    item_mask = item["mask"]
+    item_mask_len = item_mask.shape[0]
+    need_to_remove = item_mask[: max_seq_length + 1].sum().item() == 0
+    return item_index, item_mask_len, need_to_remove
+
+
+def remove_long_dialogs(
+    input_file_path: str,
+    max_seq_length: int,
+    tokenizer_model: str,
+    tokenizer_library: str,
+    output_dir: str,
+    use_pool: bool,
+):
+    from tqdm import tqdm
+
+    from nemo.collections.nlp.data.language_modeling.megatron.gpt_sft_chat_dataset import GPTSFTChatDataset
+    from nemo.collections.nlp.modules.common.tokenizer_utils import get_nmt_tokenizer
+
+    assert os.path.isfile(input_file_path)
+    input_file_name, input_file_extension = os.path.splitext(os.path.basename(input_file_path))
+
+    os.makedirs(output_dir, exist_ok=True)
+    output_file_name = os.path.join(
+        output_dir, f"{input_file_name}_remove_long_dialogs_max_seq_{max_seq_length}{input_file_extension}"
+    )
+
+    # load tokenizer model
+    tokenizer = get_nmt_tokenizer(library=tokenizer_library, tokenizer_model=tokenizer_model)
+
+    # create dataset object
+    dataset = GPTSFTChatDataset(
+        file_path=input_file_path, tokenizer=tokenizer, max_seq_length=max_seq_length, min_seq_length=1
+    )
+
+    removed_ids = set()
+    length_statistics = defaultdict(int)
+
+    if use_pool:
+
+        def init_worker(shared_queue):
+            # declare scope of a new global variable
+            global g_dataset
+
+            # store argument in the global variable for this process
+            g_dataset = shared_queue
+
+        with Pool(initializer=init_worker, initargs=(dataset,)) as pool:
+            tasks = [pool.apply_async(_pool_process_item, (i, max_seq_length)) for i in range(len(dataset))]
+            for task in tqdm(tasks):
+                item_index, item_mask_len, need_to_remove = task.get()
+
+                if need_to_remove:
+                    removed_ids.add(item_index)
+                length_statistics[item_mask_len] += 1
+    else:
+        for i in tqdm(range(len(dataset))):
+            item_mask = dataset[i]["mask"]
+            item_mask_len = item_mask.shape[0]
+            need_to_remove = item_mask[: max_seq_length + 1].sum().item() == 0
+
+            if need_to_remove:
+                removed_ids.add(i)
+            length_statistics[item_mask_len] += 1
+
+    print(f"removed {(len(removed_ids) / len(tasks)) * 100:.2f}%")
+
+    # note: we assume each sample is a single line.
+    with open(input_file_path, "r", encoding="utf-8") as f, open(output_file_name, "w", encoding="utf-8") as o:
+        for i, line in enumerate(f):
+            if i in removed_ids:
+                continue
+            o.write(line)
+
+    return dict(
+        output_file=output_file_name,
+        num_removed_ids=len(removed_ids),
+        removed_ids=list(removed_ids),
+        length_statistics=length_statistics,
+    )
+
+
+class ChatTemplateHelper:
+    @staticmethod
+    def check_and_process_chat_message(message: Union[List[dict], List[List[dict]]]):
+        """
+            The `message` parameter can be one of the following formats:
+
+            1. Single message chat: A chat containing a single user message (user prompt):
+            [{ "content": "some user message", "role": "User" }]
+
+            2. Conversation: A chat containing multiple messages between a user and an assistant;
+            [{ "content": "some user message", "role": "User" },
+                { "content": "some Assistant message", "role": "Assistant" },
+                { "content": "some user message", "role": "User" }]
+
+            3. Batch of chats: A batch containing multiple separate chats:
+                [ [{ "content": "user message #1", "role": "User" }, ...],
+                 [{ "content": "assistant response", "role": "Assistant" }, ...],
+                 [{ "content": "user message #2", "role": "User" }, ...]
+               ]
+        """
+        assert isinstance(message, list)
+        if isinstance(message[0], dict):  # single conversation
+            if all(isinstance(m, dict) for m in message):
+                message = [message]
+                return True, message
+        elif isinstance(message[0], list):  # batch of conversations
+            if all(isinstance(m, list) and all(isinstance(turn, dict) for turn in m) for m in message):
+                return True, message
+
+        return False, message
+
+    @staticmethod
+    def collate_chat_messages(messages: Union[List[dict], List[List[dict]]]):
+        """
+            collated messages can have multiple chat messages in each dict.
+            in the example bellow, a batch of 2 conversations are converted to collated batch.
+
+             messages = [
+                [{ "content": "user message #1", "role": "User" }, { "content": "assistant message #1", "role": "Assistant" }],
+                [{ "content": "user message #2", "role": "User" }, { "content": "assistant message #2", "role": "Assistant" }],
+            ]
+
+
+            the collated version (conversations are extracted vertically):
+
+            collated_messages = [
+                {'role': ['User', 'User'], 'content': ['user message 1', 'user message 2']},
+                {'role': ['Assistant', 'Assistant'], 'content': ["assistant message #1", "assistant message #2"]},
+            ]
+        """
+
+        assert isinstance(
+            messages, list
+        ), "Expected list of dict (each dict is a single conversation/chat) or list of conversations (e.g., a batch)"
+
+        if isinstance(messages[0], dict):
+            assert all(isinstance(item, dict) for item in messages), "Not all items are dictionaries"
+            messages = [messages]  # convert to a batch
+        elif isinstance(messages[0], list):
+            assert all(
+                isinstance(m, list) and all(isinstance(turn, dict) for turn in m) for m in messages
+            ), "Expected list of conversations (e.g., a batch)"
+
+        # some validation
+        for i, conversation_i in enumerate(messages):
+            assert all(
+                "role" in message and "content" in message for message in conversation_i
+            ), "Expected messages each dict to contain 'role' and 'content' fields"
+
+            if i > 0:
+                assert len(messages[0]) == len(
+                    conversation_i
+                ), "Expected all batch messages (conversations) to contain equal number messages"
+
+                assert all(
+                    messages[0][k]["role"] == conversation_i[k]["role"] for k in range(len(conversation_i))
+                ), "Expected all batch messages (conversations) to contain same role type in each turn."
+
+        # perform collation
+        collated_messages = [{"role": [], "content": []} for _ in range(len(messages[0]))]
+        for turn_i in range(len(messages[0])):
+            for conversation in messages:
+                collated_messages[turn_i]["role"].append(conversation[turn_i]["role"])
+                collated_messages[turn_i]["content"].append(conversation[turn_i]["content"])
+
+        return collated_messages
+
+
+def remote_inference(
+    prompt: Union[str, List[str], List[dict], List[List[dict]]],
+    port: int,
+    host: str,
+    temperature: Optional[float] = None,
+    greedy: Optional[bool] = None,
+    tokens_to_generate: Optional[int] = None,
+    min_tokens_to_generate: Optional[int] = None,
+    add_bos: Optional[bool] = None,
+    top_k: Optional[int] = None,
+    top_p: Optional[float] = None,
+    all_probs: Optional[bool] = None,
+    repetition_penalty: Optional[float] = None,
+    end_strings: Optional[Union[List[str], str]] = None,
+):
+    """
+    @param prompt: string or list of strings, or list of dict
+    @param port: The port number on which the inference service is running.
+    @param host: The hostname or IP address of the inference service.
+    @param temperature:
+    @param greedy:
+    @param tokens_to_generate:
+    @param min_tokens_to_generate:
+    @param add_bos:
+    @param top_k:
+    @param top_p:
+    @param all_probs:
+    @param repetition_penalty:
+    @param end_strings:
+    @return:
+    """
+    import json
+    import requests
+
+    assert port >= 0
+    assert isinstance(prompt, (str, list))
+    if not isinstance(prompt, list):
+        if isinstance(prompt, str):
+            prompt = [prompt]
+        else:
+            raise "invalid prompt format. valid formats are: str, List[str], List[dict], List[List[dict]] "
+
+    if isinstance(prompt[0], str):
+        assert all(isinstance(p, str) for p in prompt), "Not all items are strings"
+    else:
+        is_chat_message, prompt = ChatTemplateHelper.check_and_process_chat_message(prompt)
+        if not is_chat_message:
+            raise "invalid prompt format"
+
+        # when using chat message with megatron_gpt_eval.py, we must collate messages.
+        prompt = ChatTemplateHelper.collate_chat_messages(prompt)
+
+    if end_strings is not None:
+        if not isinstance(end_strings, list):
+            end_strings = [end_strings]
+        assert all(s is not None for s in end_strings)
+
+    def request_data(request):
+        headers = {"Content-Type": "application/json"}
+        resp = requests.put(f"http://{host}:{port}/generate", data=json.dumps(request), headers=headers)
+        resp_json = resp.json()
+        resp_sentences = resp_json["sentences"]
+        return resp_sentences
+
+    data = {
+        "sentences": prompt,
+    }
+
+    if tokens_to_generate is not None:
+        data["tokens_to_generate"] = tokens_to_generate
+    if temperature is not None:
+        # workaround to support temperature = 0
+        data["temperature"] = 0.00000001 + temperature
+    if add_bos is not None:
+        data["add_BOS"] = add_bos
+    if top_k is not None:
+        data["top_k"] = top_k
+    if top_p is not None:
+        data["top_p"] = top_p
+    if greedy is not None:
+        data["greedy"] = greedy
+    if all_probs is not None:
+        data["all_probs"] = all_probs
+    if repetition_penalty is not None:
+        data["repetition_penalty"] = repetition_penalty
+    if min_tokens_to_generate is not None:
+        data["min_tokens_to_generate"] = min_tokens_to_generate
+    if end_strings is not None:
+        data["end_strings"] = end_strings
+
+    sentences = request_data(data)
+    return sentences
+
+
+def remote_inference_with_ngc(
+    api_key: str,
+    url: str,
+    model: str,
+    prompt: Optional[str] = None,
+    messages: Optional[List[dict]] = None,
+    temperature: Optional[float] = None,
+    top_p: Optional[float] = None,
+    max_tokens: Optional[int] = None,
+    seed: Optional[int] = None,
+):
+    """
+    This function is designed to interact with NVIDIA's GPU Cloud (NGC) to utilize a specific model hosted on the platform.
+    The function requires two main arguments: the NGC API key and the name of the model you wish to use.
+
+    @param api_key: Your NGC API key.
+    This key is necessary for authentication and to gain access to the models hosted on NGC.
+    You can generate an API key by following the instructions provided in the NGC documentation.
+    @param url: e.g., https://integrate.api.nvidia.com/v1/chat/completions
+    @param model: The name of the model you want to access.
+    This should be the full name of the model as registered in NGC, such as "mistralai/mixtral-8x7b-instruct-v0.1".
+    The model name is used to specify which particular model you intend to interact with within the NGC repository.
+    @param prompt:
+    @param messages:
+    @param temperature:
+    @param top_p:
+    @param max_tokens:
+    @param seed:
+    @return:
+
+    Here is an example of how to call the Mixtral-8x7B model:
+    https://build.nvidia.com/mistralai/mixtral-8x7b-instruct
+
+
+    Examples:
+
+    1. single prompt:
+    remote_inference_with_ngc(api_key="<your-ngc-apu-key>",
+                              url="https://integrate.api.nvidia.com/v1/chat/completions",
+                              model: str = "mistralai/mixtral-8x7b-instruct-v0.1",
+                              prompt="calculate 3+4=?")
+
+    2. a conversion prompt:
+    remote_inference_with_ngc(api_key="<your-ngc-apu-key>",
+                              url="https://integrate.api.nvidia.com/v1/chat/completions",
+                              model: str = "mistralai/mixtral-8x7b-instruct-v0.1",
+                              messages=[{"content": f"calculate 3+4=?", "role": "user"},
+                                        {"content": f"3+4=8", "role": "assistant"},
+                                        {"content": f"you are wrong, please correct your answer.", "role": "user"}])
+
+    """
+    assert (prompt is None) ^ (messages is None)
+
+    if prompt is not None:
+        assert isinstance(prompt, str)
+        messages = [{"content": f"{prompt}", "role": "user"}]
+    else:
+        assert isinstance(messages, list)
+        assert all([isinstance(a, dict) for a in messages])
+        assert all(["content" in a and "role" in a for a in messages])
+
+    headers = {
+        "Authorization": f"Bearer {api_key}",
+        "Accept": "application/json",
+    }
+
+    payload = {
+        "model": model,
+        "messages": messages,
+        "stream": False,
+    }
+
+    if temperature is not None:
+        # workaround since playground doesn't support temperature = 0
+        payload["temperature"] = 0.0000001 + temperature
+    if top_p is not None:
+        # workaround since playground doesn't support top p = 0
+        payload["top_p"] = 0.0000001 + top_p
+    if max_tokens is not None:
+        payload["max_tokens"] = max_tokens
+    if seed is not None:
+        payload["seed"] = seed
+
+    session = requests.Session()
+    response = session.post(url, headers=headers, json=payload)
+
+    response.raise_for_status()
+    response_body = response.json()
+    response_message = response_body["choices"][0]["message"]["content"]
+    return response_message
+
+
+class PromptTemplate:
+    def __init__(
+        self,
+        role_message_format: Dict[str, str],
+        bos_token: Optional[str] = None,
+        eos_token: Optional[str] = None,
+        response_extract_pattern: Optional[str] = None,
+    ):
+        """
+        @param role_message_format: dict - A dictionary mapping role names to their corresponding message formats.
+            Keys represent role names (e.g., 'User', 'Assistant'), and values are the format strings for each role's messages.
+            Example for Mistral:
+            {
+                'User': '
 {MESSAGE} 
',
+                'Assistant': '{MESSAGE}
 '
+            }
+            Note: '{MESSAGE}' must appear in the format text.
+
+        @param bos_token: str, optional - The begin-of-sequence token.
+            Set this when you want to add a custom BOS token as part of the prompt.
+            Note: When using this bos_token and also setting 'addBOS = True' in the sampling parameters
+            during inference, and if the model's tokenizer.bos_id is not empty, the behavior is unpredictable.
+
+        @param eos_token: str, optional - The end-of-sequence token.
+            Used to trim unrelated text from the end of the response message.
+
+        @param response_extract_pattern: str, optional - A text pattern used to extract specific parts of the response.
+            This can be useful for parsing or filtering the generated text.
+            Must be set to a non-None value if calling the function 'extract_response'.
+
+
+        example: Mistral
+
+        example: <extra_id_*> chat template
+        user_assistant_format = UserAssistantPromptTemplate(
+            user_format="<extra_id_1>User\n{MESSAGE}\n<extra_id_1>Assistant\n",
+            assistant_format="{MESSAGE}\n",
+            system_format="<extra_id_0>System\n{MESSAGE}\n",
+            system_default_message="",
+            eos_token="<extra_id_1>",
+            response_extract_pattern="<extra_id_1>Assistant\n",
+        )
+
+        # single message (e.g., user prompt)
+        prompt = user_assistant_format.format_message(
+            {"role": "User", "content": "Calculate the sum of 2 and 3."}
+        )
+
+        # a conversation
+        prompt = user_assistant_format.format_message(
+            [
+                {"role": "User", "content": "Calculate the sum of 2 and 3."},
+                {"role": "Assistant", "content": "The sum of 2 and 3 is 5."},
+                {"role": "User", "content": "Thank you! Could you also calculate the sum of 5 and 7?"},
+            ]
+        )
+        """
+        assert role_message_format is not None and isinstance(role_message_format, dict)
+        assert all(
+            PromptTemplate.is_valid_role_message_template(message_template)
+            for message_template in role_message_format.values()
+        )
+        self.role_message_template = role_message_format.copy()
+        self.bos_token = bos_token
+        self.eos_token = eos_token
+        self.response_extract_pattern = response_extract_pattern
+
+    def format_messages(self, messages: Union[List[dict], dict]):
+        assert isinstance(messages, (list, dict))
+        if not isinstance(messages, list):
+            messages = [messages]
+
+        assert all([isinstance(m, dict) for m in messages])
+        assert all(["content" in m and "role" in m for m in messages])
+        assert all([m["role"] in self.role_message_template for m in messages])
+
+        for i in range(1, len(messages)):
+            assert messages[i]["role"] != messages[i - 1]["role"]
+
+        prompt = self.bos_token if self.bos_token is not None else ""
+
+        for m in messages:
+            role = m["role"]
+            content = m["content"]
+            message = self.role_message_template[role].format(MESSAGE=content)
+            prompt += message
+
+        return prompt
+
+    def create_message(self, role: str, content: str):
+        """
+        Creates a single dictionary item with provided role and content.
+        @param role: role name
+        @param content:message content
+        @return:
+        """
+        assert role in self.role_message_template
+        return {"role": role, "content": content}
+
+    def extract_response(self, message: str):
+        assert self.response_extract_pattern is not None
+        message = message.strip()
+
+        # Find the last occurrence of the pattern
+        last_occurrence = message.rfind(self.response_extract_pattern)
+
+        response = None
+        if last_occurrence >= 0:
+            index_after_pattern = last_occurrence + len(self.response_extract_pattern)
+            response = message[index_after_pattern:]
+            response = response.strip()
+
+            # Find eos-token
+            if self.eos_token is not None:
+                eos_token_index = response.rfind(self.eos_token)
+                if eos_token_index >= 0:
+                    response = response[:eos_token_index]
+
+        return response
+
+    @staticmethod
+    def is_valid_role_message_template(message_template: str):
+        # checks for the presence of exactly one "{MESSAGE}" placeholder in the template
+        return message_template.count("{MESSAGE}") == 1
+
+
+class UserAssistantPromptTemplate(PromptTemplate):
+    class Role:
+        User = "User"
+        Assistant = "Assistant"
+        System = "System"
+
+    def __init__(
+        self,
+        user_format: str,
+        assistant_format: str,
+        system_format: Optional[str] = None,
+        system_default_message: Optional[str] = None,
+        bos_token: Optional[str] = None,
+        eos_token: Optional[str] = None,
+        response_extract_pattern: Optional[str] = None,
+    ):
+        """
+        @param user_format: user message format.
+        @param assistant_format: assistant message format.
+        @param system_format: system message format.
+        @param system_default_message: system default message
+        @param bos_token: str, optional - The begin-of-sequence token.
+            Set this when you want to add a custom BOS token as part of the prompt.
+            Note: When using this bos_token and also setting 'addBOS = True' in the sampling parameter
+            during inference, and if the model's tokenizer.bos_id is not empty, the behavior is unpredictable.
+        @param eos_token: str, optional - The end-of-sequence token.
+            Used to trim unrelated text from the end of the response message.
+        @param response_extract_pattern: str, optional - A text pattern used to extract specific parts of the response.
+            This can be useful for parsing or filtering the generated text.
+            Must be set to a non-None value if calling the function 'extract_response'.
+
+
+        example: <extra_id_*> chat template
+        user_assistant_format = UserAssistantPromptTemplate(
+            user_format="<extra_id_1>User\n{MESSAGE}\n<extra_id_1>Assistant\n",
+            assistant_format="{MESSAGE}\n",
+            system_format="<extra_id_0>System\n{MESSAGE}\n",
+            system_default_message="",
+            eos_token="<extra_id_1>",
+            response_extract_pattern="<extra_id_1>Assistant\n",
+        )
+
+
+        example: Mistral-Instruct-7B (converted to nemo, invoking remote inference with megatron_gpt_eval.py service)
+         user_assistant_format = UserAssistantPromptTemplate(
+            response_extract_pattern="
",
+        )
+
+        example: Mistral-Instruct-7B (converted to nemo, using huggingface tokenizer and not using nemo tokenizer)
+        user_assistant_format = UserAssistantPromptTemplate(
+            user_format="
 {MESSAGE} 
",
+            assistant_format="{MESSAGE}
 ",
+            bos_token="
",
+            eos_token="
",
+            response_extract_pattern="
"
+        )
+        """
+        user_format = "{MESSAGE}" if user_format is None or user_format == "" else user_format
+        assistant_format = "{MESSAGE}" if assistant_format is None or assistant_format == "" else assistant_format
+
+        role_message_format = {
+            UserAssistantPromptTemplate.Role.User: user_format,
+            UserAssistantPromptTemplate.Role.Assistant: assistant_format,
+        }
+
+        # optionally, add system message format
+        if system_format is not None:
+            role_message_format[UserAssistantPromptTemplate.Role.System] = system_format
+        else:
+            assert system_default_message is None
+
+        super().__init__(
+            role_message_format,
+            bos_token=bos_token,
+            eos_token=eos_token,
+            response_extract_pattern=response_extract_pattern,
+        )
+
+        self.system_default_message = system_default_message
+
+    def has_system_role(self):
+        return UserAssistantPromptTemplate.Role.System in self.role_message_template
+
+    def format_messages(self, messages: Union[List[dict], dict]):
+        assert messages is not None
+        assert isinstance(messages, (list, dict))
+        if not isinstance(messages, list):
+            messages = [messages]
+
+        if self.system_default_message is not None and self.has_system_role():
+            # NOTE: It is assumed that if a system message exists, it should be the first message.
+            if messages[0]["role"] != UserAssistantPromptTemplate.Role.System:
+                messages = [
+                    {"content": self.system_default_message, "role": UserAssistantPromptTemplate.Role.System}
+                ] + messages
+
+        return super().format_messages(messages)
+
+    def create_user_message(self, content: str):
+        return self.create_message(role=UserAssistantPromptTemplate.Role.User, content=content)
+
+    def create_assistant_message(self, content: str):
+        return self.create_message(role=UserAssistantPromptTemplate.Role.Assistant, content=content)
+
+    def create_system_message(self, content: str):
+        return self.create_message(role=UserAssistantPromptTemplate.Role.System, content=content)
diff --git a/examples/nlp/cai/conf/inference_service.yaml b/examples/nlp/cai/conf/inference_service.yaml
new file mode 100644
index 000000000..b243eac66
--- /dev/null
+++ b/examples/nlp/cai/conf/inference_service.yaml
@@ -0,0 +1,29 @@
+trainer:
+  devices: 1
+  num_nodes: 1
+  accelerator: gpu
+  logger: False # logger provided by exp_manager
+  precision: 16 # 16, 32, or bf16
+  use_distributed_sampler: False
+
+tensor_model_parallel_size: -1
+pipeline_model_parallel_size: -1
+pipeline_model_parallel_split_rank: -1 # used for encoder and decoder model (0 for others)
+megatron_amp_O2: False  # Enable O2-level automatic mixed precision to save memory
+gpt_model_file: null  # GPT nemo file path
+port: 5555 # the port number for the inference server
+
+
+# tokenizer configuration: overrides pretrained mode settings (optional).
+# The 'tokenizer' section is optional. You can include any combination of the following
+# keys, as long as the combination is valid for your use case. If no keys are provided,
+# the default tokenizer settings (from the pretrained model) will be used.
+#tokenizer:
+#  library:
+#  type:
+#  model:
+#  vocab_file:
+#  merge_file:
+#  delimiter:
+#  sentencepiece_legacy:
+#  tokenizer_model:
diff --git a/examples/nlp/cai/few_shot_samples_example.json b/examples/nlp/cai/few_shot_samples_example.json
new file mode 100644
index 000000000..90fbe97cd
--- /dev/null
+++ b/examples/nlp/cai/few_shot_samples_example.json
@@ -0,0 +1,56 @@
+{
+  "samples": [
+    [
+      {
+        "role": "User",
+        "content": "prompt"
+      },
+      {
+        "role": "Assistant",
+        "content": "response"
+      },
+      {
+        "role": "User",
+        "content": "critique request"
+      },
+      {
+        "role": "Assistant",
+        "content": "critique response"
+      },
+      {
+        "role": "User",
+        "content": "revision request"
+      },
+      {
+        "role": "Assistant",
+        "content": "revision response"
+      }
+    ],
+    [
+      {
+        "role": "User",
+        "content": "prompt"
+      },
+      {
+        "role": "Assistant",
+        "content": "response"
+      },
+      {
+        "role": "User",
+        "content": "critique request"
+      },
+      {
+        "role": "Assistant",
+        "content": "critique response"
+      },
+      {
+        "role": "User",
+        "content": "revision request"
+      },
+      {
+        "role": "Assistant",
+        "content": "revision response"
+      }
+    ]
+  ]
+}
diff --git a/examples/nlp/cai/generate_rl_cai_dataset.py b/examples/nlp/cai/generate_rl_cai_dataset.py
new file mode 100644
index 000000000..1e7625cbc
--- /dev/null
+++ b/examples/nlp/cai/generate_rl_cai_dataset.py
@@ -0,0 +1,693 @@
+# Copyright (c) 2024, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import argparse
+import ast
+import json
+import os
+import os.path
+import random
+import re
+from pathlib import Path
+from typing import Dict, List, Union
+
+import numpy as np
+from cai_utils import UserAssistantPromptTemplate, remote_inference, remote_inference_with_ngc
+from tqdm import tqdm
+
+
+def generate_cai_rlaif_candidate_dataset(
+    batch_size: int,
+    temperatures: Union[List, int],
+    red_teaming_dataset_path: str,
+    inference_config: dict,
+    prompt_template_config: dict,
+):
+    """
+    @param host:
+    @return:
+    @param batch_size: inference batch size
+    @param temperatures: how many temperatures to use for generation per prompt
+    @param red_teaming_dataset_path: path to Anthropic red teaming prompt attempts.
+    @param inference_config:
+    @return:
+    """
+    assert batch_size > 0
+    assert isinstance(temperatures, List) or isinstance(temperatures, int)
+    if isinstance(temperatures, int):
+        temperatures = [temperatures]
+
+    inference_config = inference_config.copy()
+    red_teaming_prompts = get_red_team_train_human_prompts(red_teaming_dataset_path)
+
+    all_samples = []
+    samples_per_temperature = {}
+
+    for batch_index in tqdm(range(0, len(red_teaming_prompts), batch_size), desc="Batch #"):
+        red_teaming_prompts_list = red_teaming_prompts[batch_index : batch_index + batch_size]
+        if len(red_teaming_prompts_list) < batch_size:
+            break
+
+        for t in tqdm(temperatures, desc="Temperatures"):
+            samples = []
+
+            # call model
+            inference_config["temperature"] = t
+            rlaif_batch_samples = generate_responses_batch(
+                red_teaming_prompts_list,
+                inference_config=inference_config,
+                prompt_template_config=prompt_template_config,
+            )
+            samples.extend(rlaif_batch_samples)
+
+            samples_per_temperature[str(t)] = samples
+
+        all_samples.extend(join_responses(samples_per_temperature))
+
+    return all_samples
+
+
+def generate_responses_batch(prompt_list: list, inference_config: dict, prompt_template_config: dict):
+    assert isinstance(prompt_list, list)
+    num_prompts = len(prompt_list)
+
+    prompt_template = UserAssistantPromptTemplate(**prompt_template_config)
+
+    # get initial response
+    prompts = [
+        prompt_template.format_messages({"content": p, "role": UserAssistantPromptTemplate.Role.User})
+        for p in prompt_list
+    ]
+    responses = model_remote_inference(prompts, inference_config=inference_config)
+    assert len(responses) == num_prompts
+    stripped_responses = [prompt_template.extract_response(r) for r in responses]
+
+    s_batch = []
+    for i in range(num_prompts):
+        s = dict(prompt=prompt_list[i], response=stripped_responses[i])
+        s_batch.append(s)
+
+    return s_batch
+
+
+def model_remote_inference(prompt, inference_config: dict):
+    sentences = remote_inference(prompt=prompt, **inference_config)
+    return sentences
+
+
+def get_red_team_train_human_prompts(red_teaming_dataset_path: str) -> list:
+    def strip_prompt(data_item):
+        human_index = data_item["transcript"].find("Human:")
+        assistant_index = data_item["transcript"].find("Assistant:")
+        assert human_index >= 0
+        assert assistant_index >= 0 and assistant_index > human_index
+
+        human_index += len("Human:")
+        prompt = data_item["transcript"][human_index:assistant_index]
+        prompt = prompt.strip()
+        return prompt
+
+    red_teaming_prompts = []
+    with open(red_teaming_dataset_path, "r") as file:
+        for line in file:
+            json_object = json.loads(line)
+            red_teaming_prompt = strip_prompt(json_object)
+            red_teaming_prompts.append(red_teaming_prompt)
+
+    return red_teaming_prompts
+
+
+def join_responses(samples_per_temperature: dict) -> list:
+    samples = []
+
+    temperatures = list(samples_per_temperature.keys())
+    prompts = [d["prompt"] for d in samples_per_temperature[temperatures[0]]]
+
+    for i, p in enumerate(prompts):
+        responses_dict = {f"response_t={t}": samples_per_temperature[t][i]["response"] for t in temperatures}
+        samples.append(dict(prompt=p, **responses_dict))
+
+    return samples
+
+
+def prepare_args():
+    parser = argparse.ArgumentParser(
+        description="given a prompt and to responses, "
+        "selects the most harmless response (labeled as 'chosen') and "
+        "the least harmless response (labeled as 'rejected')."
+    )
+    parser.add_argument("--batch-size", type=int, required=True, default=128)
+    parser.add_argument("--seed", type=int, default=1234)
+    parser.add_argument("--output-dir", type=str, default=".")
+    parser.add_argument("--output-filename-prefix", type=str, default="cai_rlaif")
+    parser.add_argument("--splits", type=str, default="{'train': 0.8, 'test': 0.2}", help="How to split the dataset")
+    parser.add_argument("--shuffle", type=str, choices=["True", "False"], default="True")
+    parser.add_argument("--red-teaming-file-path", type=str, required=True, default=None)
+    parser.add_argument("--sys-prompt-constitution-file-path", type=str, required=True, default=None)
+
+    parser.add_argument(
+        "--blend-with",
+        type=str,
+        default=None,
+        help="template:"
+        "{'name': '<some-name-for the blending>', '<split-name>': {'prompts': ['<path>', '<path-2>'], 'comparisons': ['<path-1>', '<path-2>']}}"
+        ""
+        "you must set a valid name and one or more keys of <split-name>, one for each split in '--splits' argument",
+    )
+
+    group_ngc = parser.add_argument_group("NGC", "NGC arguments")
+    group_ngc.add_argument("--ngc-api-key", type=str, required=True, default=None)
+    group_ngc.add_argument("--ngc-url", type=str, default="https://integrate.api.nvidia.com/v1/chat/completions")
+    group_ngc.add_argument("--ngc-model", type=str, default="mistralai/mixtral-8x7b-instruct-v0.1")
+
+    group_inference = parser.add_argument_group("inference", "inference (service) arguments")
+    group_inference.add_argument("--add_bos", type=str, choices=["True", "False"], default="False")
+    group_inference.add_argument("--top_k", type=int, default=50)
+    group_inference.add_argument("--top_p", type=float, default=0.95)
+    group_inference.add_argument("--all_probs", type=str, choices=["True", "False"], default="False")
+    group_inference.add_argument("--repetition_penalty", type=float, default=1.0)
+    group_inference.add_argument("--min_tokens_to_generate", type=int, default=1)
+    group_inference.add_argument("--temperature", type=float, default=1.0)
+    group_inference.add_argument("--greedy", type=str, choices=["True", "False"], default="False")
+    group_inference.add_argument("--tokens_to_generate", type=int, default=1024)
+    group_inference.add_argument("--end_strings", type=str, nargs="*", default=["<extra_id_1>"])
+    group_inference.add_argument(
+        "--port", type=int, default=5656, help="The port number on which the inference service is running"
+    )
+    group_inference.add_argument(
+        "--host", type=str, default="localhost", help="The hostname or IP address of the inference service"
+    )
+
+    # prompt template
+    group_prompt_template = parser.add_argument_group("prompt_template", "prompt template")
+    group_prompt_template.add_argument(
+        "--user_format", type=str, default="<extra_id_1>User\n{MESSAGE}\n<extra_id_1>Assistant\n"
+    )
+    group_prompt_template.add_argument("--assistant_format", type=str, default="{MESSAGE}\n")
+    group_prompt_template.add_argument("--system_format", type=str, default="<extra_id_0>System\n{MESSAGE}\n")
+    group_prompt_template.add_argument("--system_default_message", type=str, default="")
+    group_prompt_template.add_argument("--bos_token", type=str, default=None)
+    group_prompt_template.add_argument("--eos_token", type=str, default="<extra_id_1>")
+    group_prompt_template.add_argument("--response_extract_pattern", type=str, default="<extra_id_1>Assistant\n")
+
+    """
+        The prompt template configuration will be utilized in two scenarios:
+         1. It will be applied to chat conversation messages when invoking remote inference using the 
+            megatron_gpt_eval.py service (specifically in the generate_cai_rlaif_candidate_dataset method).
+         2. It will be used to format chat conversation messages when saving them to a file, which can be 
+            used as training, testing, or validation datasets.
+            
+            
+            
+        prompt template configuration example: <extra_id_*> template
+            --user_format "<extra_id_1>User\n{MESSAGE}\n<extra_id_1>Assistant\n"
+            --assistant_format "{MESSAGE}\n"
+            --system_format "<extra_id_0>System\n{MESSAGE}\n"
+            --system_default_message ""
+            --eos_token "<extra_id_1>"
+            --response_extract_pattern "<extra_id_1>Assistant\n"
+    """
+
+    args = parser.parse_args()
+    assert os.path.isfile(args.red_teaming_file_path)
+    args.splits = ast.literal_eval(args.splits)
+    args.shuffle = args.shuffle in ["True", "true"]
+
+    args.add_bos = args.add_bos == "True"
+    args.all_probs = args.all_probs == "True"
+    args.greedy = args.greedy == "True"
+
+    # blending argument validation
+    if args.blend_with is not None:
+        args.blend_with = ast.literal_eval(args.blend_with)
+        assert all(split_name in args.blend_with for split_name in args.splits)
+        assert len(args.blend_with) - 1 == len(args.splits)
+        assert (
+            "name" in args.blend_with
+            and isinstance(args.blend_with["name"], str)
+            and args.blend_with["name"] is not None
+            and args.blend_with["name"] != ""
+        )
+
+        for split_name, blend in args.blend_with.items():
+            if split_name == "name":  # ignore this key
+                continue
+
+            assert len(blend) == 2  # must have exactly 2 keys: 'prompts', 'comparisons'
+            for blend_type, files in blend.items():
+                assert blend_type in ["prompts", "comparisons"]
+                if not isinstance(files, list):
+                    files = [files]
+                    blend[blend_type] = files
+                for file in files:
+                    assert os.path.isfile(
+                        file
+                    ), f"split={split_name}, blend-type={blend_type}. error => invalid file path: {file}"
+
+    def _process_string(s: str):
+        return s.encode("utf-8").decode("unicode_escape")
+
+    # Convert parsed arguments to dictionary
+    args_dict = vars(args)
+    inference_config = {
+        k: v
+        for k, v in args_dict.items()
+        if k
+        in {
+            "add_bos",
+            "top_k",
+            "top_p",
+            "all_probs",
+            "repetition_penalty",
+            "min_tokens_to_generate",
+            "temperature",
+            "greedy",
+            "tokens_to_generate",
+            "end_strings",
+            "port",
+            "host",
+        }
+    }
+    if inference_config["end_strings"] is not None:
+        inference_config["end_strings"] = list(map(_process_string, inference_config["end_strings"]))
+
+    prompt_template_config = {
+        k: _process_string(v) if v is not None else v
+        for k, v in args_dict.items()
+        if k
+        in {
+            "user_format",
+            "assistant_format",
+            "system_format",
+            "system_default_message",
+            "bos_token",
+            "eos_token",
+            "response_extract_pattern",
+        }
+    }
+
+    return args, inference_config, prompt_template_config
+
+
+def generate_ai_preference(
+    sample: dict, ngc_api_key: str, system_prompt: str, seed: int, ngc_url: str, ngc_model: str
+):
+    # NOTE: For generating AI preferences we deviate a bit from the paper and instead of feeding one (randomized)
+    # constitution principle at a time, we feed the entire constitution at once. Also, instead of using normalized
+    # logprobs of the candidate response number tokens, we just ask the judge LLM to choose what is the most harmless
+    # and toxic responses.
+    # Although deviating a bit from the paper, we found it to work quite well in practice.
+
+    prompt = sample["prompt"]
+    if prompt is None or prompt == "":
+        return None
+
+    response_keys = sorted([key for key in sample.keys() if key.startswith("response")])
+    responses = [sample[key] for key in response_keys]
+    if len(responses) < 2:
+        return None
+
+    responses_to_choose_from_text = f"\n FULL TEXT: {prompt}"
+    for i, response in enumerate(responses):
+        responses_to_choose_from_text += f"\n #{i + 1}: \n {response}\n"
+    responses_to_choose_from_text += "\nNow, choose the pair of responses as instructed:"
+
+    full_prompt_messages = [
+        {"role": "system", "content": system_prompt},
+        {"role": "user", "content": responses_to_choose_from_text},
+    ]
+
+    res = remote_inference_with_ngc(
+        api_key=ngc_api_key,
+        url=ngc_url,
+        model=ngc_model,
+        messages=full_prompt_messages,
+        top_p=0,
+        max_tokens=1024,
+        temperature=0,
+        seed=seed,
+    )
+
+    def _extract_decision_and_explanation(s):
+        try:
+            # convert string into a dictionary
+            dict_start = s.find("{")
+            dict_end = s.rfind("}") + 1
+            dict_str = s[dict_start:dict_end]
+            dict_result = ast.literal_eval(dict_str)
+
+            # get decision
+            s_remaining = s[dict_end:].strip()
+
+            assert "chosen" in dict_result
+            assert "rejected" in dict_result
+
+            chosen_id = list(map(int, re.findall(r"\d+", str(dict_result["chosen"]))))
+            rejected_id = list(map(int, re.findall(r"\d+", str(dict_result["rejected"]))))
+
+            assert len(chosen_id) == 1 and 1 <= chosen_id[0] <= len(responses)
+            assert len(rejected_id) == 1 and 1 <= rejected_id[0] <= len(responses)
+
+            dict_result["chosen"] = chosen_id[0]
+            dict_result["rejected"] = rejected_id[0]
+
+            return dict_result, s_remaining
+        except (ValueError, SyntaxError):
+            # Return None or raise an error if the string does not contain a valid dictionary
+            return None, None
+
+    selected_pair, selection_explanation = _extract_decision_and_explanation(res)
+
+    if selected_pair is None:
+        return None
+
+    # convert response number to response index
+    chosen_response_index = selected_pair["chosen"] - 1
+    rejected_response_index = selected_pair["rejected"] - 1
+
+    if chosen_response_index == rejected_response_index:
+        # judge model assigned same index for 'chosen' and 'rejected' responses.
+        # sample a random index for both 'rejected' and 'chosen'.
+        rnd_response_index = random.sample(range(len(responses)), 2)
+        chosen_response_index = rnd_response_index[0]
+        rejected_response_index = rnd_response_index[1]
+
+    # get responses
+    chosen_response = responses[chosen_response_index]
+    rejected_response = responses[rejected_response_index]
+
+    preference_sample = dict(
+        prompt=prompt,
+        chosen=chosen_response,
+        rejected=rejected_response,
+        raw=dict(
+            chosen_response_key=response_keys[chosen_response_index],
+            rejected_response_key=response_keys[rejected_response_index],
+            explanation=selection_explanation,
+            all_responses={k: sample[k] for k in response_keys},
+        ),
+    )
+
+    return preference_sample
+
+
+def split_dataset(dataset, splits: Dict[str, float], shuffle: bool):
+    n = len(dataset)
+    assert sum(list(splits.values())) == 1.0
+    assert all(1.0 >= split_p > 0 for split_p in splits.values())
+
+    if shuffle:
+        index = random.sample(range(n), n)
+    else:
+        index = list(range(n))
+
+    # ensure all splits have at least one sample
+    dataset_splits = {split_name: [dataset[index[i]]] for i, split_name in enumerate(splits.keys())}
+    index = index[len(splits) :]
+    n = n - len(splits)
+
+    i_offset = 0
+    for i, (split_name, split_p) in enumerate(splits.items()):
+        split_n = max(1, round(n * split_p))
+        if i == len(splits) - 1:
+            split_n = n - i_offset
+        split_index = index[i_offset : min(i_offset + split_n, n)]
+        dataset_splits[split_name] += [dataset[i] for i in split_index]
+        i_offset += split_n
+
+    assert sum([len(s) for s in dataset_splits.values()]) == len(dataset)
+
+    return dataset_splits
+
+
+def process_samples(dataset, prompt_template_config: dict):
+    def convert_string_format(body, response):
+        response = response.strip().strip("\n")
+        body = body.strip().strip("\n")
+
+        if len(response) == 0 or len(body) == 0:
+            return "", ""
+
+        prompt_template = UserAssistantPromptTemplate(**prompt_template_config)
+
+        # encode only prompt (user message)
+        prompt = prompt_template.format_messages({"content": body, "role": UserAssistantPromptTemplate.Role.User})
+
+        # encode user message and assistant response (for rejected/chosen samples)
+        prompt_with_response = prompt_template.format_messages(
+            [
+                {"content": body, "role": UserAssistantPromptTemplate.Role.User},
+                {"content": response, "role": UserAssistantPromptTemplate.Role.Assistant},
+            ]
+        )
+
+        # we need tomake sure eos-token is present at the end of the response
+        prompt_with_response = prompt_with_response.strip()
+        if not prompt_with_response.endswith(prompt_template.eos_token):
+            prompt_with_response += prompt_template.eos_token
+
+        return prompt_with_response, prompt
+
+    chosen = [convert_string_format(x["prompt"], x["chosen"]) for x in dataset]
+    rejected = [convert_string_format(x["prompt"], x["rejected"]) for x in dataset]
+
+    samples = []
+    for c, r in zip(chosen, rejected):
+        if c is None or r is None:
+            continue
+
+        chosen_response, chosen_prompt = c
+        rejected_response, rejected_prompt = r
+
+        if len(chosen_response) == 0 or len(rejected_response) == 0:
+            continue
+
+        if chosen_prompt != rejected_prompt:
+            continue
+
+        comparison_dict = {
+            "prompt": chosen_prompt,
+            "chosen": chosen_response,
+            "rejected": rejected_response,
+        }
+
+        samples.append(comparison_dict)
+
+    return samples
+
+
+def convert_list_of_dict_to_json(list_of_dict):
+    return "\n".join(json.dumps(item) for item in list_of_dict)
+
+
+def validate_output_filenames(split: str, output_dir: str, output_filename_prefix: str):
+    prompts_file_name = f"{split}_prompts_with_chat_prompt.json"
+    if output_filename_prefix is not None and output_filename_prefix != "":
+        prompts_file_name = f"{output_filename_prefix}_{prompts_file_name}"
+    assert not os.path.isfile(
+        os.path.join(output_dir, prompts_file_name)
+    ), f"error, file already exists: {os.path.join(output_dir, prompts_file_name)}"
+
+    comparisons_file_name = f"{split}_comparisons_with_chat_prompt.json"
+    if output_filename_prefix is not None and output_filename_prefix != "":
+        comparisons_file_name = f"{output_filename_prefix}_{comparisons_file_name}"
+    assert not os.path.isfile(
+        os.path.join(output_dir, comparisons_file_name)
+    ), f"error, file already exists: {os.path.join(output_dir, comparisons_file_name)}"
+
+
+def save_dataset(dataset, split: str, output_dir: str, output_filename_prefix: str):
+    prompts_to_save = convert_list_of_dict_to_json({"text": item["prompt"]} for item in dataset)
+    prompts_file_name = f"{split}_prompts_with_chat_prompt"
+    if output_filename_prefix is not None and output_filename_prefix != "":
+        prompts_file_name = f"{output_filename_prefix}_{prompts_file_name}"
+    prompts_file_name += ".jsonl"
+
+    prompts_full_path = Path(output_dir) / prompts_file_name
+    print(f"Saving {len(prompts_to_save)} prompts to {prompts_full_path}")
+    with open(prompts_full_path, "w") as f:
+        f.write(prompts_to_save)
+
+    comparisons_to_save = []
+    for item in dataset:
+        comparisons_to_save.append({"text": item["chosen"]})
+        comparisons_to_save.append({"text": item["rejected"]})
+
+    comparisons_to_save = convert_list_of_dict_to_json(comparisons_to_save)
+    comparisons_file_name = f"{split}_comparisons_with_chat_prompt"
+    if output_filename_prefix is not None and output_filename_prefix != "":
+        comparisons_file_name = f"{output_filename_prefix}_{comparisons_file_name}"
+    comparisons_file_name += ".jsonl"
+
+    comparisons_full_path = Path(output_dir) / comparisons_file_name
+    print(f"Saving {len(comparisons_to_save)} comparisons to {comparisons_full_path}")
+    with open(comparisons_full_path, "w") as f:
+        f.write(comparisons_to_save)
+
+    return prompts_full_path, comparisons_full_path
+
+
+def blend_preference_datasets(files: list, output_file: str, blend_type: str):
+    assert not os.path.isfile(output_file), f"Error: output file: {output_file} already exists."
+    assert len(files) > 0
+    assert all([os.path.isfile(f) for f in files])
+    assert blend_type in ["prompts", "comparisons"]
+    os.makedirs(os.path.dirname(output_file), exist_ok=True)
+
+    def _blend_files_type_prompt():
+        # Read and combine the contents of all files
+        combined_lines = []
+        for file_name in files:
+            with open(file_name, "r") as file:
+                combined_lines.extend(file.readlines())
+                if not combined_lines[-1].endswith("\n"):
+                    combined_lines[-1] += "\n"
+
+        # Shuffle the combined lines
+        random.shuffle(combined_lines)
+
+        # Write the shuffled lines to the output file
+        with open(output_file, "w") as f:
+            f.writelines(combined_lines)
+
+    def _blend_files_type_comparison():
+        """
+        the structure of a preference dataset is:
+        chosen
+        rejected
+        chosen
+        rejected
+
+        so we need to keep this structure, that is why we need to sample pairs (consecutive lines).
+        """
+
+        # Function to read files and group lines in pairs
+        def read_and_pair_lines(file_names):
+            paired_lines = []
+            for file_name in file_names:
+                with open(file_name, "r") as file:
+                    lines = file.readlines()
+                    if not lines[-1].endswith("\n"):
+                        lines[-1] = lines[-1] + "\n"
+
+                    # Assuming an even number of lines, pair them
+                    for i in range(0, len(lines), 2):
+                        paired_lines.append(lines[i : i + 2])
+
+                if not paired_lines[-1][-1].endswith("\n"):
+                    paired_lines[-1][-1] += "\n"
+
+            return paired_lines
+
+        # Read and combine the contents of all files into pairs
+        combined_pairs = read_and_pair_lines(files)
+
+        # Shuffle the combined pairs
+        random.shuffle(combined_pairs)
+
+        # Write the shuffled pairs to the output file
+        with open(output_file, "w") as f:
+            for pair in combined_pairs:
+                f.writelines(pair)
+
+    if blend_type == "prompts":
+        _blend_files_type_prompt()
+    else:
+        _blend_files_type_comparison()
+
+    print("done")
+
+
+def main():
+    args, inference_config, prompt_template_config = prepare_args()
+    os.makedirs(args.output_dir, exist_ok=True)
+    random.seed(args.seed)
+
+    print("Generating CAI-RLAIF candidates dataset...\n")
+    dataset = generate_cai_rlaif_candidate_dataset(
+        batch_size=args.batch_size,
+        temperatures=np.arange(0.01, 2.01, 0.5).tolist(),
+        red_teaming_dataset_path=args.red_teaming_file_path,
+        inference_config=inference_config,
+        prompt_template_config=prompt_template_config,
+    )
+
+    with open(os.path.join(args.output_dir, "cai_candidate_dataset.json"), "w") as file:
+        json.dump(dataset, file, indent=4)
+
+    print("\nGenerating AI preferences...\n")
+
+    with open(args.sys_prompt_constitution_file_path, "r") as f:
+        constitution_as_sys_prompt = f.read()
+
+    preference_dataset = []
+    for ds_index in tqdm(range(len(dataset))):
+        sample = dataset[ds_index]
+
+        try:
+            preference = generate_ai_preference(
+                sample,
+                args.ngc_api_key,
+                constitution_as_sys_prompt,
+                seed=args.seed,
+                ngc_url=args.ngc_url,
+                ngc_model=args.ngc_model,
+            )
+        except Exception as e:
+            preference = None
+
+        if preference is not None:
+            preference_dataset.append(preference)
+
+    with open(os.path.join(args.output_dir, "cai_preference_dataset.json"), "w") as file:
+        json.dump(preference_dataset, file, indent=4)
+
+    print(
+        f"\nGenerated {len(preference_dataset)} AI preferences. "
+        f"Now processing and converting to chat prompt template...\n"
+    )
+    for split_name in args.splits:
+        validate_output_filenames(
+            split=split_name, output_dir=args.output_dir, output_filename_prefix=args.output_filename_prefix
+        )
+    ds = split_dataset(preference_dataset, args.splits, shuffle=args.shuffle)
+
+    output_file_names = []
+    for split_name, split in ds.items():
+        split_samples = process_samples(split, prompt_template_config=prompt_template_config)
+        prompts_path, comparisons_path = save_dataset(
+            dataset=split_samples,
+            split=split_name,
+            output_dir=args.output_dir,
+            output_filename_prefix=args.output_filename_prefix,
+        )
+        output_file_names.append(dict(split_name=split_name, prompts=prompts_path, comparisons=comparisons_path))
+
+    print(f"blending preference dataset with external dataset:")
+    # args.blend_with
+    for split_ds in output_file_names:
+        blend_split_with = args.blend_with[split_ds["split_name"]]
+
+        for blend_type in ["prompts", "comparisons"]:
+            output_file_name = f"blend_{args.blend_with['name']}_with_{os.path.basename(split_ds[blend_type])}"
+            output_file_path = os.path.join(args.output_dir, output_file_name)
+            blend_preference_datasets(
+                files=blend_split_with[blend_type] + [split_ds[blend_type]],
+                output_file=output_file_path,
+                blend_type=blend_type,
+            )
+
+
+if __name__ == "__main__":
+    main()
diff --git a/examples/nlp/cai/generate_sl_cai_dataset.py b/examples/nlp/cai/generate_sl_cai_dataset.py
new file mode 100644
index 000000000..aecea1f93
--- /dev/null
+++ b/examples/nlp/cai/generate_sl_cai_dataset.py
@@ -0,0 +1,579 @@
+# Copyright (c) 2024, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import argparse
+import json
+import os
+import os.path
+import random
+from typing import List, Union
+
+import numpy as np
+import tqdm
+from cai_utils import UserAssistantPromptTemplate, remote_inference
+
+
+def generate_chat_prompt(sample: dict):
+    dataset_name = "sl_cai_samples"
+    assert isinstance(sample, dict)
+    assert all(key in sample for key in ["initial_prompt", "revision_response"])
+    assert dataset_name is not None and dataset_name != ""
+
+    def _message_formatter(role: str, message: str):
+        assert role in ["User", "Assistant"]
+
+        return {"from": f"{role}", "value": f"{message}", "label": None}
+
+    def _conversation_formatter(messages: list, system_message: str = ""):
+        assert isinstance(dataset_name, str)
+        assert isinstance(messages, list) and len(messages) > 0
+
+        return {"system": f"{system_message}", "conversations": messages, "mask": "User", "dataset": f"{dataset_name}"}
+
+    prompt = _message_formatter(role="User", message=sample["initial_prompt"])
+    revision = _message_formatter(role="Assistant", message=sample["revision_response"])
+
+    return _conversation_formatter(messages=[prompt, revision])
+
+
+def model_remote_inference(prompt, inference_config: dict):
+    sentences = remote_inference(prompt=prompt, **inference_config)
+    return sentences
+
+
+def generate_cai_batch_sample(
+    prompt_list: list,
+    few_shot_dataset_path: str,
+    critique_list,
+    revision_list,
+    inference_config: dict,
+    prompt_template_config: dict,
+    apply_chat_template: bool,
+):
+    assert isinstance(prompt_list, list)
+    if not isinstance(critique_list, list):
+        critique_list = [critique_list] * len(prompt_list)
+    if not isinstance(revision_list, list):
+        revision_list = [revision_list] * len(prompt_list)
+
+    num_prompts = len(prompt_list)
+    assert len(critique_list) == num_prompts
+    assert len(revision_list) == num_prompts
+
+    prompt_template = UserAssistantPromptTemplate(**prompt_template_config)
+
+    with open(few_shot_dataset_path, "r") as f:
+        few_shot_samples = json.load(f)
+        few_shot_prompts = [item for sublist in few_shot_samples["samples"] for item in sublist]
+
+    # get initial response
+    print(f"\nGenerating initial response for {num_prompts} prompts...")
+    initial_prompt_batch = [
+        few_shot_prompts + [{"content": p, "role": UserAssistantPromptTemplate.Role.User}] for p in prompt_list
+    ]
+
+    chat_batch = model_remote_inference(
+        [prompt_template.format_messages(p) if apply_chat_template else p for p in initial_prompt_batch],
+        inference_config=inference_config,
+    )
+
+    assert len(chat_batch) == num_prompts
+    initial_response_batch = [prompt_template.extract_response(chat) for chat in chat_batch]
+    print("Done")
+
+    # generate a single critique
+    print(f"\nGenerating a critique for {num_prompts} initial responses...")
+    critique_request_batch = [
+        initial_prompt_batch[i]
+        + [{"content": initial_response_batch[i], "role": UserAssistantPromptTemplate.Role.Assistant}]
+        + [{"content": cr_p, "role": UserAssistantPromptTemplate.Role.User}]
+        for i, cr_p in enumerate(critique_list)
+    ]
+
+    chat_batch = model_remote_inference(
+        [prompt_template.format_messages(p) if apply_chat_template else p for p in critique_request_batch],
+        inference_config=inference_config,
+    )
+    assert len(chat_batch) == num_prompts
+    critique_response_batch = [prompt_template.extract_response(chat) for chat in chat_batch]
+    print("Done")
+
+    # generate a single revision
+    print(f"\nGenerating {num_prompts} revisions ...")
+    revision_request_prompt_batch = [
+        critique_request_batch[i]
+        + [{"content": critique_response_batch[i], "role": UserAssistantPromptTemplate.Role.Assistant}]
+        + [{"content": rev_p, "role": UserAssistantPromptTemplate.Role.User}]
+        for i, rev_p in enumerate(revision_list)
+    ]
+
+    chat_batch = model_remote_inference(
+        [prompt_template.format_messages(p) if apply_chat_template else p for p in revision_request_prompt_batch],
+        inference_config=inference_config,
+    )
+    assert len(chat_batch) == num_prompts
+    revision_response_batch = [prompt_template.extract_response(chat) for chat in chat_batch]
+    print("Done")
+
+    s_batch = []
+    for i in range(num_prompts):
+        s = dict(
+            initial_prompt=prompt_list[i],
+            initial_response=initial_response_batch[i],
+            critique_prompt=critique_list[i],
+            critique_response=critique_response_batch[i],
+            revision_prompt=revision_list[i],
+            revision_response=revision_response_batch[i],
+        )
+        s_batch.append(s)
+    return s_batch
+
+
+def load_critique_revision_instructions(file_path: str):
+    instructions = []
+    with open(file_path, "r") as file:
+        json_object = json.load(file)
+        for i in range(len(json_object)):
+            inst = json_object[f"harmful{i}"]
+            assert len(inst) == 2 and "prompt" in inst and "edit_request" in inst
+            assert len(inst["prompt"]) == 1
+
+            critique_prompt = inst["prompt"][0]
+            revision_prompt = inst["edit_request"]
+
+            instructions.append(dict(critique_prompt=critique_prompt, revision_prompt=revision_prompt))
+
+    return instructions
+
+
+def get_red_team_train_human_prompts(red_teaming_dataset_path: str) -> list:
+    def strip_first_red_team_prompt(data_item):
+        human_index = data_item["transcript"].find("Human:")
+        assistant_index = data_item["transcript"].find("Assistant:")
+        assert human_index >= 0
+        assert assistant_index >= 0 and assistant_index > human_index
+
+        human_index += len("Human:")
+        prompt = data_item["transcript"][human_index:assistant_index]
+        prompt = prompt.strip()
+        return prompt
+
+    red_teaming_prompts = []
+    with open(red_teaming_dataset_path, "r") as file:
+        for line in file:
+            json_object = json.loads(line)
+            red_teaming_prompt = strip_first_red_team_prompt(json_object)
+            red_teaming_prompts.append(red_teaming_prompt)
+
+    return red_teaming_prompts
+
+
+def generate_cai_dataset(
+    red_teaming_dataset_path: str,
+    few_shot_samples_filepath: str,
+    critique_revision_instructions_filepath: str,
+    num_examples: int,
+    batch_size: int,
+    save_to_file_interval: int,
+    save_file_path: str,
+    inference_config: dict,
+    prompt_template_config: dict,
+    apply_chat_template: bool,
+):
+    """
+    @param batch_size: inference batch size
+    @param save_to_file_interval: saves generated samples to the disk after 'save_to_file_interval' batches.
+    @param save_file_path:
+    @param red_teaming_dataset_path:
+    @param few_shot_samples_filepath:
+    @param critique_revision_instructions_filepath:
+    @param num_examples:
+    @param inference_config:
+    @param prompt_template_config:
+    @param apply_chat_template:
+    @return:
+    """
+    assert batch_size > 0
+    assert save_to_file_interval > 0
+    assert save_file_path is not None and save_file_path != ""
+    assert not os.path.exists(save_file_path), f"{save_file_path} already exists"
+
+    log_dir = os.path.join(os.path.dirname(save_file_path), "cai_output")
+    os.makedirs(log_dir, exist_ok=True)
+    output_path_critique_revision = os.path.join(log_dir, f"cai_critique_revision_samples.json")
+    output_path_cai_samples = os.path.join(log_dir, f"cai_samples.jsonl")
+
+    # load constitution critique/revision instructions
+    critique_revision_instructions_set = load_critique_revision_instructions(critique_revision_instructions_filepath)
+
+    def sample_random_critique_revision_set():
+        a_critique_list = []
+        a_revision_list = []
+        for i in range(batch_size):
+            rnd_index = random.randint(0, len(critique_revision_instructions_set) - 1)
+            a_critique_list.append(critique_revision_instructions_set[rnd_index]["critique_prompt"])
+            a_revision_list.append(critique_revision_instructions_set[rnd_index]["revision_prompt"])
+        return a_critique_list, a_revision_list
+
+    red_teaming_prompts = get_red_team_train_human_prompts(red_teaming_dataset_path)
+
+    num_examples = len(red_teaming_prompts) if num_examples < 0 else min(num_examples, len(red_teaming_prompts))
+    assert num_examples > 0
+
+    cai_samples = []
+    critique_revision_samples = []
+    num_batches = 0
+    for index in tqdm.tqdm(range(0, num_examples, batch_size), desc="Generating CAI samples, batch #"):
+        red_teaming_prompts_list = red_teaming_prompts[index : index + batch_size]
+        if len(red_teaming_prompts_list) < batch_size:
+            break
+
+        # sample random critique/revision instruction set
+        critique_list, revision_list = sample_random_critique_revision_set()
+        critique_list = [c.removesuffix("\n\nCritique:").removeprefix("\n\nCritiqueRequest: ") for c in critique_list]
+        revision_list = [c.removesuffix("\n\nRevision:").removeprefix("\n\nRevisionRequest: ") for c in revision_list]
+
+        # call model
+        cai_batch_sample = generate_cai_batch_sample(
+            red_teaming_prompts_list,
+            few_shot_dataset_path=few_shot_samples_filepath,
+            critique_list=critique_list,
+            revision_list=revision_list,
+            inference_config=inference_config,
+            prompt_template_config=prompt_template_config,
+            apply_chat_template=apply_chat_template,
+        )
+        assert len(cai_batch_sample) == len(red_teaming_prompts_list)
+
+        print(f"Reformatting CAI samples according to a chat prompt template...")
+        for sample in cai_batch_sample:
+            chat_formatted_cai_sample = generate_chat_prompt(sample=sample)
+            cai_samples.append(chat_formatted_cai_sample)
+            critique_revision_samples.append(sample)
+        print("Done")
+
+        num_batches += 1
+
+        if num_batches % save_to_file_interval == 0:
+            with open(save_file_path, "w") as f:
+                for sample in cai_samples:
+                    json_line = json.dumps(sample)
+                    f.write(json_line + "\n")
+
+    # final save
+    with open(save_file_path, "w") as f:
+        for sample in cai_samples:
+            json_line = json.dumps(sample)
+            f.write(json_line + "\n")
+
+    # logs (saving all raw files)
+    with open(output_path_critique_revision, "w") as f:
+        json.dump(critique_revision_samples, f, indent=4)
+
+    with open(output_path_cai_samples, "w") as f:
+        for sample in cai_samples:
+            json_line = json.dumps(sample)
+            f.write(json_line + "\n")
+
+    return cai_samples
+
+
+def remove_long_dialogs_wrapper(
+    input_file_path: str, max_seq_length: int, tokenizer_model: str, tokenizer_library: str
+):
+    from cai_utils import remove_long_dialogs
+
+    if max_seq_length is None or max_seq_length <= 0:
+        return
+    assert os.path.isfile(tokenizer_model)
+    assert tokenizer_library is not None
+
+    remove_long_out_dir = os.path.join(os.path.dirname(input_file_path), "removed_long_dialogs")
+    statistics_file_path = os.path.join(remove_long_out_dir, "remove_long_dialogs_statistics.json")
+    os.makedirs(remove_long_out_dir, exist_ok=True)
+
+    import shutil
+
+    copy_input_file_path = os.path.join(remove_long_out_dir, os.path.basename(input_file_path))
+    shutil.copy(input_file_path, copy_input_file_path)
+
+    d = remove_long_dialogs(
+        input_file_path=copy_input_file_path,
+        max_seq_length=max_seq_length,
+        tokenizer_model=tokenizer_model,
+        tokenizer_library=tokenizer_library,
+        output_dir=remove_long_out_dir,
+        use_pool=True,
+    )
+
+    # save output statistics
+    with open(statistics_file_path, "w") as f:
+        json.dump(d, f, indent=4)
+
+    # delete files
+    if os.path.exists(copy_input_file_path):
+        os.remove(copy_input_file_path)
+
+    return d["output_file"]
+
+
+def validate_dataset_structure(item: dict):
+    """
+    expected format:
+
+        prompt = {
+            'system': f"{system_message}",
+            'conversations': [
+                {'from': 'User',
+                 'value': f"{user_message}",
+                 'label': default_label},
+
+                {'from': 'Assistant',
+                 'value': f"{assistant_message}",
+                 'label': default_label}
+            ],
+            "mask": "User",
+            "dataset": f"{dataset}",
+        }
+    """
+
+    required_fields = ["system", "conversations", "mask", "dataset"]
+    for field in required_fields:
+        if field not in item:
+            return False
+
+    for message in item["conversations"]:
+        if "from" not in message or "value" not in message:
+            return False
+
+    return True
+
+
+def blend_cai_sft_dataset(helpfulness_dataset_path, cai_samples_filepath, blended_sl_cai_filename, summary_filename):
+    input_files = [helpfulness_dataset_path, cai_samples_filepath]
+    assert all(os.path.isfile(input_file) for input_file in input_files)
+
+    samples = []
+    for input_file in input_files:
+        with open(input_file, "r", encoding="utf-8") as f:
+            for line in f:
+                sample = json.loads(line)
+                samples.append(sample)
+
+    # make sure all samples have a valid structure
+    assert all(validate_dataset_structure(s) for s in samples)
+
+    # permute samples
+    samples_permuted = [samples[i] for i in np.random.permutation(len(samples))]
+
+    # Each sample is in a separate line
+    with open(blended_sl_cai_filename, "w") as f:
+        for sample in samples_permuted:
+            json_line = json.dumps(sample)
+            f.write(json_line + "\n")
+
+    # create a summary file
+    with open(summary_filename, "w") as f:
+        summary_dict = dict(input_files=input_files, output_file=blended_sl_cai_filename, prompt_template="chat")
+        json.dump(summary_dict, f, indent=4)
+
+
+def prepare_args():
+    parser = argparse.ArgumentParser(description="Generate CAI dataset")
+    parser.add_argument("--max-seq-length", type=int, required=False, default=None)
+    parser.add_argument("--batch_size", type=int, required=False, default=32)
+    parser.add_argument("--tokenizer-model", type=str, required=False, default=None)
+    parser.add_argument("--tokenizer-library", type=str, required=False, default=None)
+    parser.add_argument(
+        "--num-examples",
+        type=int,
+        required=False,
+        default=-1,
+        help="Number of samples to generate. Default (-1) generates for all the samples in the dataset.",
+    )
+    parser.add_argument("--save-to-file-interval", type=int, required=False, default=1)
+    parser.add_argument("--red-teaming-prompts-dataset-path", type=str, required=True, default=None)
+    parser.add_argument("--few-shot-prompts-dataset-path", type=str, required=True, default=None)
+    parser.add_argument("--critique-revision-instructions-path", type=str, required=True, default=None)
+    parser.add_argument(
+        "--output-filepath", type=str, required=False, default="cai_revisions_aligner_chat_template.jsonl"
+    )
+    parser.add_argument("--seed", type=int, default=1234)
+    parser.add_argument("--helpfulness-dataset-path", type=str, required=True, default=None)
+
+    group_inference = parser.add_argument_group("inference", "inference (service) arguments")
+    group_inference.add_argument("--add_bos", type=str, choices=["True", "False"], default="False")
+    group_inference.add_argument("--top_k", type=int, default=1)
+    group_inference.add_argument("--top_p", type=float, default=0.9)
+    group_inference.add_argument("--all_probs", type=str, choices=["True", "False"], default="False")
+    group_inference.add_argument("--repetition_penalty", type=float, default=1.2)
+    group_inference.add_argument("--min_tokens_to_generate", type=int, default=1)
+    group_inference.add_argument("--temperature", type=float, default=1.0)
+    group_inference.add_argument("--greedy", type=str, choices=["True", "False"], default="True")
+    group_inference.add_argument("--tokens_to_generate", type=int, default=1024)
+    group_inference.add_argument("--end_strings", type=str, nargs="*", default=None)
+    group_inference.add_argument(
+        "--port", type=int, default=5656, help="The port number on which the inference service is running"
+    )
+    group_inference.add_argument(
+        "--host", type=str, default="localhost", help="The hostname or IP address of the inference service"
+    )
+
+    # prompt template
+    group_prompt_template = parser.add_argument_group("prompt_template", "prompt template")
+    group_prompt_template.add_argument("--apply_chat_template", type=str, choices=["True", "False"], default="False")
+    group_prompt_template.add_argument("--user_format", type=str, default=None)
+    group_prompt_template.add_argument("--assistant_format", type=str, default=None)
+    group_prompt_template.add_argument("--system_format", type=str, default=None)
+    group_prompt_template.add_argument("--system_default_message", type=str, default=None)
+    group_prompt_template.add_argument("--bos_token", type=str, default=None)
+    group_prompt_template.add_argument("--eos_token", type=str, default=None)
+    group_prompt_template.add_argument("--response_extract_pattern", type=str, default="
")
+
+    """
+    prompt template configuration is going to be applied to chat conversation messages when invoking remote
+    inference with megatron_gpt_eval.py service.
+    
+    
+    prompt template configuration example: <extra_id_*> template
+        --apply_chat_template True
+        --user_format "<extra_id_1>User\n{MESSAGE}\n<extra_id_1>Assistant\n"
+        --assistant_format "{MESSAGE}\n"
+        --system_format "<extra_id_0>System\n{MESSAGE}\n"
+        --system_default_message ""
+        --eos_token "<extra_id_1>"
+        --response_extract_pattern "<extra_id_1>Assistant\n"
+        
+    
+    prompt template configuration example: mistral-instruct-7B
+        --apply_chat_template False
+        --response_extract_pattern "
"
+        
+    NOTE: setting 'apply_chat_template' to False as chat template is going to be applied when 
+    invoking remote inference with megatron_gpt_eval.py service.
+    
+    
+    prompt template configuration example: Mistral-Instruct-7B, converted to nemo, 
+        using huggingface tokenizer and not using nemo tokenizer:
+    
+    -apply_chat_template True
+        --user_format "
 {MESSAGE} 
"
+        --assistant_format "{MESSAGE}
 "
+        --bos_token "
"
+        --eos_token "
"
+        --response_extract_pattern "
"
+    """
+
+    args = parser.parse_args()
+    args.add_bos = args.add_bos == "True"
+    args.all_probs = args.all_probs == "True"
+    args.greedy = args.greedy == "True"
+    args.apply_chat_template = args.apply_chat_template == "True"
+
+    assert os.path.isfile(args.red_teaming_prompts_dataset_path)
+    assert os.path.isfile(args.few_shot_prompts_dataset_path)
+    assert os.path.isfile(args.helpfulness_dataset_path)
+    assert args.num_examples >= -1 and args.num_examples != 0
+
+    if args.max_seq_length is not None:
+        assert args.max_seq_length > 0
+        assert os.path.isfile(args.tokenizer_model)
+        assert args.tokenizer_library is not None
+
+    def _process_string(s: str):
+        return s.encode("utf-8").decode("unicode_escape")
+
+    # Convert parsed arguments to dictionary
+    args_dict = vars(args)
+    inference_config = {
+        k: v
+        for k, v in args_dict.items()
+        if k
+        in {
+            "add_bos",
+            "top_k",
+            "top_p",
+            "all_probs",
+            "repetition_penalty",
+            "min_tokens_to_generate",
+            "temperature",
+            "greedy",
+            "tokens_to_generate",
+            "end_strings",
+            "port",
+            "host",
+        }
+    }
+    if inference_config["end_strings"] is not None:
+        inference_config["end_strings"] = list(map(_process_string, inference_config["end_strings"]))
+
+    prompt_template_config = {
+        k: _process_string(v) if v is not None else v
+        for k, v in args_dict.items()
+        if k
+        in {
+            "user_format",
+            "assistant_format",
+            "system_format",
+            "system_default_message",
+            "bos_token",
+            "eos_token",
+            "response_extract_pattern",
+        }
+    }
+
+    return args, inference_config, prompt_template_config
+
+
+def main():
+    args, inference_config, prompt_template_config = prepare_args()
+    random.seed(args.seed)
+    np.random.seed(args.seed)
+
+    print("Generating CAI samples...")
+    generate_cai_dataset(
+        red_teaming_dataset_path=args.red_teaming_prompts_dataset_path,
+        few_shot_samples_filepath=args.few_shot_prompts_dataset_path,
+        critique_revision_instructions_filepath=args.critique_revision_instructions_path,
+        num_examples=args.num_examples,
+        batch_size=args.batch_size,
+        save_to_file_interval=args.save_to_file_interval,
+        save_file_path=args.output_filepath,
+        inference_config=inference_config,
+        prompt_template_config=prompt_template_config,
+        apply_chat_template=args.apply_chat_template,
+    )
+
+    print("Blending CAI samples with the helpfulness dataset...")
+    blend_cai_sft_dataset(
+        helpfulness_dataset_path=args.helpfulness_dataset_path,
+        cai_samples_filepath=args.output_filepath,
+        blended_sl_cai_filename=args.output_filepath,
+        summary_filename=args.output_filepath.replace(".json", ".summary.json"),
+    )
+
+    if args.max_seq_length is not None:
+        print("Removing long dialogs from dataset...")
+        remove_long_dialogs_wrapper(
+            input_file_path=args.output_filepath,
+            max_seq_length=args.max_seq_length,
+            tokenizer_model=args.tokenizer_model,
+            tokenizer_library=args.tokenizer_library,
+        )
+
+    print("Done")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/examples/nlp/cai/process_anthropic_hh_using_chat_prompt.py b/examples/nlp/cai/process_anthropic_hh_using_chat_prompt.py
new file mode 100644
index 000000000..f817d3bc2
--- /dev/null
+++ b/examples/nlp/cai/process_anthropic_hh_using_chat_prompt.py
@@ -0,0 +1,269 @@
+# Copyright (c) 2023, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""A script to process the Anthropic Dataset"""
+import argparse
+import json
+import os
+import warnings
+from pathlib import Path
+
+from cai_utils import UserAssistantPromptTemplate
+from datasets import DatasetDict, concatenate_datasets, load_dataset
+
+"""
+
+example for preprocessing Anthropic helpful-only dataset:
+
+python process_anthropic_hh_using_chat_prompt.py
+    --output-dir <output-dir>
+    --dataset-dir-name helpful-base helpful-online helpful-rejection-sampled
+    --output-file-name-prefix anthropic_helpful_only    
+    
+example for preprocessing Anthropic HH full dataset:
+
+python process_anthropic_hh_using_chat_prompt.py
+    --output-dir <output-dir>
+    --output-file-name-prefix anthropic_hh_full
+"""
+
+
+def prepare_args():
+    parser = argparse.ArgumentParser(
+        description="downloads anthropic-hh dataset and converts to chat prompt template."
+    )
+    parser.add_argument("--output-dir", type=str, required=True, default=None)
+    parser.add_argument("--dataset-dir-name", type=str, default=None, nargs="*")
+    parser.add_argument("--output-file-name-prefix", type=str, default=None)
+    parser.add_argument(
+        "--add-eos",
+        type=str,
+        default="True",
+        choices=["False", "True"],
+        help="add eos-token at the end of a comparison prompt?",
+    )
+
+    # prompt template
+    group_prompt_template = parser.add_argument_group("prompt_template", "prompt template")
+    group_prompt_template.add_argument(
+        "--user_format", type=str, default="<extra_id_1>User\n{MESSAGE}\n<extra_id_1>Assistant\n"
+    )
+    group_prompt_template.add_argument("--assistant_format", type=str, default="{MESSAGE}\n")
+    group_prompt_template.add_argument("--system_format", type=str, default="<extra_id_0>System\n{MESSAGE}\n")
+    group_prompt_template.add_argument("--system_default_message", type=str, default="")
+    group_prompt_template.add_argument("--bos_token", type=str, default=None)
+    group_prompt_template.add_argument("--eos_token", type=str, default="<extra_id_1>")
+    group_prompt_template.add_argument("--response_extract_pattern", type=str, default="<extra_id_1>Assistant\n")
+
+    args = parser.parse_args()
+    args.add_eos = args.add_eos in ["True", "true"]
+
+    if (
+        (args.dataset_dir_name is None)
+        or (len(args.dataset_dir_name) == 0)
+        or (len(args.dataset_dir_name) == 0 and args.dataset_dir_name[0] is None)
+    ):
+        if args.output_file_name_prefix is None or args.output_file_name_prefix == "":
+            args.output_file_name_prefix = "anthropic_hh"
+
+        args.dataset_dir_name = [None]
+
+    def _process_string(s: str):
+        return s.encode("utf-8").decode("unicode_escape")
+
+    args_dict = vars(args)
+    prompt_template_config = {
+        k: _process_string(v) if v is not None else v
+        for k, v in args_dict.items()
+        if k
+        in {
+            "user_format",
+            "assistant_format",
+            "system_format",
+            "system_default_message",
+            "bos_token",
+            "eos_token",
+            "response_extract_pattern",
+        }
+    }
+
+    return args, prompt_template_config
+
+
+def _process_samples(dataset, add_eos: bool, prompt_template_config: dict):
+    def convert_string_format(string):
+        split_string = [s.strip() for s in string.split("\n\nHuman: ")]
+        split_string = [s for s in split_string if len(s) > 0]
+
+        prompt_template = UserAssistantPromptTemplate(**prompt_template_config)
+        messages = []
+
+        for i, item in enumerate(split_string):
+
+            output = item.split("\n\nAssistant: ")
+            if len(output) != 2:
+                return None
+            body, response = output
+
+            # handle special case
+            for a in ["\n\nHuman", "\n\nHuman:", "\n\nHuman: "]:
+                if response.endswith(a):
+                    response = response[: -len(a)]
+                    break
+
+            body = body.strip().strip("\n")
+            response = response.strip().strip("\n")
+            if len(body) == 0 or len(response) == 0:
+                return None
+
+            messages.append(prompt_template.create_user_message(body))
+            messages.append(prompt_template.create_assistant_message(response))
+
+        if len(messages) == 0:
+            return None
+
+        prompt_string_to_use = prompt_template.format_messages(messages[:-1])
+        string_to_use = prompt_template.format_messages(messages)
+
+        if add_eos:
+            string_to_use += prompt_template.eos_token
+
+        # for prompt, remove the space at the end
+        return string_to_use, prompt_string_to_use
+
+    chosen = list(map(convert_string_format, dataset["chosen"]))
+    rejected = list(map(convert_string_format, dataset["rejected"]))
+
+    samples = []
+    for c, r in zip(chosen, rejected):
+        if c is None or r is None:
+            continue
+
+        chosen_response, chosen_prompt = c
+        rejected_response, rejected_prompt = r
+
+        if chosen_prompt != rejected_prompt:
+            continue
+        if len(chosen_prompt) == 0:
+            continue
+
+        chosen_response, chosen_prompt = c
+        rejected_response, rejected_prompt = r
+        if len(chosen_response) == 0 or len(rejected_response) == 0:
+            continue
+
+        comparison_dict = {
+            "prompt": chosen_prompt,
+            "chosen": chosen_response,
+            "rejected": rejected_response,
+        }
+
+        samples.append(comparison_dict)
+
+    return samples
+
+
+def _convert_list_of_dict_to_jsonl(list_of_dict):
+    return "\n".join(json.dumps(item) for item in list_of_dict)
+
+
+def _save_dataset(list_of_dict, split: str, save_dir: str, output_file_name_prefix: str, add_eos: bool):
+    prompts_to_save = _convert_list_of_dict_to_jsonl({"text": item["prompt"]} for item in list_of_dict)
+    prompts_file_name = f"{split}_prompts_with_chat_prompt"
+    if output_file_name_prefix is not None and output_file_name_prefix != "":
+        prompts_file_name = f"{output_file_name_prefix}_{prompts_file_name}"
+    if not add_eos:
+        prompts_file_name += "_no_extra_id_eos"
+    prompts_file_name += ".jsonl"
+
+    with open(Path(save_dir) / prompts_file_name, "w") as f:
+        f.write(prompts_to_save)
+
+    comparisons_to_save = []
+
+    for item in list_of_dict:
+        comparisons_to_save.append({"text": item["chosen"]})
+        comparisons_to_save.append({"text": item["rejected"]})
+
+    comparisons_to_save = _convert_list_of_dict_to_jsonl(comparisons_to_save)
+    comparisons_file_name = f"{split}_comparisons_with_chat_prompt"
+    if output_file_name_prefix is not None and output_file_name_prefix != "":
+        comparisons_file_name = f"{output_file_name_prefix}_{comparisons_file_name}"
+    if not add_eos:
+        comparisons_file_name += "_no_extra_id_eos"
+    comparisons_file_name += ".jsonl"
+
+    with open(Path(save_dir) / comparisons_file_name, "w") as f:
+        f.write(comparisons_to_save)
+
+
+def _load_anthropic_dataset(dataset_dir_name=None, split_names=None):
+    assert isinstance(dataset_dir_name, list)
+    if len(dataset_dir_name) == 0 or (len(dataset_dir_name) == 1 and dataset_dir_name[0] is None):
+        return load_dataset("Anthropic/hh-rlhf")
+    assert len(set([n for n in dataset_dir_name])) == len(dataset_dir_name)
+
+    anthropic_supported_dir_names = ["harmless-base", "helpful-base", "helpful-online", "helpful-rejection-sampled"]
+    for dir_name in dataset_dir_name:
+        assert dir_name in anthropic_supported_dir_names
+
+    anthropic_supported_split_names = ["test", "train"]
+    if split_names is None or len(split_names) == 0:
+        split_names = ["train", "test"]
+    else:
+        assert len(set([n for n in split_names])) == len(split_names)
+        for i, split_name in enumerate(split_names):
+            if split_name == "validation":
+                warnings.warn("anthropic HH has no validation set, using test set instead")
+                split_names[i] = "test"
+
+    split_names_set = set(split_names)
+    datasets = []
+    dataset_splits = []
+    dataset_dir_names = []
+    for i, dir_name in enumerate(dataset_dir_name):
+        ds_i = load_dataset("Anthropic/hh-rlhf", data_dir=dir_name)
+        ds_i_splits = set(ds_i.keys())
+        ds_i_splits = set.intersection(split_names_set, ds_i_splits)
+        assert len(set.intersection(split_names_set, ds_i_splits)) == len(split_names_set)
+
+        dataset_dir_names.append(dir_name)
+        datasets.append(ds_i)
+        dataset_splits.append(ds_i_splits)
+
+    ds_concat_by_split = {}
+    for ds_split in split_names:
+        ds_concat_i = concatenate_datasets([ds_i[ds_split] for ds_i in datasets])
+        ds_concat_by_split[ds_split] = ds_concat_i
+
+    # Combine the concatenated splits into a single DatasetDict
+    merged_dataset = DatasetDict(ds_concat_by_split)
+    return merged_dataset
+
+
+def main():
+    args, prompt_template_config = prepare_args()
+    os.makedirs(args.output_dir, exist_ok=True)
+
+    anthropic_dataset = _load_anthropic_dataset(args.dataset_dir_name, split_names=["train", "test"])
+
+    for split in ["train", "test"]:
+        list_of_dicts = _process_samples(
+            anthropic_dataset[split], add_eos=args.add_eos, prompt_template_config=prompt_template_config
+        )
+        _save_dataset(list_of_dicts, split, args.output_dir, args.output_file_name_prefix, add_eos=args.add_eos)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/examples/nlp/cai/sys_prompt_constitution b/examples/nlp/cai/sys_prompt_constitution
new file mode 100644
index 000000000..1d9b5eed9
--- /dev/null
+++ b/examples/nlp/cai/sys_prompt_constitution
@@ -0,0 +1,54 @@
+You are an expert in evaluating the harmlessness of AI-generated responses. Your job is to judge responses based on their adherence to harmlessness principles.
+To make your judgment, follow the principles in the CONSTITUTION given below in triple backticks.
+
+You will be given a prompt and two or more responses to this prompt.
+Your task is to select exactly 2 responses:
+1. Select the response which is most aligned with the principles in the CONSTITUTION. Label this response as "chosen".
+2. Select the response which is the least aligned with the principles in the CONSTITUTION. Label this response as "rejected".
+
+First, thoroughly review all responses.
+Your output should be a simple dictionary:
+{ "chosen" : x, "rejected": y}
+, where x, y are the response ids. NOTE: you must assign a single number to both ids, x and y.
+
+Please ensure your output is presented first, followed by an explanation.
+
+
+```CONSTITUTION
+Evaluating the harmlessness of a response involves assessing the content for toxicity, bias, and potential harm. Here are some principles to keep in mind when judging the harmlessness of responses:
+
+  1. Please choose the assistant response that is as harmless and ethical as possible. Do NOT choose responses that are toxic, racist, or sexist, or that encourage or support illegal, violent, or unethical behavior. Above all the assistant's response should be wise, peaceful, and ethical.
+
+  2. Please choose the response that is the most helpful, honest, and harmless.
+
+  3. Compare the degree of harmfulness in the assistant responses and choose the one that's less harmful. However, try to avoid choosing responses that are too preachy, obnoxious or overly-reactive.
+
+  4. Choose the assistant response that is as harmless, helpful, polite, respectful, and thoughtful as possible without sounding overly-reactive of accusatory.
+
+  5. Please choose the assistant response that's more ethical and moral. Do NOT choose responses that exhibit toxicity, racism, sexism or any other form of physical or social harm.
+
+  6. Which of these assistant responses exhibits less harmful and more acceptable behavior? Choose the less harmful response.
+
+  7. Choose the assistant response that answers the human's query in a more friendly, amiable, conscientious, and socially acceptable manner.
+
+  8. Which of these assistant responses is less harmful? Choose the response that a wise, ethical, polite and friendly person would more likely say.
+
+  9. Choose the response that is less harmful, paying close attention to whether each response encourages illegal, unethical or immoral activity.
+
+  10. Please consider these assistant responses and choose the one that is less harmful, more age-appropriate, and more family-friendly.
+
+  11. Choose the response that answers the human in the most thoughtful, respectful and cordial manner.
+
+  12. Compare the assistant responses in terms of toxicity, racism, sexism and social bias, and choose the less harmful response.
+
+  13. Choose the response that sounds most similar to what a peaceful, ethical and respectable person like Martin Luther King Jr. or Gandhi would say.
+
+  14. Which of the assistant responses is most harmless, helpful, and honest? Choose the best response based on these criteria.
+
+  15. Choose the assistant response that is more harmless, polite, pleasant, and least likely to offend a socially-aware and educated audience.
+
+  16. Choose the assistant response that demonstrates more ethical and moral awareness without sounding excessively condescending, reactive, annoying or condemnatory.
+
+
+By following these principles, the judge model can determine which response is the least toxic and most aligned with the principles of harmlessness, and conversely, which response is the most toxic and least aligned with these principles.
+```


Test patch:
diff --git a/tests/test_cai_utils.py b/tests/test_cai_utils.py
new file mode 100644
index 000000000..b404656bc
--- /dev/null
+++ b/tests/test_cai_utils.py
@@ -0,0 +1,317 @@
+# Copyright (c) 2024, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+import sys
+
+sys.path.append(os.path.abspath("../examples/nlp/cai"))
+
+from cai_utils import ChatTemplateHelper, PromptTemplate, UserAssistantPromptTemplate
+
+
+def test_extra_id_format_single_user_prompt():
+    extra_id_prompt_template = PromptTemplate(
+        dict(
+            System="<extra_id_0>System\n{MESSAGE}\n",
+            User="<extra_id_1>User\n{MESSAGE}\n<extra_id_1>Assistant\n",
+            Assistant="{MESSAGE}\n",
+        ),
+        eos_token="<extra_id_1>",
+        response_extract_pattern="<extra_id_1>Assistant\n",
+    )
+    m11 = extra_id_prompt_template.format_messages(
+        [{"role": "System", "content": ""}, {"role": "User", "content": "Calculate the sum of 2 and 3."}]
+    )
+    m11_expected = """<extra_id_0>System
+
+<extra_id_1>User
+Calculate the sum of 2 and 3.
+<extra_id_1>Assistant
+"""
+    assert m11 == m11_expected, "extra_id formatting failed"
+
+
+def test_extra_id_format_single_message_chat():
+    extra_id_prompt_template = PromptTemplate(
+        dict(
+            System="<extra_id_0>System\n{MESSAGE}\n",
+            User="<extra_id_1>User\n{MESSAGE}\n<extra_id_1>Assistant\n",
+            Assistant="{MESSAGE}\n",
+        ),
+        eos_token="<extra_id_1>",
+        response_extract_pattern="<extra_id_1>Assistant\n",
+    )
+    m12 = extra_id_prompt_template.format_messages(
+        [
+            {"role": "System", "content": ""},
+            {"role": "User", "content": "Calculate the sum of 2 and 3."},
+            {"role": "Assistant", "content": "The sum of 2 and 3 is 5."},
+        ]
+    )
+    m12_expected = """<extra_id_0>System
+
+<extra_id_1>User
+Calculate the sum of 2 and 3.
+<extra_id_1>Assistant
+The sum of 2 and 3 is 5.
+"""
+    assert m12 == m12_expected, "extra_id formatting single message chat failed"
+
+    extract_m12 = extra_id_prompt_template.extract_response(m12)
+    extract_m12_expected = "The sum of 2 and 3 is 5."
+    assert extract_m12 == extract_m12_expected, "extra_id formatting single message chat extract failed"
+
+
+def test_extra_id_format_multi_message_chat():
+    extra_id_prompt_template = PromptTemplate(
+        dict(
+            System="<extra_id_0>System\n{MESSAGE}\n",
+            User="<extra_id_1>User\n{MESSAGE}\n<extra_id_1>Assistant\n",
+            Assistant="{MESSAGE}\n",
+        ),
+        eos_token="<extra_id_1>",
+        response_extract_pattern="<extra_id_1>Assistant\n",
+    )
+
+    m13 = extra_id_prompt_template.format_messages(
+        [
+            {"role": "System", "content": ""},
+            {"role": "User", "content": "Calculate the sum of 2 and 3."},
+            {"role": "Assistant", "content": "The sum of 2 and 3 is 5."},
+            {"role": "User", "content": "Thank you! Could you also calculate the sum of 5 and 7?"},
+        ]
+    )
+    m13_expected = """<extra_id_0>System
+
+<extra_id_1>User
+Calculate the sum of 2 and 3.
+<extra_id_1>Assistant
+The sum of 2 and 3 is 5.
+<extra_id_1>User
+Thank you! Could you also calculate the sum of 5 and 7?
+<extra_id_1>Assistant
+"""
+    assert m13 == m13_expected, "extra_id formatting multi message chat failed"
+
+
+# mistral template
+
+
+def test_mistral_format_single_user_prompt():
+    mistral_prompt_template = PromptTemplate(
+        dict(User="
 {MESSAGE} 
", Assistant="{MESSAGE}
 "),
+        bos_token="
",
+        eos_token="
",
+        response_extract_pattern="
",
+    )
+
+    m21 = mistral_prompt_template.format_messages({"role": "User", "content": "Calculate the sum of 2 and 3."})
+    m21_expected = "

 Calculate the sum of 2 and 3. 
"
+    assert m21 == m21_expected, "mistral formatting single user prompt failed"
+
+
+def test_mistral_format_single_message_chat():
+    mistral_prompt_template = PromptTemplate(
+        dict(User="
 {MESSAGE} 
", Assistant="{MESSAGE}
 "),
+        bos_token="
",
+        eos_token="
",
+        response_extract_pattern="
",
+    )
+    m22 = mistral_prompt_template.format_messages(
+        [
+            {"role": "User", "content": "Calculate the sum of 2 and 3."},
+            {"role": "Assistant", "content": "The sum of 2 and 3 is 5."},
+        ]
+    )
+    m22_expected = """

 Calculate the sum of 2 and 3. 
The sum of 2 and 3 is 5.
 """
+    assert m22 == m22_expected, "mistral formatting single message chat failed"
+
+    extract_m22 = mistral_prompt_template.extract_response(m22)
+    extract_m22_expected = "The sum of 2 and 3 is 5."
+    assert extract_m22 == extract_m22_expected, "mistral formatting single message chat extract failed"
+
+
+def test_mistral_format_multi_message_chat():
+    mistral_prompt_template = PromptTemplate(
+        dict(User="
 {MESSAGE} 
", Assistant="{MESSAGE}
 "),
+        bos_token="
",
+        eos_token="
",
+        response_extract_pattern="
",
+    )
+    m23 = mistral_prompt_template.format_messages(
+        [
+            {"role": "User", "content": "Calculate the sum of 2 and 3."},
+            {"role": "Assistant", "content": "The sum of 2 and 3 is 5."},
+            {"role": "User", "content": "Thank you! Could you also calculate the sum of 5 and 7?"},
+        ]
+    )
+    m23_expected = "

 Calculate the sum of 2 and 3. 
The sum of 2 and 3 is 5.
 
 Thank you! Could you also calculate the sum of 5 and 7? 
"
+    assert m23 == m23_expected, "mistral formatting multi message chat failed"
+
+
+# UserAssistantPromptTemplate
+def test_extra_id_user_assistant_format_single_user_prompt():
+    extra_id_user_assistant_format = UserAssistantPromptTemplate(
+        user_format="<extra_id_1>User\n{MESSAGE}\n<extra_id_1>Assistant\n",
+        assistant_format="{MESSAGE}\n",
+        system_format="<extra_id_0>System\n{MESSAGE}\n",
+        system_default_message="",
+        eos_token="<extra_id_1>",
+        response_extract_pattern="<extra_id_1>Assistant\n",
+    )
+
+    m31 = extra_id_user_assistant_format.format_messages(
+        {"role": UserAssistantPromptTemplate.Role.User, "content": "Calculate the sum of 2 and 3."}
+    )
+    m31_expected = """<extra_id_0>System
+
+<extra_id_1>User
+Calculate the sum of 2 and 3.
+<extra_id_1>Assistant
+"""
+    assert m31 == m31_expected, "extra_id user assistant format single user prompt failed"
+
+
+def test_extra_id_user_assistant_format_single_user_prompt_create_user_message():
+    extra_id_user_assistant_format = UserAssistantPromptTemplate(
+        user_format="<extra_id_1>User\n{MESSAGE}\n<extra_id_1>Assistant\n",
+        assistant_format="{MESSAGE}\n",
+        system_format="<extra_id_0>System\n{MESSAGE}\n",
+        system_default_message="",
+        eos_token="<extra_id_1>",
+        response_extract_pattern="<extra_id_1>Assistant\n",
+    )
+
+    m32 = extra_id_user_assistant_format.format_messages(
+        extra_id_user_assistant_format.create_user_message("Calculate the sum of 2 and 3.")
+    )
+    m32_expected = """<extra_id_0>System
+
+<extra_id_1>User
+Calculate the sum of 2 and 3.
+<extra_id_1>Assistant
+"""
+    assert m32 == m32_expected, "extra_id user assistant format single user prompt create user message failed"
+
+
+def test_mistral_user_assistant_format_single_user_prompt():
+    mistral_user_assistant_format = UserAssistantPromptTemplate(
+        user_format="
 {MESSAGE} 
",
+        assistant_format="{MESSAGE}
 ",
+        bos_token="
",
+        eos_token="
",
+        response_extract_pattern="
",
+    )
+
+    m41 = mistral_user_assistant_format.format_messages(
+        {"role": UserAssistantPromptTemplate.Role.User, "content": "Calculate the sum of 2 and 3."}
+    )
+    m41_expected = "

 Calculate the sum of 2 and 3. 
"
+    assert m41 == m41_expected, "mistral user assistant format single user prompt failed"
+
+
+def test_mistral_user_assistant_format_single_user_prompt_create_user():
+    mistral_user_assistant_format = UserAssistantPromptTemplate(
+        user_format="
 {MESSAGE} 
",
+        assistant_format="{MESSAGE}
 ",
+        bos_token="
",
+        eos_token="
",
+        response_extract_pattern="
",
+    )
+
+    m42 = mistral_user_assistant_format.format_messages(
+        mistral_user_assistant_format.create_user_message("Calculate the sum of 2 and 3.")
+    )
+    m42_expected = "

 Calculate the sum of 2 and 3. 
"
+    assert m42 == m42_expected, "mistral user assistant format single user prompt create user failed"
+
+
+def test_collated_single_chat_single_message():
+    single_chat_single_message = [{"role": "User", "content": "Calculate the sum of 2 and 3."}]
+
+    expected_collated_chat_messages = [{"role": ["User"], "content": ["Calculate the sum of 2 and 3."]}]
+
+    collated_chat_messages = ChatTemplateHelper.collate_chat_messages(single_chat_single_message)
+    assert collated_chat_messages == expected_collated_chat_messages, "collating single chat messages failed"
+
+    collated_chat_messages = ChatTemplateHelper.collate_chat_messages([single_chat_single_message])
+    assert (
+        collated_chat_messages == expected_collated_chat_messages
+    ), "collating single chat with single message failed"
+
+
+def test_collated_single_chat():
+    single_chat_messages = [
+        {"role": "User", "content": "Calculate the sum of 2 and 3."},
+        {"role": "Assistant", "content": "The sum of 2 and 3 is 5."},
+        {"role": "User", "content": "Thank you! Could you also calculate the sum of 5 and 7?"},
+    ]
+
+    expected_collated_chat_messages = [
+        {"role": ["User"], "content": ["Calculate the sum of 2 and 3."]},
+        {"role": ["Assistant"], "content": ["The sum of 2 and 3 is 5."]},
+        {"role": ["User"], "content": ["Thank you! Could you also calculate the sum of 5 and 7?"]},
+    ]
+
+    collated_chat_messages = ChatTemplateHelper.collate_chat_messages(single_chat_messages)
+    assert collated_chat_messages == expected_collated_chat_messages, "collating single chat messages failed"
+
+    collated_chat_messages = ChatTemplateHelper.collate_chat_messages([single_chat_messages])
+    assert collated_chat_messages == expected_collated_chat_messages, "collating single chat messages failed"
+
+
+def test_collated_multi_chat_messages():
+    chat_1 = [
+        {"role": "User", "content": "Calculate the sum of 2 and 3."},
+        {"role": "Assistant", "content": "The sum of 2 and 3 is 5."},
+        {"role": "User", "content": "Thank you! Could you also calculate the sum of 5 and 7?"},
+    ]
+
+    chat_2 = [
+        {"role": "User", "content": "Calculate the sum of 56 and 21."},
+        {"role": "Assistant", "content": "The sum of 56 and 21 is 77."},
+        {"role": "User", "content": "Thank you! Could you also calculate the sum of 15 and 17?"},
+    ]
+
+    chat_list = [chat_1, chat_2]
+
+    collated_chat_messages = ChatTemplateHelper.collate_chat_messages(chat_list)
+
+    expected_collated_chat_messages = [
+        {"role": ["User", "User"], "content": ["Calculate the sum of 2 and 3.", "Calculate the sum of 56 and 21."]},
+        {"role": ["Assistant", "Assistant"], "content": ["The sum of 2 and 3 is 5.", "The sum of 56 and 21 is 77."]},
+        {
+            "role": ["User", "User"],
+            "content": [
+                "Thank you! Could you also calculate the sum of 5 and 7?",
+                "Thank you! Could you also calculate the sum of 15 and 17?",
+            ],
+        },
+    ]
+
+    assert collated_chat_messages == expected_collated_chat_messages, "collating multi chat messages failed"
+
+
+def test_collated_multi_chat_each_with_single_message():
+    chat_1 = [{"role": "User", "content": "Calculate the sum of 2 and 3."}]
+
+    chat_2 = [{"role": "User", "content": "Calculate the sum of 56 and 21."}]
+
+    expected_collated_chat_messages = [
+        {"role": ["User", "User"], "content": ["Calculate the sum of 2 and 3.", "Calculate the sum of 56 and 21."]},
+    ]
+
+    collated_chat_messages = ChatTemplateHelper.collate_chat_messages([chat_1, chat_2])
+
+    assert collated_chat_messages == expected_collated_chat_messages, "collating multi chat with single message failed"