Adding exported model for chat

This commit is contained in:
Manan17 2026-02-24 01:17:09 +00:00
commit 4be677e45d
11 changed files with 169 additions and 27 deletions

View file

@ -2,6 +2,7 @@
"""
Export backend - handles model exporting in various formats
"""
import json
import logging
import os
from pathlib import Path
@ -200,6 +201,18 @@ class ExportBackend:
logger.error(traceback.format_exc())
return False, f"Failed to load checkpoint: {str(e)}"
def _write_export_metadata(self, save_directory: str):
"""Write export_metadata.json with base model info for Chat page discovery."""
try:
base_model = get_base_model_from_lora(self.current_checkpoint) if self.current_checkpoint else None
metadata = {"base_model": base_model}
metadata_path = os.path.join(save_directory, "export_metadata.json")
with open(metadata_path, "w") as f:
json.dump(metadata, f, indent=2)
logger.info(f"Wrote export metadata to {metadata_path}")
except Exception as e:
logger.warning(f"Could not write export metadata: {e}")
def export_merged_model(self,
save_directory: str,
format_type: str = "16-bit (FP16)",
@ -244,6 +257,9 @@ class ExportBackend:
self.current_tokenizer,
save_method=save_method
)
# Write export metadata so the Chat page can identify the base model
self._write_export_metadata(save_directory)
logger.info(f"Model saved successfully to {save_directory}")
# Push to hub if requested
@ -297,6 +313,9 @@ class ExportBackend:
self.current_model.save_pretrained(save_directory)
self.current_tokenizer.save_pretrained(save_directory)
# Write export metadata so the Chat page can identify the base model
self._write_export_metadata(save_directory)
logger.info(f"Model saved successfully to {save_directory}")
# Push to hub if requested

View file

@ -57,10 +57,12 @@ class ModelDetails(BaseModel):
class LoRAInfo(BaseModel):
"""LoRA adapter information"""
"""LoRA adapter or exported model information"""
display_name: str = Field(..., description="Display name for the LoRA")
adapter_path: str = Field(..., description="Path to the LoRA adapter")
adapter_path: str = Field(..., description="Path to the LoRA adapter or exported model")
base_model: Optional[str] = Field(None, description="Base model identifier")
source: Optional[str] = Field(None, description="'training' or 'exported'")
export_type: Optional[str] = Field(None, description="'lora' or 'merged' (for exports)")
class LoRAScanResponse(BaseModel):

View file

