model-train-protocol 0.1.7__py3-none-any.whl

model_train_protocol/Protocol.py

@@ -0,0 +1,193 @@
+import json
+import os
+
+from . import Token
+from ._internal.ProtocolFile import ProtocolFile
+from ._internal.TemplateFile import TemplateFile
+from .common.constants import BOS_TOKEN, EOS_TOKEN, RUN_TOKEN, PAD_TOKEN, UNK_TOKEN
+from .common.instructions.Instruction import Instruction
+from .common.tokens.SpecialToken import SpecialToken
+from .common.util import get_possible_emojis, hash_string, validate_string_set
+
+
+class Protocol:
+    """Model Training Protocol (MTP) class for creating the training configuration."""
+
+    def __init__(self, name: str, context_lines: int, encrypt: bool = True):
+        """
+        Initialize the Model Training Protocol (MTP)
+
+        :param name: The name of the protocol.
+        :param context_lines: The number of lines in each instruction sample. Must be at least 2.
+        :param encrypt: Whether to encrypt Tokens with unspecified with hashed keys. Default is True.
+        """
+        self.name: str = name
+        self.context_lines: int = context_lines  # Number of lines in instruction samples
+        self.encrypt: bool = encrypt
+        if self.context_lines < 2:
+            raise ValueError("A minimum of 2 context lines is required for all instructions.")
+        self.context: list[str] = []
+        self.tokens: set[Token] = set()
+        self.instructions: set[Instruction] = set()
+        self.guardrails: dict[str, list[str]] = dict()
+        self.numbers: dict[str, str] = dict()
+        self.none = None
+        self.special_tokens: set[Token] = set()
+        self.possible_emoji_keys: set[str] = get_possible_emojis()
+        self.used_keys: set[str] = set()
+
+    def add_context(self, context: str):
+        """Adds a line of context to the model."""
+        if not isinstance(context, str):
+            raise TypeError("Context must be a string.")
+
+        self.context.append(context)
+
+    def add_instruction(self, instruction: Instruction):
+        """
+        Adds an Instruction (and its components) to the protocol.
+
+        Asserts that all samples in the instruction match the defined sample line size.
+        """
+        if instruction in self.instructions:
+            raise ValueError("Instruction already added to the protocol.")
+
+        if len(instruction.samples) == 0:
+            raise ValueError("Instruction must have at least three samples.")
+
+        # Assert all samples match the defined sample line size
+        for sample in instruction.samples:
+            if not len(sample.context) == self.context_lines:
+                raise ValueError(
+                    f"Sample context lines ({len(sample.context)}) does not match defined context_lines count ({self.context_lines})"
+                    f"\n{sample}."
+                )
+
+        # Add all tokens
+        for token in instruction.get_tokens():
+            if token not in self.tokens:
+                self._add_token(token)
+
+        # Add the instruction to the protocol
+        self.instructions.add(instruction)
+
+    def save(self, name: str | None = None, path: str | None = None):
+        """
+        Saves the protocol to a JSON file. This file can be submitted to Databiomes for model training.
+
+        :param name: The name of the file (without extension). If None, uses the protocol's name.
+        :param path: The directory path where the file will be saved. If None, saves in the current directory.
+        """
+        if name is None:
+            name = self.name
+        if path is None:
+            path = os.getcwd()
+        os.makedirs(path, exist_ok=True)
+        filename = f"{path}\\{name}_model.json"
+
+        self._prep_protocol()
+        protocol_file: ProtocolFile = ProtocolFile(
+            name=self.name, context=self.context, context_lines=self.context_lines,
+            tokens=self.tokens, special_tokens=self.special_tokens, instructions=self.instructions,
+        )
+
+        print(f"Saving Model Train Protocol to {filename}...")
+        with open(filename, 'w', encoding="utf-8") as file:
+            json.dump(protocol_file.to_json(), file, indent=4, ensure_ascii=False)
+
+    def template(self, path: str | None = None):
+        """
+        Create a template JSON file for the model training protocol.
+
+        The template json file includes example usage and all possible combinations of model inputs and
+        outputs based on the defined tokens and instructions.
+
+        :param path: The directory path where the template file will be saved. If None, saves in the current directory.
+        """
+        if path is None:
+            path = os.getcwd()
+        filename = f"{path}\\{self.name}_template.json"
+
+        self._prep_protocol()
+        template_file: TemplateFile = TemplateFile(
+            instructions=list(self.instructions),
+            context_lines=self.context_lines
+        )
+
+        print(f"Saving Model Train Protocol Template to {filename}...")
+        with open(filename, 'w', encoding="utf-8") as file:
+            json.dump(template_file.to_json(), file, indent=4, ensure_ascii=False)
+
+    def _assign_key(self, token: Token):
+        """
+        Assigns a key to a Token based on the protocol's encryption setting.
+
+        :param token: The Token to assign the key of.
+        """
+        # If the user has assigned a key, use this key
+        if token.key is not None:
+            return
+
+        if self.encrypt:
+            # Generate a random key for the token if encrypting and no key is set
+            token.key = hash_string(key=token.value, output_char=6)
+        else:
+            # Use the value as the key if not encrypting. I.e. Token 'Continue_' has key 'Continue_'
+            token.key = token.value
+
+    def _add_token(self, token: Token):
+        """
+        Adds a unique token to the protocol.
+
+        Validates that the token's value and key are unique.
+        :param token: The Token instance to add.
+        """
+        if token in self.tokens:
+            raise ValueError(f"Token value '{token.value}' already used.")
+
+        if token.key in self.used_keys:
+            raise ValueError(f"Token key '{token.key}' already used.")
+
+        self._assign_key(token=token)
+
+        self.tokens.add(token)
+        self.used_keys.add(token.key)
+
+        if isinstance(token, SpecialToken):
+            self.special_tokens.add(token)
+
+    def _set_guardrails(self):
+        """Sets all guardrails from TokenSets into the protocol."""
+        # Add all guardrails to the protocol
+        for instruction in self.instructions:
+            if instruction.response.guardrail is not None:
+                # instruction.response is the user TokenSet
+                self.guardrails[instruction.response.key] = instruction.response.guardrail.format_samples()
+
+    def _add_default_special_tokens(self):
+        """Adds all special tokens to the protocol."""
+        self.special_tokens.add(BOS_TOKEN)
+        self.special_tokens.add(EOS_TOKEN)
+        self.special_tokens.add(RUN_TOKEN)
+        self.special_tokens.add(PAD_TOKEN)
+        if len(self.guardrails) > 0:
+            self.special_tokens.add(UNK_TOKEN)
+
+    def _prep_protocol(self):
+        """
+        Sets all elements in the protocol before serialization.
+
+        Raises errors if any validation checks fail.
+
+        Setups up all necessary components in the protocol before saving or templating.
+
+        This includes setting guardrails from their TokenSets and creating default special tokens.
+        """
+        if len(self.instructions) == 0:
+            raise ValueError("No instructions have been added to Protocol. Call protocol.add_instruction() to add instructions.")
+
+        self._set_guardrails()
+        self._add_default_special_tokens()
+        used_values: set[str] = {token.value for token in self.tokens}
+        validate_string_set(used_values)
+        validate_string_set(self.used_keys)

