utils.py

import unsloth
import os
from enum import Enum

from unsloth import to_sharegpt
from unsloth import standardize_sharegpt

import packaging.version
import torch
import transformers
from datasets import DatasetDict, load_dataset, load_from_disk
from datasets.builder import DatasetGenerationError
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    BitsAndBytesConfig,
)

from peft import LoraConfig

import json
from tqdm import tqdm
from datasets import Dataset
import torch.nn.functional as F


CHATML_CHAT_TEMPLATE_QWEN = "{% for message in messages %}\n{{'<|im_start|>' + message['role'] + '\n' + message['content'] + '<|im_end|>' + '\n'}}{% if loop.last and add_generation_prompt %}{{'<|im_start|>assistant\n' }}{% endif %}{% endfor %}"
CHATML_CHAT_TEMPLATE_LLAMA = "{% for message in messages %}\n{{'<|start_header_id|>' + message['role'] + '<|end_header_id|>' + '\n' + message['content'] + '<|eot_id|>' + '\n'}}{% if loop.last and add_generation_prompt %}{{'<|start_header_id|>assistant<|end_header_id|>\n' }}{% endif %}{% endfor %}"
# CHATML_CHAT_TEMPLATE_LLAMA = "{{- 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    {%- set date_string = \"26 Jul 2024\" %}\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 + builtin tools #}\n{{- \"<|start_header_id|>system<|end_header_id|>\\n\\n\" }}\n{%- if builtin_tools is defined or tools is not none %}\n    {{- \"Environment: ipython\\n\" }}\n{%- endif %}\n{%- if builtin_tools is defined %}\n    {{- \"Tools: \" + builtin_tools | reject('equalto', 'code_interpreter') | join(\", \") + \"\\n\\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        {%- if builtin_tools is defined and tool_call.name in builtin_tools %}\n            {{- '<|start_header_id|>assistant<|end_header_id|>\\n\\n' -}}\n            {{- \"<|python_tag|>\" + tool_call.name + \".call(\" }}\n            {%- for arg_name, arg_val in tool_call.arguments | items %}\n                {{- arg_name + '=\"' + arg_val + '\"' }}\n                {%- if not loop.last %}\n                    {{- \", \" }}\n                {%- endif %}\n                {%- endfor %}\n            {{- \")\" }}\n        {%- else  %}\n            {{- '<|start_header_id|>assistant<|end_header_id|>\\n\\n' -}}\n            {{- '{\"name\": \"' + tool_call.name + '\", ' }}\n            {{- '\"parameters\": ' }}\n            {{- tool_call.arguments | tojson }}\n            {{- \"}\" }}\n        {%- endif %}\n        {%- if builtin_tools is defined %}\n            {#- This means we're in ipython mode #}\n            {{- \"<|eom_id|>\" }}\n        {%- else %}\n            {{- \"<|eot_id|>\" }}\n        {%- endif %}\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"
DEFAULT_ZEPHYR_CHAT_TEMPLATE = "{% for message in messages %}\n{% if message['role'] == 'user' %}\n{{ '<|user|>\n' + message['content'] + eos_token }}\n{% elif message['role'] == 'system' %}\n{{ '<|system|>\n' + message['content'] + eos_token }}\n{% elif message['role'] == 'assistant' %}\n{{ '<|assistant|>\n'  + message['content'] + eos_token }}\n{% endif %}\n{% if loop.last and add_generation_prompt %}\n{{ '<|assistant|>' }}\n{% endif %}\n{% endfor %}"


class ZephyrSpecialTokens(str, Enum):
    user = "<|user|>"
    assistant = "<|assistant|>"
    system = "<|system|>"
    eos_token = "</s>"
    bos_token = "<s>"
    pad_token = "<pad>"

    @classmethod
    def list(cls):
        return [c.value for c in cls]


