INNER CODE UNIT · Python

train

HIT-SCIR/Chinese-Mixtral-8x7B · train.py:195

def train():
    model_args: ModelArguments
    data_args: DataArguments
    peft_args: PeftArguments
    training_args: TrainingArguments
    parser = transformers.HfArgumentParser((ModelArguments, DataArguments, PeftArguments, TrainingArguments))
    model_args, data_args, peft_args, training_args = parser.parse_args_into_dataclasses()

    model = AutoModelForCausalLM.from_pretrained(
        model_args.model_name_or_path,
        cache_dir=training_args.cache_dir,
        torch_dtype="auto",
        load_in_4bit=True,
        bnb_4bit_quant_type="nf4",
        attn_implementation="flash_attention_2",
    )

    tokenizer = LlamaBPEDropoutTokenizer.from_pretrained(

View source record →