/ludwig/utils/llm_utils.py

from __future__ import annotations

import copy
import logging
import tempfile
from typing import TYPE_CHECKING, Union

import torch
import torch.nn.functional as F
import transformers
from packaging import version

try:
    from bitsandbytes.nn.modules import Embedding as BnbEmbedding
except ImportError:
    BnbEmbedding = None
from transformers import AutoModelForCausalLM, TextStreamer

from ludwig.constants import IGNORE_INDEX_TOKEN_ID, LOGITS, PREDICTIONS, PROBABILITIES
from ludwig.schema.trainer import LLMTrainerConfig
from ludwig.utils.error_handling_utils import default_retry
from ludwig.utils.logging_utils import log_once
from ludwig.utils.model_utils import find_embedding_layer_with_path

if TYPE_CHECKING:
    from transformers import AutoConfig, PreTrainedModel, PreTrainedTokenizer

    from ludwig.schema.encoders.text_encoders import LLMEncoderConfig
    from ludwig.schema.model_types.llm import LLMModelConfig


logger = logging.getLogger(__name__)

transformers_436 = version.parse(transformers.__version__) >= version.parse("4.36.0")

FALLBACK_CONTEXT_LEN = 2048

_MODELS_WITH_DEVICE_MAP_AUTO_EXCLUSION = set()


@default_retry(tries=8, exceptions=OSError)
def load_pretrained_from_config(
    config_obj: LLMModelConfig | LLMEncoderConfig,
    model_config: AutoConfig | None = None,
    weights_save_path: str | None = None,
) -> PreTrainedModel:
    load_kwargs = {}
    quantization = config_obj.quantization
    if quantization and getattr(quantization, "backend", "bitsandbytes") == "bitsandbytes":
        # Apply bitsandbytes quantization configuration at model load time.
        load_kwargs["dtype"] = getattr(torch, quantization.bnb_4bit_compute_dtype)
        load_kwargs["quantization_config"] = quantization.to_bitsandbytes()
        load_kwargs["device_map"] = "auto"

        if transformers_436:
            load_kwargs["attn_implementation"] = "eager"
    else:
        # Either no quantization, or torchao — which quantizes the model *after* load, not
        # via transformers' BitsAndBytesConfig. Load in float32 by default to avoid CUBLAS
        # errors with small hidden sizes and to ensure numerical stability during training
        # without mixed-precision.
        load_kwargs["dtype"] = torch.float32

    config_modified = False
    if config_obj.model_parameters:
        # Add any model specific parameters to the load kwargs
        for param_name, param_value in config_obj.model_parameters.to_dict().items():
            # Not all parameters are supported by all models, so we only add the parameter to the load kwargs
            # if it is supported by the model.
            if param_value is None:
                continue

            if hasattr(model_config, param_name):
                if isinstance(param_value, dict):
                    # For nested dict params (e.g. rope_scaling), merge with existing
                    # config values to preserve defaults like rope_theta.
                    existing = getattr(model_config, param_name, {}) or {}
                    existing.update(param_value)
                    setattr(model_config, param_name, existing)
                    config_modified = True
                else:
                    load_kwargs[param_name] = param_value
            else:
                logger.warning(f"Parameter {param_name} is not supported by {config_obj.base_model}. Skipping.")

    # Only pass config= when we've directly modified it (e.g. rope_scaling merge).
    if config_modified:
        load_kwargs["config"] = model_config

    logger.info("Loading large language model...")
    pretrained_model_name_or_path = weights_save_path or config_obj.base_model
    trust_remote_code = getattr(config_obj, "trust_remote_code", False)
    is_multimodal = getattr(config_obj, "is_multimodal", False)
    if is_multimodal:
        # VLMs — Qwen2-VL, LLaVA, InternVL, Idefics etc. — register their own vision-to-sequence
        # head. AutoModelForVision2Seq picks the right class automatically so the vision tower,
        # projector, and LM head all come along together.
        from transformers import AutoModelForVision2Seq

        logger.info("Loading multimodal (VLM) base model via AutoModelForVision2Seq")
        model: PreTrainedModel = AutoModelForVision2Seq.from_pretrained(
            pretrained_model_name_or_path, trust_remote_code=trust_remote_code, **load_kwargs
        )
    else:
        model: PreTrainedModel = AutoModelForCausalLM.from_pretrained(
            pretrained_model_name_or_path, trust_remote_code=trust_remote_code, **load_kwargs
        )
    return model


