fusion-bench 0.2.5__py3-none-any.whl → 0.2.7__py3-none-any.whl

fusion_bench/mixins/fabric_training.py

@@ -0,0 +1,320 @@
+import itertools
+import logging
+import os
+from abc import abstractmethod
+from typing import TYPE_CHECKING, Literal, Union
+
+import torch
+from torch import Tensor, nn
+from tqdm.auto import tqdm
+
+from .lightning_fabric import LightningFabricMixin
+
+if TYPE_CHECKING:
+    from lightning.fabric.wrappers import (
+        _FabricDataLoader,
+        _FabricModule,
+        _FabricOptimizer,
+    )
+
+log = logging.getLogger(__name__)
+
+
+class FabricTrainingMixin(LightningFabricMixin):
+    """
+    This is a general purpose mixin for training a model with PyTorch Lightning.
+    """
+
+    _latest_saved_checkpoint_global_step: int = -1
+    """The global step index of the latest saved checkpoint."""
+    _expected_total_steps: int = None
+    """The expected total number of steps of the entire training."""
+    is_training: bool
+    """Whether the training is in progress. If set to False, the training will stop."""
+    epoch_idx: int
+    """The epoch index, which is the number of epochs completed."""
+    global_step_idx: int
+    """The global step index, which is the number of parameter update steps."""
+    max_epochs: int
+    """Max number of epochs of the entire training."""
+    max_steps: int
+    """Max number of parameter update steps of the entire training."""
+    max_steps_per_epoch: int
+    """Max number of parameter update steps per epoch."""
+    gradient_clip_algorithm: Literal["value", "norm"]
+    """The algorithm to clip gradients. Available options: 'value', 'norm'."""
+    gradient_clip_val: float
+    """The value to clip gradients. If None, no clipping is applied."""
+    accumulate_grad_batches: int
+    """The number of gradient accumulation steps. The effective global batch size is `the batch size per device` x `the number of devices` x `the number of gradient accumulation steps`."""
+    lr_scheduler_interval: Literal["step", "epoch"]
+    """The interval to run the learning rate scheduler. Available options: 'step', 'epoch'."""
+    lr_scheduler_frequency: int
+    """The frequency to run the learning rate scheduler."""
+    checkpoint_save_interval: Literal["step", "epoch"]
+    """The interval to save the model checkpoint. Available options: 'step', 'epoch'."""
+    checkpoint_save_frequency: int
+    """The frequency to save the model checkpoint."""
+
+    def clip_gradients_if_needed(self, model, optimizer):
+        """
+        Clips gradients if the gradient clipping value is set.
+
+        Args:
+            model (nn.Module): The model whose gradients need to be clipped.
+            optimizer (torch.optim.Optimizer): The optimizer used for training.
+        """
+        fabric = self.fabric
+
+        if self.gradient_clip_val is not None:
+            if self.gradient_clip_algorithm == "value":
+                fabric.clip_gradients(model, optimizer, clip_val=self.gradient_clip_val)
+            elif self.gradient_clip_algorithm == "norm":
+                fabric.clip_gradients(model, optimizer, max_norm=self.gradient_clip_val)
+            else:
+                raise ValueError(
+                    f"Unknown gradient clip algorithm: {self.gradient_clip_algorithm}. Available options: 'value', 'norm'"
+                )
+
+    def compute_expected_total_steps(
+        self, train_dataloader: torch.utils.data.DataLoader
+    ):
+        """
+        Computes the expected total number of steps for the entire training.
+
+        Args:
+            train_dataloader (torch.utils.data.DataLoader): The dataloader for the training data.
+        """
+        # compute expected total steps
+        self._expected_total_steps = []
+        if self.max_steps > 0:
+            self._expected_total_steps.append(self.max_steps)
+        if self.max_steps_per_epoch > 0 and self.max_epochs > 0:
+            self._expected_total_steps.append(
+                self.max_steps_per_epoch * self.max_epochs
+            )
+        if self.max_epochs > 0:
+            self._expected_total_steps.append(
+                len(train_dataloader) * self.max_epochs // self.accumulate_grad_batches
+            )
+        self._expected_total_steps = min(self._expected_total_steps)
+        log.info(f"Expected total steps: {self._expected_total_steps}")
+
+    @property
+    def expected_total_steps(self):
+        """
+        The expected total number of steps of the entire training. You need to run `compute_expected_total_steps` method to compute this value before accessing it.
+
+        Raises:
+            ValueError: If the expected total steps have not been computed.
+        """
+        if self._expected_total_steps is None:
+            raise ValueError(
+                "The expected total steps have not been computed. Run `compute_expected_total_steps` method."
+            )
+        else:
+            return self._expected_total_steps
+
+    def conditional_checkpoint_save(
+        self,
+        stage: Literal["end_of_step", "end_of_epoch", "end_of_training"],
+        *args,
+        **kwargs,
+    ):
+        """
+        Conditionally saves a checkpoint based on the current training stage.
+
+        Args:
+            stage (Literal["end_of_step", "end_of_epoch", "end_of_training"]): The current stage of training.
+        """
+        if stage == "end_of_step":
+            if (
+                self.checkpoint_save_interval == "step"
+                and (self.global_step_idx + 1) % self.checkpoint_save_frequency == 0
+            ):
+                save_path = os.path.join(
+                    self.log_dir, "checkpoints", f"step={self.global_step_idx}.ckpt"
+                )
+                self.save_checkpoint(save_path, *args, **kwargs)
+        elif stage == "end_of_epoch":
+            if (
+                self.checkpoint_save_interval == "epoch"
+                and (self.epoch_idx + 1) % self.checkpoint_save_frequency == 0
+            ):
+                save_path = os.path.join(
+                    self.log_dir, "checkpoints", f"epoch={self.epoch_idx}.ckpt"
+                )
+                self.save_checkpoint(save_path, *args, **kwargs)
+        elif stage == "end_of_training":
+            # if the checkpoint has not been saved yet, save it
+            if self.global_step_idx > self._latest_saved_checkpoint_global_step:
+                save_path = os.path.join(
+                    self.log_dir,
+                    "checkpoints",
+                    f"epoch={self.epoch_idx}_step={self.global_step_idx}.ckpt",
+                )
+                self.save_checkpoint(save_path, *args, **kwargs)
+                try:
+                    os.symlink(
+                        src=save_path,
+                        dst=os.path.join(
+                            self.log_dir, "checkpoints", "latest_model.ckpt"
+                        ),
+                        target_is_directory=os.path.isdir(save_path),
+                    )
+                except Exception as e:
+                    log.error(f"Failed to create symlink: {e}")
+        else:
+            raise ValueError(
+                f"Unknown stage: {stage}. Available options: 'end_of_step', 'end_of_epoch', 'end_of_training'"
+            )
+
+    @abstractmethod
+    def save_checkpoint(self, path, **kwargs):
+        """
+        Saves a checkpoint of the model.
+
+        Args:
+            path (str): The path where the checkpoint will be saved.
+
+        Raises:
+            NotImplementedError: If the method is not implemented.
+        """
+        raise NotImplementedError("save_checkpoint method is not implemented")
+
+    def train(
+        self,
+        model: Union[nn.Module, "_FabricModule"],
+        optimizer: Union[torch.optim.Optimizer, "_FabricOptimizer"],
+        lr_scheduler: torch.optim.lr_scheduler.LRScheduler,
+    ):
+        """
+        Trains the model.
+
+        The global batch size is `the batch size per device` x `the number of devices` x `the number of gradient accumulation steps`.
+
+        Args:
+            model (Union[nn.Module, "_FabricModule"]): The model to be trained.
+            optimizer (Union[torch.optim.Optimizer, "_FabricOptimizer"]): The optimizer used for training.
+            lr_scheduler (torch.optim.lr_scheduler.LRScheduler): The learning rate scheduler.
+        """
+        fabric = self.fabric
+        self.is_training = True
+        # number of parameter update iterations, not the number of batches
+        self.global_step_idx = 0
+        model.train()
+        optimizer.zero_grad()
+        for epoch_idx in tqdm(
+            range(self.max_epochs) if self.max_epochs > 0 else itertools.count(0),
+            "Training Epoch",
+            dynamic_ncols=True,
+            leave=False,
+            disable=not fabric.is_global_zero,
+        ):
+            self.epoch_idx = epoch_idx
+            self.train_epoch(model, optimizer, lr_scheduler)
+            # run lr_scheduler at the end of the epoch if interval is set to "epoch"
+            if (
+                self.lr_scheduler_interval == "epoch"
+                and (epoch_idx + 1) % self.lr_scheduler_frequency == 0
+            ):
+                lr_scheduler.step()
+
+            # save the model at the end of the epoch if interval is set to "epoch" and frequency is met
+            self.conditional_checkpoint_save(stage="end_of_epoch")
+
+            if not self.is_training:
+                break
+
+        optimizer.zero_grad()
+        # save the model at the end of training
+        self.conditional_checkpoint_save(stage="end_of_training")
+        return model
+
+    @abstractmethod
+    def train_epoch(
+        self,
+        model: Union[nn.Module, "_FabricModule"],
+        optimizer: Union[torch.optim.Optimizer, "_FabricOptimizer"],
+        lr_scheduler: torch.optim.lr_scheduler.LRScheduler,
+    ):
+        """
+        Trains the model for one epoch.
+
+        Args:
+            model (Union[nn.Module, "_FabricModule"]): The model to be trained.
+            optimizer (Union[torch.optim.Optimizer, "_FabricOptimizer"]): The optimizer used for training.
+            lr_scheduler (torch.optim.lr_scheduler.LRScheduler): The learning rate scheduler.
+
+        Raises:
+            NotImplementedError: If the method is not implemented.
+        """
+        raise NotImplementedError(
+            "Copy this as a template and implement your own train_epoch method"
+        )
+        fabric = self.fabric
+
+        accumulated_loss = 0
+        for step_idx, batch in enumerate(
+            pbar := tqdm(
+                self.train_dataloader,
+                desc="Training Batches",
+                dynamic_ncols=True,
+                leave=False,
+                disable=not fabric.is_global_zero,
+            )
+        ):
+            is_accumulating = (step_idx + 1) % self.accumulate_grad_batches != 0
+
+            # disable gradient synchronization if accumulating gradients across steps for improved performance
+            with fabric.no_backward_sync(self.model, enabled=is_accumulating):
+                # use_cache=True is not compatible with gradient checkpointing, so we disable it here
+                output = self.compute_loss(batch)
+                loss = output["loss"] / self.accumulate_grad_batches
+
+                fabric.backward(loss)
+                accumulated_loss += loss.item()
+
+            # 1. update the model parameters if not accumulating gradients
+            # 2. step the lr_scheduler if interval is set to "step" and frequency is met
+            # 3. save the model if interval is set to "step" and frequency is met
+            # 4. log metrics
+            # 5. increase the global step index and reset the accumulated metrics
+            if not is_accumulating:
+                self.clip_gradients_if_needed(model, optimizer)
+
+                # run lr_scheduler at the end of the step if interval is set to "step"
+                if (
+                    self.lr_scheduler_interval == "step"
+                    and (self.global_step_idx + 1) % self.lr_scheduler_frequency == 0
+                ):
+                    lr_scheduler.step()
+
+                # update the model parameters and zero the gradients
+                optimizer.step()
+                optimizer.zero_grad()
+
+                metrics = {
+                    "train/loss": accumulated_loss,
+                    "train/lr": optimizer.param_groups[0]["lr"],
+                }
+
+                fabric.log_dict(metrics, step=self.global_step_idx)
+                pbar.set_postfix(metrics)
+
+                # save the model at the end of the step if interval is set to "step" and frequency is met
+                self.conditional_checkpoint_save(stage="end_of_step")
+
+                # break if max_steps_per_epoch is set, and exit epoch
+                if (
+                    self.max_steps_per_epoch > 0
+                    and step_idx + 1 >= self.max_steps_per_epoch
+                ):
+                    break
+                # break if max_steps is set, and exit training
+                if self.max_steps > 0 and self.global_step_idx >= self.max_steps - 1:
+                    self.is_training = False
+                    break
+
+                self.global_step_idx += 1
+                accumulated_loss = 0