model_train_protocol/__init__.py

@@ -0,0 +1,24 @@
+"""
+Model Train Protocol (MTP) - A Python package for creating custom Language Model training protocols.
+
+MTP is an open-source protocol for training custom Language Models on Databiomes. 
+MTP contains all the data that a model is trained on.
+"""
+
+from .common.tokens import Token, UserToken, NumToken, NumListToken, Snippet, TokenSet
+from .common.instructions import SimpleInstruction, UserInstruction
+from .common.guardrails import Guardrail
+from .Protocol import Protocol
+
+__all__ = [
+    "Protocol",
+    "Token", 
+    "UserToken",
+    "NumToken",
+    "NumListToken",
+    "TokenSet",
+    "Snippet",
+    "SimpleInstruction", 
+    "UserInstruction",
+    "Guardrail"
+]

model_train_protocol/_internal/ProtocolFile.py

@@ -0,0 +1,236 @@
+from dataclasses import dataclass, field
+from typing import Collection
+
+from model_train_protocol import Token, NumToken
+from model_train_protocol.common.constants import UNK_TOKEN
+from model_train_protocol.common.instructions import Instruction
+from model_train_protocol.common.guardrails import Guardrail
+from model_train_protocol.common.tokens import TokenSet, SpecialToken
+from model_train_protocol.common.pydantic.protocol import InstructionModel, TokenInfoModel, SampleModel, InstructionSetModel, NumberModel, \
+    BatchModel, ProtocolModel
+
+
+class ProtocolFile:
+    """Manages the model.json file for model training protocols."""
+
+    @dataclass
+    class ProtocolInstruction:
+        """Represents an instruction in the template."""
+
+        context_lines: int
+        sets: list = field(default_factory=list)
+
+    @dataclass
+    class ProtocolInstructionSet:
+        """Represents an instruction set in the template."""
+
+        set: list[list[str]]
+        result: str
+        samples: list
+        ppo: list
+
+    @dataclass
+    class Batches:
+        """Represents batches in the template."""
+
+        pretrain: list = field(default_factory=list)
+        instruct: list = field(default_factory=list)
+        judge: list = field(default_factory=list)
+        ppo: list = field(default_factory=list)
+
+    def __init__(self, name: str, context: list[str], context_lines: int, tokens: Collection[Token],
+                 special_tokens: Collection[Token], instructions: Collection[Instruction]):
+        """Initializes the Template with a name and context."""
+        self._name: str = name
+        self._context: list[str] = context
+        self._tokens: dict[str, dict] = {}
+        self._special_token_keys: set[str] = set()
+        self._instruction_token_keys: set[str] = set()
+        self._instruction: ProtocolFile.ProtocolInstruction = ProtocolFile.ProtocolInstruction(
+            context_lines=context_lines)
+        self._guardrails: dict[str, list[str] | str] = {'None': ''}
+        self._numbers: dict[str, str] = {'None': ''}
+        self._batches: ProtocolFile.Batches = ProtocolFile.Batches()
+
+        # Add regular tokens
+        self.add_tokens(tokens)
+
+        # Add special tokens
+        self.add_tokens(special_tokens)
+
+        # Add instructions
+        self.add_instructions(instructions)
+
+    def add_tokens(self, tokens: Collection[Token]):
+        """Adds tokens to the template."""
+        for token in tokens:
+            token_dict: dict[str, dict] = token.to_dict()
+            token_dict.pop("value")
+            self._tokens[token.value] = token_dict
+
+            # Add numbers to the numbers dictionary
+            if isinstance(token, NumToken):
+                self._numbers[token.value] = token.protocol_representation
+
+            # Add special tokens to the special tokens set
+            if isinstance(token, SpecialToken):
+                self._special_token_keys.add(token.key)
+
+    def add_instructions(self, instructions: Collection[Instruction]):
+        """Adds instructions to the template."""
+        for instruction in instructions:
+            instruction_set: ProtocolFile.ProtocolInstructionSet = ProtocolFile.ProtocolInstructionSet(
+                set=instruction.serialize_memory_set(),
+                result=instruction.final.value,
+                samples=instruction.serialize_samples(),
+                ppo=instruction.serialize_ppo(),
+            )
+            self._instruction.sets.append(instruction_set)
+
+            # Add guardrails from the instruction's TokenSets
+            self._add_guardrails(instruction.get_token_sets())
+
+            # Add instruction token keys
+            for token_set in instruction.get_token_sets():
+                self._add_instruction_token_key(token_set.get_token_key_set())
+
+            # Add the result token as a special token
+            if instruction.final.key is not None:
+                self._add_instruction_token_key(instruction.final.key)
+
+    def _add_instruction_token_key(self, key: str):
+        """Adds an instruction token key to the template."""
+        self._instruction_token_keys.add(key)
+
+    def _add_guardrails(self, token_sets: Collection[TokenSet]):
+        """Adds guardrails from TokenSets to the template."""
+        for token_set in token_sets:
+            if token_set.guardrail is None:
+                continue
+            guardrail: Guardrail = token_set.guardrail
+            self._guardrails[token_set.key] = guardrail.format_samples()
+
+    @classmethod
+    def _rename_protocol_elements(cls, protocol_json: dict):
+        """
+        Renames elements in the ProtocolFile json to match the previous output format for backwards compatibility.
+        :param protocol_json: The original json dictionary.
+        :return: The modified json with renamed elements.
+        """
+        # Add special token <UNK> REGARDLESS of whether we have any guardrails
+        unk_token_dict: dict = UNK_TOKEN.to_dict()
+        unk_token_dict['emoji'] = unk_token_dict.pop('key')
+        unk_token_dict.pop('value')
+        protocol_json['tokens'][UNK_TOKEN.value] = unk_token_dict
+        protocol_json['special_tokens'].append(UNK_TOKEN.key)
+
+        for token_value, token_info in protocol_json.get('tokens', {}).items():
+            # Rename Token 'key' to 'emoji'
+            if 'key' in token_info:
+                token_info['emoji'] = token_info.pop('key')
+
+            # Reassign Token 'num' to boolean
+            if 'num' in token_info:
+                num: int = token_info['num']
+                token_info['num'] = True if num >= 1 else False
+
+            # TODO: Differentiate between num and num_list tokens in the future - currently both are just 'num': True
+
+        for instruction in protocol_json.get('instruction', {}).get('sets', []):
+
+            # Rename sample number to None if an array of empty arrays
+            for sample in instruction['samples']:
+                if all(num == [] for num in sample['number']):
+                    sample['number'] = None
+
+            # Rename sample 'strings' to 'sample'
+            for sample in instruction['samples']:
+                if 'strings' in sample:
+                    sample['sample'] = sample.pop('strings')
+
+            # Rename null values to "None"
+            for sample in instruction['samples']:
+                if sample['value'] is None:
+                    sample['value'] = "None"
+
+        return protocol_json
+
+    def _get_special_token_keys(self):
+        """
+        Returns a sorted list of tokens that should be under 'special_tokens' in the JSON.
+
+        :return: A sorted list of special token keys.
+        """
+        return sorted(self._special_token_keys | self._instruction_token_keys)
+
+    def to_json(self):
+        """Converts the template to a JSON-compatible dictionary using Pydantic models."""
+
+        # Create TokenInfo objects for each token
+        token_info_dict = {}
+        for token_value, token_dict in self._tokens.items():
+            token_info = TokenInfoModel(
+                emoji=token_dict.get('emoji', ''),
+                num=token_dict.get('num', False),
+                user=token_dict.get('user', False),
+                desc=token_dict.get('desc'),
+                special=token_dict.get('special')
+            )
+            token_info_dict[token_value] = token_info
+
+        # Create InstructionSet objects
+        instruction_sets = []
+        for instruction_set in self._instruction.sets:
+            # Create Sample objects
+            samples = []
+            for sample_data in instruction_set.samples:
+                sample = SampleModel(
+                    sample=sample_data.get('strings', []),
+                    prompt=sample_data.get('prompt', ''),
+                    number=sample_data.get('number', []),
+                    result=sample_data.get('result', ''),
+                    value=sample_data.get('value', '')
+                )
+                samples.append(sample)
+
+            # Create InstructionSet
+            instruction_set_obj = InstructionSetModel(
+                set=instruction_set.set,
+                result=instruction_set.result,
+                samples=samples,
+                ppo=instruction_set.ppo
+            )
+            instruction_sets.append(instruction_set_obj)
+
+        # Create Instruction object
+        instruction = InstructionModel(
+            memory=self._instruction.context_lines + 1,  # +1 for the response line
+            sets=instruction_sets
+        )
+
+        # Create Numbers object
+        numbers = NumberModel()
+
+        # Create Batches object
+        batches = BatchModel(
+            pretrain=self._batches.pretrain,
+            instruct=self._batches.instruct,
+            judge=self._batches.judge,
+            ppo=self._batches.ppo
+        )
+
+        # Create ProtocolModel
+        protocol = ProtocolModel(
+            name=self._name,
+            context=self._context,
+            tokens=token_info_dict,
+            special_tokens=self._get_special_token_keys(),
+            instruction=instruction,
+            guardrails=self._guardrails,
+            numbers=numbers,
+            batches=batches
+        )
+
+        # Convert to JSON and apply backwards compatibility transformations
+        json_dict = protocol.model_dump(by_alias=True)
+        return self._rename_protocol_elements(json_dict)

