Yi-34B-200K-rawrr1-LORA-DPO-experimental-r3 / yi-34b-dpo-unsloth-1.py

Upload yi-34b-dpo-unsloth-1.py

91a3383 verified 5 months ago

No virus

4.64 kB

	from unsloth import FastLanguageModel
	from datasets import Dataset, load_dataset
	from dataclasses import dataclass, field
	from typing import Dict, Optional
	import torch
	max_seq_length = 4096 # Choose any! We auto support RoPE Scaling internally!
	dtype = None # None for auto detection. Float16 for Tesla T4, V100, Bfloat16 for Ampere+
	load_in_4bit = True # Use 4bit quantization to reduce memory usage. Can be False.

	model, tokenizer = FastLanguageModel.from_pretrained(
	model_name = "...../yi-34b-200k-llamafied", # Choose ANY! eg mistralai/Mistral-7B-Instruct-v0.2
	max_seq_length = max_seq_length,
	attn_implementation="flash_attention_2",
	dtype = dtype,
	load_in_4bit = load_in_4bit,
	# token = "hf_...", # use one if using gated models like meta-llama/Llama-2-7b-hf
	)



	#@title Alignment Handbook utils
	import os
	import re
	from typing import List, Literal, Optional

	from datasets import DatasetDict, concatenate_datasets, load_dataset, load_from_disk
	from datasets.builder import DatasetGenerationError


	#DEFAULT_CHAT_TEMPLATE = "{% for message in messages %}\n{% if message['role'] == 'user' %}\n{{ '<\|user\|>\n' + message['content'] + eos_token }}\n{% elif message['role'] == 'system' %}\n{{ '<\|system\|>\n' + message['content'] + eos_token }}\n{% elif message['role'] == 'assistant' %}\n{{ '<\|assistant\|>\n' + message['content'] + eos_token }}\n{% endif %}\n{% if loop.last and add_generation_prompt %}\n{{ '<\|assistant\|>' }}\n{% endif %}\n{% endfor %}"
	tokenizer.chat_template = "{% if not add_generation_prompt is defined %}{% set add_generation_prompt = false %}{% endif %}{% for message in messages %}{{'<\|im_start\|>' + message['role'] + '\n' + message['content'] + '<\|im_end\|>' + '\n'}}{% endfor %}{% if add_generation_prompt %}{{ '<\|im_start\|>assistant\n' }}{% endif %}"

	def chatml_format(example):
	# Format system
	if len(example['system']) > 0:
	message = {"role": "system", "content": example['system']}
	system = tokenizer.apply_chat_template([message], tokenize=False)
	else:
	system = ""

	# Format instruction
	message = {"role": "user", "content": example['prompt']}
	prompt = tokenizer.apply_chat_template([message], tokenize=False, add_generation_prompt=True)

	# Format chosen answer
	chosen = example['chosen'] + "<\|im_end\|>\n"

	# Format rejected answer
	rejected = example['rejected'] + "<\|im_end\|>\n"

	return {
	"prompt": system + prompt,
	"chosen": chosen,
	"rejected": rejected,
	}

	# Load dataset
	dataset = load_dataset("adamo1139/rawrr_v1", split="train")

	import pprint
	pprint.pprint("""NOT a formatted dataset
	""")
	pprint
	pprint.pprint(dataset[25])
	pprint.pprint(dataset[26])
	pprint.pprint(dataset[27])
	pprint.pprint(dataset[28])
	pprint.pprint(dataset[29])
	# Save columns
	original_columns = dataset.column_names

	# Format dataset
	dataset = dataset.map(
	chatml_format,
	remove_columns=original_columns
	)

	# Print sample
	pprint.pprint("""formatted dataset""")
	pprint.pprint(dataset[25])
	pprint.pprint(dataset[26])
	pprint.pprint(dataset[27])
	pprint.pprint(dataset[28])
	pprint.pprint(dataset[29])


	model = FastLanguageModel.get_peft_model(
	model,
	r = 16, # Choose any number > 0 ! Suggested 8, 16, 32, 64, 128
	target_modules = ["q_proj", "k_proj", "v_proj", "o_proj",
	"gate_proj", "up_proj", "down_proj",],
	lora_alpha = 32,
	lora_dropout = 0, # Currently only supports dropout = 0
	bias = "none", # Currently only supports bias = "none"
	use_gradient_checkpointing = True,
	random_state = 3407,
	use_rslora = False, # We support rank stabilized LoRA
	loftq_config = None, # And LoftQ
	)

	from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, HfArgumentParser, TrainingArguments
	from trl import DPOTrainer

	dpo_trainer = DPOTrainer(
	model = model,
	ref_model = None,
	args = TrainingArguments(
	per_device_train_batch_size = 1,
	gradient_accumulation_steps = 16,
	warmup_ratio = 0.05,
	num_train_epochs = 1,
	learning_rate = 5e-5,
	fp16 = not torch.cuda.is_bf16_supported(),
	bf16 = torch.cuda.is_bf16_supported(),
	logging_steps = 1,
	optim = "adamw_8bit",
	weight_decay = 0.0,
	lr_scheduler_type = "linear",
	seed = 42,
	output_dir = "outputs2",
	),
	beta = 0.1,
	train_dataset = dataset,
	# eval_dataset = raw_datasets["test"],
	tokenizer = tokenizer,
	max_length = 500,
	max_prompt_length = 500,
	)
	dpo_trainer.train()
	model.save_pretrained("yi-34b-200k-rawrr_v1_unsloth_1") # Local saving