Merge branch 'nightly' into feature/canvas-lab

This commit is contained in:
Shine1i 2026-02-23 21:54:35 +01:00
commit 8739a01f56
76 changed files with 629 additions and 382 deletions

4
.gitignore vendored
View file

@ -38,6 +38,9 @@ unsloth_training_checkpoints/
.DS_Store
Thumbs.db
# Firebase
firebase-debug.log
# Other
resources/
tmp/
@ -55,3 +58,4 @@ studio/tests/
studio/backend/tests/
log_rtx.txt
log.txt
setup_leo.sh

View file

@ -24,6 +24,11 @@ echo "╔═══════════════════════
echo "║ Unsloth Studio Setup Script ║"
echo "╚══════════════════════════════════════╝"
# ── Clean up stale Unsloth compiled caches ──
rm -rf "$SCRIPT_DIR/unsloth_compiled_cache"
rm -rf "$SCRIPT_DIR/studio/backend/unsloth_compiled_cache"
rm -rf "$SCRIPT_DIR/studio/tmp/unsloth_compiled_cache"
# ── Detect Colab (like unsloth does) ──
IS_COLAB=false
keynames=$'\n'$(printenv | cut -d= -f1)

View file

@ -34,10 +34,6 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
finetune_vision_layers: false
finetune_language_layers: true
finetune_attention_modules: true
finetune_mlp_modules: true
logging:
enable_wandb: false

View file

@ -34,10 +34,6 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
finetune_vision_layers: false
finetune_language_layers: true
finetune_attention_modules: true
finetune_mlp_modules: true
logging:
enable_wandb: false

View file

@ -35,10 +35,6 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
finetune_vision_layers: false
finetune_language_layers: true
finetune_attention_modules: true
finetune_mlp_modules: true
logging:
enable_wandb: false

View file

@ -35,10 +35,6 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
finetune_vision_layers: false
finetune_language_layers: true
finetune_attention_modules: true
finetune_mlp_modules: true
logging:
enable_wandb: false

View file

@ -33,10 +33,6 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
finetune_vision_layers: false
finetune_language_layers: true
finetune_attention_modules: true
finetune_mlp_modules: true
logging:
enable_wandb: false

View file

@ -34,10 +34,6 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
finetune_vision_layers: false
finetune_language_layers: true
finetune_attention_modules: true
finetune_mlp_modules: true
logging:
enable_wandb: false

View file

@ -35,10 +35,6 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
finetune_vision_layers: false
finetune_language_layers: true
finetune_attention_modules: true
finetune_mlp_modules: true
logging:
enable_wandb: false

View file

@ -29,7 +29,7 @@ lora:
- "all-linear"
use_rslora: false
use_loftq: false
finetune_vision_layers: false
finetune_vision_layers: true
finetune_language_layers: true
finetune_attention_modules: true
finetune_mlp_modules: true

View file

@ -29,7 +29,7 @@ lora:
- "all-linear"
use_rslora: false
use_loftq: false
finetune_vision_layers: false
finetune_vision_layers: true
finetune_language_layers: true
finetune_attention_modules: true
finetune_mlp_modules: true

View file

@ -29,7 +29,7 @@ lora:
- "all-linear"
use_rslora: false
use_loftq: false
finetune_vision_layers: false
finetune_vision_layers: true
finetune_language_layers: true
finetune_attention_modules: true
finetune_mlp_modules: true

View file

@ -35,10 +35,6 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
finetune_vision_layers: false
finetune_language_layers: true
finetune_attention_modules: true
finetune_mlp_modules: true
logging:
enable_wandb: false

View file

@ -35,10 +35,6 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
finetune_vision_layers: false
finetune_language_layers: true
finetune_attention_modules: true
finetune_mlp_modules: true
logging:
enable_wandb: false

View file

@ -37,10 +37,6 @@ lora:
- "shared_mlp.output_linear"
use_rslora: false
use_loftq: false
finetune_vision_layers: false
finetune_language_layers: true
finetune_attention_modules: true
finetune_mlp_modules: true
logging:
enable_wandb: false

View file

@ -37,10 +37,6 @@ lora:
- "shared_mlp.output_linear"
use_rslora: false
use_loftq: false
finetune_vision_layers: false
finetune_language_layers: true
finetune_attention_modules: true
finetune_mlp_modules: true
logging:
enable_wandb: false

View file

@ -34,10 +34,6 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
finetune_vision_layers: false
finetune_language_layers: true
finetune_attention_modules: true
finetune_mlp_modules: true
logging:
enable_wandb: false

View file

@ -35,10 +35,6 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
finetune_vision_layers: false
finetune_language_layers: true
finetune_attention_modules: true
finetune_mlp_modules: true
logging:
enable_wandb: false

View file

@ -35,10 +35,6 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
finetune_vision_layers: false
finetune_language_layers: true
finetune_attention_modules: true
finetune_mlp_modules: true
logging:
enable_wandb: false

View file

@ -34,10 +34,6 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
finetune_vision_layers: false
finetune_language_layers: true
finetune_attention_modules: true
finetune_mlp_modules: true
logging:
enable_wandb: false

View file

@ -34,10 +34,6 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
finetune_vision_layers: false
finetune_language_layers: true
finetune_attention_modules: true
finetune_mlp_modules: true
logging:
enable_wandb: false

View file

@ -34,10 +34,6 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
finetune_vision_layers: false
finetune_language_layers: true
finetune_attention_modules: true
finetune_mlp_modules: true
logging:
enable_wandb: false

View file

@ -34,10 +34,6 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
finetune_vision_layers: false
finetune_language_layers: true
finetune_attention_modules: true
finetune_mlp_modules: true
logging:
enable_wandb: false

View file

@ -30,10 +30,6 @@ lora:
- "v_proj"
use_rslora: false
use_loftq: false
finetune_vision_layers: false
finetune_language_layers: true
finetune_attention_modules: true
finetune_mlp_modules: true
logging:
enable_wandb: false

View file

@ -35,7 +35,7 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
finetune_vision_layers: false
finetune_vision_layers: true
finetune_language_layers: true
finetune_attention_modules: true
finetune_mlp_modules: true

View file

@ -34,10 +34,6 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
finetune_vision_layers: false
finetune_language_layers: true
finetune_attention_modules: true
finetune_mlp_modules: true
logging:
enable_wandb: false