def to_device(
    model: PreTrainedModel,
    device: str | torch.DeviceObjType,
    config_obj: LLMModelConfig,
    curr_device: torch.DeviceObjType,
) -> tuple[PreTrainedModel, torch.DeviceObjType]:
    """Move an LLM to the requested device, accounting for sharding and adapters.

    Args:
        model: Pretrained model to put on device
        config_obj: LLM config
        curr_device: The current device that the model is on

    Returns:
        `model` moved to `device`
    """
    device = torch.device(device)

    if device.type == curr_device.type:
        log_once(f"Model already on device'{device}'.")
        return model, device
    else:
        log_once(f"Moving LLM from '{curr_device}' to '{device}'.")

    model_kwargs = {}
    num_gpus = torch.cuda.device_count()
    if device == torch.device("cuda") and num_gpus > 1:
        # TODO: make this configurable in the future. These parameters are from FastChat:
        # https://github.com/lm-sys/FastChat/blob/0e958b852a14f4bef5f0e9d7a5e7373477329cf2/fastchat/serve/inference.py#L90
        # TODO: Wrap device_map="auto" in a try-except block since it may not be supported for all models (E.g. BertLMHead)
        # We don't add quantization here (float16 or bfloat16) since we may not always want to quantize. We should
        # make quantization configurable in the future via the trainer config.
        model_kwargs.update(
            {
                "low_cpu_mem_usage": True,
                "max_memory": dict.fromkeys(range(num_gpus), "13GiB"),
            }
        )

        if config_obj.base_model not in _MODELS_WITH_DEVICE_MAP_AUTO_EXCLUSION:
            model_kwargs["device_map"] = "auto"

        if config_obj.quantization:
            model_kwargs["quantization_config"] = config_obj.quantization.to_bitsandbytes()

        # we save and reload the weights to ensure that they can be sharded across the GPUs using `from_pretrained`
        with tempfile.TemporaryDirectory() as tmpdir:
            model.save_pretrained(tmpdir)

            if config_obj.adapter:
                model = AutoModelForCausalLM.from_pretrained(
                    config_obj.base_model,
                    trust_remote_code=getattr(config_obj, "trust_remote_code", False),
                    **model_kwargs,
                )

                # Leave this import inline to support a minimal install of Ludwig
                from peft import PeftModel

                model = PeftModel.from_pretrained(
                    model,
                    tmpdir,
                    torch_dtype=torch.float16,
                )
            else:
                model = AutoModelForCausalLM.from_pretrained(
                    tmpdir,
                    trust_remote_code=getattr(config_obj, "trust_remote_code", False),
                    **model_kwargs,
                )
    else:
        model = model.to(device)

    return model, device


def _load_peft_config(pretrained_adapter_weights: str):
    """Load a PeftConfig, fixing known compatibility issues with newer PEFT versions."""
    import json

    from huggingface_hub import hf_hub_download
    from peft import PeftConfig

    config_file = hf_hub_download(pretrained_adapter_weights, "adapter_config.json")
    with open(config_file) as f:
        config_dict = json.load(f)

    # AdaLoRA requires total_step > 0 in newer PEFT versions, but pretrained
    # configs may have total_step=None.
    if config_dict.get("peft_type") == "ADALORA" and not config_dict.get("total_step"):
        config_dict["total_step"] = 10000

    return PeftConfig.from_peft_type(**config_dict)


def initialize_adapter(model: PreTrainedModel, config_obj: "LLMModelConfig") -> Union["PeftModel", PreTrainedModel]:  # noqa F821
    """Wrap a pretrained model with a PEFT model for fine-tuning.

    Dispatches to the multi-adapter path when ``config_obj.adapters`` is set (several
    named adapters registered on the same base, optional weighted merge, runtime
    switching via ``set_adapter``) and to the single-adapter path for ``config_obj.adapter``.
    The two fields are mutually exclusive at the schema layer.

    Args:
         model: Pretrained model to fine-tune with an adapter.
         config_obj: LLM config

    Returns:
        ``model`` wrapped in a PEFT model if an adapter config was provided, otherwise
        ``model`` unaltered.
    """
    if getattr(config_obj, "adapters", None) is not None:
        return _initialize_multi_adapters(model, config_obj)

    if config_obj.adapter:
        if config_obj.adapter.pretrained_adapter_weights:
            # Load pretrained adapter weights if specified.
            logger.info(f"Using pretrained adapter weights: {config_obj.adapter.pretrained_adapter_weights}")

            # Leave this import inline to support a minimal install of Ludwig
            from peft import MODEL_TYPE_TO_PEFT_MODEL_MAPPING, PeftConfig  # noqa

            peft_config = _load_peft_config(config_obj.adapter.pretrained_adapter_weights)

            model = MODEL_TYPE_TO_PEFT_MODEL_MAPPING[peft_config.task_type].from_pretrained(
                model, config_obj.adapter.pretrained_adapter_weights, config=peft_config
            )
        else:
            # Leave this import inline to support a minimal install of Ludwig
            from peft import get_peft_model, TaskType

            # If no pretrained adapter is provided, we want to load untrained weights into the model
            peft_config = config_obj.adapter.to_config(
                task_type=TaskType.CAUSAL_LM, tokenizer_name_or_path=config_obj.base_model
            )

            model = get_peft_model(model, peft_config)

    return model