class ChatmlSpecialTokensLlama(str, Enum):
    # For Qwen 2.5 14b instruct
    # user = "<|im_start|>user"
    # assistant = "<|im_start|>assistant"
    # system = "<|im_start|>system"
    # eos_token = "<|im_end|>"
    # bos_token = "<s>"
    # pad_token = "<pad>"

    # For llama 3.1 8b instruct
    user = "<|start_header_id|>user<|end_header_id|>"
    assistant = "<|start_header_id|>assistant<|end_header_id|>"
    system = "<|start_header_id|>system<|end_header_id|>"
    eos_token = "<|eot_id|>"
    bos_token = "<|begin_of_text|>"
    pad_token = "<pad>"

    # prefix = "<prefix>"
    # answer = "<answer>"
    # knowledge_0 = "<knowledge_0>"
    # knowledge_1 = "<knowledge_1>"
    # knowledge_2 = "<knowledge_2>"
    # knowledge_3 = "<knowledge_3>"
    # knowledge_4 = "<knowledge_4>"
    # knowledge_5 = "<knowledge_5>"
    # knowledge_6 = "<knowledge_6>"
    # knowledge_7 = "<knowledge_7>"

    # reason_0 = "<reason_0>"
    # reason_1 = "<reason_1>"
    # reason_2 = "<reason_2>"
    # reason_3 = "<reason_3>"
    # reason_4 = "<reason_4>"
    # reason_5 = "<reason_5>"
    # reason_6 = "<reason_6>"
    # reason_7 = "<reason_7>"
    
    @classmethod
    def list(cls):
        return [c.value for c in cls]

class ChatmlSpecialTokensQwen(str, Enum):
    # For Qwen 2.5 14b instruct
    user = "<|im_start|>user"
    assistant = "<|im_start|>assistant"
    system = "<|im_start|>system"
    eos_token = "<|im_end|>"
    bos_token = "<s>"
    pad_token = "<pad>"

    # For llama 3.1 8b instruct
    # user = "<|start_header_id|>user<|end_header_id|>"
    # assistant = "<|start_header_id|>assistant<|end_header_id|>"
    # system = "<|start_header_id|>system<|end_header_id|>"
    # eos_token = "<|eot_id|>"
    # bos_token = "<|begin_of_text|>"
    # pad_token = "<pad>"

    # prefix = "<prefix>"
    # answer = "<answer>"
    # knowledge_0 = "<knowledge_0>"
    # knowledge_1 = "<knowledge_1>"
    # knowledge_2 = "<knowledge_2>"
    # knowledge_3 = "<knowledge_3>"
    # knowledge_4 = "<knowledge_4>"
    # knowledge_5 = "<knowledge_5>"
    # knowledge_6 = "<knowledge_6>"
    # knowledge_7 = "<knowledge_7>"

    # reason_0 = "<reason_0>"
    # reason_1 = "<reason_1>"
    # reason_2 = "<reason_2>"
    # reason_3 = "<reason_3>"
    # reason_4 = "<reason_4>"
    # reason_5 = "<reason_5>"
    # reason_6 = "<reason_6>"
    # reason_7 = "<reason_7>"
    
    @classmethod
    def list(cls):
        return [c.value for c in cls]

