1# -*- coding: utf-8 -*-
2import os
3from unsloth import FastModel
4import torch
5from trl import SFTConfig, SFTTrainer
6from teich import mask_data, prepare_data
7
8MAX_SEQ_LEN = 32768
9MODEL_NAME = os.environ.get("MODEL_NAME", "qwen/Qwen3.5-9B")
10OUTPUT_DIR = os.environ.get("OUTPUT_DIR", "outputs/qwen-tool-sft")
11HUB_REPO_ID = os.environ.get("HUB_REPO_ID", "armand0e/Qwen3.5-9B-Opus-Agent")
12HF_TOKEN = os.environ.get("HF_TOKEN", "")
13
14model, tokenizer = FastModel.from_pretrained(
15 model_name=MODEL_NAME,
16 max_seq_length=MAX_SEQ_LEN,
17 load_in_4bit=False,
18 load_in_8bit=False,
19 full_finetuning=False,
20)
21
22model = FastModel.get_peft_model(
23 model,
24 finetune_vision_layers = False, # Turn off for just text!
25 finetune_language_layers = True, # Should leave on!
26 finetune_attention_modules = True, # Attention good for GRPO
27 finetune_mlp_modules = True, # Should leave on always!
28
29 r = 32, # Larger = higher accuracy, but might overfit
30 lora_alpha = 64, # Recommended alpha == r at least
31 lora_dropout = 0,
32 bias = "none",
33 random_state = 3407,
34)
35
36train_dataset = prepare_data(
37 {
38 "chat": {
39 "source": "TeichAI/claude-4.5-opus-high-reasoning-250x"
40 },
41 "opus-agent": {
42 "source": "armand0e/badlogicgames-pi-mono-opus-filtered",
43 },
44 },
45 tokenizer,
46 split="train",
47 hf_token=HF_TOKEN,
48 chat_template_kwargs={"enable_thinking": True},
49 max_length=MAX_SEQ_LEN,
50 drop_oversized_examples=True,
51 trim_oversized_followups=True,
52 tokenize=True,
53 strict=True,
54)
55
56trainer = SFTTrainer(
57 model=model,
58 tokenizer=tokenizer,
59 train_dataset=train_dataset,
60 eval_dataset=None,
61 args=SFTConfig(
62 dataset_text_field="text",
63 dataset_num_proc=1,
64 max_length=MAX_SEQ_LEN,
65 packing=False,
66 per_device_train_batch_size=1,
67 gradient_accumulation_steps=8,
68 warmup_steps= 5,
69 num_train_epochs=2,
70 learning_rate=2e-5,
71 logging_steps=1,
72 save_steps=100,
73 save_total_limit=3,
74 optim="adamw_8bit",
75 weight_decay=0.01,
76 max_grad_norm=0.3,
77 lr_scheduler_type="linear",
78 output_dir=OUTPUT_DIR,
79 seed=3407,
80 report_to="none",
81 ),
82)
83
84trainer = mask_data(
85 trainer,
86 tokenizer=tokenizer,
87 train_on_reasoning=False,
88 train_on_final_answers=True,
89 train_on_tools=True,
90)
91
92print(trainer.train_dataset.preview())
93
94trainer_stats = trainer.train(resume_from_checkpoint=False)
95
96model.push_to_hub(f"{HUB_REPO_ID}-LoRA", token=HF_TOKEN)
97tokenizer.push_to_hub(f"{HUB_REPO_ID}-LoRA", token=HF_TOKEN)
98
99model.push_to_hub_merged(HUB_REPO_ID, tokenizer, save_method="merged_16bit", token=HF_TOKEN)
