damBruh commited on
Commit
d892042
·
verified ·
1 Parent(s): b06c237

Upload train_brand_voice.py with huggingface_hub

Browse files
Files changed (1) hide show
  1. train_brand_voice.py +65 -0
train_brand_voice.py ADDED
@@ -0,0 +1,65 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # /// script
2
+ # dependencies = ["trl>=0.12.0", "peft>=0.7.0", "trackio", "datasets", "transformers", "accelerate", "bitsandbytes"]
3
+ # ///
4
+
5
+ from datasets import load_dataset
6
+ from peft import LoraConfig
7
+ from trl import SFTTrainer, SFTConfig
8
+ import trackio
9
+ import os
10
+
11
+ # Get HF token from environment
12
+ hf_token = os.environ.get("HF_TOKEN")
13
+
14
+ print("Loading SkyyRose brand voice dataset...")
15
+ dataset = load_dataset("damBruh/skyyrose-brand-voice-training", split="train", token=hf_token)
16
+ print(f"Dataset loaded: {len(dataset)} examples")
17
+
18
+ # LoRA configuration for efficient fine-tuning
19
+ peft_config = LoraConfig(
20
+ r=16,
21
+ lora_alpha=32,
22
+ lora_dropout=0.05,
23
+ target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
24
+ bias="none",
25
+ task_type="CAUSAL_LM",
26
+ )
27
+
28
+ # Training configuration
29
+ training_args = SFTConfig(
30
+ output_dir="skyyrose-brand-voice-llm",
31
+ num_train_epochs=3,
32
+ per_device_train_batch_size=2,
33
+ gradient_accumulation_steps=4,
34
+ learning_rate=2e-4,
35
+ warmup_ratio=0.1,
36
+ weight_decay=0.01,
37
+ max_grad_norm=1.0,
38
+ logging_steps=1,
39
+ save_strategy="epoch",
40
+ save_total_limit=2,
41
+ gradient_checkpointing=True,
42
+ push_to_hub=True,
43
+ hub_model_id="damBruh/skyyrose-brand-voice-llm",
44
+ hub_token=hf_token,
45
+ report_to="trackio",
46
+ run_name="skyyrose-brand-voice-qwen2.5",
47
+ seed=42,
48
+ )
49
+
50
+ print("Initializing SFTTrainer with Qwen/Qwen2.5-1.5B-Instruct...")
51
+ trainer = SFTTrainer(
52
+ model="Qwen/Qwen2.5-1.5B-Instruct",
53
+ train_dataset=dataset,
54
+ peft_config=peft_config,
55
+ args=training_args,
56
+ )
57
+
58
+ print("Starting training...")
59
+ trainer.train()
60
+
61
+ print("Saving and pushing to Hub...")
62
+ trainer.save_model()
63
+ trainer.push_to_hub()
64
+
65
+ print("SUCCESS: Model available at huggingface.co/damBruh/skyyrose-brand-voice-llm")