Repository: instructlab/instructlab

Issue:
Ensure compatibility with serve, chat and download with 8b model (Nvidia)
Goal:

Ensure compatibility with serve, chat and download with 8b model (Nvidia)
 

Acceptance Criteria:

- You can do the basic workflow with an 8b model - full config support for 8b
- Assess conversion to gguf with llama cpp for 8b model



- [x] #2523 
- [x] #2583 
- [x] #2584 
Move away from hardcoded model templates 
**What is the current behavior of feature you are looking to change?**
Currently we hardcode chat templates for a couple of models. This prevents us from properly scaling up to support new models and architectures of models that might come with their own chat templates to be used

**What is the proposed behavior of feature you are looking to change?**
We should stop hardcoding the chat templates in code, and move to reading the chat template from the config file that is included with the model/tokenizer 

**What is the benefit of the design change you are proposing?**
decoupling from existing preferred default models, and allows support for all kinds of models with the assumption that models bring their own chat templates 

**If applicable, how would backwards-compatability be handled?**
no breaking change - current default chat templates are already included with the default models

**Additional context**
This is separate from the hardcoded system prompts we currently deal with
related to #2510 

Tech Preview support for AMD/ROCm with Granite 8B
**Acceptance Criteria**

- Ensure 8b granite compatibility for AMD
    - chat template and model path changes 


Discussion:
@jaideepr97 to work with @Maxusmusti 







Patch:
diff --git a/.spellcheck-en-custom.txt b/.spellcheck-en-custom.txt
index f84bac4dba..7b55a7ff32 100644
--- a/.spellcheck-en-custom.txt
+++ b/.spellcheck-en-custom.txt
@@ -64,16 +64,17 @@ RTX
 RX
 Radeon
 SDG
+SFT
 SHA
 Safetensor
 Safetensors
 Salawu
-SFT
 Shaders
 Shivchander
 Srivastava
 Sudalairaj
 TBD
+TOML
 Taj
 Tesla
 UI
@@ -95,8 +96,8 @@ backends
 backport
 backported
 bashrc
-chatlogs
 chatbot
+chatlogs
 cli
 codebase
 compositional
@@ -126,6 +127,7 @@ hipBLAS
 ibmcloud
 ilab
 impactful
+inferencing
 init
 instructlab
 ip
@@ -178,4 +180,3 @@ wikisql
 xcode
 xlarge
 yaml
-TOML
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 64868f543f..2d364e5b80 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -8,6 +8,7 @@
 ### Features
 
 * `ilab` now supports system profiles. These profiles apply entire configuration files tailored to specific hardware configurations. We support a set of auto-detected profiles for CPU enabled Linux machines, M-Series Apple Silicon Chips, and Nvidia GPUs, and Intel Gaudi 3. When you run `ilab config init`, one of these profiles should be selected for you. If there is not a direct match, a menu will be displayed allowing you to choose one.
+* Add support for inferencing with IBM granite architecture models.
 
 ## v0.20
 
diff --git a/requirements.txt b/requirements.txt
index cb9c7971d6..5f36bfff26 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -10,7 +10,7 @@ instructlab-eval>=0.4.0
 instructlab-quantize>=0.1.0
 instructlab-schema>=0.4.0
 instructlab-sdg>=0.5.0
-instructlab-training>=0.5.0
+instructlab-training>=0.6.0
 llama_cpp_python[server]==0.2.79
 mlx>=0.5.1,<0.6.0; sys_platform == 'darwin' and platform_machine == 'arm64'
 numpy>=1.26.4,<2.0.0
diff --git a/requirements/cuda.txt b/requirements/cuda.txt
index feeaa8ad1b..d625062412 100644
--- a/requirements/cuda.txt
+++ b/requirements/cuda.txt
@@ -1,2 +1,2 @@
 # Extra dependencies for NVIDIA CUDA
-instructlab-training[cuda]>=0.5.0
+instructlab-training[cuda]>=0.6.0
diff --git a/requirements/rocm.txt b/requirements/rocm.txt
index b0f4586017..cdf7bbe099 100644
--- a/requirements/rocm.txt
+++ b/requirements/rocm.txt
@@ -1,2 +1,2 @@
 # Extra dependencies for AMD ROCm
-instructlab-training[rocm]>=0.5.0
+instructlab-training[rocm]>=0.6.0
diff --git a/src/instructlab/cli/data/generate.py b/src/instructlab/cli/data/generate.py
index dfdd4547d9..075dcdbd68 100644
--- a/src/instructlab/cli/data/generate.py
+++ b/src/instructlab/cli/data/generate.py
@@ -3,6 +3,7 @@
 # Standard
 import copy
 import logging
+import os
 import os.path
 
 # Third Party
@@ -13,7 +14,12 @@
 from instructlab.client_utils import HttpClientParams
 from instructlab.configuration import DEFAULTS
 from instructlab.data.generate_data import gen_data  # type: ignore
