File size: 4,244 Bytes
e6b028b 658e879 e6b028b 1dd9bc4 e6b028b 1dd9bc4 e6b028b 1dd9bc4 e6b028b 658e879 e6b028b 658e879 e6b028b 658e879 e6b028b 658e879 e6b028b 658e879 e6b028b 658e879 e6b028b 658e879 e6b028b | 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 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 | # /// script
# dependencies = ["trl>=0.12.0", "peft>=0.7.0", "trackio", "datasets", "accelerate", "torch", "bitsandbytes"]
# ///
from datasets import load_dataset, concatenate_datasets
from peft import LoraConfig, PeftModel
from transformers import AutoModelForCausalLM, AutoTokenizer
from trl import SFTTrainer, SFTConfig
import trackio
import torch
print("Loading base model and merging SFT adapter...")
# Load the base Qwen model
base_model = AutoModelForCausalLM.from_pretrained(
"Qwen/Qwen3-0.6B",
torch_dtype=torch.bfloat16,
device_map="auto",
trust_remote_code=True
)
# Load and merge the existing SFT adapter to preserve code capabilities
model = PeftModel.from_pretrained(base_model, "chaddy81/qwen3-0.6b-multicode-sft")
model = model.merge_and_unload() # Merge adapter into base model
print("SFT adapter merged successfully!")
# Load tokenizer
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-0.6B", trust_remote_code=True)
if tokenizer.pad_token is None:
tokenizer.pad_token = tokenizer.eos_token
print("Loading datasets...")
# Load the validated Tailwind dataset (already in messages format)
tailwind_ds = load_dataset("summykai/tailwind-v4-sft-mix-001", split="train")
print(f"Tailwind dataset: {len(tailwind_ds)} examples")
# Load the high-quality UI generation dataset
uigen_ds = load_dataset("smirki/UIGEN-T1.1-TAILWIND", split="train")
print(f"UIGEN dataset: {len(uigen_ds)} examples")
# Format UIGEN dataset to messages format
def format_uigen_to_messages(example):
# Extract code from markdown code blocks if present
answer = example["answer"]
if answer.startswith("```html"):
answer = answer[7:] # Remove ```html
if answer.endswith("```"):
answer = answer[:-3] # Remove trailing ```
# Include the reasoning as part of a richer response
response = answer.strip()
return {
"messages": [
{"role": "user", "content": example["question"]},
{"role": "assistant", "content": response}
]
}
print("Formatting UIGEN dataset...")
uigen_formatted = uigen_ds.map(
format_uigen_to_messages,
remove_columns=uigen_ds.column_names
)
# Combine datasets
print("Combining datasets...")
combined_dataset = concatenate_datasets([tailwind_ds, uigen_formatted])
combined_dataset = combined_dataset.shuffle(seed=42)
print(f"Total training examples: {len(combined_dataset)}")
# Create train/eval split
dataset_split = combined_dataset.train_test_split(test_size=0.05, seed=42)
print(f"Train: {len(dataset_split['train'])}, Eval: {len(dataset_split['test'])}")
print("Initializing trainer...")
trainer = SFTTrainer(
model=model, # Use the merged model
processing_class=tokenizer,
train_dataset=dataset_split["train"],
eval_dataset=dataset_split["test"],
peft_config=LoraConfig(
r=16, # Reduced from 32 to save memory
lora_alpha=32,
target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
lora_dropout=0.05,
bias="none",
task_type="CAUSAL_LM"
),
args=SFTConfig(
output_dir="qwen3-0.6b-design-sft",
push_to_hub=True,
hub_model_id="chaddy81/qwen3-0.6b-design-sft",
num_train_epochs=3,
per_device_train_batch_size=1, # Reduced from 2
gradient_accumulation_steps=16, # Increased to maintain effective batch size
learning_rate=2e-4,
lr_scheduler_type="cosine",
warmup_ratio=0.1,
max_length=2048, # Reduced from 4096 to save memory
logging_steps=10,
eval_strategy="no", # Skip eval during training to save memory
save_strategy="steps",
save_steps=300,
bf16=True,
gradient_checkpointing=True,
gradient_checkpointing_kwargs={"use_reentrant": False}, # More memory efficient
optim="adamw_8bit", # 8-bit optimizer to save memory
report_to="trackio",
run_name="qwen3-design-sft-v2",
)
)
print("Starting training...")
trainer.train()
print("Pushing to Hub...")
trainer.push_to_hub()
print("Training complete!")
|