feat: fixed smart trainer

Author: xEcho1337

brain4j-core/src/main/java/org/brain4j/core/training/advanced/SmartTrainer.java

@@ -33,58 +33,63 @@ public void addListener(TrainListener listener) {
     public void abort() {
         this.running = false;
     }
-
-    public void start(Model model, ListDataSource dataSource, double lossThreshold) {
-        this.start = System.nanoTime();
-        this.running = true;
-        this.epoches = 0;
-
-        this.listeners.forEach(listener -> listener.register(this, model));
-
-        while (running && evaluation.loss() > lossThreshold) {
-            iterate(model, dataSource, Integer.MAX_VALUE);
-        }
-
-        this.running = false;
-        this.end = System.nanoTime();
-    }
-
+    
     public void step(Model model, ListDataSource dataSource, int totalEpoches) {
         long start = System.nanoTime();
         this.listeners.forEach(listener -> listener.onEpochStarted(epoches, totalEpoches, start));
-
+        
         model.fit(dataSource);
-
+        
         long took = System.nanoTime() - start;
         this.listeners.forEach(listener -> listener.onEpochCompleted(epoches, totalEpoches, took));
     }
-
-    public void startFor(Model model, ListDataSource dataSource, int epochesAmount) {
+    
+    private void setupPreTraining(Model model, ListDataSource trainSource, ListDataSource evalSource) {
+        if (trainSource == null || evalSource == null) {
+            throw new IllegalArgumentException("Training source and evaluation source are required and cannot be null!");
+        }
+        
         this.start = System.nanoTime();
         this.running = true;
         this.epoches = 0;
-
+        
+        this.evaluation = model.evaluate(evalSource);
         this.listeners.forEach(listener -> listener.register(this, model));
+    }
+    
+    public void start(Model model, ListDataSource trainSource, ListDataSource evalSource, double lossThreshold) {
+        setupPreTraining(model, trainSource, evalSource);
+        
+        while (running && evaluation.loss() > lossThreshold) {
+            iterate(model, trainSource, evalSource, Integer.MAX_VALUE);
+        }
 
+        this.running = false;
+        this.end = System.nanoTime();
+    }
+    
+    public void startFor(Model model, ListDataSource trainSource, ListDataSource evalSource, int epochesAmount) {
+        setupPreTraining(model, trainSource, evalSource);
+        
         for (int i = 0; i < epochesAmount && running; i++) {
-            iterate(model, dataSource, epochesAmount);
+            iterate(model, trainSource, evalSource, epochesAmount);
         }
 
         this.running = false;
         this.end = System.nanoTime();
     }
 
-    private void iterate(Model model, ListDataSource dataSource, int totalEpoches) {
-        step(model, dataSource, totalEpoches);
+    private void iterate(Model model, ListDataSource trainSource, ListDataSource evalSource, int totalEpoches) {
+        step(model, trainSource, totalEpoches);
 
         this.epoches++;
 
         if (epoches % evaluateEvery == 0) {
             long start = System.nanoTime();
-            EvaluationResult result = model.evaluate(dataSource);
+            this.evaluation = model.evaluate(evalSource);
             long took = System.nanoTime() - start;
 
-            this.listeners.forEach(listener -> listener.onEvaluated(dataSource, result, epoches, took));
+            this.listeners.forEach(listener -> listener.onEvaluated(evalSource, evaluation, epoches, took));
         }
     }
 

brain4j-core/src/main/java/org/brain4j/core/training/advanced/TrainListener.java

@@ -27,10 +27,29 @@ public void onEpochCompleted(int epoch, int totalEpoches, long took) {
      */
     public void onEpochStarted(int epoch, int totalEpoches, long start) {
     }
-
+    
+    /**
+     * Called after an evaluation has completed.
+     * @param dataSource the data source where the evaluation has been done
+     * @param evaluation the evaluation result
+     * @param epoch the current epoch
+     * @param took how many nanoseconds it took to evaluate
+     */
     public void onEvaluated(ListDataSource dataSource, EvaluationResult evaluation, int epoch, long took) {
     }
-
+    
+    /**
+     * Called when the loss gets increased
+     * @param loss the current loss
+     * @param previousLoss the previous loss
+     */
     public void onLossIncreased(double loss, double previousLoss) {
     }
+    
+    /**
+     * Aborts the training session, delegates to {@link SmartTrainer#abort()}
+     */
+    public void abort() {
+        trainer.abort();
+    }
 }