SFT training code

sft_train.py

@@ -0,0 +1,152 @@
+# %%
+from datasets import Dataset, load_dataset
+import pandas as pd
+from transformers import AutoTokenizer, AutoModelForCausalLM, DataCollatorForSeq2Seq, TrainingArguments, Trainer
+from peft import LoraConfig, TaskType, get_peft_model, AutoPeftModelForCausalLM
+
+# %%
+df = pd.read_csv('data/riddles_data.csv')
+df =df.sample(frac = 1)
+#df = df[:1000]
+
+# %%
+df.describe()
+
+# %%
+ds = Dataset.from_pandas(df)
+
+# %%
+ds[:3]
+
+# %%
+llm_model_name="Qwen/Qwen1.5-0.5B-Chat"
+model = AutoModelForCausalLM.from_pretrained(llm_model_name)
+tokenizer = AutoTokenizer.from_pretrained(llm_model_name,trust_remote_code=True, pad_token='<|endoftext|>')
+
+tokenizer
+
+
+# %%
+def process_func(example):
+    MAX_LENGTH = 512 
+    input_ids, attention_mask, labels = [], [], []
+    instruction = tokenizer(f"<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n<|im_start|>user\n猜谜语：\n谜面：{example['riddle']}\n\n谜底是什么？<|im_end|>\n<|im_start|>assistant\n", add_special_tokens=False)  # add_special_tokens 不在开头加 special_tokens
+    response = tokenizer(f"谜底是：{example['label']}", add_special_tokens=False)
+    input_ids = instruction["input_ids"] + response["input_ids"] + [tokenizer.pad_token_id]
+    attention_mask = instruction["attention_mask"] + response["attention_mask"] + [1]
+    labels = [-100] * len(instruction["input_ids"]) + response["input_ids"] + [tokenizer.pad_token_id]  
+    if len(input_ids) > MAX_LENGTH:  # 做一个截断
+        input_ids = input_ids[:MAX_LENGTH]
+        attention_mask = attention_mask[:MAX_LENGTH]
+        labels = labels[:MAX_LENGTH]
+        print (f"{tokenizer.decode(input_ids)} Too Long")
+    return {
+        "input_ids": input_ids,
+        "attention_mask": attention_mask,
+        "labels": labels
+    }
+
+# %%
+tokenized_id = ds.map(process_func, remove_columns=ds.column_names)
+tokenized_id
+
+# %%
+tokenizer.decode(tokenized_id[0]['input_ids'])
+
+# %%
+tokenizer.decode(list(filter(lambda x: x != -100, tokenized_id[1]["labels"])))
+
+# %%
+config = LoraConfig(
+    task_type=TaskType.CAUSAL_LM, 
+    target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
+    #target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
+    inference_mode=False,
+    r=32,
+    lora_alpha=32,
+    lora_dropout=0.05
+)
+
+
+# %%
+model = get_peft_model(model, config)
+config
+
+# %%
+model.print_trainable_parameters()
+
+# %%
+args = TrainingArguments(
+    output_dir="./Qwen1.5_0.5B_Chat_sft_full",
+    logging_steps=10,
+    num_train_epochs=2,
+    save_steps=10,
+    learning_rate=1e-4,
+    save_on_each_node=True,
+    fp16=False
+)
+
+# %%
+trainer = Trainer(
+    model=model,
+    args=args,
+    train_dataset=tokenized_id,
+    data_collator=DataCollatorForSeq2Seq(tokenizer=tokenizer, padding=True),
+)
+
+trainer.train(resume_from_checkpoint=True)
+
+# %%
+trainer.save_model("./qwen_sft_full")
+
+# %%
+llm_model_name="Qwen/Qwen1.5-0.5B-Chat"
+
+#model = AutoModelForCausalLM.from_pretrained(llm_model_name)
+# # Load PEFT model on CPU
+model = AutoPeftModelForCausalLM.from_pretrained(
+    "Qwen1.5_0.5B_Chat_sft_full_ckpt_200_ok/checkpoint-210",
+    #low_cpu_mem_usage=True,
+)
+# # Merge LoRA and base model and save
+#merged_model = model.merge_and_unload()
+#merged_model.save_pretrained("./qwen_sft",safe_serialization=False, max_shard_size="2GB")
+
+tokenizer = AutoTokenizer.from_pretrained(llm_model_name,trust_remote_code=True, pad_token='<|endoftext|>')
+
+
+# %%
+prompt = "谜面：一生受用（猜一字）\n谜底是什么？请解释。"
+messages = [
+    {"role": "system", "content": "You are a helpful assistant."},
+    {"role": "user", "content": prompt}
+]
+text = tokenizer.apply_chat_template(
+    messages,
+    tokenize=False,
+    add_generation_prompt=True
+)
+
+print(text)
+model_inputs = tokenizer([text], return_tensors="pt").to("cpu")
+
+generated_ids = model.generate(
+        model_inputs.input_ids,
+        max_new_tokens=128,
+        do_sample=False,
+        top_p=0.0
+)
+
+generated_ids = [
+    output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
+]
+
+response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
+
+# %%
+response 
+
+# %%
+
+
+