File size: 1,117 Bytes
730b6df |
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 |
from datasets import load_dataset
from transformers import AutoTokenizer, AutoModelForCausalLM, TrainingArguments
from trl import DPOTrainer
# --------------------
# 1. USTAW MODEL
# --------------------
model_name = "mistralai/Mistral-7B-Instruct-v0.2" # możesz zmienić
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
model_name,
load_in_4bit=True,
device_map="auto"
)
# --------------------
# 2. WCZYTAJ DANE
# --------------------
dataset = load_dataset("json", data_files="dpo_data.jsonl")["train"]
# --------------------
# 3. ARGUMENTY TRENINGU
# --------------------
training_args = TrainingArguments(
output_dir="./dpo_output",
per_device_train_batch_size=2,
gradient_accumulation_steps=4,
num_train_epochs=2,
learning_rate=5e-6,
bf16=True,
logging_steps=10,
save_steps=500,
)
# --------------------
# 4. START DPO
# --------------------
trainer = DPOTrainer(
model=model,
ref_model=None,
tokenizer=tokenizer,
train_dataset=dataset,
beta=0.1,
args=training_args,
)
trainer.train()
|