// A class organizes a research experiment around the native checked fit API.
// The class instance is local; only supported model/data buffers reach workers.
import aner.dataset
import aner.nn
import aner.ml
import aner.parallel

class ResearchClassifier {
    let model: ClassifierPipeline
    let config: TrainingConfig

    fn train(training: Dataset, execution: ParallelContext) {
        return parallel.fit(model: model, data: training,
            config: config, execution: execution)
    }
}

let parts = dataset.iris().partition(validation: 0.2, test: 0.2, seed: 42)
let experiment = ResearchClassifier(
    model: ml.classifier(nn.mlp(4, List(8), 3, 42),
        standardize: true, loss: "cross_entropy", optimizer: "adam"),
    config: ml.training(epochs: 40, batch: 32, rate: 0.01, seed: 42)
        .validation(parts.validation()))
let fitted = experiment.train(parts.train(),
    parallel.device(device: "cpu", dtype: "Float64"))
print(fitted.evaluate(parts.test()).summary())