From 53f7f4f158d62844c7df940c0a44ceaa516e7909 Mon Sep 17 00:00:00 2001
From: imagineer99 <samleejackson0@gmail.com>
Date: Tue, 24 Feb 2026 08:16:45 +0000
Subject: [PATCH] feat: improve training config UX and remove unused logging
 options

---
 .../studio/sections/training-section.tsx      | 175 ++++++------------
 .../src/features/studio/tour/steps/save.tsx   |   4 +-
 .../src/features/training/lib/yaml-config.ts  |   7 -
 3 files changed, 59 insertions(+), 127 deletions(-)

diff --git a/studio/frontend/src/features/studio/sections/training-section.tsx b/studio/frontend/src/features/studio/sections/training-section.tsx
index dfcd72632e..3f342ad68c 100644
--- a/studio/frontend/src/features/studio/sections/training-section.tsx
+++ b/studio/frontend/src/features/studio/sections/training-section.tsx
@@ -2,13 +2,11 @@ import { SectionCard } from "@/components/section-card";
 import { Button } from "@/components/ui/button";
 import { ChartContainer } from "@/components/ui/chart";
 import type { ChartConfig } from "@/components/ui/chart";
-import { Checkbox } from "@/components/ui/checkbox";
 import {
-  Collapsible,
-  CollapsibleContent,
-  CollapsibleTrigger,
-} from "@/components/ui/collapsible";
-import { Input } from "@/components/ui/input";
+  Tooltip,
+  TooltipContent,
+  TooltipTrigger,
+} from "@/components/ui/tooltip";
 import {
   parseYamlConfig,
   serializeConfigToYaml,
@@ -17,14 +15,13 @@ import {
 } from "@/features/training";
 import {
   Archive04Icon,
-  ArrowDown01Icon,
   ChartAverageIcon,
   CleanIcon,
   CloudUploadIcon,
   Rocket01Icon,
 } from "@hugeicons/core-free-icons";
 import { HugeiconsIcon } from "@hugeicons/react";
-import { useRef, useState } from "react";
+import { useRef } from "react";
 import { toast } from "sonner";
 import { CartesianGrid, Line, LineChart, XAxis, YAxis } from "recharts";
 
@@ -44,7 +41,6 @@ const placeholderData = [
 export function TrainingSection() {
   const store = useTrainingConfigStore();
   const { isStarting, startError, startTrainingRun } = useTrainingActions();
-  const [logOpen, setLogOpen] = useState(false);
   const isIncompatible =
     !store.isVisionModel && store.isDatasetMultimodal === true;
   const fileInputRef = useRef<HTMLInputElement>(null);
@@ -79,7 +75,13 @@ export function TrainingSection() {
     const url = URL.createObjectURL(blob);
     const a = document.createElement("a");
     a.href = url;
-    a.download = "training-config.yaml";
+
+    const model = (store.selectedModel ?? "model").split("/").pop();
+    const method = store.trainingMethod ?? "qlora";
+    const dataset = (store.dataset ?? "dataset").split("/").pop();
+    const timestamp = new Date().toISOString().replace(/[:T]/g, "-").slice(0, 19);
+    a.download = `${model}_${method}_${dataset}_${timestamp}.yaml`;
+
     a.click();
     URL.revokeObjectURL(url);
   };
@@ -163,36 +165,52 @@ export function TrainingSection() {
         )}
 
         {/* Upload / Save / Reset */}
+        <p className="text-xs text-muted-foreground">Training Config</p>
         <div className="grid grid-cols-3 gap-2">
-          <Button
-            variant="outline"
-            size="sm"
-            className="cursor-pointer"
-            onClick={() => fileInputRef.current?.click()}
-          >
-            <HugeiconsIcon icon={CloudUploadIcon} className="size-3.5" />
-            Upload
-          </Button>
-          <Button
-            data-tour="studio-save"
-            variant="outline"
-            size="sm"
-            className="cursor-pointer"
-            onClick={handleSaveConfig}
-          >
-            <HugeiconsIcon icon={Archive04Icon} className="size-3.5" />
-            Save
-          </Button>
-          <Button
-            variant="outline"
-            size="sm"
-            className="cursor-pointer"
-            onClick={handleResetConfig}
-            disabled={!store.selectedModel}
-          >
-            <HugeiconsIcon icon={CleanIcon} className="size-3.5" />
-            Reset
-          </Button>
+          <Tooltip>
+            <TooltipTrigger asChild>
+              <Button
+                variant="outline"
+                size="sm"
+                className="cursor-pointer"
+                onClick={() => fileInputRef.current?.click()}
+              >
+                <HugeiconsIcon icon={CloudUploadIcon} className="size-3.5" />
+                Upload
+              </Button>
+            </TooltipTrigger>
+            <TooltipContent>Load a saved YAML config</TooltipContent>
+          </Tooltip>
+          <Tooltip>
+            <TooltipTrigger asChild>
+              <Button
+                data-tour="studio-save"
+                variant="outline"
+                size="sm"
+                className="cursor-pointer"
+                onClick={handleSaveConfig}
+              >
+                <HugeiconsIcon icon={Archive04Icon} className="size-3.5" />
+                Save
+              </Button>
+            </TooltipTrigger>
+            <TooltipContent>Download current config as YAML</TooltipContent>
+          </Tooltip>
+          <Tooltip>
+            <TooltipTrigger asChild>
+              <Button
+                variant="outline"
+                size="sm"
+                className="cursor-pointer"
+                onClick={handleResetConfig}
+                disabled={!store.selectedModel}
+              >
+                <HugeiconsIcon icon={CleanIcon} className="size-3.5" />
+                Reset
+              </Button>
+            </TooltipTrigger>
+            <TooltipContent>Reset to model defaults</TooltipContent>
+          </Tooltip>
         </div>
         <input
           ref={fileInputRef}
@@ -201,85 +219,6 @@ export function TrainingSection() {
           className="hidden"
           onChange={handleFileUpload}
         />
-
-        {/* Logging */}
-        <Collapsible open={logOpen} onOpenChange={setLogOpen}>
-          <CollapsibleTrigger className="flex w-full cursor-pointer items-center gap-1.5 text-xs text-muted-foreground">
-            <HugeiconsIcon
-              icon={ArrowDown01Icon}
-              className={`size-3.5 transition-transform ${logOpen ? "rotate-180" : ""}`}
-            />
-            Logging
-          </CollapsibleTrigger>
-          <CollapsibleContent className="mt-3 flex flex-col gap-3">
-            {/* W&B */}
-            <div className="flex items-center gap-2">
-              <Checkbox
-                id="wandb"
-                checked={store.enableWandb}
-                onCheckedChange={(v) => store.setEnableWandb(!!v)}
-              />
-              <label
-                htmlFor="wandb"
-                className="text-xs cursor-pointer text-muted-foreground"
-              >
-                Weights & Biases
-              </label>
-            </div>
-            {store.enableWandb && (
-              <div className="flex flex-col gap-2 pl-6">
-                <Input
-                  placeholder="W&B API Token"
-                  type="password"
-                  value={store.wandbToken}
-                  onChange={(e) => store.setWandbToken(e.target.value)}
-                />
-                <Input
-                  placeholder="Project name"
-                  value={store.wandbProject}
-                  onChange={(e) => store.setWandbProject(e.target.value)}
-                />
-              </div>
-            )}
-
-            {/* TensorBoard */}
-            <div className="flex items-center gap-2">
-              <Checkbox
-                id="tensorboard"
-                checked={store.enableTensorboard}
-                onCheckedChange={(v) => store.setEnableTensorboard(!!v)}
-              />
-              <label
-                htmlFor="tensorboard"
-                className="text-xs cursor-pointer text-muted-foreground"
-              >
-                TensorBoard
-              </label>
-            </div>
-            {store.enableTensorboard && (
-              <div className="flex flex-col gap-2 pl-6">
-                <Input
-                  placeholder="Log directory"
-                  value={store.tensorboardDir}
-                  onChange={(e) => store.setTensorboardDir(e.target.value)}
-                />
-                <div className="flex items-center justify-between">
-                  <span className="text-xs font-medium text-muted-foreground">
-                    Log frequency
-                  </span>
-                  <Input
-                    type="number"
-                    value={store.logFrequency}
-                    onChange={(e) =>
-                      store.setLogFrequency(Number(e.target.value))
-                    }
-                    className="w-24"
-                  />
-                </div>
-              </div>
-            )}
-          </CollapsibleContent>
-        </Collapsible>
         </div>
       </SectionCard>
     </div>
diff --git a/studio/frontend/src/features/studio/tour/steps/save.tsx b/studio/frontend/src/features/studio/tour/steps/save.tsx
index d137aecf41..d93548fc3d 100644
--- a/studio/frontend/src/features/studio/tour/steps/save.tsx
+++ b/studio/frontend/src/features/studio/tour/steps/save.tsx
@@ -6,8 +6,8 @@ export const studioSaveStep: TourStep = {
   title: "Save config",
   body: (
     <>
-      Save configs that worked. Re-running the same baseline makes it obvious
-      if a change helped (or if you just got lucky).
+      Save your training config as a YAML file. Re-running the same baseline
+      makes it obvious if a change helped (or if you just got lucky).
     </>
   ),
 };
diff --git a/studio/frontend/src/features/training/lib/yaml-config.ts b/studio/frontend/src/features/training/lib/yaml-config.ts
index 594122fa7d..3bc113f173 100644
--- a/studio/frontend/src/features/training/lib/yaml-config.ts
+++ b/studio/frontend/src/features/training/lib/yaml-config.ts
@@ -75,13 +75,6 @@ export function serializeConfigToYaml(
       lr_scheduler_type: state.lrSchedulerType,
     },
     lora,
-    logging: {
-      enable_wandb: state.enableWandb,
-      wandb_project: state.wandbProject,
-      enable_tensorboard: state.enableTensorboard,
-      tensorboard_dir: state.tensorboardDir,
-      log_frequency: state.logFrequency,
-    },
   };
 
   return yaml.dump(config, { lineWidth: -1, noRefs: true });