From 0f2239ca2359ec66cb5f2a894411f0fcb4dfcfd7 Mon Sep 17 00:00:00 2001
From: Jaret Burkett <jaretburkett@gmail.com>
Date: Sun, 31 Aug 2025 16:58:27 -0600
Subject: [PATCH] Add force sample toggle to the ui

---
 ui/src/app/jobs/new/SimpleJob.tsx | 31 +++++++++++++++++++++++++++++--
 ui/src/app/jobs/new/jobConfig.ts  |  1 +
 ui/src/docs.tsx                   | 11 +++++++++++
 ui/src/types.ts                   |  1 +
 4 files changed, 42 insertions(+), 2 deletions(-)

diff --git a/ui/src/app/jobs/new/SimpleJob.tsx b/ui/src/app/jobs/new/SimpleJob.tsx
index 2d9d9268..fb42599a 100644
--- a/ui/src/app/jobs/new/SimpleJob.tsx
+++ b/ui/src/app/jobs/new/SimpleJob.tsx
@@ -809,7 +809,28 @@ export default function SimpleJob({
                       label="Skip First Sample"
                       className="pt-4"
                       checked={jobConfig.config.process[0].train.skip_first_sample || false}
-                      onChange={value => setJobConfig(value, 'config.process[0].train.skip_first_sample')}
+                      onChange={value => {
+                        setJobConfig(value, 'config.process[0].train.skip_first_sample');
+                        // cannot do both, so disable the other
+                        if (value){
+                          setJobConfig(false, 'config.process[0].train.force_first_sample');
+                        }
+                      }}
+                    />
+                  </div>
+                  <div>
+                    <Checkbox
+                      label="Force First Sample"
+                      className="pt-1"
+                      checked={jobConfig.config.process[0].train.force_first_sample || false}
+                      docKey={'train.force_first_sample'}
+                      onChange={value => {
+                        setJobConfig(value, 'config.process[0].train.force_first_sample');
+                        // cannot do both, so disable the other
+                        if (value){
+                          setJobConfig(false, 'config.process[0].train.skip_first_sample');
+                        }
+                      }}
                     />
                   </div>
                   <div>
@@ -817,7 +838,13 @@ export default function SimpleJob({
                       label="Disable Sampling"
                       className="pt-1"
                       checked={jobConfig.config.process[0].train.disable_sampling || false}
-                      onChange={value => setJobConfig(value, 'config.process[0].train.disable_sampling')}
+                      onChange={value => {
+                        setJobConfig(value, 'config.process[0].train.disable_sampling');
+                        // cannot do both, so disable the other
+                        if (value){
+                          setJobConfig(false, 'config.process[0].train.force_first_sample');
+                        }
+                      }}
                     />
                   </div>
                 </FormGroup>
diff --git a/ui/src/app/jobs/new/jobConfig.ts b/ui/src/app/jobs/new/jobConfig.ts
index df257bb9..40dad0be 100644
--- a/ui/src/app/jobs/new/jobConfig.ts
+++ b/ui/src/app/jobs/new/jobConfig.ts
@@ -75,6 +75,7 @@ export const defaultJobConfig: JobConfig = {
             ema_decay: 0.99,
           },
           skip_first_sample: false,
+          force_first_sample: false,
           disable_sampling: false,
           dtype: 'bf16',
           diff_output_preservation: false,
diff --git a/ui/src/docs.tsx b/ui/src/docs.tsx
index 6b414488..cadf0e01 100644
--- a/ui/src/docs.tsx
+++ b/ui/src/docs.tsx
@@ -155,6 +155,17 @@ const docs: { [key: string]: ConfigDoc } = {
       </>
     ),
   },
+  'train.force_first_sample': {
+    title: 'Force First Sample',
+    description: (
+      <>
+        This option will force the trainer to generate samples when it starts. The trainer will normally only generate a first sample
+        when nothing has been trained yet, but will not do a first sample when resuming from an existing checkpoint. This option
+        forces a first sample every time the trainer is started. This can be useful if you have changed sample prompts and want to see
+        the new prompts right away.
+      </>
+    ),
+  },
 };
 
 export const getDoc = (key: string | null | undefined): ConfigDoc | null => {
diff --git a/ui/src/types.ts b/ui/src/types.ts
index 5034d4f9..eae7e17f 100644
--- a/ui/src/types.ts
+++ b/ui/src/types.ts
@@ -117,6 +117,7 @@ export interface TrainConfig {
     weight_decay: number;
   };
   skip_first_sample: boolean;
+  force_first_sample: boolean;
   disable_sampling: boolean;
   diff_output_preservation: boolean;
   diff_output_preservation_multiplier: number;