def _initialize_multi_adapters(model: PreTrainedModel, config_obj: "LLMModelConfig") -> "PeftModel":  # noqa F821
    """Attach several named PEFT adapters to ``model`` and (optionally) a merged one.

    PEFT's public multi-adapter surface:
      * ``get_peft_model(base, cfg, adapter_name=...)`` creates a PeftModel with one named
        adapter. We use the first configured adapter here as the anchor.
      * ``peft_model.add_adapter(adapter_name, cfg)`` registers additional adapters on the
        same PeftModel.
      * ``peft_model.add_weighted_adapter(source_names, weights, name, combination_type,
        density)`` produces a new adapter by combining existing ones. Used when
        ``adapters.merge`` is set.
      * ``peft_model.set_adapter(name)`` designates the default adapter. Only one adapter
        is active at a time — users switch explicitly at inference, or ask for the merged
        adapter as the default.
    """
    from peft import get_peft_model, TaskType  # imported inline for minimal installs

    adapters_cfg = config_obj.adapters
    items = list(adapters_cfg.adapters.items())  # insertion order, validated non-empty by schema
    first_name, first_cfg = items[0]

    if not hasattr(first_cfg, "to_config"):
        # The schema stores entries as raw dicts; re-materialize them into adapter configs via
        # the adapter registry so each has a working `to_config()` method.
        items = [(name, _materialize_adapter_config(cfg)) for name, cfg in items]
        first_name, first_cfg = items[0]

    first_peft_config = first_cfg.to_config(task_type=TaskType.CAUSAL_LM, tokenizer_name_or_path=config_obj.base_model)
    model = get_peft_model(model, first_peft_config, adapter_name=first_name)

    for name, adapter_cfg in items[1:]:
        peft_config = adapter_cfg.to_config(task_type=TaskType.CAUSAL_LM, tokenizer_name_or_path=config_obj.base_model)
        model.add_adapter(name, peft_config)

    if adapters_cfg.merge is not None:
        merge = adapters_cfg.merge
        weights = merge.weights if merge.weights is not None else [1.0] * len(merge.sources)
        kwargs = {
            "adapters": merge.sources,
            "weights": weights,
            "adapter_name": merge.name,
            "combination_type": merge.combination_type,
        }
        if merge.combination_type in ("ties", "dare_linear", "dare_ties", "magnitude_prune"):
            kwargs["density"] = merge.density
        model.add_weighted_adapter(**kwargs)
        logger.info(
            "Merged adapters %s (weights=%s) via %s into %r",
            merge.sources,
            weights,
            merge.combination_type,
            merge.name,
        )

    active = adapters_cfg.active or first_name
    model.set_adapter(active)
    logger.info("Registered adapters: %s (active=%s)", [n for n, _ in items], active)

    return model


def _materialize_adapter_config(raw):
    """Turn a raw dict from ``adapters.adapters`` into a BaseAdapterConfig instance."""
    from ludwig.schema.llms.peft import adapter_registry

    if hasattr(raw, "to_config"):
        return raw
    if not isinstance(raw, dict):
        raise TypeError(f"Expected dict adapter config, got {type(raw).__name__}")
    adapter_type = raw.get("type")
    if adapter_type is None:
        raise ValueError("Adapter config is missing required `type` field.")
    if adapter_type not in adapter_registry:
        raise ValueError(f"Unknown adapter type {adapter_type!r}. Known: {list(adapter_registry.keys())}")
    return adapter_registry[adapter_type].model_validate(raw)