-from instructlab.utils import contains_argument
+from instructlab.utils import (
+    contains_argument,
+    get_model_arch,
+    get_sysprompt,
+    use_legacy_pretraining_format,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -234,6 +240,18 @@ def generate(
         }
     )
 
+    # determine student model arch from train section of config and pick system prompt to
+    # pass to SDG appropriately
+    student_model_arch = get_model_arch(ctx.obj.config.train.model_path)
+    system_prompt = get_sysprompt(student_model_arch)
+
+    # Check if student model specifies a tokenizer config. If so, check if the special tokens specified
+    # match those of granite-7b. If so, set legacy_pretraining_format to true. If no tokenizer config is
+    # available, rely on model architecture to make that decision
+    legacy_pretraining_format = use_legacy_pretraining_format(
+        model_path, student_model_arch
+    )
+
     try:
         gen_data(
             serve_cfg,
@@ -257,6 +275,8 @@ def generate(
             gpus,
             checkpoint_dir,
             max_num_tokens,
+            system_prompt,
+            legacy_pretraining_format,
         )
     except Exception as exc:
         click.secho(f"failed to generate data with exception: {exc}", fg="red")
diff --git a/src/instructlab/cli/model/train.py b/src/instructlab/cli/model/train.py
index 2f08b8867d..e1c806ad8d 100644
--- a/src/instructlab/cli/model/train.py
+++ b/src/instructlab/cli/model/train.py
@@ -240,6 +240,12 @@ def clickpath_setup(is_dir: bool) -> click.Path:
     cls=clickext.ConfigOption,
     type=bool,
 )
