ssdavid commited on
Commit
ef69abc
·
verified ·
1 Parent(s): 7871aae

Upload train.py with huggingface_hub

Browse files
Files changed (1) hide show
  1. train.py +77 -0
train.py ADDED
@@ -0,0 +1,77 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # /// script
2
+ # dependencies = ["trl>=0.12.0", "peft>=0.7.0", "trackio", "datasets"]
3
+ # ///
4
+ """
5
+ Obsidian Bases SLM Training Script
6
+ Fine-tunes Qwen 3 0.6B to generate .base files from natural language.
7
+ """
8
+
9
+ import os
10
+ from datasets import load_dataset
11
+ from peft import LoraConfig
12
+ from trl import SFTTrainer, SFTConfig
13
+ import trackio
14
+
15
+ # Load dataset from Hub
16
+ dataset = load_dataset("ssdavid/obsidian-bases-query-v1", split="train")
17
+
18
+ # Format for SFT - convert to messages format
19
+ def format_example(example):
20
+ return {
21
+ "messages": [
22
+ {"role": "user", "content": example["instruction"]},
23
+ {"role": "assistant", "content": example["output"]}
24
+ ]
25
+ }
26
+
27
+ dataset = dataset.map(format_example)
28
+
29
+ # Split for evaluation
30
+ dataset_split = dataset.train_test_split(test_size=0.1, seed=42)
31
+
32
+ # LoRA config for efficient fine-tuning
33
+ peft_config = LoraConfig(
34
+ r=16,
35
+ lora_alpha=32,
36
+ lora_dropout=0.05,
37
+ target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
38
+ bias="none",
39
+ task_type="CAUSAL_LM"
40
+ )
41
+
42
+ # Training config
43
+ training_args = SFTConfig(
44
+ output_dir="obsidian-bases-slm",
45
+ push_to_hub=True,
46
+ hub_model_id="ssdavid/obsidian-bases-slm",
47
+ num_train_epochs=3,
48
+ per_device_train_batch_size=4,
49
+ gradient_accumulation_steps=4,
50
+ learning_rate=2e-4,
51
+ warmup_ratio=0.1,
52
+ logging_steps=10,
53
+ eval_strategy="steps",
54
+ eval_steps=50,
55
+ save_strategy="steps",
56
+ save_steps=100,
57
+ max_length=512,
58
+ report_to="trackio",
59
+ project="obsidian-bases-slm",
60
+ run_name="qwen3-0.6b-bases-v1",
61
+ )
62
+
63
+ # Create trainer
64
+ trainer = SFTTrainer(
65
+ model="Qwen/Qwen3-0.6B",
66
+ train_dataset=dataset_split["train"],
67
+ eval_dataset=dataset_split["test"],
68
+ peft_config=peft_config,
69
+ args=training_args,
70
+ )
71
+
72
+ # Train
73
+ trainer.train()
74
+
75
+ # Push to Hub
76
+ trainer.push_to_hub()
77
+ print("Training complete! Model pushed to ssdavid/obsidian-bases-slm")