def get_context_len(model_config: AutoConfig):
    """Determines the maximum length of the context (input + output tokens) based on the provided model
    configuration.

    Args:
        model_config (AutoConfig): The model configuration object containing information about the model's properties.

    Returns:
        int: The maximum context length, which can be derived from the model configuration. If no relevant attribute
             is found, the default value of 2048 is returned.

    This function examines the provided model configuration object to identify the attribute that specifies the maximum
    context length. It checks for attributes in the following order of preference:
    1. 'max_sequence_length': If this attribute is present in the model configuration, its value is returned.
    2. 'max_position_embeddings': If 'max_sequence_length' is not found but 'max_position_embeddings' is present, its
       value is returned.
    3. 'n_positions': If neither 'max_sequence_length' nor 'max_position_embeddings' are found, and 'n_positions' is
       present, its value is returned.
    4. Default: If none of the relevant attributes are present, the function returns a default value of 2048.

    Note:
    - The maximum context length is important for defining the size of input and output sequences in a model.

    Example Usage:
    >>> config = AutoConfig.from_pretrained("bert-base-uncased")
    >>> context_len = get_context_len(config)
    >>> print(context_len)
    512
    """
    if hasattr(model_config, "max_sequence_length"):
        return model_config.max_sequence_length
    elif hasattr(model_config, "max_position_embeddings"):
        return model_config.max_position_embeddings
    elif hasattr(model_config, "n_positions"):
        return model_config.n_positions
    else:
        return FALLBACK_CONTEXT_LEN


def has_padding_token(input_tensor: torch.Tensor, tokenizer: PreTrainedTokenizer):
    """Checks if the input tensor contains any padding tokens.

    Args:
        input_tensor (torch.Tensor): The input tensor.
        tokenizer (PreTrainedTokenizer): The tokenizer used to encode the input.

    Returns:
        bool: True if the input tensor contains any padding tokens, False otherwise.

    Example:
        >>> import torch
        >>> from transformers import PreTrainedTokenizer
        >>> tokenizer = PreTrainedTokenizer.from_pretrained('bert-base-uncased')
        >>> input_sentence = "This is an example sentence."
        >>> input_ids = tokenizer.encode(input_sentence, add_special_tokens=True)
        >>> padded_input_ids = torch.nn.functional.pad(input_ids, (0, 10 - len(input_ids)))
        >>> has_padding = has_padding_token(padded_input_ids, tokenizer)
        >>> has_padding
        True
    """
    if input_tensor.dim() == 1:
        return torch.any(input_tensor == tokenizer.pad_token_id).item()
    elif input_tensor.dim() == 2:
        return torch.any(input_tensor == tokenizer.pad_token_id, dim=-1).item()
    else:
        raise ValueError(
            f"Input tensor must be 1D (single sequence) or 2D (batch of sequences), got {input_tensor.dim()}D tensor."
        )


def remove_left_padding(input_ids_sample: torch.Tensor, tokenizer: PreTrainedTokenizer):
    """Removes left padding and other tokens until the first BOS token from the input_ids tensor.

    Args:
        input_ids_sample (torch.Tensor): The input tensor with padding and other tokens.
        tokenizer (PreTrainedTokenizer): The tokenizer used to encode the input.

    Returns:
        torch.Tensor: The input tensor without left padding and other tokens until the first BOS token.

    Example:
        >>> import torch
        >>> from transformers import PreTrainedTokenizer
        >>> tokenizer = PreTrainedTokenizer.from_pretrained('bert-base-uncased')
        >>> input_sentence = "This is an example sentence."
        >>> input_ids = tokenizer.encode(input_sentence, add_special_tokens=True)
        >>> padded_input_ids = torch.nn.functional.pad(input_ids, (10 - len(input_ids), 0))
        >>> input_ids_no_padding = remove_left_padding(padded_input_ids, tokenizer)
        >>> input_ids_no_padding
        tensor([[1, 2, 3]])
    """
    # Remove all PAD tokens
    pad_idxs = torch.where(input_ids_sample == tokenizer.pad_token_id)[0]  # all PAD token locations
    input_ids_no_padding = input_ids_sample
    if len(pad_idxs) != 0:
        pad_idx = pad_idxs[-1]  # get last PAD token location
        input_ids_no_padding = input_ids_sample[pad_idx + 1 :]

    # Start from the first BOS token
    bos_idxs = torch.where(input_ids_no_padding == tokenizer.bos_token_id)[0]  # all BOS token locations
    if len(bos_idxs) != 0:
        bos_idx = bos_idxs[0]  # get first BOS token location
    else:
        bos_idx = 0

    input_ids_no_bos = input_ids_no_padding[bos_idx:].unsqueeze(0)
    return input_ids_no_bos