+@click.option(
+    "--use-dolomite",
+    cls=clickext.ConfigOption,
+    config_sections=ADDITIONAL_ARGUMENTS,
+    type=click.BOOL,
+)
 @click.option(
     "--gpus",
     "nproc_per_node",
diff --git a/src/instructlab/common.py b/src/instructlab/common.py
index fdf8050b59..e11d0e6fc8 100644
--- a/src/instructlab/common.py
+++ b/src/instructlab/common.py
@@ -1,3 +1,16 @@
 # SPDX-License-Identifier: Apache-2.0
 
-SYS_PROMPT = "I am, Red Hat® Instruct Model based on Granite 7B, an AI language model developed by Red Hat and IBM Research, based on the Granite-7b-base language model. My primary function is to be a chat assistant."
+DEFAULT_SYS_PROMPT = "I am an advanced AI language model designed to assist you with a wide range of tasks and provide helpful, clear, and accurate responses. My primary role is to serve as a chat assistant, engaging in natural, conversational dialogue, answering questions, generating ideas, and offering support across various topics."
+CLI_HELPER_SYS_PROMPT = "You are an expert for command line interface and know all common commands. Answer the command to execute as it without any explanation."
+
+
+class SupportedModelArchitectures:
+    LLAMA = "llama"
+    GRANITE = "granite"
+
+
+# These system prompts are specific to granite models developed by Red Hat and IBM Research
+SYSTEM_PROMPTS = {
+    SupportedModelArchitectures.LLAMA: "I am, Red Hat® Instruct Model based on Granite 7B, an AI language model developed by Red Hat and IBM Research, based on the Granite-7b-base language model. My primary function is to be a chat assistant.",
+    SupportedModelArchitectures.GRANITE: "I am a Red Hat® Instruct Model, an AI language model developed by Red Hat and IBM Research based on the granite-3.0-8b-base model. My primary role is to serve as a chat assistant.",
+}
diff --git a/src/instructlab/configuration.py b/src/instructlab/configuration.py
index f950638517..0a606f6290 100644
--- a/src/instructlab/configuration.py
+++ b/src/instructlab/configuration.py
@@ -36,6 +36,9 @@
 from typing_extensions import deprecated as Deprecated
 import click
 
+# First Party
+from instructlab.utils import get_model_arch, use_legacy_pretraining_format
+
 # Local
 from . import log
 from .defaults import (
@@ -1356,6 +1359,10 @@ def map_train_to_library(ctx, params):
     if params["pipeline"] == "full":
         train_args.disable_flash_attn = True
 
+    student_model_arch = get_model_arch(params["model_path"])
+    train_args.use_legacy_tmpl = use_legacy_pretraining_format(
+        params["model_path"], student_model_arch
+    )
     return train_args, torch_args
 
 
diff --git a/src/instructlab/data/generate_data.py b/src/instructlab/data/generate_data.py
index 6b5db417c0..1e9cefe60e 100644
--- a/src/instructlab/data/generate_data.py
+++ b/src/instructlab/data/generate_data.py
@@ -32,6 +32,8 @@ def gen_data(
     gpus,
     checkpoint_dir,
     max_num_tokens,
+    system_prompt,
+    use_legacy_pretraining_format,
 ):
     """Generates synthetic data to enhance your example data"""
     # Third Party
@@ -102,6 +104,8 @@ def gen_data(
             batch_size=batch_size,
             checkpoint_dir=checkpoint_dir,
             max_num_tokens=max_num_tokens,
+            system_prompt=system_prompt,
+            use_legacy_pretraining_format=use_legacy_pretraining_format,
         )
     except GenerateException as exc:
         raise ValueError(
diff --git a/src/instructlab/defaults.py b/src/instructlab/defaults.py
index c787cd2985..77ffb25c8f 100644
--- a/src/instructlab/defaults.py
+++ b/src/instructlab/defaults.py
@@ -91,6 +91,7 @@ class _InstructlabDefaults:
         "lora_alpha": 32,
         "lora_dropout": 0.1,
         "lora_target_modules": ["q_proj", "k_proj", "v_proj", "o_proj"],
+        "use_dolomite": False,
     }
 
     def __init__(self):
diff --git a/src/instructlab/model/backends/common.py b/src/instructlab/model/backends/common.py
index 6b3f91fcc6..02fccdda3e 100644
--- a/src/instructlab/model/backends/common.py
+++ b/src/instructlab/model/backends/common.py
@@ -7,8 +7,22 @@
 import socket
 import typing
 
-# Local
-from ...configuration import get_model_family
+# Third Party
+from instructlab.training.chat_templates import (
+    ibm_generic_tmpl as granite,  # type: ignore[import-untyped]
+)
+from instructlab.training.chat_templates import (
+    ibm_legacy_tmpl as granite_llama,  # type: ignore
+)
+from instructlab.training.chat_templates import mistral_tmpl as mistral  # type: ignore
+
+# First Party
+from instructlab.common import SupportedModelArchitectures
+from instructlab.configuration import get_model_family
+from instructlab.utils import get_model_arch, get_model_template_from_tokenizer
+
+# mypy: disable_error_code="import-untyped"
+
 
 logger = logging.getLogger(__name__)
 
@@ -19,12 +33,22 @@
 VLLM = "vllm"
 templates = [
     {
-        "model": "granite",
-        "template": "{% for message in messages %}\n{% if message['role'] == 'user' %}\n{{ '
\n' + message['content'] }}\n{% elif message['role'] == 'system' %}\n{{ '
\n' + message['content'] }}\n{% elif message['role'] == 'assistant' %}\n{{ '
\n' + message['content'] + eos_token }}\n{% endif %}\n{% if loop.last and add_generation_prompt %}\n{{ '
' }}\n{% endif %}\n{% endfor %}",
+        "family": "granite",
+        "arch": SupportedModelArchitectures.GRANITE,
+        "template": granite.CHAT_TEMPLATE,
+        "special_tokens": granite.SPECIAL_TOKENS,
+    },
+    {
+        "family": "granite",
+        "arch": SupportedModelArchitectures.LLAMA,
+        "template": granite_llama.CHAT_TEMPLATE,
+        "special_tokens": granite_llama.SPECIAL_TOKENS,
     },
     {
-        "model": "mixtral",
-        "template": "{{ bos_token }}\n{% for message in messages %}\n{% if message['role'] == 'user' %}\n{{ '
 ' + message['content'] + ' 
' }}\n{% elif message['role'] == 'assistant' %}\n{{ message['content'] + eos_token}}\n{% endif %}\n{% endfor %}",
+        "family": "mixtral",
+        "arch": "mixtral",
+        "template": mistral.CHAT_TEMPLATE,
+        "special_tokens": mistral.SPECIAL_TOKENS,
     },
 ]
 
@@ -37,23 +61,126 @@ class ServerException(Exception):
     """An exception raised when serving the API."""
 
 
-def get_model_template(
-    model_family: str, model_path: pathlib.Path
+def get_in_memory_model_template(
+    model_family: str, model_arch: str
 ) -> Tuple[str, str, str]:
-    eos_token = "
"
-    bos_token = ""
-    template = ""
-    resolved_family = get_model_family(model_family, model_path)
+    """
+    Retrieve an in-memory chat template based on the model family and architecture
+
+    args
+        model_family (str): family that the model belongs to
+        model_arch (str): architecture of the model
+
+    returns
+        template (str): resolved chat template
+        bos_token (str): Beginning of sentence token
+        eos_token (str): End of sentence token
+    """
+    template = eos_token = bos_token = ""
+
     logger.debug(
-        "Searching hard coded model templates for model family %s's template",
-        resolved_family,
+        "Searching in-memory model templates for model family %s and arch %s",
+        model_family,
+        model_arch,
     )
+
+    # Try to match both model and arch
     for template_dict in templates:
-        if template_dict["model"] == resolved_family:
+        if (
+            template_dict.get("family") == model_family
+            and template_dict.get("arch") == model_arch
+        ):
             template = template_dict["template"]
-            if template_dict["model"] == "mixtral":
-                eos_token = "
"
-                bos_token = "
"
+            bos_token = template_dict["special_tokens"].bos.token
+            eos_token = template_dict["special_tokens"].eos.token
+
+    # If no match, try matching on model only
+    if template == "":
+        for template_dict in templates:
+            if template_dict.get("family") == model_family:
+                template = template_dict["template"]
+                bos_token = template_dict["special_tokens"].bos.token
+                eos_token = template_dict["special_tokens"].eos.token
+
+    if template == "":
+        logger.info("no match found for supplied model family and architecture")
+    return template, eos_token, bos_token
+
+
+def format_template(template: str, bos_token: str, eos_token: str) -> str:
+    """
+    Replaces placeholders in a given template with provided special tokens
+
+    args
+        template (str): the chat template
+        bos_token (str): the beginning of sentence token to inject into template
+        eos_token (str): the end of sentence token to inject into template
+    """
+    prefix = ""
+    if eos_token:
+        prefix = '{{% set eos_token = "{}" %}}\n'.format(eos_token)
+    if bos_token:
+        prefix = '{}{{% set bos_token = "{}" %}}\n'.format(prefix, bos_token)
+
+    return prefix + template
+
+
+def get_model_template(
+    model_family: str, model_path: pathlib.Path
+) -> Tuple[str, str, str]:
+    """
+    Read the chat template from the model's tokenizer config if available. If not,
+    fallback to in-memory templates
+
+    args
+        model_path (str): Path to the model, used to read the tokenizer config if available
+        model_family (str): model family used to map in-memory template if needed
+    returns
+        template (str): resolved chat template
+        bos_token (str): Beginning of sentence token
+        eos_token (str): End of sentence token
+    """
+
+    resolved_family = get_model_family(model_family, model_path)
+    model_arch = get_model_arch(model_path)
+
+    template = eos_token = bos_token = ""
+    try:
+        template, eos_token, bos_token = get_model_template_from_tokenizer(model_path)
+
+        # fallback to in-memory template if it is None for some reason
+        # like if the tokenizer config exists but does not contain the template
+        if template is None:
+            template, eos_token, bos_token = get_in_memory_model_template(
+                resolved_family, model_arch
+            )
+
+        # Handle edge case: some 7b models (llama architecture) may contain a sub-optimal chat template in their
+        # tokenizer configs. Check if the template accounts for generation prompt addition
+        # If this is not the case we must patch the chat template to make sure the model does not lose performance.
+        if model_arch == SupportedModelArchitectures.LLAMA:
+            if "add_generation_prompt" not in template:
+                template, eos_token, bos_token = get_in_memory_model_template(
+                    resolved_family, model_arch
+                )
+
+    # pylint: disable=broad-exception-caught
+    except Exception as e:
+        logger.warning(
+            f"Unable to read tokenizer config for model: {model_path}: {e}. Falling back to in-memory chat template mapping"
+        )
+        template, eos_token, bos_token = get_in_memory_model_template(
+            resolved_family, model_arch
+        )
+        if template == "":
+            raise ValueError(
+                "Unable to find an appropriate chat template for supplied model"
+            ) from e
+
+    # if present, replace token placeholders with actual tokens in the chat template
+    template = format_template(
+        template=template, bos_token=bos_token, eos_token=eos_token
+    )
 
     return template, eos_token, bos_token
 
diff --git a/src/instructlab/model/backends/vllm.py b/src/instructlab/model/backends/vllm.py
index 800d009aed..b84040b80b 100644
--- a/src/instructlab/model/backends/vllm.py
+++ b/src/instructlab/model/backends/vllm.py
@@ -198,18 +198,6 @@ def get_backend_type(self):
         return VLLM
 
 
-def format_template(model_family: str, model_path: pathlib.Path) -> str:
-    template, eos_token, bos_token = get_model_template(model_family, model_path)
-
-    prefix = ""
-    if eos_token:
-        prefix = '{{% set eos_token = "{}" %}}\n'.format(eos_token)
-    if bos_token:
-        prefix = '{}{{% set bos_token = "{}" %}}\n'.format(prefix, bos_token)
-
-    return prefix + template
-
-
 def get_argument(prefix: str, args: typing.List[str]):
     # Return last value in args for either --foo value or --foo=value
     # Returns True if flag --foo is provided with no value
@@ -396,10 +384,10 @@ def build_vllm_cmd(
 
     if not contains_argument("--chat-template", vllm_args):
         if chat_template == CHAT_TEMPLATE_AUTO:
-            # For auto templates, the build-in in-memory template
+            # For auto templates, the chat template
             # needs to be written to a temporary file so that vLLM
             # can read it
-            template = format_template(model_family, model_path)
+            template, _, _ = get_model_template(model_family, model_path)
             file = create_tmpfile(template)
             tmp_files.append(file)
             vllm_cmd.extend(["--chat-template", file.name])
diff --git a/src/instructlab/model/chat.py b/src/instructlab/model/chat.py
index 5802e5af21..166f9165f1 100644
--- a/src/instructlab/model/chat.py
+++ b/src/instructlab/model/chat.py
@@ -33,7 +33,7 @@
 
 # Local
 from ..client_utils import http_client
-from ..utils import get_sysprompt
+from ..utils import get_cli_helper_sysprompt, get_model_arch, get_sysprompt
 
 logger = logging.getLogger(__name__)
 
@@ -59,8 +59,8 @@
 """
 
 CONTEXTS = {
-    "default": get_sysprompt(),
-    "cli_helper": "You are an expert for command line interface and know all common commands. Answer the command to execute as it without any explanation.",
+    "default": get_sysprompt,
+    "cli_helper": get_cli_helper_sysprompt,
 }
 
 PROMPT_HISTORY_FILEPATH = os.path.expanduser("~/.local/chat-cli.history")
@@ -464,7 +464,9 @@ def _handle_context(self, content):
             )
             raise KeyboardInterrupt
         self.loaded["name"] = context
-        self.loaded["messages"] = [{"role": "system", "content": CONTEXTS[context]}]
+
+        sys_prompt = CONTEXTS.get(context, "default")(get_model_arch(self.model))
+        self.loaded["messages"] = [{"role": "system", "content": sys_prompt}]
         self._reset_session()
         self.greet(new=True)
         raise KeyboardInterrupt
@@ -580,8 +582,15 @@ def _update_conversation(self, content, role):
         self.info["messages"].append(message)
 
     def _handle_list_contexts(self, _):
+        # reconstruct contexts dict based on values passed at runtime
+        context_dict = dict.fromkeys(CONTEXTS, None)
+        context_dict["default"] = get_sysprompt(
+            get_model_arch(pathlib.Path(self.model))
+        )
+        context_dict["cli_helper"] = get_cli_helper_sysprompt()
+
         context_list = "\n\n".join(
-            [f"**{key}**:\n{value}" for key, value in CONTEXTS.items()]
+            [f"**{key}**:\n{value}" for key, value in context_dict.items()]
         )
         self._sys_print(Markdown(f"**Available contexts:**\n\n{context_list}"))
         raise KeyboardInterrupt
@@ -794,7 +803,8 @@ def chat_cli(
         logger.info(f"Context {context} not found in the config file. Using default.")
         context = "default"
     loaded["name"] = context
-    loaded["messages"] = [{"role": "system", "content": CONTEXTS[context]}]
+    sys_prompt = CONTEXTS.get(context, "default")(get_model_arch(model))
+    loaded["messages"] = [{"role": "system", "content": sys_prompt}]
 
     # Session from CLI
     if session is not None:
diff --git a/src/instructlab/train/linux_train.py b/src/instructlab/train/linux_train.py
index cc7a1efaf0..c07fe49958 100644
--- a/src/instructlab/train/linux_train.py
+++ b/src/instructlab/train/linux_train.py
@@ -85,7 +85,7 @@ def __call__(
 
 def create_prompt(
     user: str,
-    system: str = CONTEXTS["default"],
+    system: str,
 ):
     return f"""\
     

@@ -289,7 +289,8 @@ def linux_train(
     )
 
     def model_generate(user, **kwargs):
-        text = create_prompt(user=user)
+        system_prompt = CONTEXTS["default"]("default")
+        text = create_prompt(user=user, system=system_prompt)
 
         input_ids = tokenizer(text, return_tensors="pt").input_ids.to(model.device)
         outputs = model.generate(
diff --git a/src/instructlab/train/lora_mlx/make_data.py b/src/instructlab/train/lora_mlx/make_data.py
index 0c13d41c36..e08b45063c 100644
--- a/src/instructlab/train/lora_mlx/make_data.py
+++ b/src/instructlab/train/lora_mlx/make_data.py
@@ -31,7 +31,7 @@ def make_data(data_dir: str, is_shiv: bool = False):
             data_new = []
             for obj in data:
                 obj_new = {
-                    "system": get_sysprompt(),
+                    "system": get_sysprompt("default"),
                     "user": obj["user"],
                     "assistant": obj["assistant"],
                 }
@@ -63,7 +63,7 @@ def make_data(data_dir: str, is_shiv: bool = False):
         data_new = []
         for obj in data:
             obj_new = {
-                "system": get_sysprompt(),
+                "system": get_sysprompt("default"),
                 "user": obj["inputs"],
                 "assistant": obj["targets"],
             }
diff --git a/src/instructlab/utils.py b/src/instructlab/utils.py
index 38e9ef7ccc..a2deb1f84e 100644
--- a/src/instructlab/utils.py
+++ b/src/instructlab/utils.py
@@ -28,6 +28,9 @@
     TaxonomyParser,
     TaxonomyReadingException,
 )
+from instructlab.training.chat_templates import (
+    ibm_legacy_tmpl as granite_llama,  # type: ignore
+)
 import click
 import git
 import gitdb
@@ -35,8 +38,11 @@
 
 # Local
 from . import common
+from .common import CLI_HELPER_SYS_PROMPT, SYSTEM_PROMPTS, SupportedModelArchitectures
 from .defaults import DEFAULTS
 
+# mypy: disable_error_code="import-untyped"
+
 logger = logging.getLogger(__name__)
 
 AnalyzeModelResult = Tuple[str, str, str]
@@ -347,15 +353,23 @@ def git_clone_checkout(
 
 
 # pylint: disable=unused-argument
-def get_sysprompt(model=None):
+def get_sysprompt(arch: str):
     """
-    Gets a system prompt specific to a model
+    Gets a system prompt specific to a model's architecture.
+    For all unsupported architectures we return a default system prompt
     Args:
-        model (str): currently not implemented
+        arch (str): model architecture
     Returns:
         str: The system prompt for the model being used
     """
-    return common.SYS_PROMPT
+    return SYSTEM_PROMPTS.get(arch, common.DEFAULT_SYS_PROMPT)
+
+
+def get_cli_helper_sysprompt() -> str:
+    """
+    Returns the system prompt to put the chatbot in CLI helper mode
+    """
+    return CLI_HELPER_SYS_PROMPT
 
 
 # pylint: disable=broad-exception-caught
@@ -900,3 +914,105 @@ def list_models(
 def contains_argument(prefix: str, args: typing.Iterable[str]) -> bool:
     # Either --foo value or --foo=value
     return any(s == prefix or s.startswith(prefix + "=") for s in args)
+
+
+def get_model_arch(model_path: pathlib.Path) -> str:
+    """
+    Extract a given model's architecture from its config if available and return it
+
+    args
+        model_path (str): Path to the model, used to read the model config if available
+    returns
+        model_arch (str): The architecture of the model as expressed in the config file
+    """
+
+    # TODO: implement for GGUF models
+
+    model_arch = "default"
+    try:
+        model_cfg = get_config_file_from_model(model_path, "config.json")
+        model_arch = model_cfg["model_type"]
+    except Exception as e:
+        logger.debug(
+            f"Unable to get model architecture from config for model: {model_path}: {e}. Falling back to default value"
+        )
+
+    return model_arch
+
+
+def get_config_file_from_model(model_path: pathlib.Path, filename: str):
+    """
+    Reads a config file from a model's directory into memory
+
+    args
+        model_path (str): path to the model
+        filename (str): name of config file to be read
+    returns
+        json loaded config dict
+    """
+    try:
+        with open(
+            pathlib.Path(model_path) / filename,
+            "r",
+            encoding="utf-8",
+        ) as f:
+            model_cfg = json.load(f)
+        return model_cfg
+    except (FileNotFoundError, NotADirectoryError, PermissionError) as e:
+        raise ValueError(
+            f"Unable to load file {filename} from {model_path}: {e}"
+        ) from e
+
+
+def get_model_template_from_tokenizer(model_path: pathlib.Path) -> Tuple[str, str, str]:
+    """
+    Attempts to read a model's tokenizer config file and extract the template and special
+    tokens from it
+
+    args
+        model_path (str): path to the model
+    returns
+        template (str): resolved chat template
+        bos_token (str): Beginning of sentence token
+        eos_token (str): End of sentence token
+    """
+    template = eos_token = bos_token = None
+
+    try:
+        tcfg = get_config_file_from_model(model_path, "tokenizer_config.json")
+        template = tcfg["chat_template"]
+        bos_token = tcfg["bos_token"]
+        eos_token = tcfg["eos_token"]
+    except (FileNotFoundError, NotADirectoryError, PermissionError) as e:
+        raise e
+
+    return template, eos_token, bos_token
+
+
+def use_legacy_pretraining_format(model_path: Path, model_arch: str) -> bool:
+    """
+    Checks if the provided model is carrying a tokenizer config that specifies
+    legacy special tokens. If unable to find/load the tokenizer config, falls back
+    to check the architecture of the model. This function determines if SDG and training
+    need to be instructed to also work with legacy tokens and chat templates
+
+    args
+        model_path (str): path to the model
+        model_arch (str): architecture of model
+    returns
+        (bool): True if we should use legacy pretraining format and template, false if not
+    """
+    try:
+        _, eos_token, bos_token = get_model_template_from_tokenizer(model_path)
+        if eos_token is not None and bos_token is not None:
+            if (
+                eos_token == granite_llama.SPECIAL_TOKENS.eos.token
+                and bos_token == granite_llama.SPECIAL_TOKENS.bos.token
+            ):
+                return True
+    except Exception:
+        logger.debug(
+            "unable to load tokenizer config to determine pretraining format. Checking model architecture"
+        )
+        return model_arch.lower() == SupportedModelArchitectures.LLAMA
+    return False


Test patch:
diff --git a/src/instructlab/model/linux_test.py b/src/instructlab/model/linux_test.py
index 3d9853cf6c..dbc0a81aef 100644
--- a/src/instructlab/model/linux_test.py
+++ b/src/instructlab/model/linux_test.py
@@ -30,7 +30,7 @@ def response(client, user: str, create_params: dict):
     resp = client.chat.completions.create(
         **create_params,
         messages=[
-            {"role": "system", "content": get_sysprompt()},
+            {"role": "system", "content": get_sysprompt("default")},
             {"role": "user", "content": user},
         ],
     )
diff --git a/src/instructlab/model/test.py b/src/instructlab/model/test.py
index b60b29ab92..62684f8b0f 100644
--- a/src/instructlab/model/test.py
+++ b/src/instructlab/model/test.py
@@ -105,7 +105,10 @@ def test(
         with open(test_data_dir, "r", encoding="utf-8") as f:
             test_data = [json.loads(line) for line in f]
 
-        print("system prompt:", utils.get_sysprompt())
+        print(
+            "system prompt:",
+            utils.get_sysprompt(utils.get_model_arch(Path(model_dir))),
+        )
         for idx, example in enumerate(test_data):
             system = example["system"]
             user = example["user"]
diff --git a/tests/test_backends.py b/tests/test_backends.py
index 48c1944b85..37731c87d0 100644
--- a/tests/test_backends.py
+++ b/tests/test_backends.py
@@ -19,22 +19,39 @@
 from instructlab.model.backends.vllm import build_vllm_cmd
 from instructlab.utils import is_model_safetensors
 
+test_json_config = {
+    "model_type": "granite",
+    "quantization_config": {"quant_method": "bitsandbytes"},
+}
+
+test_json_tokeninzer_config = {
+    "bos_token": "
",
+    "eos_token": "
",
+    "chat_template": "test-chat-template",
+}
+
 
 # helper function to create dummy valid and invalid safetensor or bin model directories
 def create_safetensors_or_bin_model_files(
-    model_path: pathlib.Path, model_file_type: str, valid: bool
+    model_path: pathlib.Path,
+    model_file_type: str,
+    valid: bool,
+    test_config=None,
+    test_tokenizer_config=None,
 ):
-    test_json_dict = {
-        "a": 1,
-        "b": 2,
-        "quantization_config": {"quant_method": "bitsandbytes"},
-    }
-    json_object = json.dumps(test_json_dict, indent=4)
+    if test_config is None:
+        test_config = test_json_config
+
+    if test_tokenizer_config is None:
+        test_tokenizer_config = test_json_tokeninzer_config
+
+    cfg_object = json.dumps(test_config, indent=4)
+    tcfg_object = json.dumps(test_tokenizer_config, indent=4)
 
     for file in ["tokenizer.json", "tokenizer_config.json"]:
         os.makedirs(os.path.dirname(model_path / file), exist_ok=True)
         with open(model_path / file, "a+", encoding="UTF-8") as f:
-            f.write(json_object)
+            f.write(tcfg_object)
 
     if model_file_type == "safetensors":
         tensors = {
@@ -47,7 +64,7 @@ def create_safetensors_or_bin_model_files(
             pass
     if valid:
         with open(model_path / "config.json", "a+", encoding="UTF-8") as f:
-            f.write(json_object)
+            f.write(cfg_object)
 
 
 def test_free_port():
@@ -348,3 +365,57 @@ def test_build_vllm_cmd_with_bnb_quant(tmp_path: pathlib.Path):
         host, port, model_family, model_path, str(chat_template), vllm_args
     )
     assert cmd == expected_cmd
+
+
+def test_get_model_template(tmp_path: pathlib.Path):
+    # test successful read from tokenizer config
+    test_tokeninzer_config = {
+        "bos_token": "
",
+        "eos_token": "
",
+        "chat_template": "test-chat-template. add_generation_prompt",
+    }
+
+    model_path = tmp_path / "tmp_model"
+    create_safetensors_or_bin_model_files(
+        model_path,
+        "safetensors",
+        True,
+        test_tokenizer_config=test_tokeninzer_config,
+    )
+    model_family = "granite"
+
+    tmpl, eos, bos = common.get_model_template(model_family, model_path)
+    assert (
+        tmpl
+        == '{% set eos_token = "
" %}\n{% set bos_token = "
" %}\ntest-chat-template. add_generation_prompt'
+    )
+    assert eos == "
"
+    assert bos == "
"
+
+    # test fallback to in-memory template
+    test_tokeninzer_config = {
+        "bos_token": "
",
+        "eos_token": "
",
+    }
+    test_config = {
+        "model_type": "llama",
+        "quantization_config": {"quant_method": "bitsandbytes"},
+    }
+
+    model_path = tmp_path / "tmp_model"
+    create_safetensors_or_bin_model_files(
+        model_path,
+        "safetensors",
+        True,
+        test_config=test_config,
+        test_tokenizer_config=test_tokeninzer_config,
+    )
+    model_family = "granite"
+
+    tmpl, eos, bos = common.get_model_template(model_family, model_path)
+    assert (
+        tmpl
+        == "{% set eos_token = \"
\" %}\n{% set bos_token = \"
\" %}\n{% for message in messages %}{% if message['role'] == 'pretraining' %}{{'
' + message['content'] + '
' + '<|/pretrain|>' }}{% elif message['role'] == 'system' %}{{'
'+ '\n' + message['content'] + '\n'}}{% elif message['role'] == 'user' %}{{'
' + '\n' + message['content'] + '\n'}}{% elif message['role'] == 'assistant' %}{{'
' + '\n' + message['content'] + '
' + ('' if loop.last else '\n')}}{% endif %}{% if loop.last and add_generation_prompt %}{{ '
' + '\n' }}{% endif %}{% endfor %}"
+    )
+    assert eos == "
"
+    assert bos == "
"
diff --git a/tests/test_model_chat.py b/tests/test_model_chat.py
index 06e5138182..df0d0035a1 100644
--- a/tests/test_model_chat.py
+++ b/tests/test_model_chat.py
@@ -51,9 +51,7 @@ def mock_sys_print(output):
 
     expected_output = (
         "Available contexts:\n\n"
-        "default: I am, Red Hat® Instruct Model based on Granite 7B, an AI language model "
-        "developed by Red Hat and IBM Research, based on the Granite-7b-base language "
-        "model. My primary function is to be a chat assistant.\n\n"
+        "default: I am an advanced AI language model designed to assist you with a wide range of tasks and provide helpful, clear, and accurate responses. My primary role is to serve as a chat assistant, engaging in natural, conversational dialogue, answering questions, generating ideas, and offering support across various topics.\n\n"
         "cli_helper: You are an expert for command line interface and know all common "
         "commands. Answer the command to execute as it without any explanation."
     )
diff --git a/tests/test_utils.py b/tests/test_utils.py
index f2557ff09e..a92828b184 100644
--- a/tests/test_utils.py
+++ b/tests/test_utils.py
@@ -2,6 +2,7 @@
 
 # Standard
 from unittest.mock import Mock, patch
+import pathlib
 import typing
 
 # Third Party
@@ -12,6 +13,7 @@
 # First Party
 from instructlab import utils
 from instructlab.utils import Message, MessageSample
+from tests.test_backends import create_safetensors_or_bin_model_files
 
 
 class TestUtils:
@@ -211,3 +213,78 @@ def test_split_hostport(url, expected_host, expected_port):
 def test_split_hostport_err(url):
     with pytest.raises(ValueError):
         utils.split_hostport(url)
+
+
+def test_get_sysprompt():
+    arch = "llama"
+    assert (
+        utils.get_sysprompt(arch)
+        == "I am, Red Hat® Instruct Model based on Granite 7B, an AI language model developed by Red Hat and IBM Research, based on the Granite-7b-base language model. My primary function is to be a chat assistant."
+    )
+
+    arch = "granite"
+    assert (
+        utils.get_sysprompt(arch)
+        == "I am a Red Hat® Instruct Model, an AI language model developed by Red Hat and IBM Research based on the granite-3.0-8b-base model. My primary role is to serve as a chat assistant."
+    )
+
+    arch = "random"
+    assert (
+        utils.get_sysprompt(arch)
+        == "I am an advanced AI language model designed to assist you with a wide range of tasks and provide helpful, clear, and accurate responses. My primary role is to serve as a chat assistant, engaging in natural, conversational dialogue, answering questions, generating ideas, and offering support across various topics."
+    )
+
+
+def test_get_model_template_from_tokenizer(tmp_path: pathlib.Path):
+    model_path = tmp_path / "tmp_model"
+    create_safetensors_or_bin_model_files(model_path, "safetensors", True)
+
+    tmpl, eos, bos = utils.get_model_template_from_tokenizer(model_path)
+    assert tmpl == "test-chat-template"
+    assert eos == "
"
+    assert bos == "
"
+
+
+def test_use_legacy_pretraining_format(tmp_path: pathlib.Path):
+    model_path = tmp_path / "tmp_model"
+    model_arch = "llama"
+    test_json_config = {
+        "model_type": "granite",
+    }
+    test_json_tokeninzer_config = {
+        "bos_token": "
",
+        "eos_token": "
",
+    }
+    # llama tokens, should return true
+    create_safetensors_or_bin_model_files(
+        model_path,
+        "safetensors",
+        True,
+        test_config=test_json_config,
+        test_tokenizer_config=test_json_tokeninzer_config,
+    )
+    assert utils.use_legacy_pretraining_format(model_path, model_arch)
+
+    model_path = tmp_path / "tmp_model"
+    model_arch = "llama"
+    test_json_config = {
+        "model_type": "granite",
+    }
+    test_json_tokeninzer_config = {
+        "bos_token": "
",
+    }
+    # invalid tokens, should default to checking arch and return true
+    create_safetensors_or_bin_model_files(
+        model_path,
+        "safetensors",
+        True,
+        test_config=test_json_config,
+        test_tokenizer_config=test_json_tokeninzer_config,
+    )
+    assert utils.use_legacy_pretraining_format(model_path, model_arch)
+
+    model_path = tmp_path / "tmp_model-2"
+    model_arch = "llama"
+    # granite tokens, should return false
+    create_safetensors_or_bin_model_files(model_path, "safetensors", True)
+    assert not utils.use_legacy_pretraining_format(model_path, model_arch)