Spaces:
Runtime error
Runtime error
File size: 1,555 Bytes
a365d48 | 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 | import os
from datasets import load_dataset
from transformers import AutoTokenizer, AutoModelForCausalLM, TrainingArguments
from trl import SFTTrainer
from peft import LoraConfig
BASE_MODEL = os.environ.get("BASE_MODEL", "Qwen/Qwen2.5-0.5B-Instruct")
def main():
ds = load_dataset("json", data_files="data/sft.jsonl")["train"]
tok = AutoTokenizer.from_pretrained(BASE_MODEL, use_fast=True)
if tok.pad_token is None:
tok.pad_token = tok.eos_token
model = AutoModelForCausalLM.from_pretrained(
BASE_MODEL,
device_map="auto",
torch_dtype="auto",
)
peft_cfg = LoraConfig(
r=16, lora_alpha=32, lora_dropout=0.05,
bias="none", task_type="CAUSAL_LM",
target_modules=["q_proj","k_proj","v_proj","o_proj","up_proj","down_proj","gate_proj"]
)
args = TrainingArguments(
output_dir="adapter_sft",
per_device_train_batch_size=2,
gradient_accumulation_steps=8,
learning_rate=2e-4,
num_train_epochs=1,
logging_steps=20,
save_steps=200,
fp16=True,
report_to="none"
)
trainer = SFTTrainer(
model=model,
tokenizer=tok,
train_dataset=ds,
peft_config=peft_cfg,
max_seq_length=1024,
args=args,
packing=False,
dataset_text_field=None, # because we use "messages"
)
trainer.train()
trainer.save_model("adapter_sft")
tok.save_pretrained("adapter_sft")
print("Saved adapter_sft/")
if __name__ == "__main__":
main()
|