def add_left_padding(input_ids, max_length, pad_value=0):
    """Adds left padding to the input_ids tensor.

    Args:
        input_ids (torch.Tensor): The input tensor.
        max_length (int): The maximum length of the tensor after padding.
        pad_value (int, optional): The value used for padding. Defaults to 0.

    Returns:
        torch.Tensor: The input_ids tensor with left padding.

    Example:
        >>> input_ids = torch.tensor([1, 2, 3])
        >>> max_length = 5
        >>> padded_tensor = add_left_padding(input_ids, max_length)
        >>> padded_tensor
        tensor([0, 0, 1, 2, 3])
    """
    padding = torch.tensor([pad_value] * (max_length - input_ids.shape[0]), dtype=torch.int64, device=input_ids.device)
    return torch.cat((padding, input_ids), dim=-1)


def create_attention_mask(input_ids: torch.Tensor, tokenizer: PreTrainedTokenizer):
    """Creates an attention mask for the input_ids tensor. This also sets the last padding token ID to 1 if it
    exists.

    Args:
        input_ids (torch.Tensor): The input tensor.
        tokenizer (PreTrainedTokenizer): The tokenizer used to encode the input.

    Returns:
        torch.Tensor: The attention mask tensor.

    Example:
        >>> import torch # noqa
        >>> from transformers import PreTrainedTokenizer
        >>> tokenizer = PreTrainedTokenizer.from_pretrained('bert-base-uncased')
        >>> input_sentence = "This is an example sentence."
        >>> input_ids = tokenizer.encode(input_sentence, add_special_tokens=True)
        >>> attention_mask = create_attention_mask(input_ids, tokenizer)
        >>> attention_mask
        tensor([1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1])
    """
    attention_mask = input_ids != tokenizer.pad_token_id
    # Last token may not be padding if we've already hit the max sequence length
    if not attention_mask[-1]:
        # last token is padding, always attended to even if it is padding
        attention_mask[-1] = 1
    attention_mask = attention_mask.to(torch.int64)
    return attention_mask


def find_last_matching_index(tensor_a: torch.Tensor, tensor_b: torch.Tensor):
    """Returns the last index of `tensor_a` that matches `tensor_b`. Specifically, this checks whether the tensor_b
    is in the last tensor_b.shape[0] elements of tensor_a.

    Args:
        tensor_a (torch.Tensor): The first tensor.
        tensor_b (torch.Tensor): The second tensor.

    Returns:
        int: The last index of `tensor_a` that matches `tensor_b`. Returns -1 if there is no matching index.

    Example:
        >>> import torch
        >>> tensor_a = torch.tensor([1, 2, 3, 4, 5, 6, 7, 8])
        >>> tensor_b = torch.tensor([6, 7, 8])
        >>> last_matching_index = find_last_matching_index(tensor_a, tensor_b)
        >>> last_matching_index
        5
    """
    last_index = -1

    tensor_a_length = tensor_a.shape[0]
    tensor_b_length = tensor_b.shape[0]

    # Get the last tensor_b_length elements of tensor_a.
    tensor_a_truncated = tensor_a[-tensor_b_length:]

    # Find the last matching index.
    for i in range(tensor_b_length):
        if torch.equal(tensor_a_truncated[i:], tensor_b[: tensor_b_length - i]):
            last_index = tensor_a_length - tensor_b_length + i
            break

    return last_index


