Yi-34B-200K-AEZAKMI-RAW-2301-LoRA / yi-34b-aezakmi-sft-1-hf.py

Upload yi-34b-aezakmi-sft-1-hf.py

5cd4da8 verified 5 months ago

No virus

5.56 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 = "/run/.../yi-34b-rawrr-dpo-2-unsloth", # 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['instruction']}
	prompt = tokenizer.apply_chat_template([message], tokenize=False, add_generation_prompt=True)

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


	return {
	"text": system + prompt + response,
	}

	# Load dataset
	#dataset = load_dataset("adamo1139/AEZAKMI_v2", split="train")
	dataset = load_dataset("json", data_files="/run/..../datasets/aezakmi_v2/aezakmi_v2.jsonl", 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 transformers.utils import logging
	from trl import SFTTrainer

	sft_trainer = SFTTrainer(
	model = model,
	tokenizer = tokenizer,
	train_dataset = dataset,
	dataset_text_field = "text",
	max_seq_length = 2200,
	packing=True,
	args = TrainingArguments(
	evaluation_strategy = "no",
	per_device_train_batch_size = 1,
	gradient_accumulation_steps = 1,
	num_train_epochs = 1.4,
	warmup_steps = 100,
	learning_rate = 0.00006,
	fp16 = not torch.cuda.is_bf16_supported(),
	bf16 = torch.cuda.is_bf16_supported(),
	logging_steps = 1,
	output_dir = "outputs3",
	optim = "adamw_8bit",
	weight_decay = 0.0,
	lr_scheduler_type = "cosine",
	lr_scheduler_kwargs = {
	"num_cycles" : 0.3,
	},
	seed = 42,
	save_strategy = "steps",
	save_steps = 1000,
	save_total_limit = 10,
	),
	)

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


	sft_trainer.train()
	model.save_pretrained("yi-34b-200k-aezakmi-raw-unsloth-2") # Local saving