View file

@ -34,10 +34,6 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
finetune_vision_layers: false
finetune_language_layers: true
finetune_attention_modules: true
finetune_mlp_modules: true
logging:
enable_wandb: false

View file

@ -34,10 +34,6 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
finetune_vision_layers: false
finetune_language_layers: true
finetune_attention_modules: true
finetune_mlp_modules: true
logging:
enable_wandb: false

View file

@ -33,10 +33,6 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
finetune_vision_layers: false
finetune_language_layers: true
finetune_attention_modules: true
finetune_mlp_modules: true
logging:
enable_wandb: false

View file

@ -30,10 +30,6 @@ lora:
- "v_proj"
use_rslora: false
use_loftq: false
finetune_vision_layers: false
finetune_language_layers: true
finetune_attention_modules: true
finetune_mlp_modules: true
logging:
enable_wandb: false

View file

@ -35,10 +35,6 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
finetune_vision_layers: false
finetune_language_layers: true
finetune_attention_modules: true
finetune_mlp_modules: true
logging:
enable_wandb: false

View file

@ -34,10 +34,6 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
finetune_vision_layers: false
finetune_language_layers: true
finetune_attention_modules: true
finetune_mlp_modules: true
logging:
enable_wandb: false

View file

@ -1,17 +1,17 @@
# Model defaults for unsloth/Phi-4-14B-Instruct
# Based on Phi_4-Conversational.ipynb
# Also applies to: unsloth/Phi-4-14B-Instruct-unsloth-bnb-4bit, microsoft/Phi-4-14B-Instruct, unsloth/Phi-4-14B-Instruct-bnb-4bit
# Model defaults for unsloth/GLM-4.7-Flash
# Based on GLM_Flash_A100(80GB).py
# Also applies to: unsloth/GLM-4.7-Flash-unsloth-bnb-4bit, unsloth/GLM-4.7-Flash-bnb-4bit, THUDM/GLM-4.7-Flash
training:
max_seq_length: 2048
# num_epochs: 4
num_epochs: 0
learning_rate: 2e-4
batch_size: 2
gradient_accumulation_steps: 4
batch_size: 4
gradient_accumulation_steps: 2
warmup_steps: 5
max_steps: 30
save_steps: 30
max_steps: 60
save_steps: 60
weight_decay: 0.001
random_seed: 3407
packing: false
@ -21,7 +21,7 @@ training:
lr_scheduler_type: "linear"
lora:
lora_r: 16
lora_r: 8
lora_alpha: 16
lora_dropout: 0.0
target_modules:
@ -32,12 +32,9 @@ lora:
- "gate_proj"
- "up_proj"
- "down_proj"
- "out_proj"
use_rslora: false
use_loftq: false
finetune_vision_layers: false
finetune_language_layers: true
finetune_attention_modules: true
finetune_mlp_modules: true
logging:
enable_wandb: false
@ -46,3 +43,8 @@ logging:
tensorboard_dir: "runs"
log_frequency: 10
inference:
temperature: 0.7
top_p: 0.8
top_k: 20

View file

@ -29,10 +29,6 @@ lora:
- "all-linear"
use_rslora: false
use_loftq: false
finetune_vision_layers: false
finetune_language_layers: true
finetune_attention_modules: true
finetune_mlp_modules: true
logging:
enable_wandb: false

View file

@ -37,10 +37,6 @@ lora:
- "out_proj"
use_rslora: false
use_loftq: false
finetune_vision_layers: false
finetune_language_layers: true
finetune_attention_modules: true
finetune_mlp_modules: true
logging:
enable_wandb: false

View file

@ -33,10 +33,6 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
finetune_vision_layers: false
finetune_language_layers: true
finetune_attention_modules: true
finetune_mlp_modules: true
logging:
enable_wandb: false

View file

@ -35,10 +35,6 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
finetune_vision_layers: false
finetune_language_layers: true
finetune_attention_modules: true
finetune_mlp_modules: true
logging:
enable_wandb: false

View file

@ -34,10 +34,6 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
finetune_vision_layers: false
finetune_language_layers: true
finetune_attention_modules: true
finetune_mlp_modules: true
logging:
enable_wandb: false

View file

@ -29,10 +29,6 @@ lora:
- "v_proj"
use_rslora: false
use_loftq: false
finetune_vision_layers: false
finetune_language_layers: true
finetune_attention_modules: true
finetune_mlp_modules: true
logging:
enable_wandb: false

View file

@ -34,10 +34,6 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
finetune_vision_layers: false
finetune_language_layers: true
finetune_attention_modules: true
finetune_mlp_modules: true
logging:
enable_wandb: false

View file

@ -34,10 +34,6 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
finetune_vision_layers: false
finetune_language_layers: true
finetune_attention_modules: true
finetune_mlp_modules: true
logging:
enable_wandb: false

View file

@ -35,10 +35,6 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
finetune_vision_layers: false
finetune_language_layers: true
finetune_attention_modules: true
finetune_mlp_modules: true
logging:
enable_wandb: false

View file

@ -0,0 +1,51 @@
# Model defaults for imdatta0/tiny_qwen3_moe_2.8B_0.7B
# Based on TinyQwen3_MoE.py
# Dummy model of qwen3moe architecture created to fit in T4
# MoE model - includes gate_up_proj for MoE layers
training:
max_seq_length: 2048
# num_epochs: 4
num_epochs: 0
learning_rate: 2e-4
batch_size: 1
gradient_accumulation_steps: 1
warmup_steps: 5
max_steps: 50
save_steps: 50
weight_decay: 0.001
random_seed: 3407
packing: false
train_on_completions: true
gradient_checkpointing: "unsloth"
optim: "adamw_8bit"
lr_scheduler_type: "linear"
lora:
lora_r: 32
lora_alpha: 64
lora_dropout: 0.0
target_modules:
- "q_proj"
- "k_proj"
- "v_proj"
- "o_proj"
- "gate_proj"
- "up_proj"
- "down_proj"
- "gate_up_proj"
use_rslora: false
use_loftq: false
logging:
enable_wandb: false
wandb_project: "llm-finetuning"
enable_tensorboard: false
tensorboard_dir: "runs"
log_frequency: 10
inference:
temperature: 0.6
top_k: 20
top_p: 0.95

View file