def pad_target_tensor_for_fine_tuning(
    targets: dict[str, torch.Tensor],
    predictions: dict[str, torch.Tensor],
    model_inputs: torch.Tensor,
    of_name: str,
) -> dict[str, torch.Tensor]:
    """Pad and adjust target tensors for fine-tuning LLMS models.

    This function is used to pad and adjust the target tensors with IGNORE_INDEX_TOKEN_ID based on the model inputs and
    predictions during the fine-tuning process of Language Models. Here's what this function does:
        1. If none of the tokens from the target were in the model inputs, we create a tensor of the length of model
            inputs with value IGNORE_INDEX_TOKEN_IDs. This ignores this row from affecting loss.
        2. If the target tokens were entirely inside the model inputs, we want to pad all the tokens in model_inputs
            coming from the input with IGNORE_INDEX_TOKEN_IDs and leave the target tokens as is. This ensures that all
            of the target tokens are used during loss computation.
        3. In the scenario that only some part of the target tokens were in the model inputs, we want to pad the model
            inputs until that point and only leave the partial tokens of the target as is. This ensures that we will
            only compute loss on the target tokens that were in the model inputs.

    Args:
        targets (Dict[str, torch.Tensor]): A dictionary containing the target tensors.
        predictions (Dict[str, torch.Tensor]): A dictionary containing the predicted tensors.
        model_inputs (torch.Tensor): The input tensor passed into the model's forward pass.
        of_name (str): The name of the target tensor to be padded and adjusted.

    Returns:
        Dict[str, torch.Tensor]: A dictionary containing the updated target
        dictionaries.
    """
    target_length = targets.get(of_name).size()[1]
    prediction_length = predictions[of_name].get(PREDICTIONS).size()[1]

    if target_length == prediction_length:
        return targets

    updated_targets = []
    for idx, target in enumerate(targets[of_name]):
        # Remove any leading IGNORE_INDEX_TOKEN_IDs in the target that were temporarily added for alignment
        end_index = (target != IGNORE_INDEX_TOKEN_ID).nonzero()[0]
        target = target[end_index:]
        target_device = target.device

        # See if any part of the target was in the tensor passed into the model's forward pass
        last_matching_index = find_last_matching_index(model_inputs[idx], target)

        # If the last matching index is -1, it means that the input tensor passed into the model was truncated
        # and did not contain the target tensor. In this case, we need to truncate the target tensors as well
        # and just set it to a tensor of IGNORE_INDEX_TOKEN_ID so that we don't compute loss on this target tensor.
        if last_matching_index == -1:
            updated_targets.append(torch.full((prediction_length,), IGNORE_INDEX_TOKEN_ID).to(device=target_device))

        # If the last matching index is not -1, it means that the input tensor passed into the model was not
        # truncated and contained either a part of the target tensor or the entire target tensor. In this case,
        # we need to set the target tensor to the part of the target tensor that was passed into the model while
        # also padding it to the correct length with IGNORE_INDEX_TOKEN_ID.
        else:
            padding = torch.full((last_matching_index,), IGNORE_INDEX_TOKEN_ID).to(device=target_device)
            updated_targets.append(torch.cat((padding, target), dim=-1)[:prediction_length])

    targets[of_name] = torch.stack(updated_targets).to(device=targets.get(of_name).device, dtype=torch.int64)

    return targets


def generate_merged_ids(
    input_ids: torch.tensor,
    target_ids: torch.tensor,
    tokenizer: PreTrainedTokenizer,
    max_sequence_length: int | None = None,
):
    """Generate merged input and target IDs tensor.

    This function merges the input_ids and target_ids together to create a unified tensor
    to pass into the model. It also returns attention masks for the merged tensors.

    Args:
        input_ids (torch.Tensor): The input IDs tensor.
        target_ids (torch.Tensor or None): The target IDs tensor or None.
        max_sequence_length (int or None): The maximum sequence length to pad or truncate to.
        tokenizer (PreTrainedTokenizer): The tokenizer used to encode the input_ids and target_ids.

    Returns:
        torch.Tensor: The merged input and target IDs tensor.
        torch.Tensor: The attention masks for the merged tensor.
    """
    merged_input_and_targets = []
    lengths = []

    eos_tensor = torch.tensor([tokenizer.eos_token_id]).to(target_ids[0].device)

    # Merge input_ids and target_ids by concatenating them together.
    # We remove the left padding from both input_ids and target_ids before concatenating them.
    for input_id_sample, target_id_sample in zip(input_ids, target_ids):
        input_id_sample_no_padding = remove_left_padding(input_id_sample, tokenizer)[0]
        target_id_sample_no_padding = remove_left_padding(target_id_sample, tokenizer)[0]
        target_id_sample_no_padding = torch.cat((target_id_sample_no_padding, eos_tensor), dim=-1)

        merged_sample_ids = torch.cat((input_id_sample_no_padding, target_id_sample_no_padding), dim=-1)
        # If the merged tensor is longer than the maximum sequence length, we truncate it.
        if max_sequence_length and merged_sample_ids.shape[0] > max_sequence_length:
            merged_sample_ids = merged_sample_ids[:max_sequence_length]

        merged_input_and_targets.append(merged_sample_ids)
        lengths.append(merged_sample_ids.shape[0])

    # Since we remove the left padding from the target_ids, the merged input_ids and target_ids
    # may not have the same lengths. We need to align them to the same length by adding left padding
    # and generate an attention mask for just the part of the input that is not padding.
    max_length = max(lengths)
    attention_masks = []
    for i, merged_sample_ids in enumerate(merged_input_and_targets):
        merged_input_and_targets[i] = add_left_padding(merged_sample_ids, max_length)
        attention_masks.append(create_attention_mask(merged_input_and_targets[i], tokenizer))

    return torch.stack(merged_input_and_targets), torch.stack(attention_masks)


