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

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

# %% ../../../nbs/hf.transformers.llama.ipynb 3
def prepare_llama2_for_training(tokenizer, model):
    tokenizer.pad_token = tokenizer.eos_token
    tokenizer.padding_side = "right"
    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 = "left"
    model.config.use_cache = True

# %% ../../../nbs/hf.transformers.llama.ipynb 4
def chat2text(example, key="messages", tokenizer=None):
    if tokenizer is None:
        from transformers import AutoTokenizer

        tokenizer = AutoTokenizer.from_pretrained("NousResearch/Llama-2-7b-chat-hf")
    text = tokenizer.apply_chat_template(example[key], tokenize=False)
    return {"text": text}


# %% ../../../nbs/hf.transformers.llama.ipynb 5
def _remove_special_tokens(text):
    for token in ["</s>", "<s>[INST]", "[/INST]", "<<SYS>>",  "<</SYS>>"]:
        text = text.replace(token, "")
    return text.strip()

def _extract_messages_from_conv(conv):
    inp, out = conv.split(" [/INST] ")
    inp_messages = inp.split(r" <</SYS>> ")
    if len(inp_messages) == 1:
        system_message, user_message = "", inp_messages[0]
    else:
        system_message, user_message = inp_messages
    system_message = _remove_special_tokens(system_message)
    user_message = _remove_special_tokens(user_message)
    assistant_message = _remove_special_tokens(out)
    if system_message:
        yield {'role': 'system', 'content': system_message}
    yield {'role': 'user', 'content': user_message}
    if assistant_message:
        yield {'role': 'assistant', 'content': assistant_message}

def _split_into_interactions(prompt):
    for conv in prompt.split("</s>"):
        conv = conv.strip()
        if not conv:
            continue
        yield conv

def text2chat(text: str):
    messages = [message for conv in _split_into_interactions(text) for message in _extract_messages_from_conv(conv)]
    return {"messages": messages}