special_tokens_list = []
new_token_dict = []
initialize_tokens = []
def prepare_prefix_prompt(num_prefix=2, num_knowledge=4):
    global special_tokens_list
    global new_token_dict
    global initialize_tokens
    
    # new_token_dict = {'prefix': '', 'answer': '', 'knowledge_0': '', 'knowledge_1': '', 'knowledge_2': '', 'knowledge_3': '', 'knowledge_4': '', 'knowledge_5': '', 'knowledge_6': '', 'knowledge_7': '', 'reason_0': '', 'reason_1': '', 'reason_2': '', 'reason_3': '', 'reason_4': '', 'reason_5': '', 'reason_6': '', 'reason_7': '', 'reason_8': ''}
    new_token_dict = {'prefix': '', 'answer': '', 'knowledge_0': '', 'knowledge_1': '', 'knowledge_2': '', 'knowledge_3': '', 'knowledge_4': '', 'knowledge_5': '', 'knowledge_6': '', 'knowledge_7': '','reason_0': '', 'reason_1': '', 'reason_2': '', 'reason_3': '', 'reason_4': '', 'reason_5': '', 'reason_6': '', 'reason_7': ''}
    # since we are doing sft, if we add <question> as the prefix of question sentence, I don't know whether we can update the <question> embedding.
    # so here, we using the original embedding of <quesiton> instead of regarding it as a new token.
    special_tokens_list = []
    initialize_tokens = []
    for k in new_token_dict:
        text = ''
        if k == 'prefix':
            num_tokens = num_prefix
        else:
            num_tokens = num_knowledge
        for i in range(num_tokens):
            if k == 'prefix' or k == 'answer' or k == 'query':
                # token_name = f'<{k}_{i}>'
                token_name = f'<{k}>'
                special_tokens_list.append(token_name)
                initialize_tokens.append(k)
                text += ' ' + token_name
                # text += token_name
            elif k in ['reason_0','reason_1','reason_2','reason_3','reason_4','reason_5','reason_6','reason_7','knowledge_0','knowledge_1','knowledge_2','knowledge_3','knowledge_4','knowledge_5','knowledge_6','knowledge_7']:
                token_name = f"<{k}>"
                # special_tokens_list.append(token_name)
                # initialize_words_list.append(k.split('_')[0].strip())
                text += ' ' + token_name
                # text += token_name
            else:
                print(f"wrong token...{k}")
        if k in ['reason_0','reason_1','reason_2','reason_3','reason_4','reason_5','reason_6','reason_7','knowledge_0','knowledge_1','knowledge_2','knowledge_3','knowledge_4','knowledge_5','knowledge_6','knowledge_7']:
            token_name = f"<{k}>"
            special_tokens_list.append(token_name)
            initialize_tokens.append(k.split('_')[0].strip())
        new_token_dict[k] = text

    # print(f"special_tokens_list: {special_tokens_list}")
    # print(f"new_token_dict: {new_token_dict}")
    # print(f"initialize_tokens: {initialize_tokens}")

def prepare_prefix_prompt_disentangle(num_prefix=2, num_knowledge=4):
    global special_tokens_list
    global new_token_dict
    global initialize_tokens
    
    # new_token_dict = {'prefix': '', 'answer': '', 'knowledge_0': '', 'knowledge_1': '', 'knowledge_2': '', 'knowledge_3': '', 'knowledge_4': '', 'knowledge_5': '', 'knowledge_6': '', 'knowledge_7': '', 'reason_0': '', 'reason_1': '', 'reason_2': '', 'reason_3': '', 'reason_4': '', 'reason_5': '', 'reason_6': '', 'reason_7': '', 'reason_8': ''}
    new_token_dict = {'prefix': '', 'answer': '', 'knowledge': '', 'reason': ''}
    # since we are doing sft, if we add <question> as the prefix of question sentence, I don't know whether we can update the <question> embedding.
    # so here, we using the original embedding of <quesiton> instead of regarding it as a new token.
    special_tokens_list = []
    initialize_tokens = []
    for k in new_token_dict:
        text = ''
        if k == 'prefix':
            num_tokens = num_prefix
        else:
            num_tokens = num_knowledge
        for i in range(num_tokens):
            if k == 'prefix' or k == 'answer':
                token_name = f'<{k}_{i}>'
                # token_name = f'<{k}>'
                special_tokens_list.append(token_name)
                initialize_tokens.append(k)
                text += ' ' + token_name
                # text += token_name
            elif k in ['reason', 'knowledge']:
                # token_name = f"<{k}>"
                token_name = f'<{k}_{i}>'
                special_tokens_list.append(token_name)
                initialize_tokens.append(k)
                text += ' ' + token_name
                # text += token_name
            else:
                print(f"wrong token...{k}")
        # if k in ['reason_0','reason_1','reason_2','reason_3','reason_4','reason_5','reason_6','reason_7','knowledge_0','knowledge_1','knowledge_2','knowledge_3','knowledge_4','knowledge_5','knowledge_6','knowledge_7']:
        #     token_name = f"<{k}>"
        #     special_tokens_list.append(token_name)
        #     initialize_tokens.append(k.split('_')[0].strip())
        new_token_dict[k] = text