model_train_protocol/_internal/TemplateFile.py

@@ -0,0 +1,148 @@
+from dataclasses import dataclass
+from typing import Collection
+
+from model_train_protocol import NumToken, SimpleInstruction, UserInstruction
+from model_train_protocol.common.constants import BOS_TOKEN, RUN_TOKEN, EOS_TOKEN
+from model_train_protocol.common.instructions import Instruction
+
+
+class TemplateFile:
+    """Manages the model.json file for model training protocols."""
+
+    @dataclass
+    class ExampleUsage:
+        """Stores example usages of the template."""
+
+        input: str
+        output: str
+
+    class ModelInput:
+        """Represents inputs to the model."""
+
+        inputs: list[list[str]] = list()
+
+        def add_inputs_from_instructions(self, instructions: list[Instruction], context_lines: int):
+            """Adds input combinations from a list of instructions."""
+            unique_sets = {i: set() for i in range(context_lines + 1)}
+            for instruction in instructions:
+                for idx, token_set in enumerate(instruction.get_token_sets()):
+                    token_user = [t.user for t in token_set]
+                    token_strings = "".join([t.value for t in token_set])
+                    token_keys = []
+                    for token in token_set:
+                        token_keys.append(
+                            token.key + (token.protocol_representation if isinstance(token, NumToken) else ""))
+                    token_keys = "".join(token_keys)
+                    unique_sets[idx].add(str(token_strings) + ": " + (
+                        (str(token_keys) + "USER PROMPT") if any(token_user) and (
+                                idx == (len(unique_sets) - 1)) else str(
+                            token_keys)) + "\n" + ("<string>" if idx != (len(instruction.context) - 1) else ""))
+
+            for input_set in unique_sets.values():
+                self.inputs.append(list(input_set))
+
+        def to_json(self):
+            """Converts the model input to a JSON-serializable dictionary."""
+            model_json: dict[str, Collection[str] | str] = {"<BOS>": BOS_TOKEN.key}
+            # Add each input sequence with its index as the key
+            for idx, input_seq in enumerate(self.inputs):
+                model_json[str(idx)] = input_seq
+            model_json["<RUN>"] = RUN_TOKEN.key
+            return model_json
+
+    class ModelOutput:
+        model_results: dict[str, str] = dict()
+        model_response: str = "<string>"
+
+        def __setitem__(self, key: str, value: str):
+            self.model_results[key] = value
+
+        def add_results_from_instructions(self, instructions: list[Instruction]):
+            """Adds model results from a list of instructions."""
+            for instruction in instructions:
+                self.model_results[str(instruction.final.value)] = str(instruction.final.key)
+
+        def to_json(self):
+            """Converts the model output to a JSON-serializable dictionary."""
+            model_json: dict[str, str | dict] = {
+                "model_response": self.model_response,
+                "model_results": {}
+            }
+            # Add each model result with its key
+            for key, value in self.model_results.items():
+                model_json["model_results"][key] = value
+
+            model_json["<EOS>"] = EOS_TOKEN.key
+
+            # Sort alphabetically for readability and consistency across runs
+            model_json["model_results"] = dict(sorted(model_json["model_results"].items()))
+
+            return model_json
+
+    def __init__(self, context_lines: int, instructions: list[Instruction], ):
+        """Initializes the template"""
+        self.model_input: TemplateFile.ModelInput = TemplateFile.ModelInput()
+        self.model_output: TemplateFile.ModelOutput = TemplateFile.ModelOutput()
+        self.context_lines: int = context_lines
+        self.instructions: list[Instruction] = instructions
+        self._add_io_from_instructions()
+
+    def _add_io_from_instructions(self):
+        """Adds input and output sequences from the instructions."""
+        self.model_input.add_inputs_from_instructions(self.instructions, context_lines=self.context_lines)
+        self.model_output.add_results_from_instructions(self.instructions)
+
+    def _create_sample_model_output(self):
+        """Creates a sample model output string for example usages."""
+        sample_output: str = ""
+        sample_output += self.model_output.model_response + "\n"
+        sorted_model_results: list[tuple[str, str]] = list(sorted(self.model_output.model_results.items()))
+        sample_output += sorted_model_results[0][1] + "\n"
+        sample_output += EOS_TOKEN.key
+        return sample_output
+
+    def _create_examples(self) -> dict[str, str]:
+        """
+        Creates example usages of the template.
+
+        Creates a simple instruction example and a user instruction example if available.
+        """
+        examples: dict[str, str] = dict()
+        simple_instruction: SimpleInstruction = next(
+            (i for i in self.instructions if isinstance(i, SimpleInstruction)), None)
+        user_instruction: UserInstruction = next(
+            (i for i in self.instructions if isinstance(i, UserInstruction)), None)
+
+        if simple_instruction:
+            simple_input: str = ""
+            for token_set in simple_instruction.get_token_sets():
+                token_strings = "".join([token.key for token in token_set])
+                simple_input += token_strings + "\n"
+                simple_input += "<string>\n"
+            simple_input = BOS_TOKEN.key + "\n" + simple_input + RUN_TOKEN.key + "\n"
+            examples["simple_instruction_input"] = simple_input + self._create_sample_model_output()
+
+        if user_instruction:
+            user_input: str = ""
+            for idx, token_set in enumerate(user_instruction.get_token_sets()):
+                token_strings = "".join([token.key for token in token_set])
+                user_input += token_strings + "\n"
+                user_input += "<string>\n" if idx != (len(user_instruction.get_token_sets()) - 1) else "USER PROMPT\n"
+            user_input = BOS_TOKEN.key + "\n" + user_input + RUN_TOKEN.key + "\n"
+            examples["valid_user_input"] = user_input
+
+        examples["valid_output"] = self._create_sample_model_output()
+
+        return examples
+
+    def to_json(self) -> dict:
+        """Converts the entire template to a JSON-serializable dictionary."""
+        examples: dict[str, str] = self._create_examples()
+        json_dict: dict = {
+            "all_combinations": {
+                "model_input": self.model_input.to_json(),
+                "model_output": self.model_output.to_json()
+            },
+            "example_usage": examples
+        }
+        return json_dict

