Adding exported model for chat
This commit is contained in:
parent
77b0978d5f
commit
4be677e45d
11 changed files with 169 additions and 27 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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',
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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>
|
||||
))
|
||||
)}
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -410,6 +410,8 @@ export function ChatPage(): ReactElement {
|
|||
name: lora.name,
|
||||
baseModel: lora.baseModel,
|
||||
updatedAt: lora.updatedAt,
|
||||
source: lora.source,
|
||||
exportType: lora.exportType,
|
||||
})),
|
||||
[lorasFromStore],
|
||||
);
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
};
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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 {
|
||||
|
|
|
|||
|
|
@ -33,4 +33,6 @@ export interface ChatLoraSummary {
|
|||
name: string;
|
||||
baseModel: string;
|
||||
updatedAt?: number;
|
||||
source?: "training" | "exported";
|
||||
exportType?: "lora" | "merged";
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue