import { Order } from 'blockly/javascript';
/**
 * @category Datasets
 * @module Mnist
 */
/**
 * Loads the MNIST dataset with specified training and test set sizes.
 *
 * @param {Number} trainingSize - The size of training dataset
 * @param {Number} testSize - The size of test dataset
 * @returns {String} load_mnist_dataset(trainingSize, testSize)
 */
function loadMnist(block, generator) {
    const trainingSize = generator.valueToCode(block, 'TRAINING', Order.NONE) || '';
    const testSize = generator.valueToCode(block, 'TEST', Order.NONE) || '';
    const loadMnistDataset = generator.provideFunction_('load_mnist_dataset', `function ${generator.FUNCTION_NAME_PLACEHOLDER_}(training_set_size,test_set_size) {
      const {training, test} = mnist.set(training_set_size, test_set_size);
      
      const x_train = tf.tensor(training.map((item) => item.input))
      const y_train = tf.tensor(training.map((item) => item.output))
      const x_test = tf.tensor(test.map((item) => item.input))
      const y_test = tf.tensor(test.map((item) => item.output))

      return [x_train, y_train, x_test, y_test]
}`);
    return [`${loadMnistDataset}(${trainingSize},${testSize})`, Order.VOID];
}
/**
 * Loads a pre-trained MNIST model from a specified URL.
 *
 * @returns {String} await tf.loadLayersModel('https://raw.githubusercontent.com/Cheeetah97/mnist_pretrained_model/main/model.json')
 */
function loadMnistModel(block, generator) {
    return [
        `await tf.loadLayersModel('https://raw.githubusercontent.com/Cheeetah97/mnist_pretrained_model/main/model.json');`,
        Order.VOID,
    ];
}
export const datasetBlockGenerator = {
    load_mnist: loadMnist,
    load_model: loadMnistModel,
};
//# sourceMappingURL=generators.js.map