fusion_bench/mixins/lightning_fabric.py

@@ -1,6 +1,7 @@
+import functools
 import logging
 import os
-from typing import TYPE_CHECKING, Any, Optional, TypeVar, List
+from typing import TYPE_CHECKING, Any, List, Optional, TypeVar
 
 import lightning as L
 import torch
@@ -8,11 +9,12 @@ from lightning.fabric.loggers import TensorBoardLogger
 from lightning.fabric.utilities.rank_zero import rank_zero_only
 from omegaconf import DictConfig, OmegaConf
 
-from fusion_bench.utils.instantiate import instantiate
 from fusion_bench.utils import import_object
+from fusion_bench.utils.instantiate import instantiate
 
 if TYPE_CHECKING:
     import lightning.fabric.loggers.tensorboard
+    from lightning.fabric.strategies import FSDPStrategy
 
 log = logging.getLogger(__name__)
 
@@ -32,6 +34,13 @@ def get_policy(*args: str) -> set:
     return {import_object(arg) for arg in args}
 
 
+def get_size_based_auto_wrap_policy(*args, **kwargs):
+    from torch.distributed.fsdp.wrap import size_based_auto_wrap_policy
+
+    policy = functools.partial(size_based_auto_wrap_policy, *args, **kwargs)
+    return policy
+
+
 class LightningFabricMixin:
     """
     A mixin class for integrating Lightning Fabric into a project.

fusion_bench/modelpool/__init__.py

@@ -16,6 +16,7 @@ _import_structure = {
         "HuggingFaceGPT2ClassificationPool",
         "GPT2ForSequenceClassificationPool",
     ],
+    "seq_classification_lm": ["SeqenceClassificationModelPool"],
 }
 
 
@@ -31,6 +32,7 @@ if TYPE_CHECKING:
     from .nyuv2_modelpool import NYUv2ModelPool
     from .PeftModelForSeq2SeqLM import PeftModelForSeq2SeqLMPool
     from .seq2seq_lm import Seq2SeqLMPool
+    from .seq_classification_lm import SeqenceClassificationModelPool
 
 else:
     sys.modules[__name__] = LazyImporter(

fusion_bench/modelpool/causal_lm/__init__.py

@@ -1,2 +1,2 @@
 # flake8: noqa F401
-from .causal_lm import CausalLMBackbonePool, CausalLMPool
+from .causal_lm import CausalLMBackbonePool, CausalLMPool, load_peft_causal_lm

fusion_bench/modelpool/causal_lm/causal_lm.py

@@ -3,6 +3,7 @@ import os
 from copy import deepcopy
 from typing import Any, Optional, TypeAlias, Union, cast  # noqa: F401
 
+import peft
 from omegaconf import DictConfig, flag_override
 from torch import nn
 from torch.nn.modules import Module
@@ -23,28 +24,6 @@ log = logging.getLogger(__name__)
 CausalLM: TypeAlias = Union[LlamaForCausalLM, MistralForCausalLM, Any]
 
 
-def config_priority_get(priority_config, general_config, key, default):
-    """
-    Retrieve a configuration value with priority.
-
-    This function retrieves the value associated with `key` from `priority_config` if it exists.
-    If the key is not found in `priority_config`, it retrieves the value from `general_config`.
-    If the key is not found in either configuration, it returns the provided `default` value.
-
-    Args:
-        priority_config (dict): The configuration dictionary with higher priority.
-        general_config (dict): The general configuration dictionary.
-        key (str): The key to look up in the configuration dictionaries.
-        default: The default value to return if the key is not found in either configuration.
-
-    Returns:
-        The value associated with `key` from `priority_config` or `general_config`, or the `default` value if the key is not found.
-    """
-    if key in priority_config:
-        return priority_config[key]
-    return general_config.get(key, default)
-
-
 class CausalLMPool(BaseModelPool):
     _config_mapping = BaseModelPool._config_mapping | {
         "_tokenizer": "tokenizer",
@@ -138,3 +117,23 @@ class CausalLMBackbonePool(CausalLMPool):
             model_name_or_config, *args, **kwargs
         )
         return model.model.layers
+
+
+def load_peft_causal_lm(
+    base_model_path: str,
+    peft_model_path: str,
+    torch_dtype: str = "bfloat16",
+    is_trainable: bool = True,
+    merge_and_unload: bool = False,
+):
+    base_model = LlamaForCausalLM.from_pretrained(
+        base_model_path, torch_dtype=torch_dtype
+    )
+    model = peft.PeftModel.from_pretrained(
+        base_model,
+        peft_model_path,
+        is_trainable=is_trainable,
+    )
+    if merge_and_unload:
+        model = model.merge_and_unload()
+    return model

fusion_bench/modelpool/seq_classification_lm/__init__.py

@@ -0,0 +1,2 @@
+from .reward_model import create_reward_model_from_pretrained
+from .seq_classification_lm import SeqenceClassificationModelPool

fusion_bench/modelpool/seq_classification_lm/reward_model.py

@@ -0,0 +1,15 @@
+from transformers import AutoModelForSequenceClassification
+
+
+def create_reward_model_from_pretrained(pretrained_model_name_or_path: str, **kwargs):
+    """
+    Create a reward model for reward modeling (RLHF).
+
+    Args:
+        pretrained_model_name_or_path (str): The name or path of the pretrained model.
+        **kwargs: Additional keyword arguments passed to the model class.
+    """
+    model = AutoModelForSequenceClassification.from_pretrained(
+        pretrained_model_name_or_path, num_labels=1, **kwargs
+    )
+    return model

fusion_bench/modelpool/seq_classification_lm/seq_classification_lm.py

@@ -0,0 +1,98 @@
+import logging
+import os
+from copy import deepcopy
+from typing import TYPE_CHECKING, Any, Optional, TypeAlias, Union, cast  # noqa: F401
+
+from omegaconf import DictConfig, flag_override
+from transformers import PreTrainedModel, PreTrainedTokenizer
+from typing_extensions import override
+
+from fusion_bench.modelpool import BaseModelPool
+from fusion_bench.utils import instantiate
+from fusion_bench.utils.dtype import parse_dtype
+
+if TYPE_CHECKING:
+    from transformers import LlamaForSequenceClassification
+
+log = logging.getLogger(__name__)
+
+
+class SeqenceClassificationModelPool(BaseModelPool):
+
+    def __init__(
+        self,
+        models,
+        *,
+        tokenizer: Optional[DictConfig],
+        model_kwargs: Optional[DictConfig] = None,
+        **kwargs,
+    ):
+        super().__init__(models, **kwargs)
+        # process `model_kwargs`
+        self._tokenizer = tokenizer
+        self._model_kwargs = model_kwargs
+        if self._model_kwargs is None:
+            self._model_kwargs = DictConfig({})
+        with flag_override(self._model_kwargs, "allow_objects", True):
+            if hasattr(self._model_kwargs, "torch_dtype"):
+                self._model_kwargs.torch_dtype = parse_dtype(
+                    self._model_kwargs.torch_dtype
+                )
+
+    @override
+    def load_model(
+        self,
+        model_name_or_config: str | DictConfig,
+        *args,
+        **kwargs,
+    ) -> Union[PreTrainedModel, "LlamaForSequenceClassification"]:
+        model_kwargs = deepcopy(self._model_kwargs)
+        model_kwargs.update(kwargs)
+        if isinstance(model_name_or_config, str):
+            log.info(f"Loading model: {model_name_or_config}", stacklevel=2)
+        return super().load_model(model_name_or_config, *args, **model_kwargs)
+
+    def load_tokenizer(self, *args, **kwargs) -> PreTrainedTokenizer:
+        assert self._tokenizer is not None, "Tokenizer is not defined in the config"
+        log.info("Loading tokenizer.", stacklevel=2)
+        tokenizer = instantiate(self._tokenizer, *args, **kwargs)
+        return tokenizer
+
+    @override
+    def save_model(
+        self,
+        model: PreTrainedModel,
+        path: str,
+        push_to_hub: bool = False,
+        model_dtype: Optional[str] = None,
+        save_tokenizer: bool = False,
+        tokenizer_kwargs=None,
+        **kwargs,
+    ):
+        """
+        Save the model to the specified path.
+
+        Args:
+            model (PreTrainedModel): The model to be saved.
+            path (str): The path where the model will be saved.
+            push_to_hub (bool, optional): Whether to push the model to the Hugging Face Hub. Defaults to False.
+            save_tokenizer (bool, optional): Whether to save the tokenizer along with the model. Defaults to False.
+            **kwargs: Additional keyword arguments passed to the `save_pretrained` method.
+        """
+        path = os.path.expanduser(path)
+        if save_tokenizer:
+            if tokenizer_kwargs is None:
+                tokenizer_kwargs = {}
+            # load the tokenizer
+            tokenizer = self.load_tokenizer(**tokenizer_kwargs)
+            tokenizer.save_pretrained(
+                path,
+                push_to_hub=push_to_hub,
+            )
+        if model_dtype is not None:
+            model.to(dtype=parse_dtype(model_dtype))
+        model.save_pretrained(
+            path,
+            push_to_hub=push_to_hub,
+            **kwargs,
+        )

fusion_bench/models/chat_templates/__init__.py

@@ -0,0 +1 @@
+from .load_tokenizer import chat_template_mapping, load_tokenizer_with_chat_template

fusion_bench/models/chat_templates/llama_3_Instruct.py

@@ -0,0 +1 @@
+CHAT_TEMPLATE = '{{- bos_token }}\n{%- if custom_tools is defined %}\n    {%- set tools = custom_tools %}\n{%- endif %}\n{%- if not tools_in_user_message is defined %}\n    {%- set tools_in_user_message = true %}\n{%- endif %}\n{%- if not date_string is defined %}\n    {%- if strftime_now is defined %}\n        {%- set date_string = strftime_now("%d %b %Y") %}\n    {%- else %}\n        {%- set date_string = "26 Jul 2024" %}\n    {%- endif %}\n{%- endif %}\n{%- if not tools is defined %}\n    {%- set tools = none %}\n{%- endif %}\n\n{#- This block extracts the system message, so we can slot it into the right place. #}\n{%- if messages[0][\'role\'] == \'system\' %}\n    {%- set system_message = messages[0][\'content\']|trim %}\n    {%- set messages = messages[1:] %}\n{%- else %}\n    {%- set system_message = "" %}\n{%- endif %}\n\n{#- System message #}\n{{- "<|start_header_id|>system<|end_header_id|>\\n\\n" }}\n{%- if tools is not none %}\n    {{- "Environment: ipython\\n" }}\n{%- endif %}\n{{- "Cutting Knowledge Date: December 2023\\n" }}\n{{- "Today Date: " + date_string + "\\n\\n" }}\n{%- if tools is not none and not tools_in_user_message %}\n    {{- "You have access to the following functions. To call a function, please respond with JSON for a function call." }}\n    {{- \'Respond in the format {"name": function name, "parameters": dictionary of argument name and its value}.\' }}\n    {{- "Do not use variables.\\n\\n" }}\n    {%- for t in tools %}\n        {{- t | tojson(indent=4) }}\n        {{- "\\n\\n" }}\n    {%- endfor %}\n{%- endif %}\n{{- system_message }}\n{{- "<|eot_id|>" }}\n\n{#- Custom tools are passed in a user message with some extra guidance #}\n{%- if tools_in_user_message and not tools is none %}\n    {#- Extract the first user message so we can plug it in here #}\n    {%- if messages | length != 0 %}\n        {%- set first_user_message = messages[0][\'content\']|trim %}\n        {%- set messages = messages[1:] %}\n    {%- else %}\n        {{- raise_exception("Cannot put tools in the first user message when there\'s no first user message!") }}\n{%- endif %}\n    {{- \'<|start_header_id|>user<|end_header_id|>\\n\\n\' -}}\n    {{- "Given the following functions, please respond with a JSON for a function call " }}\n    {{- "with its proper arguments that best answers the given prompt.\\n\\n" }}\n    {{- \'Respond in the format {"name": function name, "parameters": dictionary of argument name and its value}.\' }}\n    {{- "Do not use variables.\\n\\n" }}\n    {%- for t in tools %}\n        {{- t | tojson(indent=4) }}\n        {{- "\\n\\n" }}\n    {%- endfor %}\n    {{- first_user_message + "<|eot_id|>"}}\n{%- endif %}\n\n{%- for message in messages %}\n    {%- if not (message.role == \'ipython\' or message.role == \'tool\' or \'tool_calls\' in message) %}\n        {{- \'<|start_header_id|>\' + message[\'role\'] + \'<|end_header_id|>\\n\\n\'+ message[\'content\'] | trim + \'<|eot_id|>\' }}\n    {%- elif \'tool_calls\' in message %}\n        {%- if not message.tool_calls|length == 1 %}\n            {{- raise_exception("This model only supports single tool-calls at once!") }}\n        {%- endif %}\n        {%- set tool_call = message.tool_calls[0].function %}\n        {{- \'<|start_header_id|>assistant<|end_header_id|>\\n\\n\' -}}\n        {{- \'{"name": "\' + tool_call.name + \'", \' }}\n        {{- \'"parameters": \' }}\n        {{- tool_call.arguments | tojson }}\n        {{- "}" }}\n        {{- "<|eot_id|>" }}\n    {%- elif message.role == "tool" or message.role == "ipython" %}\n        {{- "<|start_header_id|>ipython<|end_header_id|>\\n\\n" }}\n        {%- if message.content is mapping or message.content is iterable %}\n            {{- message.content | tojson }}\n        {%- else %}\n            {{- message.content }}\n        {%- endif %}\n        {{- "<|eot_id|>" }}\n    {%- endif %}\n{%- endfor %}\n{%- if add_generation_prompt %}\n    {{- \'<|start_header_id|>assistant<|end_header_id|>\\n\\n\' }}\n{%- endif %}\n'

fusion_bench/models/chat_templates/load_tokenizer.py

@@ -0,0 +1,43 @@
+import logging
+
+from transformers import AutoTokenizer
+
+from .llama_3_Instruct import CHAT_TEMPLATE as LLAMA_3_INSTRUCT_CHAT_TEMPLATE
+
+chat_template_mapping = {"llama_3_instruct": LLAMA_3_INSTRUCT_CHAT_TEMPLATE}
+
+log = logging.getLogger(__name__)
+
+
+def load_tokenizer_with_chat_template(
+    pretrained_model_name_or_path: str,
+    model_family: str,
+    overwrite_chat_template: bool = True,
+    **kwargs,
+):
+    """
+    Load the tokenizer for Llama 3 model.
+
+    Args:
+        pretrained_model_name_or_path (str): The name or path of the pretrained model.
+        model_family (str): The model family.
+        **kwargs: Additional keyword arguments passed to the tokenizer class.
+    """
+    assert (
+        model_family in chat_template_mapping
+    ), f"Model family {model_family} not found. Available model families: {chat_template_mapping.keys()}"
+
+    tokenizer = AutoTokenizer.from_pretrained(
+        pretrained_model_name_or_path,
+        **kwargs,
+    )
+
+    if tokenizer.chat_template is None:
+        tokenizer.chat_template = chat_template_mapping[model_family]
+    else:
+        if overwrite_chat_template:
+            log.warning("Overwriting the chat template with the default chat template.")
+            tokenizer.chat_template = chat_template_mapping[model_family]
+        else:
+            log.warning("Chat template already exists. Skipping overwriting.")
+    return tokenizer