def load_data(data_path, data_type = "train", data_format='reason'):
    global special_tokens_list
    global new_token_dict
    global initialize_tokens

    datas = []
    print(f"load original data...")

    with open(data_path, 'r', encoding='utf-8') as file:
        datas = json.load(file)

    # TODO: add options to strategyQA in advance
    if 'StrategyQA' in data_path:
        StrategyQA_option_flag = True
    else:
        StrategyQA_option_flag = False

    formatted_data = []
    
    if data_format == 'think-reason':
        prepare_prefix_prompt(num_prefix=1, num_knowledge=1)
    elif data_format == 'disentangle':
        prepare_prefix_prompt_disentangle(num_prefix=2, num_knowledge=4)
    else:
        prepare_prefix_prompt(num_prefix=2, num_knowledge=4)
    

    for i in tqdm(range(len(datas)), total=len(datas)):
        # prompt = f"<GSM8K>{datas[i]['prompt']}"
        try:
            data = datas[i]
            question = data['question']
            answer = data['answer']
            cot_steps = data['cot_steps']
            if data_format == 'think' or data_format == 'think-reason':
                think_steps = data['think']
            else:
                think_steps = ''

            if data_format == 'reason' or data_format == 'disentangle':
                tmp_sol = ''
                for j in range(len(cot_steps)):
                    cot_step = cot_steps[j]

                    label = cot_step['requirement'][1:-1]
                    rest_of_step = cot_step['content']

                    added_tokens = new_token_dict['prefix'] + new_token_dict[label]
                    new_cot_step = added_tokens + " " + rest_of_step + '\n '

                    tmp_sol += new_cot_step

                new_answer = new_token_dict['prefix'] + new_token_dict['answer'] + ' The answer is: ' + answer + '\n '

                sol = tmp_sol + new_answer
                
            elif data_format == 'direct':
                sol = 'The answer is: ' + answer + '\n'

            elif data_format == 'cot' or data_format == 'react':
                sol = cot_steps + '\n' + 'The answer is: ' + answer + '\n'

            elif data_format == 'think':
                sol = '<think>' + '\n' + think_steps + '\n' + '</think>' + '\n' + cot_steps + '\n' + 'The answer is: ' + answer + '\n'

            elif data_format == 'think-reason':
                tmp_sol = ''
                
                for j in range(len(cot_steps)):
                    cot_step = cot_steps[j]

                    label = cot_step['requirement'][1:-1]
                    rest_of_step = cot_step['content']

                    # added_tokens = new_token_dict['prefix'] + new_token_dict[label]
                    # new_cot_step = added_tokens + " " + rest_of_step + '\n '
                    added_tokens = new_token_dict[label]
                    new_cot_step = added_tokens + " " + rest_of_step + " " + f"</{label}>" + '\n '

                    tmp_sol += new_cot_step

                # sol = '<think>' + '\n' + think_steps + '\n' + '</think>' + '\n' + tmp_sol + '\n' + new_token_dict['prefix'] + new_token_dict['answer'] + ' The answer is: ' + answer + '\n'
                sol = '<think>' + '\n' + think_steps + '\n' + '</think>' + '\n' + tmp_sol + '\n' + 'The answer is: ' + answer + '\n'

            else:
                print(f"no valid data format found")

            if data_type == "test":
                # prompt = '<question> <question>' + ' ' + question + '\n' + new_token_dict['prefix']
                if data_format == 'react':
                    prompt = question
                else:
                    prompt = 'Question:' + ' ' + question
            else:
                if data_format == 'react':
                    prompt = question
                else:
                    prompt = 'Question:' + ' ' + question
            
            if StrategyQA_option_flag:
                prompt += " Options: True, False"

            # print(prompt)
            # print(sol)

            # print('*' * 100)
            # print()

            if data_format == 'reason' or data_format == 'disentangle':
                instruction = 'Answer the following single-choice question step by step using the supporting facts in your knowledge.'
            elif data_format == 'direct':
                instruction = ''
            elif data_format == 'cot':
                instruction = ''
            elif data_format == 'react':
                instruction = """Solve a multiple choice question answering task with interleaving Thought, Action, Observation steps. Thought can reason about the current situation, and Action can be three types: 
(1) Search[entity], which searches the exact entity on Wikipedia and returns the first paragraph if it exists. If not, it will return some similar entities to search.
(2) Lookup[keyword], which returns the next sentence containing keyword in the current passage.
(3) Finish[answer], which returns the answer (A, B, C, D, etc.) and finishes the task.""".strip()
            elif data_format == 'think':
                instruction = ''
            elif data_format == 'think-reason':
                instruction = ''
            else:
                print(f"no valid data format found")

            formatted_data.append({
                # "instruction": 'Answer the following question True or False step by step using the supporting facts in your knowledge.',
                "instruction": instruction,
                "input": prompt,
                "output": sol
            })
            # print(formatted_data[-1])
        except:
            print(i)

    return formatted_data