@ -34,10 +34,6 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
finetune_vision_layers: false
finetune_language_layers: true
finetune_attention_modules: true
finetune_mlp_modules: true
logging:
enable_wandb: false

View file

@ -34,10 +34,6 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
finetune_vision_layers: false
finetune_language_layers: true
finetune_attention_modules: true
finetune_mlp_modules: true
logging:
enable_wandb: false

View file

@ -34,10 +34,6 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
finetune_vision_layers: false
finetune_language_layers: true
finetune_attention_modules: true
finetune_mlp_modules: true
logging:
enable_wandb: false

View file

@ -34,10 +34,6 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
finetune_vision_layers: false
finetune_language_layers: true
finetune_attention_modules: true
finetune_mlp_modules: true
logging:
enable_wandb: false

View file

@ -35,10 +35,6 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
finetune_vision_layers: false
finetune_language_layers: true
finetune_attention_modules: true
finetune_mlp_modules: true
logging:
enable_wandb: false

View file

@ -34,10 +34,6 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
finetune_vision_layers: false
finetune_language_layers: true
finetune_attention_modules: true
finetune_mlp_modules: true
logging:
enable_wandb: false

View file

@ -35,10 +35,6 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
finetune_vision_layers: false
finetune_language_layers: true
finetune_attention_modules: true
finetune_mlp_modules: true
logging:
enable_wandb: false

View file

@ -35,10 +35,6 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
finetune_vision_layers: false
finetune_language_layers: true
finetune_attention_modules: true
finetune_mlp_modules: true
logging:
enable_wandb: false

View file

@ -35,10 +35,6 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
finetune_vision_layers: false
finetune_language_layers: true
finetune_attention_modules: true
finetune_mlp_modules: true
logging:
enable_wandb: false

View file

@ -0,0 +1,51 @@
# Model defaults for unsloth/Qwen3-30B-A3B-Instruct-2507
# Based on Qwen3_MoE.py
# Also applies to: Qwen/Qwen3-30B-A3B-Instruct-2507, unsloth/Qwen3-30B-A3B-Instruct-2507-bnb-4bit
# MoE model - includes gate_up_proj for MoE layers
training:
max_seq_length: 2048
# num_epochs: 4
num_epochs: 0
learning_rate: 2e-4
batch_size: 1
gradient_accumulation_steps: 1
warmup_steps: 5
max_steps: 50
save_steps: 50
weight_decay: 0.001
random_seed: 3407
packing: false
train_on_completions: true
gradient_checkpointing: "unsloth"
optim: "adamw_8bit"
lr_scheduler_type: "linear"
lora:
lora_r: 32
lora_alpha: 64
lora_dropout: 0.0
target_modules:
- "q_proj"
- "k_proj"
- "v_proj"
- "o_proj"
- "gate_proj"
- "up_proj"
- "down_proj"
- "gate_up_proj"
use_rslora: false
use_loftq: false
logging:
enable_wandb: false
wandb_project: "llm-finetuning"
enable_tensorboard: false
tensorboard_dir: "runs"
log_frequency: 10
inference:
temperature: 0.6
top_k: 20
top_p: 0.95

View file

@ -35,10 +35,6 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
finetune_vision_layers: false
finetune_language_layers: true
finetune_attention_modules: true
finetune_mlp_modules: true
logging:
enable_wandb: false

View file

@ -35,10 +35,6 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
finetune_vision_layers: false
finetune_language_layers: true
finetune_attention_modules: true
finetune_mlp_modules: true
logging:
enable_wandb: false

View file

@ -35,10 +35,6 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
finetune_vision_layers: false
finetune_language_layers: true
finetune_attention_modules: true
finetune_mlp_modules: true
logging:
enable_wandb: false

View file

@ -184,6 +184,10 @@ class InferenceBackend:
# Clear GPU memory cache
clear_gpu_cache()
# Remove stale compiled cache so the next model gets a fresh one
from utils.cache_cleanup import clear_unsloth_compiled_cache
clear_unsloth_compiled_cache()
logger.info(f"Model '{model_name}' successfully unloaded.")
return True
except Exception as e:

View file

