# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/ml.llama.ipynb.

# %% auto 0
__all__ = ['LLAMA2_CHAT_PROMPT_TEMPLATE', 'prepare_llama2_for_training', 'prepare_llama2_for_inference']

# %% ../../nbs/ml.llama.ipynb 3
LLAMA2_CHAT_PROMPT_TEMPLATE = """<s>[INST] <<SYS>>
{system_prompt}
<</SYS>>

{user_message} [/INST]"""

# %% ../../nbs/ml.llama.ipynb 4
def prepare_llama2_for_training(tokenizer, model):
    tokenizer.pad_token = tokenizer.eos_token
    tokenizer.padding_side = "right"  # Fix weird overflow issue with fp16 training
    model.config.pretraining_tp = 1
    model.config.use_cache = False

def prepare_llama2_for_inference(tokenizer, model):
    tokenizer.pad_token = tokenizer.eos_token
    tokenizer.padding_side = "right"  # Fix weird overflow issue with fp16 training
    model.config.use_cache = True
