Merge pull request #241 from unslothai/feature/adding-exported-models-for-chat

Adding exported model for chat

studio/backend/core/export/export.py

@@ -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

studio/backend/core/inference/inference.py

@@ -6,8 +6,10 @@ from unsloth.chat_templates import get_chat_template
 from transformers import TextStreamer
 from peft import PeftModel, PeftModelForCausalLM
 
+import json
 import sys
 import torch
+from pathlib import Path
 from typing import Optional, Union, Generator, Tuple
 from utils.models import ModelConfig, get_base_model_from_lora
 from utils.paths import is_model_cached
@@ -112,7 +114,18 @@ class InferenceBackend:
                 # In that case, load the real processor from the base model.
                 from transformers import ProcessorMixin
                 if not (isinstance(processor, ProcessorMixin) or hasattr(processor, "image_processor")):
+                    # For LoRA adapters, use the base model. For local merged exports,
+                    # read export_metadata.json to find the original base model.
                     processor_source = config.base_model if config.is_lora else config.identifier
+                    if not config.is_lora and config.is_local:
+                        _meta_path = Path(config.path) / "export_metadata.json"
+                        try:
+                            if _meta_path.exists():
+                                _meta = json.loads(_meta_path.read_text())
+                                if _meta.get("base_model"):
+                                    processor_source = _meta["base_model"]
+                        except Exception:
+                            pass
                     logger.warning(
                         f"FastVisionModel returned {type(processor).__name__} (no image_processor) "
                         f"for '{model_name}' — loading proper processor from '{processor_source}'"

studio/backend/models/models.py

@@ -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):

studio/backend/routes/models.py

@@ -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(

studio/backend/utils/models/__init__.py

@@ -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',

studio/backend/utils/models/model_config.py

@@ -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.

studio/frontend/src/components/assistant-ui/model-selector.tsx

@@ -62,7 +62,7 @@ function ModelSelectorTrigger({
         className={cn(
           "flex items-center gap-2 transition-colors",
           variant === "outline" &&
-            "rounded-full border border-border/60 hover:bg-accent",
+          "rounded-full border border-border/60 hover:bg-accent",
           variant === "ghost" && "rounded-md hover:bg-accent",
           variant === "muted" && "rounded-md bg-muted hover:bg-muted/80",
           size === "sm" && "h-8 px-3 text-xs",
@@ -183,9 +183,20 @@ export function ModelSelector({
       all.set(model.id, model);
     }
     for (const lora of loraModels) {
+      // Strip "/ suffix" from display name (e.g. "foo_123/foo" → "foo_123")
+      const displayName = lora.name.includes("/")
+        ? lora.name.split("/")[0].trim()
+        : lora.name;
+      // Show type tag instead of base model name
+      const isExported = lora.source === "exported";
+      const isMerged = lora.exportType === "merged";
+      const tag = isExported
+        ? isMerged ? "Merged · Exported" : "LoRA"
+        : "LoRA";
       all.set(lora.id, {
         ...lora,
-        description: lora.baseModel || lora.description,
+        name: displayName,
+        description: tag,
       });
     }
     return all;

studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx

@@ -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>
             ))
           )}

studio/frontend/src/components/assistant-ui/model-selector/types.ts

@@ -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;
 }
 

studio/frontend/src/features/chat/chat-page.tsx

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

studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts

@@ -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,
   };
 }
 

studio/frontend/src/features/chat/types/api.ts

@@ -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 {

studio/frontend/src/features/chat/types/runtime.ts

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