@ -7,7 +7,7 @@ import os
os.environ["TOKENIZERS_PARALLELISM"] = "false"
import torch
from utils.hardware import clear_gpu_cache
from utils.hardware import clear_gpu_cache, safe_num_proc
torch._dynamo.config.recompile_limit = 64
from unsloth import FastLanguageModel, FastVisionModel, is_bfloat16_supported
from unsloth.chat_templates import get_chat_template
@ -126,6 +126,10 @@ class UnslothTrainer:
print("\nClearing GPU memory before training...")
clear_gpu_cache()
# Remove stale compiled cache so the new model gets a fresh one
from utils.cache_cleanup import clear_unsloth_compiled_cache
clear_unsloth_compiled_cache()
# Detect if this is a vision model AND dataset is multimodal
# A vision-capable model with a text-only dataset should use FastLanguageModel
self.is_vlm = is_vision_model(model_name) and is_dataset_multimodal
@ -707,7 +711,7 @@ class UnslothTrainer:
"output_dir": output_dir,
"report_to": ["wandb"] if training_args.get('enable_wandb', False) else "none",
"include_num_input_tokens_seen": True, # Enable token counting
"dataset_num_proc": max(1, os.cpu_count() // 4),
"dataset_num_proc": safe_num_proc(max(1, os.cpu_count() // 4)),
}
# Add warmup parameter - use warmup_ratio if provided, otherwise warmup_steps
@ -792,7 +796,7 @@ class UnslothTrainer:
trainer_kwargs = {
"model": self.model,
"train_dataset": dataset['dataset'],
"processing_class": self.tokenizer.tokenizer,
"processing_class": self.tokenizer,
"data_collator": data_collator,
"args": SFTConfig(**config_args),
}
@ -867,9 +871,44 @@ class UnslothTrainer:
self.trainer,
instruction_part=instruction_part,
response_part=response_part,
num_proc=config_args.get("dataset_num_proc", max(1, os.cpu_count() // 4)),
num_proc=config_args.get("dataset_num_proc", safe_num_proc(max(1, os.cpu_count() // 4))),
)
print("Train on responses only configured successfully\n")
# ── Safety net: check if all samples were filtered out ──
# Unsloth's train_on_responses_only masks non-response
# tokens with -100. If max_seq_length is too short and the
# response portion gets truncated away, EVERY sample ends
# up with all labels == -100 and Unsloth removes them,
# leaving 0 usable training samples.
filtered_len = len(self.trainer.train_dataset)
original_len = len(dataset["dataset"])
dropped = original_len - filtered_len
drop_pct = round(100 * dropped / original_len, 1) if original_len > 0 else 0
if filtered_len == 0 or drop_pct > 30:
max_seq = training_args.get('max_seq_length', 2048)
error_msg = (
f"{dropped}/{original_len} samples ({drop_pct}%) "
f"were dropped after applying 'train on responses "
f"only' — only {filtered_len} remain. This usually "
f"means max_seq_length ({max_seq}) is too short "
f"and the response portion is being truncated "
f"away. Try increasing max_seq_length (e.g. 8192) "
f"or disabling 'Train on completions'."
)
logger.error(error_msg)
self._update_progress(error=error_msg, is_training=False)
return
if dropped > 0:
print(
f"⚠️ {dropped}/{original_len} samples "
f"({drop_pct}%) were dropped (all labels "
f"masked). {filtered_len} samples remain.\n"
)
print(f"Post-filter dataset size: {filtered_len} samples\n")
except Exception as e:
logger.warning(f"Failed to apply train on responses only: {e}")
train_on_responses_enabled = False
@ -951,7 +990,7 @@ class UnslothTrainer:
progress_callback = ProgressCallback(self)
self.trainer.add_callback(progress_callback)
num_samples = len(dataset["dataset"])
num_samples = len(self.trainer.train_dataset)
batch_size = training_args.get('batch_size', 2)
grad_accum = training_args.get('gradient_accumulation_steps', 4)
num_epochs = training_args.get('num_epochs', 3)

View file

@ -3,7 +3,6 @@ Main FastAPI application for Unsloth UI Backend
"""
import os
import secrets
import shutil
from contextlib import asynccontextmanager
from fastapi import FastAPI
@ -27,12 +26,15 @@ from auth import storage
from utils.hardware import detect_hardware, get_device, DeviceType
import utils.hardware.hardware as _hw_module
UNSLOTH_CACHE_DIR = Path(__file__).parent / "unsloth_compiled_cache"
from utils.cache_cleanup import clear_unsloth_compiled_cache
@asynccontextmanager
async def lifespan(app: FastAPI):
"""Startup: detect hardware, print setup token if needed. Shutdown: clean up compiled cache."""
# Clean up any stale compiled cache from previous runs
clear_unsloth_compiled_cache()
# Detect hardware first — sets DEVICE global used everywhere
detect_hardware()
@ -60,7 +62,7 @@ async def lifespan(app: FastAPI):
yield
# Cleanup
_hw_module.DEVICE = None
shutil.rmtree(UNSLOTH_CACHE_DIR, ignore_errors=True)
clear_unsloth_compiled_cache()
# Create FastAPI app

View file

@ -0,0 +1,30 @@
"""
Utility for cleaning up the Unsloth compiled cache directory.
The unsloth_compiled_cache is created by unsloth_zoo/compiler.py during
FastModel.from_pretrained() and contains model-type-specific compiled Python
files. It should be cleared between model loads to avoid stale artefacts.
"""
import shutil
import logging
from pathlib import Path
logger = logging.getLogger(__name__)
# Possible locations where unsloth_compiled_cache may appear
_BACKEND_DIR = Path(__file__).resolve().parent.parent # studio/backend
_PROJECT_ROOT = _BACKEND_DIR.parent.parent # repo root
_CACHE_DIRS = [
_BACKEND_DIR / "unsloth_compiled_cache",
_PROJECT_ROOT / "unsloth_compiled_cache",
_PROJECT_ROOT / "studio" / "tmp" / "unsloth_compiled_cache",
]
def clear_unsloth_compiled_cache() -> None:
"""Remove every known unsloth_compiled_cache directory (idempotent)."""
for cache_dir in _CACHE_DIRS:
if cache_dir.exists():
logger.info(f"Removing unsloth compiled cache: {cache_dir}")
shutil.rmtree(cache_dir, ignore_errors=True)

View file

@ -283,9 +283,11 @@ def apply_chat_template_to_dataset(
}
if not isinstance(dataset, IterableDataset):
from multiprocessing import cpu_count
from utils.hardware import safe_num_proc
if num_proc is None or type(num_proc) is not int:
num_proc = max(1, cpu_count() // 3)
num_proc = safe_num_proc()
else:
num_proc = safe_num_proc(num_proc)
dataset_map_kwargs['num_proc'] = num_proc
dataset_map_kwargs['desc'] = "Applying template to Alpaca format"
@ -347,9 +349,11 @@ def apply_chat_template_to_dataset(
}
if not isinstance(dataset, IterableDataset):
from multiprocessing import cpu_count
from utils.hardware import safe_num_proc
if num_proc is None or type(num_proc) is not int:
num_proc = max(1, cpu_count() // 3)
num_proc = safe_num_proc()
else:
num_proc = safe_num_proc(num_proc)
dataset_map_kwargs['num_proc'] = num_proc
dataset_map_kwargs['desc'] = f"Applying chat template to {final_format}"

View file

@ -107,10 +107,12 @@ def standardize_chat_format(
}
if not isinstance(dataset, IterableDataset):
from multiprocessing import cpu_count
from utils.hardware import safe_num_proc
if num_proc is None or type(num_proc) is not int:
num_proc = max(1, cpu_count() // 3)
num_proc = safe_num_proc()
else:
num_proc = safe_num_proc(num_proc)
dataset_map_kwargs['num_proc'] = num_proc
dataset_map_kwargs['desc'] = "Standardizing chat format"
@ -173,10 +175,12 @@ def convert_chatml_to_alpaca(dataset, batch_size=1000, num_proc=None):
}
if not isinstance(dataset, IterableDataset):
from multiprocessing import cpu_count
from utils.hardware import safe_num_proc
if num_proc is None or type(num_proc) is not int:
num_proc = max(1, cpu_count() // 3)
num_proc = safe_num_proc()
else:
num_proc = safe_num_proc(num_proc)
dataset_map_kwargs['num_proc'] = num_proc
dataset_map_kwargs['desc'] = "Converting ChatML to Alpaca format"
@ -221,10 +225,12 @@ def convert_alpaca_to_chatml(dataset, batch_size=1000, num_proc=None):
}
if not isinstance(dataset, IterableDataset):
from multiprocessing import cpu_count
from utils.hardware import safe_num_proc
if num_proc is None or type(num_proc) is not int:
num_proc = max(1, cpu_count() // 3)
num_proc = safe_num_proc()
else:
num_proc = safe_num_proc(num_proc)
dataset_map_kwargs['num_proc'] = num_proc
dataset_map_kwargs['desc'] = "Converting Alpaca to ChatML format"

View file

@ -383,7 +383,13 @@ TEMPLATE_TO_MODEL_MAPPER = {
"unsloth/yi-34b-chat-bnb-4bit",
"01-ai/Yi-6B-Chat",
"01-ai/Yi-34B-Chat",
)
),
"glm": (
"unsloth/GLM-4.7-Flash-unsloth-bnb-4bit",
"unsloth/GLM-4.7-Flash",
"THUDM/GLM-4.7-Flash",
"unsloth/GLM-4.7-Flash-bnb-4bit",
),
}
MODEL_TO_TEMPLATE_MAPPER = {}
@ -506,4 +512,9 @@ TEMPLATE_TO_RESPONSES_MAPPER = {
"instruction": "<|im_start|>user\n",
"response": "<|im_start|>assistant\n",
},
"glm": {
"instruction": "[gMASK]<sop><|user|>",
"response": "<|assistant|><think>",
},
}

View file

@ -13,6 +13,8 @@ from .hardware import (
get_gpu_summary,
get_package_versions,
get_gpu_utilization,
get_physical_gpu_count,
safe_num_proc,
)
__all__ = [
@ -27,4 +29,6 @@ __all__ = [
'get_gpu_summary',
'get_package_versions',
'get_gpu_utilization',
'get_physical_gpu_count',
'safe_num_proc',
]

View file

@ -385,3 +385,69 @@ def get_gpu_utilization() -> Dict[str, Any]:
"power_limit_w": power_limit,
"power_utilization_pct": power_pct,
}
# ========== Multi-GPU Detection & Safe num_proc ==========
_physical_gpu_count: Optional[int] = None
def get_physical_gpu_count() -> int:
"""
Return the number of physical NVIDIA GPUs on the machine.
Uses ``nvidia-smi -L`` which is NOT affected by CUDA_VISIBLE_DEVICES,
so it always reflects the true hardware count.
Result is cached after the first call.
"""
global _physical_gpu_count
if _physical_gpu_count is not None:
return _physical_gpu_count
try:
import subprocess
result = subprocess.run(
["nvidia-smi", "-L"],
capture_output=True, text=True, timeout=5,
)
if result.returncode == 0 and result.stdout.strip():
_physical_gpu_count = len(result.stdout.strip().splitlines())
else:
_physical_gpu_count = 1
except Exception:
_physical_gpu_count = 1
return _physical_gpu_count
def safe_num_proc(desired: Optional[int] = None) -> int:
"""
Return a safe ``num_proc`` for ``dataset.map()`` calls.
On multi-GPU machines the NVIDIA driver spawns extra background threads,
making ``os.fork()`` prone to deadlocks when many workers are created.
This helper caps ``num_proc`` to 4 on such machines.
On single-GPU (or CPU-only) machines the original value is returned
unchanged.
Args:
desired: The num_proc you *want*. If None, auto-computes from
``os.cpu_count()``.
Returns:
A safe integer 1.
"""
import os
if desired is None or not isinstance(desired, int):
desired = max(1, os.cpu_count() // 3)
if get_physical_gpu_count() > 1:
capped = min(4, desired)
print(
f"⚙️ Multi-GPU detected ({get_physical_gpu_count()} GPUs) — "
f"capping num_proc {desired}{capped} to avoid fork deadlocks"
)
return capped
return desired

View file

@ -26,25 +26,22 @@ import {
type FC,
type PropsWithChildren,
useEffect,
useMemo,
useState,
} from "react";
import { useShallow } from "zustand/shallow";
const useFileSrc = (file: File | undefined): string | undefined => {
const objectUrl = useMemo(
() => (file ? URL.createObjectURL(file) : undefined),
[file],
);
const [objectUrl, setObjectUrl] = useState<string | undefined>(undefined);
useEffect(() => {
if (!objectUrl) {
return undefined;
if (!file) {
setObjectUrl(undefined);
return;
}
return () => {
URL.revokeObjectURL(objectUrl);
};
}, [objectUrl]);
const url = URL.createObjectURL(file);
setObjectUrl(url);
return () => URL.revokeObjectURL(url);
}, [file]);
return objectUrl;
};

View file

@ -151,12 +151,9 @@ export function HubModelPicker({
const metricsById = useMemo(
() =>
new Map(
results.map((result) => [
result.id,
result.totalParams
? formatCompact(result.totalParams)
: `${formatCompact(result.downloads)}`,
]),
results
.filter((result) => result.totalParams)
.map((result) => [result.id, formatCompact(result.totalParams!)]),
),
[results],
);
@ -167,11 +164,7 @@ export function HubModelPicker({
{ est: number; status: VramFitStatus | null; detail: string | null }
>();
for (const r of results) {
const detail = r.totalParams
? formatCompact(r.totalParams)
: r.downloads != null
? `${formatCompact(r.downloads)}`
: null;
const detail = r.totalParams ? formatCompact(r.totalParams) : null;
if (r.totalParams) {
const est = estimateLoadingVram(r.totalParams, "qlora");
const status = gpu.available

View file

@ -218,9 +218,9 @@ function InlineSidebar({
className={cn(
"bg-sidebar text-sidebar-foreground h-full overflow-hidden rounded-2xl corner-squircle transition-[width] duration-200 ease-linear",
!collapsed &&
(side === "left"
? "border-r border-0 border-sidebar-border"
: "border-l border-0 border-sidebar-border"),
(side === "left"
? "border-r border-0 border-sidebar-border"
: "border-l border-0 border-sidebar-border"),
collapsed ? "w-0" : "w-(--sidebar-width)",
)}
>
@ -301,9 +301,18 @@ export function ChatPage(): ReactElement {
const handleCheckpointChange = useCallback(
(value: string, meta?: { isLora: boolean }) => {
void selectModel({ id: value, isLora: meta?.isLora });
const currentCheckpoint =
useChatRuntimeStore.getState().params.checkpoint;
if (!value || value === currentCheckpoint) return;
setView({ mode: "single", newThreadNonce: crypto.randomUUID() });
void (async () => {
if (currentCheckpoint) {
await ejectModel();
}
await selectModel({ id: value, isLora: meta?.isLora });
})();
},
[selectModel],
[selectModel, ejectModel],
);
const handleEject = useCallback(() => {
void ejectModel();
@ -349,6 +358,41 @@ export function ChatPage(): ReactElement {
setViewBeforeCompare(null);
}, [viewBeforeCompare]);
const handleThreadSelect = useCallback(
(nextView: ChatView) => {
setView(nextView);
const threadId =
nextView.mode === "single" ? nextView.threadId : undefined;
const pairId =
nextView.mode === "compare" ? nextView.pairId : undefined;
void (async () => {
let thread: import("./types").ThreadRecord | undefined;
if (threadId) {
thread = await db.threads.get(threadId);
} else if (pairId) {
thread = await db.threads
.where("pairId")
.equals(pairId)
.first();
}
const threadModelId = thread?.modelId;
if (!threadModelId) return;
const currentCheckpoint =
useChatRuntimeStore.getState().params.checkpoint;
if (threadModelId === currentCheckpoint) return;
if (currentCheckpoint) {
await ejectModel();
}
await selectModel({ id: threadModelId });
})();
},
[ejectModel, selectModel],
);
const models = useMemo<ModelOption[]>(
() =>
modelsFromStore.map((model) => ({
@ -475,99 +519,99 @@ export function ChatPage(): ReactElement {
return (
<div className="h-[calc(100dvh-4rem)] bg-background overflow-hidden">
<GuidedTour {...tour.tourProps} />
<SidebarProvider
defaultOpen={true}
open={sidebarOpen}
onOpenChange={setSidebarOpen}
className="!min-h-0 h-full w-full max-w-7xl mx-auto px-2 sm:px-4"
style={
{
"--sidebar-width": "14rem",
"--sidebar-width-icon": "3rem",
} as CSSProperties
}
>
<InlineSidebar>
<ThreadSidebar
view={view}
onSelect={setView}
onNewThread={handleNewThread}
onNewCompare={handleNewCompare}
showCompare={canCompare}
/>
</InlineSidebar>
<GuidedTour {...tour.tourProps} />
<SidebarProvider
defaultOpen={true}
open={sidebarOpen}
onOpenChange={setSidebarOpen}
className="!min-h-0 h-full w-full max-w-7xl mx-auto px-2 sm:px-4"
style={
{
"--sidebar-width": "14rem",
"--sidebar-width-icon": "3rem",
} as CSSProperties
}
>
<InlineSidebar>
<ThreadSidebar
view={view}
onSelect={handleThreadSelect}
onNewThread={handleNewThread}
onNewCompare={handleNewCompare}
showCompare={canCompare}
/>
</InlineSidebar>
<div className="flex min-h-0 min-w-0 flex-1 flex-col">
<div className="flex h-11 shrink-0 items-center px-1.5 sm:px-2">
<div className="flex items-center gap-1">
<SidebarTrigger />
<TopBarActions
onNewThread={handleNewThread}
onNewCompare={handleNewCompare}
showCompare={canCompare}
/>
<ModelSelector
models={models}
loraModels={loraModels}
value={inferenceParams.checkpoint}
onValueChange={handleCheckpointChange}
onEject={handleEject}
variant="ghost"
open={modelSelectorOpen}
onOpenChange={handleModelSelectorOpenChange}
triggerDataTour="chat-model-selector"
contentDataTour="chat-model-selector-popover"
className="max-w-[62vw] sm:max-w-none"
/>
{loadingModel ? (
<div
className="flex items-center gap-1.5 text-muted-foreground"
title={`Loading ${loadingModel.displayName}. This may include downloading.`}
>
<Spinner className="size-3.5 shrink-0" />
<span className="text-xs">
Downloading model
</span>
</div>
) : null}
</div>
{modelsError && (
<div className="ml-2 text-xs text-destructive truncate max-w-[28rem]">
{modelsError}
<div className="flex min-h-0 min-w-0 flex-1 flex-col">
<div className="flex h-11 shrink-0 items-center px-1.5 sm:px-2">
<div className="flex items-center gap-1">
<SidebarTrigger />
<TopBarActions
onNewThread={handleNewThread}
onNewCompare={handleNewCompare}
showCompare={canCompare}
/>
<ModelSelector
models={models}
loraModels={loraModels}
value={inferenceParams.checkpoint}
onValueChange={handleCheckpointChange}
onEject={handleEject}
variant="ghost"
open={modelSelectorOpen}
onOpenChange={handleModelSelectorOpenChange}
triggerDataTour="chat-model-selector"
contentDataTour="chat-model-selector-popover"
className="max-w-[62vw] sm:max-w-none"
/>
{loadingModel ? (
<div
className="flex items-center gap-1.5 text-muted-foreground"
title={`Loading ${loadingModel.displayName}. This may include downloading.`}
>
<Spinner className="size-3.5 shrink-0" />
<span className="text-xs">
Downloading model
</span>
</div>
) : null}
</div>
{modelsError && (
<div className="ml-2 text-xs text-destructive truncate max-w-[28rem]">
{modelsError}
</div>
)}
<div className="flex-1" />
<button
type="button"
onClick={() => setSettingsOpen((o) => !o)}
className="flex h-9 w-9 items-center justify-center rounded-md text-muted-foreground transition-colors hover:bg-accent hover:text-foreground"
title="Inference settings"
data-tour="chat-settings"
>
<HugeiconsIcon icon={Settings04Icon} className="size-5" />
</button>
</div>
{view.mode === "single" ? (
<SingleContent
key={view.threadId ?? view.newThreadNonce ?? "new"}
threadId={view.threadId}
newThreadNonce={view.newThreadNonce}
/>
) : (
<CompareContent key={view.pairId} pairId={view.pairId} />
)}
<div className="flex-1" />
<button
type="button"
onClick={() => setSettingsOpen((o) => !o)}
className="flex h-9 w-9 items-center justify-center rounded-md text-muted-foreground transition-colors hover:bg-accent hover:text-foreground"
title="Inference settings"
data-tour="chat-settings"
>
<HugeiconsIcon icon={Settings04Icon} className="size-5" />
</button>
</div>
{view.mode === "single" ? (
<SingleContent
key={view.threadId ?? view.newThreadNonce ?? "new"}
threadId={view.threadId}
newThreadNonce={view.newThreadNonce}
/>
) : (
<CompareContent key={view.pairId} pairId={view.pairId} />
)}
</div>
<ChatSettingsPanel
open={settingsOpen}
params={inferenceParams}
onParamsChange={setInferenceParams}
autoTitle={autoTitle}
onAutoTitleChange={setAutoTitle}
/>
</SidebarProvider>
<ChatSettingsPanel
open={settingsOpen}
params={inferenceParams}
onParamsChange={setInferenceParams}
autoTitle={autoTitle}
onAutoTitleChange={setAutoTitle}
/>
</SidebarProvider>
</div>
);
}

View file

@ -19,6 +19,20 @@ db.version(2)
})
.upgrade((tx) => tx.table("messages").clear());
db.version(3)
.stores({
threads: "id, modelType, pairId, archived, createdAt",
messages: "id, threadId, createdAt",
})
.upgrade((tx) =>
tx
.table("threads")
.toCollection()
.modify((thread) => {
if (!thread.modelId) thread.modelId = "";
}),
);
export { db };
export function useLiveQuery<T>(

View file

@ -165,8 +165,10 @@ export function useChatModelRuntime() {
setLoadingModel({ id: modelId, displayName });
try {
async function performLoad(): Promise<void> {
if (params.checkpoint) {
await unloadModel({ model_path: params.checkpoint });
const currentCheckpoint =
useChatRuntimeStore.getState().params.checkpoint;
if (currentCheckpoint) {
await unloadModel({ model_path: currentCheckpoint });
}
const loadResponse = await loadModel({

View file

@ -35,6 +35,14 @@ const DEFAULT_SUGGESTIONS = [
"Format a comparison of 3 databases as a markdown table with pros and cons",
];
type TitleResponse = {
choices?: Array<{
message?: {
content?: string;
};
}>;
};
class VisionImageAdapter implements AttachmentAdapter {
accept = "image/jpeg,image/png,image/webp,image/gif";
@ -216,7 +224,7 @@ async function generateTitleWithModel(payload: {
}),
});
const body = (await response.json().catch(() => null)) as any;
const body = (await response.json().catch(() => null)) as TitleResponse | null;
if (!response.ok) return null;
const raw: string | undefined = body?.choices?.[0]?.message?.content;
if (!raw) return null;
@ -233,27 +241,42 @@ function fallbackTitleFromUserText(userText: string): string {
return cleaned.slice(0, max) + (cleaned.length > max ? "..." : "");
}
function cloneContent(content: ThreadMessage["content"]): ThreadMessage["content"] {
return Array.isArray(content)
? JSON.parse(JSON.stringify(content))
: [];
}
function cloneAttachments(
attachments: readonly CompleteAttachment[] | undefined,
): readonly CompleteAttachment[] {
if (!Array.isArray(attachments)) {
return [];
}
return JSON.parse(JSON.stringify(attachments));
}
function toThreadMessage(m: MessageRecord): ThreadMessage {
const base = {
id: m.id,
createdAt: new Date(m.createdAt),
content:
Array.isArray(m.content) && m.content.length > 0
? m.content
: [{ type: "text" as const, text: "" }],
};
const content =
Array.isArray(m.content) && m.content.length > 0
? cloneContent(m.content)
: [{ type: "text" as const, text: "" }];
if (m.role === "user") {
return {
...base,
id: m.id,
createdAt: new Date(m.createdAt),
role: "user" as const,
attachments: [],
content: content as Extract<ThreadMessage, { role: "user" }>["content"],
attachments: cloneAttachments(m.attachments),
metadata: { custom: {} },
};
}
return {
...base,
id: m.id,
createdAt: new Date(m.createdAt),
role: "assistant" as const,
content: content as Extract<ThreadMessage, { role: "assistant" }>["content"],
status: { type: "complete" as const, reason: "unknown" as const },
metadata: {
custom: (m.metadata as Record<string, unknown>) ?? {},
@ -300,10 +323,13 @@ function createDexieAdapter(
},
async initialize(threadId: string) {
const currentModelId =
useChatRuntimeStore.getState().params.checkpoint ?? "";
await db.threads.add({
id: threadId,
title: "New Chat",
modelType,
modelId: currentModelId,
pairId,
archived: false,
createdAt: Date.now(),
@ -441,9 +467,9 @@ function ThreadHistoryProvider({
async append({ message }: ExportedMessageRepositoryItem) {
const { remoteId } = await aui.threadListItem().initialize();
const content = Array.isArray(message.content)
? JSON.parse(JSON.stringify(message.content))
: [];
const content = cloneContent(message.content);
const attachments =
message.role === "user" ? cloneAttachments(message.attachments) : [];
const custom = message.metadata?.custom;
const existing = await db.messages.get(message.id);
const createdAt =
@ -455,6 +481,7 @@ function ThreadHistoryProvider({
threadId: remoteId,
role: message.role,
content,
...(attachments.length > 0 && { attachments }),
...(custom && Object.keys(custom).length > 0 && { metadata: custom }),
createdAt,
});

View file

@ -8,6 +8,7 @@ export interface ThreadRecord {
id: string;
title: string;
modelType: ModelType;
modelId?: string;
pairId?: string;
archived: boolean;
createdAt: number;
@ -18,6 +19,7 @@ export interface MessageRecord {
threadId: string;
role: import("@assistant-ui/react").ThreadMessage["role"];
content: import("@assistant-ui/react").ThreadMessage["content"];
attachments?: import("@assistant-ui/react").ThreadMessage["attachments"];
metadata?: Record<string, unknown>;
createdAt: number;
}

View file

@ -244,10 +244,6 @@ export function ModelSelectionStep() {
<span className="text-xs text-muted-foreground shrink-0">
{sizeLabel}
</span>
) : r?.downloads != null ? (
<span className="text-[10px] text-muted-foreground shrink-0">
{formatCompact(r.downloads)}
</span>
) : null}
</ComboboxItem>
);

View file

@ -210,11 +210,7 @@ export function ModelSection() {
{ est: number; status: VramFitStatus | null; detail: string | null }
>();
for (const r of hfResults) {
const detail = r.totalParams
? formatCompact(r.totalParams)
: r.downloads != null
? `\u2193${formatCompact(r.downloads)}`
: null;
const detail = r.totalParams ? formatCompact(r.totalParams) : null;
if (r.totalParams) {
const est = estimateLoadingVram(r.totalParams, method);
const status = gpu.available

View file

@ -30,7 +30,7 @@ import {
ZapIcon,
} from "@hugeicons/core-free-icons";
import { HugeiconsIcon } from "@hugeicons/react";
import { useState, type ReactElement, type ReactNode } from "react";
import { useEffect, useState, type ReactElement, type ReactNode } from "react";
import { Link, useNavigate } from "@tanstack/react-router";
import { useShallow } from "zustand/react/shallow";
import { useGpuUtilization } from "@/hooks";
@ -83,6 +83,13 @@ export function ProgressSection(): ReactElement {
const { stopTrainingRun } = useTrainingActions();
const gpu = useGpuUtilization(runtime.isTrainingRunning);
const [stopDialogOpen, setStopDialogOpen] = useState(false);
const [stopRequested, setStopRequested] = useState(false);
useEffect(() => {
if (!runtime.isTrainingRunning) {
setStopRequested(false);
}
}, [runtime.isTrainingRunning]);
const pct =
runtime.totalSteps > 0
@ -209,11 +216,12 @@ export function ProgressSection(): ReactElement {
data-tour="studio-training-stop"
variant="destructive"
size="sm"
className="h-7 cursor-pointer px-3 text-xs"
className={`h-7 px-3 text-xs ${stopRequested ? "cursor-not-allowed opacity-60" : "cursor-pointer"}`}
onClick={() => setStopDialogOpen(true)}
disabled={!runtime.isTrainingRunning}
disabled={!runtime.isTrainingRunning || stopRequested}
>
<HugeiconsIcon icon={StopIcon} className="size-3" /> Stop
<HugeiconsIcon icon={StopIcon} className="size-3" />
{stopRequested ? "Stopping…" : "Stop"}
</Button>
<AlertDialogContent overlayClassName="bg-background/40 supports-backdrop-filter:backdrop-blur-[1px]">
<AlertDialogHeader>
@ -226,12 +234,24 @@ export function ProgressSection(): ReactElement {
<AlertDialogCancel>Continue Training</AlertDialogCancel>
<AlertDialogAction
variant="destructive"
onClick={() => void stopTrainingRun(false)}
onClick={() => {
setStopRequested(true);
setStopDialogOpen(false);
void stopTrainingRun(false).then((ok) => {
if (!ok) setStopRequested(false);
});
}}
>
Cancel Training
</AlertDialogAction>
<AlertDialogAction
onClick={() => void stopTrainingRun(true)}
onClick={() => {
setStopRequested(true);
setStopDialogOpen(false);
void stopTrainingRun(true).then((ok) => {
if (!ok) setStopRequested(false);
});
}}
>
Stop and Save
</AlertDialogAction>

View file

@ -1,6 +1,6 @@
import type { PipelineType } from "@huggingface/hub";
import { listModels } from "@huggingface/hub";
import { useCallback } from "react";
import { useCallback, useMemo } from "react";
import { useHfPaginatedSearch } from "./use-hf-paginated-search";
export interface HfModelResult {
@ -64,6 +64,53 @@ function mapModel(raw: unknown): HfModelResult | null {
};
}
/** Number of unsloth results to pull up-front before yielding general results. */
const UNSLOTH_PREFETCH = 20;
/**
* Creates a merged async generator that yields unsloth-owned models first,
* then general results (with deduplication).
*/
async function* mergedModelIterator(
query: string,
task?: PipelineType,
accessToken?: string,
): AsyncGenerator<unknown> {
const common = {
additionalFields: ["safetensors", "tags"] as ("safetensors" | "tags")[],
fetch: withPopularitySort,
...(accessToken ? { credentials: { accessToken } } : {}),
};
// Fire both iterators immediately (parallel network requests on first pull)
const unslothIter = listModels({
search: { query, owner: "unsloth", ...(task ? { task } : {}) },
...common,
});
const generalIter = listModels({
search: { query, ...(task ? { task } : {}) },
...common,
});
// Phase 1: pull & yield unsloth models first
const seen = new Set<string>();
let count = 0;
for await (const model of unslothIter) {
const m = model as { name?: string };
if (m.name) seen.add(m.name);
yield model;
count++;
if (count >= UNSLOTH_PREFETCH) break;
}
// Phase 2: yield general results, skipping already-seen unsloth models
for await (const model of generalIter) {
const m = model as { name?: string };
if (m.name && seen.has(m.name)) continue;
yield model;
}
}
export function useHfModelSearch(
query: string,
options?: { task?: PipelineType; accessToken?: string },
@ -71,18 +118,36 @@ export function useHfModelSearch(
const { task, accessToken } = options ?? {};
const createIter = useCallback(
() =>
listModels({
search: {
...(query.trim() ? { query } : { owner: "unsloth" }),
...(task ? { task } : {}),
},
additionalFields: ["safetensors", "tags"],
fetch: withPopularitySort,
...(accessToken ? { credentials: { accessToken } } : {}),
}) as AsyncGenerator<unknown>,
() => {
const trimmed = query.trim();
if (!trimmed) {
// No query → show default unsloth models
return listModels({
search: { owner: "unsloth", ...(task ? { task } : {}) },
additionalFields: ["safetensors", "tags"],
fetch: withPopularitySort,
...(accessToken ? { credentials: { accessToken } } : {}),
}) as AsyncGenerator<unknown>;
}
// Dual-query: unsloth first, then general
return mergedModelIterator(trimmed, task, accessToken) as AsyncGenerator<unknown>;
},
[query, task, accessToken],
);
return useHfPaginatedSearch(createIter, mapModel);
const search = useHfPaginatedSearch(createIter, mapModel);
// Secondary sort guarantee: unsloth models always float to the top
const results = useMemo(
() =>
[...search.results].sort((a, b) => {
const aFirst = a.id.startsWith("unsloth/") ? 0 : 1;
const bFirst = b.id.startsWith("unsloth/") ? 0 : 1;
return aFirst - bFirst;
}),
[search.results],
);
return { ...search, results };
}