def create_datasets(tokenizer, data_args, training_args, apply_chat_template=False):
    # def preprocess(samples):
    #     batch = []
    #     for conversation in samples["messages"]:
    #         batch.append(tokenizer.apply_chat_template(conversation, tokenize=False))
    #     return {"content": batch}

    def preprocess(samples):
        batch = []
        for conversation in samples["conversations"]:
            batch.append(tokenizer.apply_chat_template(conversation, tokenize=False))
        return {"content": batch}
        # convos = examples["conversations"]
        # texts = [tokenizer.apply_chat_template(convo, tokenize = False, add_generation_prompt = False) for convo in convos]
        # return { "text" : texts, }

    raw_datasets = DatasetDict()

    dataset_set = data_args.dataset_name.split(',')

    for split in data_args.splits.split(","):
        # try:
        #     # Try first if dataset on a Hub repo
        #     dataset = load_dataset(data_args.dataset_name, split=split)
        # except DatasetGenerationError:
        #     # If not, check local dataset
        #     # dataset = load_from_disk(os.path.join(data_args.dataset_name, split))

        data_format = data_args.data_format
        if data_format not in ["direct", "reason", "cot", "think", "think-reason", "disentangle", "react"]:
            raise ValueError(f"data_format must be either 'direct', 'reason', 'cot', 'disentangle', 'think', 'think-reason' or 'react', got {data_format}")

        if "train" in split:
            train_data = load_data(dataset_set[0], data_type='train', data_format=data_format)
            # 将train_data封装为load_dataset的格式
            train_dataset = Dataset.from_dict({
                "instruction": [item["instruction"] for item in train_data],
                "input": [item["input"] for item in train_data],
                "output": [item["output"] for item in train_data]
            })

            if data_format == 'reason':
                merged_prompt = "{instruction}\n{input}"
            elif data_format == 'direct':
                merged_prompt = "{instruction}{input}"
            elif data_format == 'cot':
                merged_prompt = "{instruction}{input}"
            elif data_format == 'think':
                merged_prompt = "{instruction}{input}"
            elif data_format == 'think-reason':
                merged_prompt = "{instruction}{input}"
            elif data_format == 'disentangle':
                merged_prompt = "{instruction}\n{input}"
            elif data_format == 'react':    
                merged_prompt = "{instruction}\n{input}"
            else:
                print(f"no valid data format found")

            train_dataset = to_sharegpt(
                train_dataset,
                # merged_prompt = "{instruction}[[\nYour input is:\n{input}]]",
                merged_prompt = merged_prompt,
                output_column_name = "output",
                conversation_extension = 1, # Select more to handle longer conversations
            )

            train_dataset = standardize_sharegpt(train_dataset)

            raw_datasets["train"] = train_dataset
        elif "test" in split:
            eval_data = load_data(dataset_set[1], data_type='test', data_format=data_format)

            # 将eval_data封装为load_dataset的格式
            eval_dataset = Dataset.from_dict({
                "instruction": [item["instruction"] for item in eval_data],
                "input": [item["input"] for item in eval_data],
                "output": [item["output"] for item in eval_data]
            })

            if data_format == 'reason':
                merged_prompt = "{instruction}\n{input}"
            elif data_format == 'direct':
                merged_prompt = "{instruction}{input}"
            elif data_format == 'cot':
                merged_prompt = "{instruction}{input}"
            elif data_format == 'think':
                merged_prompt = "{instruction}{input}"
            elif data_format == 'think-reason':
                merged_prompt = "{instruction}{input}"
            elif data_format == 'disentangle':
                merged_prompt = "{instruction}\n{input}"
            elif data_format == 'react':
                merged_prompt = "{instruction}\n{input}"
            else:
                print(f"no valid data format found")

            eval_dataset = to_sharegpt(
                eval_dataset,
                # merged_prompt = "{instruction}[[\nYour input is:\n{input}]]",
                merged_prompt = merged_prompt,
                output_column_name = "output",
                conversation_extension = 1, # Select more to handle longer conversations
            )

            eval_dataset = standardize_sharegpt(eval_dataset)
            raw_datasets["test"] = eval_dataset

            
        else:
            raise ValueError(f"Split type {split} not recognized as one of test or train.")

    if apply_chat_template:
        raw_datasets = raw_datasets.map(
            preprocess,
            batched=True,
            # remove_columns=raw_datasets["train"].column_names,
        )

    train_data = raw_datasets["train"]
    valid_data = raw_datasets["test"]
    print(f"Size of the train set: {len(train_data)}. Size of the validation set: {len(valid_data)}")
    print(f"A sample of train dataset: {train_data[0]}")

    return train_data, valid_data


