Upload app.py

app.py

@@ -114,12 +114,15 @@ class SKAModel(nn.Module):
             self.output_history = []
 
 
-def get_mnist_per_class(samples_per_class):
+def get_mnist_per_class(samples_per_class, data_seed=0):
     """Select N samples per class from MNIST, return dict of {digit: images}."""
     targets = mnist_dataset.targets.numpy()
+    rng = np.random.RandomState(data_seed)
     digit_images = {}
     for digit in range(10):
-        indices = np.where(targets == digit)[0][:samples_per_class]
+        all_indices = np.where(targets == digit)[0]
+        rng.shuffle(all_indices)
+        indices = all_indices[:samples_per_class]
         images_list = []
         for idx in indices:
             img, _ = mnist_dataset[idx]
@@ -128,7 +131,7 @@ def get_mnist_per_class(samples_per_class):
     return digit_images
 
 
-def run_ska_per_class(neurons_str, K, tau, samples_per_class):
+def run_ska_per_class(neurons_str, K, tau, samples_per_class, data_seed):
     # Parse layer sizes
     try:
         layer_sizes = [int(x.strip()) for x in neurons_str.split(",")]
@@ -137,10 +140,11 @@ def run_ska_per_class(neurons_str, K, tau, samples_per_class):
 
     K = int(K)
     samples_per_class = int(samples_per_class)
+    data_seed = int(data_seed)
     learning_rate = tau / K
 
     # Get data per class
-    digit_images = get_mnist_per_class(samples_per_class)
+    digit_images = get_mnist_per_class(samples_per_class, data_seed)
 
     # Run SKA separately for each digit
     all_entropy_histories = {}
@@ -197,6 +201,7 @@ with gr.Blocks(title="SKA Per-Class Entropy Explorer") as demo:
             k_slider = gr.Slider(1, 200, value=50, step=1, label="K (forward steps)")
             tau_slider = gr.Slider(0.25, 0.75, value=0.5, step=0.01, label="Learning budget τ (τ = η.K)")
             samples_slider = gr.Slider(1, 100, value=100, step=1, label="Samples per class")
+            seed_slider = gr.Slider(0, 99, value=0, step=1, label="Data seed (shuffle samples)")
             run_btn = gr.Button("Run SKA Per-Class", variant="primary")
 
             gr.Markdown("---")
@@ -218,7 +223,7 @@ with gr.Blocks(title="SKA Per-Class Entropy Explorer") as demo:
 
     run_btn.click(
         fn=run_ska_per_class,
-        inputs=[neurons_input, k_slider, tau_slider, samples_slider],
+        inputs=[neurons_input, k_slider, tau_slider, samples_slider, seed_slider],
         outputs=[plot_entropy],
     )