def _get_decoded_targets_and_predictions(
    targets: dict[str, torch.Tensor],
    predictions: dict[str, dict[str, torch.Tensor]],
    tokenizer: PreTrainedTokenizer,
    of_name: str,
):
    """Returns the decoded targets and predictions, accounting for IGNORE_INDEX_TOKEN_ID."""
    target_tensor = targets[of_name]
    pred_tensor = predictions[of_name][PREDICTIONS]
    # Ensure targets and predictions are on the same device
    if target_tensor.device != pred_tensor.device:
        target_tensor = target_tensor.to(pred_tensor.device)
    sanitized_targets = torch.where(target_tensor != IGNORE_INDEX_TOKEN_ID, target_tensor, tokenizer.pad_token_id)
    sanitized_predictions = torch.where(
        pred_tensor != IGNORE_INDEX_TOKEN_ID,
        pred_tensor,
        tokenizer.pad_token_id,
    )
    decoded_targets = tokenizer.batch_decode(sanitized_targets, skip_special_tokens=True)
    decoded_predictions = tokenizer.batch_decode(sanitized_predictions, skip_special_tokens=True)
    return decoded_targets, decoded_predictions


def get_realigned_target_and_prediction_tensors_for_inference(
    targets: dict[str, torch.Tensor],
    predictions: dict[str, dict[str, torch.Tensor]],
    of_name: str,
    tokenizer: PreTrainedTokenizer,
    pad_value: int | None = None,
) -> tuple[dict[str, torch.Tensor], dict[str, torch.Tensor]]:
    """Realigns the target tensor with the predictions.

    This is necessary for text metrics that require the target and prediction to be of the same length.

    Args:
        targets: The target tensor.
        predictions: The prediction tensor.
        of_name: The output feature's name.
        tokenizer: The HF tokenizer.
        pad_direction: The direction to pad the tensors. Can be 'left' or 'right'.
            Defaults to 'right'.

    Returns:
        Tuple of realigned (targets, decoded_targets, predictions, decoded_predictions).
        - targets is a map of feature name -> tensor of token ids.
        - predictions is a map from output feature name -> map of tensors with the following items:
            - "predictions": tensor of token ids.
            - "probabilities": tensor of probabilities.
            - "logits": tensor of logits.
    """
    target_length = targets.get(of_name).size()[1]
    prediction_length = predictions[of_name].get(PREDICTIONS).size()[1]

    if target_length == prediction_length:
        return targets, predictions

    if not pad_value:
        pad_value = tokenizer.pad_token_id if tokenizer.pad_token_id is not None else tokenizer.eos_token_id

    zeros_to_add = (
        target_length - prediction_length if target_length > prediction_length else prediction_length - target_length
    )

    # We don't want to modify the original targets and predictions tensors, so we create a copy of them.
    _targets = copy.deepcopy(targets)
    _predictions = copy.deepcopy(predictions)

    # Align target and prediction tensors for text to text metric computation
    if target_length > prediction_length:
        # Pad the predictions.
        _predictions[of_name][PREDICTIONS] = F.pad(
            _predictions[of_name][PREDICTIONS], (0, zeros_to_add), value=pad_value
        ).to(torch.int64)

        _predictions[of_name][PROBABILITIES] = F.pad(_predictions[of_name][PROBABILITIES], (0, 0, 0, zeros_to_add)).to(
            torch.float32
        )

        _predictions[of_name][LOGITS] = F.pad(_predictions[of_name][LOGITS], (0, 0, 0, zeros_to_add)).to(torch.float32)
    else:
        _targets[of_name] = F.pad(_targets[of_name], (0, zeros_to_add), value=pad_value).to(torch.int64)

    return _targets, _predictions