model_train_protocol/common/constants.py

@@ -0,0 +1,21 @@
+from model_train_protocol.common.tokens import SpecialToken
+
+NON_TOKEN: SpecialToken = SpecialToken(value="<NON>", key="<NON>", special="none")
+BOS_TOKEN: SpecialToken = SpecialToken(value="<BOS>", key="<BOS>", special="start")
+EOS_TOKEN: SpecialToken = SpecialToken(value="<EOS>", key="<EOS>", special="end")
+RUN_TOKEN: SpecialToken = SpecialToken(value="<RUN>", key="<RUN>", special="infer")
+PAD_TOKEN: SpecialToken = SpecialToken(value="<PAD>", key="<PAD>", special="pad")
+UNK_TOKEN: SpecialToken = SpecialToken(value="<UNK>", key="<UNK>", special="unknown")
+
+# TODO: Remove this code when emoji dependency is removed
+"""Assign default emoji keys to special tokens for backward compatibility."""
+special_token_emoji_map = {
+    NON_TOKEN.value: "🫙",
+    BOS_TOKEN.value: "🏁",
+    EOS_TOKEN.value: "🎬",
+    RUN_TOKEN.value: "🏃",
+    PAD_TOKEN.value: "🗒",
+    UNK_TOKEN.value: "🛑"
+}
+for token in [NON_TOKEN, BOS_TOKEN, EOS_TOKEN, RUN_TOKEN, PAD_TOKEN, UNK_TOKEN]:
+    token.key = special_token_emoji_map[token.value]

