Lora微调大模型

最新推荐文章于 2024-09-03 17:37:15 发布

axb你猜我猜不猜

最新推荐文章于 2024-09-03 17:37:15 发布

阅读量105

点赞数 2

文章标签：人工智能深度学习

本文链接：https://blog.csdn.net/weixin_44763000/article/details/140741306

版权

import os
import warnings
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
os.environ['TRANSFORMERS_OFFLINE'] = '1'

warnings.filterwarnings("ignore")

from transformers import (AutoModelForCausalLM,
                          AutoTokenizer,
                          TrainingArguments,
                          Trainer,
                          DataCollatorForSeq2Seq)

from peft import (LoraConfig, 
                  get_peft_model, 
                  TaskType)

from datasets import load_dataset

_tokenizer = AutoTokenizer.from_pretrained("Qwen2-0___5B")
_model = AutoModelForCausalLM.from_pretrained("Qwen2-0___5B")

# for name ,param in _model.named_parameters():
#     print(name)

_dataset = load_dataset("json",data_files="data.json",split="train")

def preprocess_dataset(example):
    MAX_LENGTH = 256
    _input_ids, _attention_mask, _labels = [], [], []
    _instruction = _tokenizer(f"User: {example['instruction']}Assistant: ",add_special_tokens=False)
    _response = _tokenizer(example["output"] + _tokenizer.eos_token,add_special_tokens=False)
    _input_ids = _instruction["input_ids"] + _response["input_ids"]
    _attention_mask = _instruction["attention_mask"] + _response["attention_mask"]
    _labels = [-100] * len(_instruction["input_ids"]) + _response["input_ids"]
    if len(_input_ids) > MAX_LENGTH:
        _input_ids = _input_ids[:MAX_LENGTH]
        _attention_mask = _attention_mask[:MAX_LENGTH]
        _labels = _labels[:MAX_LENGTH]
    return {
        "input_ids": _input_ids,
        "attention_mask": _attention_mask,
        "labels": _labels
    }


_dataset = _dataset.map(preprocess_dataset,remove_columns=_dataset.column_names)
_dataset = _dataset.shuffle()


config = LoraConfig(task_type=TaskType.CAUSAL_LM, 
                    target_modules="all-linear",
                    r=8,
                    lora_alpha=16)


_model = get_peft_model(_model, config)

# model.print_trainable_parameters()

_training_args = TrainingArguments(
    output_dir="checkpoints/lora",
    per_device_train_batch_size=10,
    gradient_accumulation_steps=2,
    num_train_epochs=100,
    save_steps=300
)

trainer = Trainer(
    model=_model,
    args=_training_args,
    train_dataset=_dataset,
    data_collator=DataCollatorForSeq2Seq(tokenizer=_tokenizer,padding=True)
)

trainer.train()

from transformers import (
                          AutoModelForCausalLM,
                          AutoTokenizer,
                          pipeline)
import peft
_model = AutoModelForCausalLM.from_pretrained("Qwen2-0___5B")
_tokenizer = AutoTokenizer.from_pretrained("Qwen2-0___5B")
_peft_model = peft.peft_model.PeftModelForCausalLM.from_pretrained(model =_model,model_id = "/home/easyai/Desktop/ai/ai/checkpoints/lora/checkpoint-7500")
pipe = pipeline("text-generation",model = _peft_model,tokenizer=_tokenizer)
ipt = f"User:你是谁？ Anwser:"
print(pipe(ipt))