@ -18,6 +18,7 @@ from auth.authentication import get_current_subject
try:
from utils.models import (
scan_trained_loras,
scan_exported_models,
load_model_defaults,
get_base_model_from_lora,
is_vision_model,
@ -32,6 +33,7 @@ except ImportError:
sys.path.insert(0, str(parent_backend))
from utils.models import (
scan_trained_loras,
scan_exported_models,
load_model_defaults,
get_base_model_from_lora,
is_vision_model,
@ -289,35 +291,45 @@ async def get_model_config(
@router.get("/loras")
async def scan_loras(
outputs_dir: str = Query(default="./outputs", description="Directory to scan for LoRA adapters"),
exports_dir: str = Query(default="./exports", description="Directory to scan for exported models"),
current_subject: str = Depends(get_current_subject),
):
"""
Scan for trained LoRA adapters in the outputs directory.
This endpoint wraps the backend scan_trained_loras function.
Scan for trained LoRA adapters and exported models.
Returns both training outputs (from outputs_dir) and exported models
(from exports_dir) in a single list, distinguished by source field.
"""
try:
# Call backend scan function
trained_loras = scan_trained_loras(outputs_dir=outputs_dir)
# Convert to LoRAInfo objects
lora_list = []
# Scan training outputs
trained_loras = scan_trained_loras(outputs_dir=outputs_dir)
for display_name, adapter_path in trained_loras:
# Get base model if available
base_model = get_base_model_from_lora(adapter_path)
lora_info = LoRAInfo(
lora_list.append(LoRAInfo(
display_name=display_name,
adapter_path=adapter_path,
base_model=base_model
)
lora_list.append(lora_info)
base_model=base_model,
source="training",
))
# Scan exported models (merged, LoRA, base — skips GGUF)
exported = scan_exported_models(exports_dir=exports_dir)
for display_name, model_path, export_type, base_model in exported:
lora_list.append(LoRAInfo(
display_name=display_name,
adapter_path=model_path,
base_model=base_model,
source="exported",
export_type=export_type,
))
return LoRAScanResponse(
loras=lora_list,
outputs_dir=outputs_dir
)
except Exception as e:
logger.error(f"Error scanning LoRAs: {e}", exc_info=True)
raise HTTPException(

View file

@ -5,6 +5,7 @@ from .model_config import (
ModelConfig,
is_vision_model,
scan_trained_loras,
scan_exported_models,
load_model_defaults,
get_base_model_from_lora,
load_model_config,
@ -17,6 +18,7 @@ __all__ = [
'ModelConfig',
'is_vision_model',
'scan_trained_loras',
'scan_exported_models',
'load_model_defaults',
'get_base_model_from_lora',
'load_model_config',

View file

@ -465,6 +465,90 @@ def scan_trained_loras(outputs_dir: str = "./outputs") -> List[Tuple[str, str]]:
logger.error(f"Error scanning outputs folder: {e}")
return []
def scan_exported_models(exports_dir: str = "./exports") -> List[Tuple[str, str, str, Optional[str]]]:
"""
Scan exports folder for exported models (merged, LoRA, base).
Skips GGUF-only exports (not loadable by Unsloth inference backend).
The exports directory is two levels deep: {run}/{checkpoint}/
Returns:
List of tuples: [(display_name, model_path, export_type, base_model), ...]
export_type: "lora" | "merged"
"""
results = []
exports_path = Path(exports_dir)
if not exports_path.exists():
return results
try:
for run_dir in exports_path.iterdir():
if not run_dir.is_dir():
continue
for checkpoint_dir in run_dir.iterdir():
if not checkpoint_dir.is_dir():
continue
adapter_config = checkpoint_dir / "adapter_config.json"
config_file = checkpoint_dir / "config.json"
has_weights = (
any(checkpoint_dir.glob("*.safetensors"))
or any(checkpoint_dir.glob("*.bin"))
)
has_gguf = any(checkpoint_dir.glob("*.gguf"))
base_model = None
export_type = None
if adapter_config.exists():
export_type = "lora"
try:
cfg = json.loads(adapter_config.read_text())
base_model = cfg.get("base_model_name_or_path")
except Exception:
pass
elif config_file.exists() and has_weights:
export_type = "merged"
# Read base model from export_metadata.json (written at export time)
export_meta = checkpoint_dir / "export_metadata.json"
try:
if export_meta.exists():
meta = json.loads(export_meta.read_text())
base_model = meta.get("base_model")
except Exception:
pass
elif has_gguf:
# GGUF-only — not loadable by current inference backend
continue
else:
continue
# Fallback: read base model from the original training run's
# adapter_config.json in ./outputs/{run_name}/
if not base_model:
outputs_adapter_cfg = Path("./outputs") / run_dir.name / "adapter_config.json"
try:
if outputs_adapter_cfg.exists():
cfg = json.loads(outputs_adapter_cfg.read_text())
base_model = cfg.get("base_model_name_or_path")
except Exception:
pass
display_name = f"{run_dir.name} / {checkpoint_dir.name}"
model_path = str(checkpoint_dir)
results.append((display_name, model_path, export_type, base_model))
logger.debug(f"Found exported model: {display_name} ({export_type})")
results.sort(key=lambda x: Path(x[1]).stat().st_mtime, reverse=True)
logger.info(f"Found {len(results)} exported models in {exports_dir}")
return results
except Exception as e:
logger.error(f"Error scanning exports folder: {e}")
return []
def get_base_model_from_lora(lora_path: str) -> Optional[str]:
"""
Read the base model name from a LoRA adapter's config.

View file

@ -365,15 +365,26 @@ export function LoraModelPicker({
<div key={baseModel}>
{index > 0 ? <div className="my-1" /> : null}
<ListLabel>{baseModel}</ListLabel>
{adapters.map((adapter) => (
<ModelRow
key={adapter.id}
label={adapter.name}
meta="LoRA"
selected={value === adapter.id}
onClick={() => onSelect(adapter.id, { source: "lora", isLora: true })}
/>
))}
{adapters.map((adapter) => {
const isExported = adapter.source === "exported";
const isMerged = adapter.exportType === "merged";
const tag = isExported
? isMerged ? "Merged" : "LoRA"
: "LoRA";
const meta = isExported ? `${tag} · Exported` : tag;
return (
<ModelRow
key={adapter.id}
label={adapter.name}
meta={meta}
selected={value === adapter.id}
onClick={() => onSelect(adapter.id, {
source: isExported ? "exported" : "lora",
isLora: !isMerged,
})}
/>
);
})}
</div>
))
)}

View file

@ -10,10 +10,12 @@ export interface ModelOption {
export interface LoraModelOption extends ModelOption {
baseModel?: string;
updatedAt?: number;
source?: "training" | "exported";
exportType?: "lora" | "merged";
}
export interface ModelSelectorChangeMeta {
source: "hub" | "lora";
source: "hub" | "lora" | "exported";
isLora: boolean;
}

View file

@ -410,6 +410,8 @@ export function ChatPage(): ReactElement {
name: lora.name,
baseModel: lora.baseModel,
updatedAt: lora.updatedAt,
source: lora.source,
exportType: lora.exportType,
})),
[lorasFromStore],
);

View file

@ -68,6 +68,8 @@ function toLoraSummary(lora: {
display_name: string;
adapter_path: string;
base_model?: string | null;
source?: "training" | "exported" | null;
export_type?: "lora" | "merged" | null;
}): ChatLoraSummary {
const idTail = lora.adapter_path.split("/").filter(Boolean).at(-1) ?? "";
const updatedAt =
@ -78,6 +80,8 @@ function toLoraSummary(lora: {
name: stripTrailingEpoch(lora.display_name),
baseModel: lora.base_model || "Unknown base model",
updatedAt,
source: lora.source ?? undefined,
exportType: lora.export_type ?? undefined,
};
}

View file

@ -14,6 +14,8 @@ export interface BackendLoraInfo {
display_name: string;
adapter_path: string;
base_model?: string | null;
source?: "training" | "exported" | null;
export_type?: "lora" | "merged" | null;
}
export interface ListLorasResponse {

View file

@ -33,4 +33,6 @@ export interface ChatLoraSummary {
name: string;
baseModel: string;
updatedAt?: number;
source?: "training" | "exported";
exportType?: "lora" | "merged";
}