初始化项目,由ModelHub XC社区提供模型
Model: vicgalle/Humanish-Roleplay-Llama-3.1-8B Source: Original Platform
This commit is contained in:
146
train_human.py
Normal file
146
train_human.py
Normal file
@@ -0,0 +1,146 @@
|
||||
# %%
|
||||
from transformers import AutoTokenizer, AutoModelForCausalLM
|
||||
|
||||
from datasets import load_dataset, Dataset
|
||||
|
||||
from trl import DPOTrainer, DPOConfig
|
||||
from peft import LoraConfig
|
||||
from peft import prepare_model_for_kbit_training
|
||||
import torch
|
||||
|
||||
import pandas as pd
|
||||
|
||||
# %%
|
||||
dataset = load_dataset("Undi95/Weyaxi-humanish-dpo-project-noemoji")["train"]
|
||||
|
||||
model_name = "Undi95/Meta-Llama-3.1-8B-Claude-bf16"
|
||||
|
||||
tokenizer = AutoTokenizer.from_pretrained(model_name)
|
||||
tokenizer.padding_side = "right"
|
||||
|
||||
tokenizer.pad_token = tokenizer.eos_token
|
||||
|
||||
# %%
|
||||
tokenizer.chat_template = "{% set loop_messages = messages %}{% for message in loop_messages %}{% set content = '<|start_header_id|>' + message['role'] + '<|end_header_id|>\n\n'+ message['content'] | trim + '<|eot_id|>' %}{% if loop.index0 == 0 %}{% set content = bos_token + content %}{% endif %}{{ content }}{% endfor %}{{ '<|start_header_id|>assistant<|end_header_id|>\n\n' }}"
|
||||
|
||||
|
||||
# %%
|
||||
dataset2 = load_dataset("ResplendentAI/NSFW_RP_Format_DPO")['train']
|
||||
|
||||
# %%
|
||||
dataset = dataset.to_pandas(
|
||||
)
|
||||
dataset2 = dataset2.to_pandas()
|
||||
|
||||
dataset = Dataset.from_pandas(pd.concat([dataset.sample(400), dataset2]).sample(frac=1))
|
||||
|
||||
# %%
|
||||
def template_prompt(system, prompt):
|
||||
if system is None:
|
||||
messages = [
|
||||
{"role": "user", "content": prompt},
|
||||
]
|
||||
else:
|
||||
messages = [
|
||||
{
|
||||
"role": "system",
|
||||
"content": system,
|
||||
},
|
||||
{"role": "user", "content": prompt},
|
||||
]
|
||||
prompt = tokenizer.apply_chat_template(
|
||||
messages, tokenize=False, add_generation_prompt=False
|
||||
)
|
||||
return prompt
|
||||
|
||||
|
||||
def template_answer(answer):
|
||||
messages = [
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": answer,
|
||||
},
|
||||
]
|
||||
answer = tokenizer.apply_chat_template(
|
||||
messages, tokenize=False, add_generation_prompt=False
|
||||
)
|
||||
return answer
|
||||
|
||||
# %%
|
||||
# create new columns
|
||||
dataset = dataset.map(
|
||||
lambda x: {
|
||||
"prompt": template_prompt(None, x["prompt"]).replace("<|start_header_id|>assistant<|end_header_id|>\n\n", "")
|
||||
}, # change this according to the dataset!!!
|
||||
)
|
||||
|
||||
# %%
|
||||
dataset = dataset.map(
|
||||
lambda x: {"chosen": template_answer(x["chosen"]).replace('<|begin_of_text|>', '').replace('><|start_header_id|>assistant<|end_header_id|>\n\n', '>')},
|
||||
)
|
||||
dataset = dataset.map(
|
||||
lambda x: {"rejected": template_answer(x["rejected"]).replace('<|begin_of_text|>', '').replace('><|start_header_id|>assistant<|end_header_id|>\n\n', '>')},
|
||||
)
|
||||
|
||||
# %%
|
||||
dataset[0]
|
||||
|
||||
# %%
|
||||
# LoRA configuration
|
||||
peft_config = LoraConfig(
|
||||
r=16,
|
||||
lora_alpha=32,
|
||||
lora_dropout=0.05,
|
||||
bias="none",
|
||||
task_type="CAUSAL_LM",
|
||||
target_modules=[
|
||||
"k_proj",
|
||||
"gate_proj",
|
||||
"v_proj",
|
||||
"up_proj",
|
||||
"q_proj",
|
||||
"o_proj",
|
||||
"down_proj",
|
||||
],
|
||||
)
|
||||
|
||||
# Model to fine-tune
|
||||
model = AutoModelForCausalLM.from_pretrained(
|
||||
model_name,
|
||||
torch_dtype=torch.float16,
|
||||
load_in_4bit=True,
|
||||
device_map="auto",
|
||||
)
|
||||
model.config.use_cache = False
|
||||
|
||||
|
||||
model.gradient_checkpointing_enable()
|
||||
model = prepare_model_for_kbit_training(model)
|
||||
|
||||
# %%
|
||||
output_name = f"checkpoints/exp_human_{model_name}"
|
||||
|
||||
training_args = DPOConfig(
|
||||
per_device_train_batch_size=1,
|
||||
gradient_accumulation_steps=4,
|
||||
num_train_epochs=1,
|
||||
gradient_checkpointing=True,
|
||||
output_dir=output_name,
|
||||
logging_steps=1,
|
||||
max_steps=50
|
||||
)
|
||||
|
||||
trainer = DPOTrainer(
|
||||
model,
|
||||
ref_model=None,
|
||||
train_dataset=dataset,
|
||||
tokenizer=tokenizer,
|
||||
args=training_args,
|
||||
peft_config=peft_config,
|
||||
)
|
||||
|
||||
trainer.train()
|
||||
|
||||
trainer.save_model(output_name)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user