# -*- coding: utf-8 -*-import osfrom unsloth import FastModelimport torchfrom trl import SFTConfig, SFTTrainerfrom teich import mask_data, prepare_data MAX_SEQ_LEN = 32768MODEL_NAME = os.environ.get("MODEL_NAME", "qwen/Qwen3.5-9B")OUTPUT_DIR = os.environ.get("OUTPUT_DIR", "outputs/qwen-tool-sft")HUB_REPO_ID = os.environ.get("HUB_REPO_ID", "armand0e/Qwen3.5-9B-Opus-Agent")HF_TOKEN = os.environ.get("HF_TOKEN", "") model, tokenizer = FastModel.from_pretrained( model_name=MODEL_NAME, max_seq_length=MAX_SEQ_LEN, load_in_4bit=False, load_in_8bit=False, full_finetuning=False,) model = FastModel.get_peft_model( model, finetune_vision_layers = False, # Turn off for just text! finetune_language_layers = True, # Should leave on! finetune_attention_modules = True, # Attention good for GRPO finetune_mlp_modules = True, # Should leave on always! r = 32, # Larger = higher accuracy, but might overfit lora_alpha = 64, # Recommended alpha == r at least lora_dropout = 0, bias = "none", random_state = 3407,) train_dataset = prepare_data( { "chat": { "source": "TeichAI/claude-4.5-opus-high-reasoning-250x" }, "opus-agent": { "source": "armand0e/badlogicgames-pi-mono-opus-filtered", }, }, tokenizer, split="train", hf_token=HF_TOKEN, chat_template_kwargs={"enable_thinking": True}, max_length=MAX_SEQ_LEN, drop_oversized_examples=True, trim_oversized_followups=True, tokenize=True, strict=True,) trainer = SFTTrainer( model=model, tokenizer=tokenizer, train_dataset=train_dataset, eval_dataset=None, args=SFTConfig( dataset_text_field="text", dataset_num_proc=1, max_length=MAX_SEQ_LEN, packing=False, per_device_train_batch_size=1, gradient_accumulation_steps=8, warmup_steps= 5, num_train_epochs=2, learning_rate=2e-5, logging_steps=1, save_steps=100, save_total_limit=3, optim="adamw_8bit", weight_decay=0.01, max_grad_norm=0.3, lr_scheduler_type="linear", output_dir=OUTPUT_DIR, seed=3407, report_to="none", ),) trainer = mask_data( trainer, tokenizer=tokenizer, train_on_reasoning=False, train_on_final_answers=True, train_on_tools=True,) print(trainer.train_dataset.preview()) trainer_stats = trainer.train(resume_from_checkpoint=False) model.push_to_hub(f"{HUB_REPO_ID}-LoRA", token=HF_TOKEN)tokenizer.push_to_hub(f"{HUB_REPO_ID}-LoRA", token=HF_TOKEN) model.push_to_hub_merged(HUB_REPO_ID, tokenizer, save_method="merged_16bit", token=HF_TOKEN)