From e132dbae765a78bea1ddd773413822cb034e1838 Mon Sep 17 00:00:00 2001
From: Jaret Burkett <jaretburkett@gmail.com>
Date: Thu, 15 Jan 2026 08:03:31 -0700
Subject: [PATCH] Add number of repeats for a dataset in the ui

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

diff --git a/ui/src/app/jobs/new/SimpleJob.tsx b/ui/src/app/jobs/new/SimpleJob.tsx
index 82d8969..5db650c 100644
--- a/ui/src/app/jobs/new/SimpleJob.tsx
+++ b/ui/src/app/jobs/new/SimpleJob.tsx
@@ -810,6 +810,14 @@ export default function SimpleJob({
                         onChange={value => setJobConfig(value, `config.process[0].datasets[${i}].network_weight`)}
                         placeholder="eg. 1.0"
                       />
+                      <NumberInput
+                        label="Num Repeats"
+                        value={dataset.num_repeats || 1}
+                        className="pt-2"
+                        onChange={value => setJobConfig(value, `config.process[0].datasets[${i}].num_repeats`)}
+                        placeholder="eg. 1"
+                        docKey={'dataset.num_repeats'}
+                      />
                     </div>
                     <div>
                       <TextInput
diff --git a/ui/src/app/jobs/new/jobConfig.ts b/ui/src/app/jobs/new/jobConfig.ts
index ca78393..66ff196 100644
--- a/ui/src/app/jobs/new/jobConfig.ts
+++ b/ui/src/app/jobs/new/jobConfig.ts
@@ -16,6 +16,7 @@ export const defaultDatasetConfig: DatasetConfig = {
   num_frames: 1,
   flip_x: false,
   flip_y: false,
+  num_repeats: 1,
 };
 
 export const defaultSliderConfig: SliderConfig = {
diff --git a/ui/src/docs.tsx b/ui/src/docs.tsx
index 82fc0d1..c82d800 100644
--- a/ui/src/docs.tsx
+++ b/ui/src/docs.tsx
@@ -306,6 +306,17 @@ const docs: { [key: string]: ConfigDoc } = {
       </>
     ),
   },
+  'dataset.num_repeats': {
+    title: 'Num Repeats',
+    description: (
+      <>
+        Number of Repeats will allow you to repeate the items in a dataset multiple times. This is useful when you are using multiple
+        datasets and want to balance the number of samples from each dataset. For instance, if you have a small dataset of 10 images 
+        and a large dataset of 100 images, you can set the small dataset to have 10 repeats to effectively make it 100 images, making
+        the two datasets occour equally during training.
+      </>
+    ),
+  },
 };
 
 export const getDoc = (key: string | null | undefined): ConfigDoc | null => {
diff --git a/ui/src/types.ts b/ui/src/types.ts
index be4af89..07f1e1d 100644
--- a/ui/src/types.ts
+++ b/ui/src/types.ts
@@ -103,6 +103,7 @@ export interface DatasetConfig {
   fps?: number;
   flip_x: boolean;
   flip_y: boolean;
+  num_repeats?: number;
   control_path_1?: string | null;
   control_path_2?: string | null;
   control_path_3?: string | null;