def create_and_prepare_model(args, data_args, training_args):
    global special_tokens_list
    global initialize_tokens

    if args.use_unsloth:
        from unsloth import FastLanguageModel
    bnb_config = None
    quant_storage_dtype = None

    if (
        torch.distributed.is_available()
        and torch.distributed.is_initialized()
        and torch.distributed.get_world_size() > 1
        and args.use_unsloth
    ):
        raise NotImplementedError("Unsloth is not supported in distributed training")

    if args.use_4bit_quantization:
        compute_dtype = getattr(torch, args.bnb_4bit_compute_dtype)
        quant_storage_dtype = getattr(torch, args.bnb_4bit_quant_storage_dtype)

        bnb_config = BitsAndBytesConfig(
            load_in_4bit=args.use_4bit_quantization,
            bnb_4bit_quant_type=args.bnb_4bit_quant_type,
            bnb_4bit_compute_dtype=compute_dtype,
            bnb_4bit_use_double_quant=args.use_nested_quant,
            bnb_4bit_quant_storage=quant_storage_dtype,
        )

        if compute_dtype == torch.float16 and args.use_4bit_quantization:
            major, _ = torch.cuda.get_device_capability()
            if major >= 8:
                print("=" * 80)
                print("Your GPU supports bfloat16, you can accelerate training with the argument --bf16")
                print("=" * 80)
        elif args.use_8bit_quantization:
            bnb_config = BitsAndBytesConfig(load_in_8bit=args.use_8bit_quantization)

    if args.use_unsloth:
        # Load model
        model, _ = FastLanguageModel.from_pretrained(
            model_name=args.model_name_or_path,
            max_seq_length=data_args.max_seq_length,
            dtype=None,
            load_in_4bit=args.use_4bit_quantization,
        )
    else:
        torch_dtype = (
            quant_storage_dtype if quant_storage_dtype and quant_storage_dtype.is_floating_point else torch.float32
        )
        model = AutoModelForCausalLM.from_pretrained(
            args.model_name_or_path,
            quantization_config=bnb_config,
            trust_remote_code=True,
            attn_implementation="flash_attention_2" if args.use_flash_attn else "eager",
            torch_dtype=torch_dtype,
            # device_map="auto",
        )

    peft_config = None
    chat_template = None
    if args.use_peft_lora and not args.use_unsloth:
        peft_config = LoraConfig(
            lora_alpha=args.lora_alpha,
            lora_dropout=args.lora_dropout,
            r=args.lora_r,
            bias="none",
            task_type="CAUSAL_LM",
            target_modules=args.lora_target_modules.split(",")
            if args.lora_target_modules != "all-linear"
            else args.lora_target_modules,
            # modules_to_save=[
            #     "lm_head",
            #     "embed_tokens"
            # ],
            # trainable_token_indices = {'embed_tokens': tokenizer.convert_tokens_to_ids(special_tokens.list())},
        )

    special_tokens = None
    chat_template = None
    if args.chat_template_format == "chatml":
        special_tokens = ChatmlSpecialTokensLlama if 'llama' in args.model_name_or_path else ChatmlSpecialTokensQwen
        # chat_template = DEFAULT_CHATML_CHAT_TEMPLATE
        chat_template = CHATML_CHAT_TEMPLATE_LLAMA if 'llama' in args.model_name_or_path else CHATML_CHAT_TEMPLATE_QWEN
    elif args.chat_template_format == "zephyr":
        special_tokens = ZephyrSpecialTokens
        chat_template = DEFAULT_ZEPHYR_CHAT_TEMPLATE

    print(f"chat_template: {chat_template}")

    if special_tokens is not None:

        # print(f"special tokens: {special_tokens.list()}")
        tokenizer = AutoTokenizer.from_pretrained(
            args.model_name_or_path,
            pad_token=special_tokens.pad_token.value,
            bos_token=special_tokens.bos_token.value,
            eos_token=special_tokens.eos_token.value,
            additional_special_tokens=special_tokens.list(),
            trust_remote_code=True,
        )
        tokenizer.chat_template = chat_template

        # make embedding resizing configurable?
        # Transformers 4.46.0+ defaults uses mean_resizing by default, which fails with QLoRA + FSDP because the
        # embedding could be on meta device, therefore, we set mean_resizing=False in that case (i.e. the status quo
        # ante). See https://github.com/huggingface/accelerate/issues/1620.
        uses_transformers_4_46 = packaging.version.parse(transformers.__version__) >= packaging.version.parse("4.46.0")
        uses_fsdp = os.environ.get("ACCELERATE_USE_FSDP").lower() == "true"
        if (bnb_config is not None) and uses_fsdp and uses_transformers_4_46:
            model.resize_token_embeddings(len(tokenizer), pad_to_multiple_of=8, mean_resizing=False)
        else:
            model.resize_token_embeddings(len(tokenizer), pad_to_multiple_of=8)
    else:
        tokenizer = AutoTokenizer.from_pretrained(args.model_name_or_path, trust_remote_code=True)
        tokenizer.pad_token = tokenizer.eos_token


    # initialize_tokens = tokenizer.convert_tokens_to_ids(initialize_tokens)
    # n_tokens = len(special_tokens_list)
    # model.n_tokens = n_tokens
    # model.prompt_tokens = None

    # orig_vocab_size = model.get_input_embeddings().weight.size(0)
    # initialize_tokens = torch.tensor(initialize_tokens, 
    #                                 dtype=torch.long, device=model.device)
    
    # model.config.vocab_size = orig_vocab_size + n_tokens

    # model.set_input_embeddings(InputEmbedding(
    #     model.get_input_embeddings(), n_tokens, initialize_tokens))
    # model.set_output_embeddings(OutputEmbedding(
    #     model.get_output_embeddings(), n_tokens, initialize_tokens))


    if args.use_unsloth:
        # Do model patching and add fast LoRA weights
        model = FastLanguageModel.get_peft_model(
            model,
            lora_alpha=args.lora_alpha,
            lora_dropout=args.lora_dropout,
            r=args.lora_r,
            target_modules=args.lora_target_modules.split(",")
            if args.lora_target_modules != "all-linear"
            else args.lora_target_modules,
            use_gradient_checkpointing=training_args.gradient_checkpointing,
            random_state=training_args.seed,
            max_seq_length=data_args.max_seq_length,
        )

    return model, peft_config, tokenizer