# -*- coding: utf-8 -*-

# Copyright 2023 XuMing(xuming624@qq.com) and The HuggingFace Inc. team. All rights reserved.

#

# Licensed under the Apache License, Version 2.0 (the "License");

# you may not use this file except in compliance with the License.

# You may obtain a copy of the License at

#

# http://www.apache.org/licenses/LICENSE-2.0

#

# Unless required by applicable law or agreed to in writing, software

# distributed under the License is distributed on an "AS IS" BASIS,

# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

# See the License for the specific language governing permissions and

# limitations under the License.

"""

Fine-tuning the library models for causal language modeling (GPT, GPT-2, CTRL, ...) on a text file or a dataset.

part of this code is adapted from https://github.com/shibing624/textgen

"""

import math

import os

from dataclasses import dataclass, field

from glob import glob

from typing import List, Sequence, Optional, Dict

import torch

from datasets import load_dataset

from loguru import logger

from peft import LoraConfig, TaskType, get_peft_model, PeftModel, prepare_model_for_int8_training

from transformers import (

AutoConfig,

BloomForCausalLM,

AutoModel,

AutoModelForCausalLM,

LlamaTokenizer,

LlamaForCausalLM,

BloomTokenizerFast,

AutoTokenizer,

HfArgumentParser,

Trainer,

TrainingArguments,

set_seed,

BitsAndBytesConfig,

deepspeed,

)

from transformers.deepspeed import is_deepspeed_zero3_enabled

from transformers.trainer import TRAINING_ARGS_NAME

from transformers.trainer_pt_utils import LabelSmoother

MODEL_CLASSES = {

}

IGNORE_INDEX = LabelSmoother.ignore_index

@dataclass

class ModelArguments:

raise ValueError("You must specify a valid model_max_length >= 256 to run training.")

@dataclass

class DataTrainingArguments:

)

@dataclass

class PeftArguments(TrainingArguments):

qlora: bool = field(default=False, metadata={"help": "Whether to use qlora"})

class CastOutputToFloat(torch.nn.Sequential):

return super().forward(x).to(torch.float32)

class SeparatorStyle:

PHOENIX = 9

@dataclass

class Conversation:

}

# A global registry for all conversation templates

conv_templates: Dict[str, Conversation] = {}

def register_conv_template(template: Conversation, override: bool = False):

conv_templates[template.name] = template

# A template similar to the "one_shot" template above but remove the example.

register_conv_template(

Conversation(

name="zero_shot",

system="A chat between a curious human and an artificial intelligence assistant. "

"The assistant gives helpful, detailed, and polite answers to the human's questions.",

roles=("Human", "Assistant"),

messages=[],

offset=0,

sep_style=SeparatorStyle.ADD_COLON_SINGLE,

sep="\n### ",

stop_str="###",

)

)

# Vicuna v1.1 template

register_conv_template(

Conversation(

name="vicuna",

system="A chat between a curious user and an artificial intelligence assistant. "

"The assistant gives helpful, detailed, and polite answers to the user's questions.",

roles=("USER", "ASSISTANT"),

messages=[],

offset=0,

sep_style=SeparatorStyle.ADD_COLON_TWO,

sep=" ",

sep2="</s>",

)

)

# Alpaca default template

register_conv_template(

Conversation(

name="alpaca",

system="Below is an instruction that describes a task. Write a response that appropriately completes the request.",

roles=("### Instruction", "### Response"),

messages=[],

offset=0,

sep_style=SeparatorStyle.ADD_COLON_TWO,

sep="\n\n",

sep2="</s>",

)

)

# ChatGLM default template

register_conv_template(

Conversation(

name="chatglm",

system="",

roles=("问", "答"),

messages=[],

offset=0,

sep_style=SeparatorStyle.CHATGLM,

sep="\n",

sep2="</s>",

)

)

# ChatGLM2 default template

register_conv_template(

Conversation(

name="chatglm2",

system="",

roles=("问", "答"),

messages=[],

offset=0,

sep_style=SeparatorStyle.CHATGLM,

sep="\n\n",

sep2="</s>",

)

)

# Phoenix default template

register_conv_template(

Conversation(

name="phoenix",

system="A chat between a curious human and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the human's questions.\n\n",

roles=("Human", "Assistant"),

messages=[],

offset=0,

sep_style=SeparatorStyle.PHOENIX,

sep="</s>",

)

)

# Internlm-chat template

register_conv_template(

Conversation(

name="internlm-chat",

system="A chat between a curious <|User|> and an <|Bot|>. The <|Bot|> gives helpful, detailed, and polite answers to the <|User|>'s questions.\n\n",

roles=("<|User|>", "<|Bot|>"),

messages=[],

offset=0,

sep_style=SeparatorStyle.CHATINTERN,

sep="<eoh>",

sep2="<eoa>",

stop_token_ids=[1, 103028],

stop_str="<|User|>",

)

)

# StarChat template

register_conv_template(

Conversation(

name="starchat",

system="<system>\n",

roles=("<|user|>", "<|assistant|>"),

messages=[],

offset=0,

sep_style=SeparatorStyle.CHATML,

sep="<|end|>",

stop_token_ids=[0, 49155],

stop_str="<|end|>",

)

)

# Baichuan-13B-Chat template

register_conv_template(

# source: https://huggingface.co/baichuan-inc/Baichuan-13B-Chat/blob/f5f47be2adbbdceb784f334d6fa1ca2c73e65097/modeling_baichuan.py#L507

# https://huggingface.co/baichuan-inc/Baichuan-13B-Chat/blob/main/generation_config.json

Conversation(

name="baichuan-chat",

system="",

roles=(" <reserved_102> ", " <reserved_103> "),

messages=[],

offset=0,

sep_style=SeparatorStyle.NO_COLON_TWO,

sep="",

sep2="</s>",

stop_token_ids=[2, 195],

)

)

# llama2 template

# reference: https://github.com/facebookresearch/llama/blob/cfc3fc8c1968d390eb830e65c63865e980873a06/llama/generation.py#L212

register_conv_template(

Conversation(

name="llama-2",

system="<s>[INST] <<SYS>>\nYou are a helpful, respectful and honest assistant. Always answer as helpfully as possible, while being safe. "

"Your answers should not include any harmful, unethical, racist, sexist, toxic, dangerous, or illegal content. "

"Please ensure that your responses are socially unbiased and positive in nature.\n\n"

"If a question does not make any sense, or is not factually coherent, explain why instead of answering something not correct. "

"If you don't know the answer to a question, please don't share false information.\n<</SYS>>\n\n",

roles=("[INST]", "[/INST]"),

messages=[],

offset=0,

sep_style=SeparatorStyle.LLAMA2,

sep=" ",

sep2=" </s><s>",

stop_token_ids=[2],

)

)

def get_conv_template(name: str) -> Conversation:

"""Get a conversation template."""

return conv_templates[name].copy()

class SavePeftModelTrainer(Trainer):

self.model.save_pretrained(output_dir)

def save_model(output_dir, model, tokenizer, args):

torch.save(args, os.path.join(output_dir, TRAINING_ARGS_NAME))

def print_trainable_parameters(model):

)

def find_all_linear_names(peft_model, int4=False, int8=False):

return sorted(lora_module_names)

def main():

trainer.save_metrics("eval", metrics)

if __name__ == "__main__":

main()