def update_embedding_layer(model: AutoModelForCausalLM, config_obj: LLMTrainerConfig) -> AutoModelForCausalLM:
    """Updates the embedding layer of the model to use the 8-bit embedding layer from bitsandbytes.nn.modules.

    This is necessary when using 8-bit optimizers from bitsandbytes.
    See: https://github.com/TimDettmers/bitsandbytes#tldr
    """
    # If we're using an 8-bit optimizer, we need to replace the embedding layer with a custom embedding layer from
    # bnb.nn.modules.Embedding.
    if hasattr(config_obj, "optimizer") and config_obj.optimizer.is_8bit:
        embedding_layer, module_path = find_embedding_layer_with_path(model)
        if embedding_layer is None:
            raise ValueError(
                "Could not find an embedding layer in the model. This is required when using 8-bit optimizers"
                "  since a custom 8-bit embedding layer is used in place of the original embedding layer."
            )

        # Initialize the BNB embedding layer with the same parameters and weights as the original embedding layer.
        bnb_embedding = BnbEmbedding(
            num_embeddings=embedding_layer.num_embeddings,
            embedding_dim=embedding_layer.embedding_dim,
            padding_idx=embedding_layer.padding_idx,
            max_norm=embedding_layer.max_norm,
            norm_type=embedding_layer.norm_type,
            scale_grad_by_freq=embedding_layer.scale_grad_by_freq,
            sparse=embedding_layer.sparse,
            _weight=embedding_layer.weight,
            device=model.device,
        )

        # Update the model's original embedding layer to use the BNB embedding layer using the module_path
        # returned by find_embedding_layer_with_path.
        module_path = module_path.split(".")
        module = model
        for module_name in module_path[:-1]:
            module = getattr(module, module_name)
        setattr(module, module_path[-1], bnb_embedding)

        # Set the get input embeddings lambda function to return the BNB embedding layer
        model.get_input_embeddings = lambda: bnb_embedding

        logger.info("Updated the pretrained embedding layer to use the embedding layer from bitsandbytes.")

    return model


def create_text_streamer(tokenizer: PreTrainedTokenizer) -> TextStreamer:
    """Creates a TextStreamer object for streaming text to stdout during generation."""
    return TextStreamer(tokenizer=tokenizer, skip_prompt=True)


def generate_merged_ids_packed(
    input_ids: torch.tensor,
    target_ids: torch.tensor,
    tokenizer,
    max_sequence_length: int | None = None,
    max_sequences_per_pack: int = 8,
):
    """Generate merged IDs with sequence packing for throughput improvement.

    Instead of padding each sequence to the same length, packs multiple short
    sequences into a single batch entry with block-diagonal attention masks
    to prevent cross-sequence attention.

    Args:
        input_ids: [batch, input_len] token IDs for prompts
        target_ids: [batch, target_len] token IDs for completions
        tokenizer: HuggingFace tokenizer
        max_sequence_length: Maximum pack length
        max_sequences_per_pack: Maximum sequences per pack

    Returns:
        packed_ids: [num_packs, max_seq_len] packed token IDs
        packed_attention_mask: [num_packs, max_seq_len, max_seq_len] block-diagonal attention
    """
    from ludwig.utils.sequence_packing import pack_sequences

    # First, merge each input+target pair into a single sequence (without padding)
    merged_sequences = []
    eos_tensor = torch.tensor([tokenizer.eos_token_id]).to(target_ids[0].device)
    pad_token_id = tokenizer.pad_token_id if tokenizer.pad_token_id is not None else 0

    for input_id_sample, target_id_sample in zip(input_ids, target_ids):
        input_no_pad = remove_left_padding(input_id_sample, tokenizer)[0]
        target_no_pad = remove_left_padding(target_id_sample, tokenizer)[0]
        target_no_pad = torch.cat((target_no_pad, eos_tensor), dim=-1)
        merged = torch.cat((input_no_pad, target_no_pad), dim=-1)
        if max_sequence_length and merged.shape[0] > max_sequence_length:
            merged = merged[:max_sequence_length]
        merged_sequences.append(merged)

    # Create dummy attention masks (all ones, actual masking done by block-diagonal)
    attention_masks = [torch.ones(seq.shape[0]) for seq in merged_sequences]

    # Pack sequences using greedy bin packing
    packed_ids, packed_attn, _, _ = pack_sequences(
        merged_sequences,
        attention_masks,
        max_length=max_sequence_length or max(s.shape[0] for s in merged_sequences),
        pad_token_id=pad_token_id,
        max_sequences_per_pack=max_sequences_per_pack,
    )

    return packed_ids, packed_attn