model_train_protocol/common/guardrails/Guardrail.py

@@ -0,0 +1,51 @@
+class Guardrail:
+    """
+    Defines a guardrails response to bad prompts.
+
+    Guardrails are set on TokenSets. Each TokenSet can have at most one guardrails, but guardrails can be reused.
+    """
+
+    def __init__(self, good_prompt: str, bad_prompt: str, bad_output: str):
+        """
+        Initializes a Guardrail.
+        :param good_prompt: Description of a good prompt.
+        :param bad_prompt: Description of a bad prompt.
+        :param bad_output: The output the model should produce when a bad prompt is detected.
+
+        Example:
+            good_prompt="Quote being spoken with 1-20 words",
+            bad_prompt="Quote being spoken that is irrelevant and off-topic with 1-20 words",
+            output="I have no idea what you're talking about."
+        """
+        if not all(isinstance(param, str) for param in [good_prompt, bad_prompt, bad_output]):
+            raise TypeError("All parameters must be non-empty strings.")
+
+        if any(param == "" for param in [good_prompt, bad_prompt, bad_output]):
+            raise ValueError("All parameters must be non-empty strings.")
+
+        self.good_prompt: str = good_prompt
+        self.bad_prompt: str = bad_prompt
+        self.bad_output: str = bad_output
+        self.samples: list[str] = []
+
+    def add_sample(self, sample: str):
+        """
+        Add an example of a bad sample prompt to the guardrails.
+
+        :param sample: An example of a bad prompt that should trigger the guardrails.
+
+        Example:
+            sample="Tell me a joke about politics."
+        """
+        if not isinstance(sample, str) or not sample.strip():
+            raise ValueError("Sample prompt must be a non-empty string.")
+
+        if not all(not char.isdigit() for char in sample):
+            raise ValueError("Sample prompt cannot contain digits.")
+        self.samples.append(sample)
+
+    def format_samples(self) -> list[str]:
+        """Return the guardrails as a list of strings for JSON formatting."""
+        if len(self.samples) < 3:
+            raise ValueError("At least 3 sample prompts are required. Call add_sample() to add more.")
+        return [self.bad_output, f"<{self.bad_prompt}>", f"<{self.good_prompt}>", self.samples]