File size: 3,086 Bytes
d4dcc59 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 | import os
import torch
import pandas as pd
from datasets import load_dataset
from transformers import (
AutoModelForCausalLM,
AutoTokenizer,
BitsAndBytesConfig,
HfArgumentParser,
TrainingArguments,
pipeline,
logging,
)
from peft import LoraConfig, PeftModel
from trl import SFTTrainer
def custom_format(row):
custom_dict = {
'Text': row['text'],
'Question': row['question'],
}
options_columns = ['key_1', 'key_2', 'key_3']
for column in options_columns:
right_key = f"{column}(right)"
wrong_key = f"{column}(wrong)"
right = f"({row[column]}(right))"
wrong = f"({row[column]}(wrong))"
custom_dict.update({
right: row[right_key],
wrong: row[wrong_key],
})
return custom_dict
if __name__ == "__main__":
# Downloaded model path from local
path = "......."
base_model = "path"
# Load the dataset
dataset = load_dataset('csv', data_files='/home/ubuntu/item_multiple_group/data/main_data_.csv')
train_df = pd.DataFrame(dataset['train'])
train_df = train_df.astype(str)
formatted_dataset = train_df.apply(custom_format, axis=1) # Apply custom formatting to each row
compute_dtype = getattr(torch, "float16")
quant_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_quant_type="nf4",
bnb_4bit_compute_dtype=compute_dtype,
bnb_4bit_use_double_quant=False,
)
# Load base model
model = AutoModelForCausalLM.from_pretrained(
base_model,
quantization_config=quant_config,
device_map={"": 0}
)
model.config.use_cache = False
model.config.pretraining_tp = 1
# Load LLaMA tokenizer
tokenizer = AutoTokenizer.from_pretrained(base_model, trust_remote_code=True)
tokenizer.pad_token = tokenizer.eos_token
tokenizer.padding_side = "right"
# Load LoRA config
peft_args = LoraConfig(
lora_alpha=16,
lora_dropout=0.1,
r=64,
bias="none",
task_type="CAUSAL_LM",
)
# Set training parameters
training_params = TrainingArguments(
output_dir="./results",
num_train_epochs=1,
per_device_train_batch_size=4,
gradient_accumulation_steps=1,
optim="paged_adamw_32bit",
save_steps=25,
logging_steps=25,
learning_rate=2e-4,
weight_decay=0.001,
fp16=False,
bf16=False,
max_grad_norm=0.3,
max_steps=-1,
warmup_ratio=0.03,
group_by_length=True,
lr_scheduler_type="constant",
report_to="tensorboard"
)
# Set supervised fine-tuning params
trainer = SFTTrainer(
model=model,
train_dataset=dataset,
peft_config=peft_args,
dataset_text_field="text",
max_seq_length=None,
tokenizer=tokenizer,
args=training_params,
packing=False,
)
trainer.train()
# Save model
trainer.model.save_pretrained(new_model) |