ArchibaldAI commited on
Commit
5585f60
·
verified ·
1 Parent(s): 7cde1fa

Upload train_intent_router.py with huggingface_hub

Browse files
Files changed (1) hide show
  1. train_intent_router.py +63 -0
train_intent_router.py ADDED
@@ -0,0 +1,63 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # /// script
2
+ # dependencies = ["trl>=0.12.0", "peft>=0.7.0", "datasets", "trackio", "accelerate", "bitsandbytes"]
3
+ # ///
4
+
5
+ from datasets import load_dataset
6
+ from peft import LoraConfig
7
+ from trl import SFTTrainer, SFTConfig
8
+ import json
9
+
10
+ # Load dataset
11
+ dataset = load_dataset("ArchibaldAI/agent-intent-router")
12
+
13
+ print(f"Train: {len(dataset['train'])} examples")
14
+ print(f"Test: {len(dataset['test'])} examples")
15
+ print(f"Sample: {dataset['train'][0]}")
16
+
17
+ # Model: SmolLM2-360M - tiny, fast, perfect for classification
18
+ model_name = "HuggingFaceTB/SmolLM2-360M-Instruct"
19
+ output_name = "ArchibaldAI/agent-intent-router-v1"
20
+
21
+ # LoRA config - lightweight fine-tuning
22
+ peft_config = LoraConfig(
23
+ r=16,
24
+ lora_alpha=32,
25
+ lora_dropout=0.05,
26
+ target_modules=["q_proj", "v_proj", "k_proj", "o_proj"],
27
+ bias="none",
28
+ task_type="CAUSAL_LM",
29
+ )
30
+
31
+ # Training config
32
+ training_args = SFTConfig(
33
+ output_dir="./intent-router",
34
+ push_to_hub=True,
35
+ hub_model_id=output_name,
36
+ num_train_epochs=5,
37
+ per_device_train_batch_size=8,
38
+ per_device_eval_batch_size=8,
39
+ learning_rate=2e-4,
40
+ warmup_ratio=0.1,
41
+ logging_steps=10,
42
+ eval_strategy="epoch",
43
+ save_strategy="epoch",
44
+ load_best_model_at_end=True,
45
+ metric_for_best_model="eval_loss",
46
+ report_to="trackio",
47
+ run_name="intent-router-v1",
48
+ max_length=256, # Short sequences for classification
49
+ bf16=True,
50
+ )
51
+
52
+ # Train
53
+ trainer = SFTTrainer(
54
+ model=model_name,
55
+ train_dataset=dataset["train"],
56
+ eval_dataset=dataset["test"],
57
+ peft_config=peft_config,
58
+ args=training_args,
59
+ )
60
+
61
+ trainer.train()
62
+ trainer.push_to_hub()
63
+ print(f"\n✅ Model pushed to {output_name}")