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))