Spaces:

elvis-hf
/

cnn_visualizer

Running

cnn_visualizer / frontends /react /src /NetworkVisualizer.tsx

move training loss curve under Train tab

76f51ae 3 months ago

15.5 kB

import { useEffect, useState, useRef } from "react";
import * as tf from "@tensorflow/tfjs";
import Plot from "react-plotly.js";

import { MnistData } from "./mnist.js";
import { Cnn, train } from "./train.ts";
import type { TrainController, RunInfo, OptimizerParams } from "./train.ts";
import { Button, InputField, Tabs, Dropdown, Card } from "@elvis/ui";
import InfoViewer from "./InfoViewer.tsx";

const DEFAULT_ARCHITECTURE = `[conv2d filters=8 kernel=11 
stride=1 padding=1 activation=relu]
[maxpool size=2 stride=2]
[flatten]
[dense units=10 activation=softmax]`;

const isFirefox = navigator.userAgent.toLowerCase().includes("firefox");
await tf.setBackend(isFirefox ? "cpu" : "webgl");
await tf.ready();

export default function NetworkVisualizer() {
  const [dataset, setDataset] = useState<MnistData | null>(null);

  useEffect(() => {
    const loadData = async () => {
      const data = new MnistData();
      await data.load();
      setDataset(data);
      console.log("dataset loaded");
    }

    loadData();
  }, []);

  // architecture states
  const [architecture, setArchitecture] = useState(DEFAULT_ARCHITECTURE);
  const [optimizerType, setOptimizerType] = useState<string>('adam');
  const [optimizerParams, setOptimizerParams] = useState<OptimizerParams>({
    learningRate: '0.001',
    beta1: '0.9',
    beta2: '0.999',
    epsilon: '1e-8',
    batchSize: '32',
    epochs: '5',
  });
  const modelRef = useRef<Cnn | null>(null);
  const optimizerRef = useRef<tf.Optimizer | null>(null);

  function handleArchitectureChange(newArchitecture: string) {
    if (isTraining) {
      alert('Cannot change architecture while training is in progress.');
    } else {
      setArchitecture(newArchitecture);
    }
  }

  function handleOptimizerChange(newOptimizerType: string, newOptimizerParams: OptimizerParams) {
    if (isTraining) {
      alert('Cannot change optimizer settings while training is in progress.');
    } else {
      setOptimizerType(newOptimizerType);
      setOptimizerParams(newOptimizerParams);
    }
  }

  function handleSampleIndexChange() {
    trainController.current.sampleIndex += 1;
    refreshSampleInfo();
  }

 
  function resetModel() {
    if (!dataset) return;

    if (modelRef.current) {
      modelRef.current.dispose();
    }

    const cnn = new Cnn(architecture, dataset.numInputChannels);
    modelRef.current = cnn;
  }

  function resetOptimizer() {
    if (optimizerType === 'adam') {
      const learningRate = parseFloat(optimizerParams.learningRate);
      const beta1 = parseFloat(optimizerParams.beta1 || "0.9");
      const beta2 = parseFloat(optimizerParams.beta2 || "0.999");
      const epsilon = parseFloat(optimizerParams.epsilon || "1e-8");

      if (Number.isNaN(learningRate) || learningRate <= 0) {
        alert('Invalid learning rate for Adam optimizer.');
        return;
      }
      if (Number.isNaN(beta1) || beta1 < 0) {
        alert('Invalid beta1 for Adam optimizer.');
        return;
      }
      if (Number.isNaN(beta2) || beta2 < 0) {
        alert('Invalid beta2 for Adam optimizer.');
        return;
      }
      if (Number.isNaN(epsilon) || epsilon <= 0) {
        alert('Invalid epsilon for Adam optimizer.');
        return;
      }

      const opt = tf.train.adam(learningRate, beta1, beta2, epsilon);
      if (optimizerRef.current) {
        optimizerRef.current.dispose();
      }
      optimizerRef.current = opt;
    } else if (optimizerType === 'sgd') {
      const learningRate = parseFloat(optimizerParams.learningRate);
      if (Number.isNaN(learningRate) || learningRate <= 0) {
        alert('Invalid learning rate for SGD optimizer.');
        return;
      }

      const opt = tf.train.sgd(learningRate);
      if (optimizerRef.current) {
        optimizerRef.current.dispose();
      }
      optimizerRef.current = opt;
    } else {
      alert(`Unsupported optimizer type: ${optimizerType}`);
    }
  }

  // Reset the model only when its architecture or input shape changes.
  useEffect(() => {
    resetModel();
  }, [architecture, dataset]);

  // Optimizer settings can change without discarding the trained model.
  useEffect(() => {
    resetOptimizer();
  }, [optimizerType, optimizerParams]);

  // training states
  const [isTraining, setIsTraining] = useState<boolean>(false);
  const isTrainingRef = useRef<boolean>(false);
  const lossesRef = useRef<Array<number>>([]);
  const trainController = useRef<TrainController>({
    isPaused: false,
    stopRequested: false,
    sampleIndex: 0,
  });
  const infoRef = useRef<Array<RunInfo>>([]);

  function refreshSampleInfo() {
    if (!dataset || !modelRef.current) return;

    let nextInfo: RunInfo[] = [];
    const cleanupToken = tf.tidy(() => {
      const sample = dataset.getTestSample(trainController.current.sampleIndex);
      const { output, info } = modelRef.current!.forwardWithInfo(
        sample.xs.reshape([
          1,
          dataset.imageSize,
          dataset.imageSize,
          dataset.numInputChannels,
        ]) as tf.Tensor4D,
      );
      const probs = tf.softmax(output);
      info.push({
        type: 'output',
        output: probs.dataSync(),
        shape: probs.shape,
      });
      nextInfo = info;
      return tf.scalar(0);
    });
    cleanupToken.dispose();
    infoRef.current = nextInfo;
    updateTick();
  }

  useEffect(() => {
    refreshSampleInfo();
  }, [architecture, dataset]);

  // render timing
  const [, setTick] = useState<number>(0);

  function updateTick() {
    setTick((tick) => tick + 1);
  }

  async function startTraining() {
    if (!modelRef.current || !dataset || !optimizerRef.current || isTrainingRef.current) {
      return;
    }

    isTrainingRef.current = true;
    setIsTraining(true);
    trainController.current.isPaused = false;
    trainController.current.stopRequested = false;

    const batchSize = parseFloat(optimizerParams.batchSize);
    if (Number.isNaN(batchSize) || batchSize <= 0) {
      alert('Invalid batch size.');
      isTrainingRef.current = false;
      setIsTraining(false);
      return;
    }
    const epochs = parseFloat(optimizerParams.epochs);
    if (Number.isNaN(epochs) || epochs <= 0) {
      alert('Invalid number of epochs.');
      isTrainingRef.current = false;
      setIsTraining(false);
      return;
    }

    let lastTickUpdate = 0;
    if (!modelRef.current) return;
    if (!optimizerRef.current) return;

    try {
      await train(
        dataset,
        modelRef.current,
        optimizerRef.current,
        batchSize,
        epochs,
        trainController.current,
        (_epoch, _batch, loss, info) => {
          // lossesRef.current.push({ epoch, batch, loss });
          lossesRef.current.push(loss);
          console.log(loss);

          infoRef.current = info;

          // update tick every 50ms
          const now = performance.now();
          if (now - lastTickUpdate > 50) {
            lastTickUpdate = now;
            updateTick();
          }
        },
      );
    } finally {
      const wasStopped = trainController.current.stopRequested;
      isTrainingRef.current = false;
      setIsTraining(false);
      trainController.current.isPaused = false;
      trainController.current.stopRequested = false;
      if (!wasStopped) {
        alert('Training finished.');
      }
    }
  }

  function handleTrain() {
    if (isTrainingRef.current) {
      trainController.current.isPaused = false;
      return;
    }
    console.log('Starting training...');
    startTraining();
  }

  function handlePauseTraining() {
    console.log('Pausing training...');
    trainController.current.isPaused = true;
  }

  function handleStopTraining() {
    console.log('Stopping training...');
    trainController.current.stopRequested = true;
    trainController.current.isPaused = false;
  }

  async function waitUntilNotTraining() {
    return new Promise<void>((resolve) => {
      function check() {
        if (!isTrainingRef.current) {
          resolve();
        } else {
          requestAnimationFrame(check);
        }
      }
      check();
    });
  }

  async function handleResetTraining() {
    console.log('Resetting training...');
    handleStopTraining();

    await waitUntilNotTraining();
    console.log('Training stopped. Resetting model.');
    lossesRef.current = [];
    infoRef.current = [];
    resetModel();
    resetOptimizer();
    refreshSampleInfo();
  }

  return (
    <div className="flex flex-col min-h-0 h-full gap-3">
      <p className="text-sm text-gray-700">
        Configure and train a convolutional neural network (CNN) to classify handwritten digits from the MNIST dataset.
      </p>
      <div className="grid grid-cols-[2fr_1fr] min-h-0 flex-1 gap-12">
        <TrainingViewer 
          infoRef={infoRef}
          handleSampleIndexChange={handleSampleIndexChange}
        />
        <Sidebar 
          architecture={architecture}
          onArchitectureChange={handleArchitectureChange}
          optimizerType={optimizerType}
          optimizerParams={optimizerParams}
          onOptimizerChange={handleOptimizerChange}
          isTraining={isTraining}
          lossesRef={lossesRef}
          onTrain={handleTrain}
          onPauseTraining={handlePauseTraining}
          onResetTraining={handleResetTraining}
        />
      </div>
    </div>
  );
}

interface TrainingViewerProps {
  infoRef: React.RefObject<Array<RunInfo>>;
  handleSampleIndexChange: () => void;
}
function TrainingViewer({
  infoRef,
  handleSampleIndexChange,
}: TrainingViewerProps) {
  return (
    <div className="flex flex-col h-full overflow-auto gap-4 w-full min-h-0">
      <InfoViewer info={infoRef.current} onSampleIndexChange={handleSampleIndexChange} />
    </div>
  )
}

interface SidebarProps {
  architecture: string;
  onArchitectureChange: (newArchitecture: string) => void;

  optimizerType: string;
  optimizerParams: OptimizerParams;
  onOptimizerChange: (newOptimizerType: string, newOptimizerParams: OptimizerParams) => void;
  isTraining: boolean;
  lossesRef: React.RefObject<Array<number>>;

  onTrain: () => void;
  onPauseTraining: () => void;
  onResetTraining: () => void;
}

type TrainingInputRowProps = {
  label: string;
  value: string | undefined;
  onChange: (value: string) => void;
};

function TrainingInputRow({ label, value, onChange }: TrainingInputRowProps) {
  return (
    <div className="grid grid-cols-[7rem_minmax(0,1fr)] items-center gap-2">
      <span className="text-sm text-right text-gray-700">{label}</span>
      <div className="[&_label]:hidden [&_input]:py-1">
        <InputField label="" value={value ?? ""} onChange={onChange} />
      </div>
    </div>
  );
}

function Sidebar({
  architecture,
  onArchitectureChange,
  optimizerType,
  optimizerParams,
  onOptimizerChange,
  isTraining,
  lossesRef,
  onTrain,
  onPauseTraining,
  onResetTraining,
}: SidebarProps) {
  const tabs = ["Architecture", "Train"];
  const [activeTab, setActiveTab] = useState<string>(tabs[0]);

  const [architectureDraft, setArchitectureDraft] = useState<string>(architecture);

  return (
    <Card className="flex flex-col h-full p-4 gap-2 overflow-auto">
      <Tabs tabs={tabs} activeTab={activeTab} onChange={setActiveTab} />
      <div className="flex flex-col p-4 gap-4 h-full overflow-auto">
        { isFirefox && (
          <p className="text-red-500">
            Warning: This demo may be quite slow on Firefox.
          </p>
        )}
        { activeTab === "Architecture" && (
          <>
            <InputField 
              label="Architecture" 
              value={architectureDraft}
              onChange={setArchitectureDraft}
              rows={15}
            />
            <Button
              label="Apply architecture"
              onClick={() => onArchitectureChange(architectureDraft)}
            />
          </>
        )}
        { activeTab === "Train" && (
          <div className="flex flex-col gap-2">
            <p className="text-sm">Training {isTraining ? "in progress" : "not in progress"}</p>
            <Plot
              data={[
                {
                  x: lossesRef.current.map((_, i) => i),
                  y: lossesRef.current,
                  mode: 'lines',
                  type: 'scatter',
                },
              ]}
              layout={{
                xaxis: { title: { text: 'Training steps' } },
                yaxis: { title: { text: 'Train loss' } },
                margin: { t: 15, r: 15, b: 40, l: 45 },
              }}
              className="w-full h-[220px]"
              config={{ responsive: true }}
            />
            <div className="grid grid-cols-[7rem_minmax(0,1fr)] items-center gap-2">
              <span className="text-sm text-right text-gray-700">Optimizer</span>
              <div className="[&_label]:hidden [&_select]:py-1">
                <Dropdown
                  label=""
                  options={["sgd", "adam"]}
                  activeOption={optimizerType}
                  onChange={(newOptimizerType) => onOptimizerChange(newOptimizerType, optimizerParams)}
                />
              </div>
            </div>
            { optimizerType === 'sgd' && (
              <TrainingInputRow
                label="Learning Rate"
                value={optimizerParams.learningRate}
                onChange={(newLearningRate) => onOptimizerChange(optimizerType, {...optimizerParams, learningRate: newLearningRate})}
              />
            )}
            { optimizerType === 'adam' && (
              <>
                <TrainingInputRow
                  label="Learning Rate"
                  value={optimizerParams.learningRate}
                  onChange={(newLearningRate) => onOptimizerChange(optimizerType, {...optimizerParams, learningRate: newLearningRate})}
                />
                <TrainingInputRow
                  label="Beta 1"
                  value={optimizerParams.beta1}
                  onChange={(newBeta1) => onOptimizerChange(optimizerType, {...optimizerParams, beta1: newBeta1})}
                />
                <TrainingInputRow
                  label="Beta 2"
                  value={optimizerParams.beta2}
                  onChange={(newBeta2) => onOptimizerChange(optimizerType, {...optimizerParams, beta2: newBeta2})}
                />
                <TrainingInputRow
                  label="Epsilon"
                  value={optimizerParams.epsilon}
                  onChange={(newEpsilon) => onOptimizerChange(optimizerType, {...optimizerParams, epsilon: newEpsilon})}
                />
              </>
            )}
            <TrainingInputRow
              label="Batch Size"
              value={optimizerParams.batchSize}
              onChange={(newBatchSize) => onOptimizerChange(optimizerType, {...optimizerParams, batchSize: newBatchSize})}
            />
            <TrainingInputRow
              label="Epochs"
              value={optimizerParams.epochs}
              onChange={(newEpochs) => onOptimizerChange(optimizerType, {...optimizerParams, epochs: newEpochs})}
            />
            <div className="grid grid-cols-3 gap-2 mt-2 [&_button]:px-2 [&_button]:py-1 [&_button]:text-xs">
              <Button label="Train" onClick={onTrain} />
              <Button label="Pause" onClick={onPauseTraining} />
              <Button label="Reset" onClick={onResetTraining} />
            </div>
          </div>
        )}
      </div>
    </Card>
  )
}

	import { useEffect, useState, useRef } from "react";
	import * as tf from "@tensorflow/tfjs";
	import Plot from "react-plotly.js";

	import { MnistData } from "./mnist.js";
	import { Cnn, train } from "./train.ts";
	import type { TrainController, RunInfo, OptimizerParams } from "./train.ts";
	import { Button, InputField, Tabs, Dropdown, Card } from "@elvis/ui";
	import InfoViewer from "./InfoViewer.tsx";

	const DEFAULT_ARCHITECTURE = `[conv2d filters=8 kernel=11
	stride=1 padding=1 activation=relu]

	[maxpool size=2 stride=2]

	[flatten]

	[dense units=10 activation=softmax]`;

	const isFirefox = navigator.userAgent.toLowerCase().includes("firefox");
	await tf.setBackend(isFirefox ? "cpu" : "webgl");
	await tf.ready();

	export default function NetworkVisualizer() {
	const [dataset, setDataset] = useState<MnistData \| null>(null);

	useEffect(() => {
	const loadData = async () => {
	const data = new MnistData();
	await data.load();
	setDataset(data);
	console.log("dataset loaded");
	}

	loadData();
	}, []);

	// architecture states
	const [architecture, setArchitecture] = useState(DEFAULT_ARCHITECTURE);
	const [optimizerType, setOptimizerType] = useState<string>('adam');
	const [optimizerParams, setOptimizerParams] = useState<OptimizerParams>({
	learningRate: '0.001',
	beta1: '0.9',
	beta2: '0.999',
	epsilon: '1e-8',
	batchSize: '32',
	epochs: '5',
	});
	const modelRef = useRef<Cnn \| null>(null);
	const optimizerRef = useRef<tf.Optimizer \| null>(null);

	function handleArchitectureChange(newArchitecture: string) {
	if (isTraining) {
	alert('Cannot change architecture while training is in progress.');
	} else {
	setArchitecture(newArchitecture);
	}
	}

	function handleOptimizerChange(newOptimizerType: string, newOptimizerParams: OptimizerParams) {
	if (isTraining) {
	alert('Cannot change optimizer settings while training is in progress.');
	} else {
	setOptimizerType(newOptimizerType);
	setOptimizerParams(newOptimizerParams);
	}
	}

	function handleSampleIndexChange() {
	trainController.current.sampleIndex += 1;
	refreshSampleInfo();
	}


	function resetModel() {
	if (!dataset) return;

	if (modelRef.current) {
	modelRef.current.dispose();
	}

	const cnn = new Cnn(architecture, dataset.numInputChannels);
	modelRef.current = cnn;
	}

	function resetOptimizer() {
	if (optimizerType === 'adam') {
	const learningRate = parseFloat(optimizerParams.learningRate);
	const beta1 = parseFloat(optimizerParams.beta1 \|\| "0.9");
	const beta2 = parseFloat(optimizerParams.beta2 \|\| "0.999");
	const epsilon = parseFloat(optimizerParams.epsilon \|\| "1e-8");

	if (Number.isNaN(learningRate) \|\| learningRate <= 0) {
	alert('Invalid learning rate for Adam optimizer.');
	return;
	}
	if (Number.isNaN(beta1) \|\| beta1 < 0) {
	alert('Invalid beta1 for Adam optimizer.');
	return;
	}
	if (Number.isNaN(beta2) \|\| beta2 < 0) {
	alert('Invalid beta2 for Adam optimizer.');
	return;
	}
	if (Number.isNaN(epsilon) \|\| epsilon <= 0) {
	alert('Invalid epsilon for Adam optimizer.');
	return;
	}

	const opt = tf.train.adam(learningRate, beta1, beta2, epsilon);
	if (optimizerRef.current) {
	optimizerRef.current.dispose();
	}
	optimizerRef.current = opt;
	} else if (optimizerType === 'sgd') {
	const learningRate = parseFloat(optimizerParams.learningRate);
	if (Number.isNaN(learningRate) \|\| learningRate <= 0) {
	alert('Invalid learning rate for SGD optimizer.');
	return;
	}

	const opt = tf.train.sgd(learningRate);
	if (optimizerRef.current) {
	optimizerRef.current.dispose();
	}
	optimizerRef.current = opt;
	} else {
	alert(`Unsupported optimizer type: ${optimizerType}`);
	}
	}

	// Reset the model only when its architecture or input shape changes.
	useEffect(() => {
	resetModel();
	}, [architecture, dataset]);

	// Optimizer settings can change without discarding the trained model.
	useEffect(() => {
	resetOptimizer();
	}, [optimizerType, optimizerParams]);

	// training states
	const [isTraining, setIsTraining] = useState<boolean>(false);
	const isTrainingRef = useRef<boolean>(false);
	const lossesRef = useRef<Array<number>>([]);
	const trainController = useRef<TrainController>({
	isPaused: false,
	stopRequested: false,
	sampleIndex: 0,
	});
	const infoRef = useRef<Array<RunInfo>>([]);

	function refreshSampleInfo() {
	if (!dataset \|\| !modelRef.current) return;

	let nextInfo: RunInfo[] = [];
	const cleanupToken = tf.tidy(() => {
	const sample = dataset.getTestSample(trainController.current.sampleIndex);
	const { output, info } = modelRef.current!.forwardWithInfo(
	sample.xs.reshape([
	1,
	dataset.imageSize,
	dataset.imageSize,
	dataset.numInputChannels,
	]) as tf.Tensor4D,
	);
	const probs = tf.softmax(output);
	info.push({
	type: 'output',
	output: probs.dataSync(),
	shape: probs.shape,
	});
	nextInfo = info;
	return tf.scalar(0);
	});
	cleanupToken.dispose();
	infoRef.current = nextInfo;
	updateTick();
	}

	useEffect(() => {
	refreshSampleInfo();
	}, [architecture, dataset]);

	// render timing
	const [, setTick] = useState<number>(0);

	function updateTick() {
	setTick((tick) => tick + 1);
	}

	async function startTraining() {
	if (!modelRef.current \|\| !dataset \|\| !optimizerRef.current \|\| isTrainingRef.current) {
	return;
	}

	isTrainingRef.current = true;
	setIsTraining(true);
	trainController.current.isPaused = false;
	trainController.current.stopRequested = false;

	const batchSize = parseFloat(optimizerParams.batchSize);
	if (Number.isNaN(batchSize) \|\| batchSize <= 0) {
	alert('Invalid batch size.');
	isTrainingRef.current = false;
	setIsTraining(false);
	return;
	}
	const epochs = parseFloat(optimizerParams.epochs);
	if (Number.isNaN(epochs) \|\| epochs <= 0) {
	alert('Invalid number of epochs.');
	isTrainingRef.current = false;
	setIsTraining(false);
	return;
	}

	let lastTickUpdate = 0;
	if (!modelRef.current) return;
	if (!optimizerRef.current) return;

	try {
	await train(
	dataset,
	modelRef.current,
	optimizerRef.current,
	batchSize,
	epochs,
	trainController.current,
	(_epoch, _batch, loss, info) => {
	// lossesRef.current.push({ epoch, batch, loss });
	lossesRef.current.push(loss);
	console.log(loss);

	infoRef.current = info;

	// update tick every 50ms
	const now = performance.now();
	if (now - lastTickUpdate > 50) {
	lastTickUpdate = now;
	updateTick();
	}
	},
	);
	} finally {
	const wasStopped = trainController.current.stopRequested;
	isTrainingRef.current = false;
	setIsTraining(false);
	trainController.current.isPaused = false;
	trainController.current.stopRequested = false;
	if (!wasStopped) {
	alert('Training finished.');
	}
	}
	}

	function handleTrain() {
	if (isTrainingRef.current) {
	trainController.current.isPaused = false;
	return;
	}
	console.log('Starting training...');
	startTraining();
	}

	function handlePauseTraining() {
	console.log('Pausing training...');
	trainController.current.isPaused = true;
	}

	function handleStopTraining() {
	console.log('Stopping training...');
	trainController.current.stopRequested = true;
	trainController.current.isPaused = false;
	}

	async function waitUntilNotTraining() {
	return new Promise<void>((resolve) => {
	function check() {
	if (!isTrainingRef.current) {
	resolve();
	} else {
	requestAnimationFrame(check);
	}
	}
	check();
	});
	}

	async function handleResetTraining() {
	console.log('Resetting training...');
	handleStopTraining();

	await waitUntilNotTraining();
	console.log('Training stopped. Resetting model.');
	lossesRef.current = [];
	infoRef.current = [];
	resetModel();
	resetOptimizer();
	refreshSampleInfo();
	}

	return (
	<div className="flex flex-col min-h-0 h-full gap-3">
	<p className="text-sm text-gray-700">
	Configure and train a convolutional neural network (CNN) to classify handwritten digits from the MNIST dataset.
	</p>
	<div className="grid grid-cols-[2fr_1fr] min-h-0 flex-1 gap-12">
	<TrainingViewer
	infoRef={infoRef}
	handleSampleIndexChange={handleSampleIndexChange}
	/>
	<Sidebar
	architecture={architecture}
	onArchitectureChange={handleArchitectureChange}

	optimizerType={optimizerType}
	optimizerParams={optimizerParams}
	onOptimizerChange={handleOptimizerChange}
	isTraining={isTraining}
	lossesRef={lossesRef}

	onTrain={handleTrain}
	onPauseTraining={handlePauseTraining}
	onResetTraining={handleResetTraining}
	/>
	</div>
	</div>
	);
	}

	interface TrainingViewerProps {
	infoRef: React.RefObject<Array<RunInfo>>;
	handleSampleIndexChange: () => void;
	}
	function TrainingViewer({
	infoRef,
	handleSampleIndexChange,
	}: TrainingViewerProps) {
	return (
	<div className="flex flex-col h-full overflow-auto gap-4 w-full min-h-0">
	<InfoViewer info={infoRef.current} onSampleIndexChange={handleSampleIndexChange} />
	</div>
	)
	}

	interface SidebarProps {
	architecture: string;
	onArchitectureChange: (newArchitecture: string) => void;

	optimizerType: string;
	optimizerParams: OptimizerParams;
	onOptimizerChange: (newOptimizerType: string, newOptimizerParams: OptimizerParams) => void;
	isTraining: boolean;
	lossesRef: React.RefObject<Array<number>>;

	onTrain: () => void;
	onPauseTraining: () => void;
	onResetTraining: () => void;
	}

	type TrainingInputRowProps = {
	label: string;
	value: string \| undefined;
	onChange: (value: string) => void;
	};

	function TrainingInputRow({ label, value, onChange }: TrainingInputRowProps) {
	return (
	<div className="grid grid-cols-[7rem_minmax(0,1fr)] items-center gap-2">
	<span className="text-sm text-right text-gray-700">{label}</span>
	<div className="[&_label]:hidden [&_input]:py-1">
	<InputField label="" value={value ?? ""} onChange={onChange} />
	</div>
	</div>
	);
	}

	function Sidebar({
	architecture,
	onArchitectureChange,
	optimizerType,
	optimizerParams,
	onOptimizerChange,
	isTraining,
	lossesRef,
	onTrain,
	onPauseTraining,
	onResetTraining,
	}: SidebarProps) {
	const tabs = ["Architecture", "Train"];
	const [activeTab, setActiveTab] = useState<string>(tabs[0]);

	const [architectureDraft, setArchitectureDraft] = useState<string>(architecture);

	return (
	<Card className="flex flex-col h-full p-4 gap-2 overflow-auto">
	<Tabs tabs={tabs} activeTab={activeTab} onChange={setActiveTab} />
	<div className="flex flex-col p-4 gap-4 h-full overflow-auto">
	{ isFirefox && (
	<p className="text-red-500">
	Warning: This demo may be quite slow on Firefox.
	</p>
	)}

	{ activeTab === "Architecture" && (
	<>
	<InputField
	label="Architecture"
	value={architectureDraft}
	onChange={setArchitectureDraft}
	rows={15}
	/>
	<Button
	label="Apply architecture"
	onClick={() => onArchitectureChange(architectureDraft)}
	/>
	</>
	)}

	{ activeTab === "Train" && (
	<div className="flex flex-col gap-2">
	<p className="text-sm">Training {isTraining ? "in progress" : "not in progress"}</p>
	<Plot
	data={[
	{
	x: lossesRef.current.map((_, i) => i),
	y: lossesRef.current,
	mode: 'lines',
	type: 'scatter',
	},
	]}
	layout={{
	xaxis: { title: { text: 'Training steps' } },
	yaxis: { title: { text: 'Train loss' } },
	margin: { t: 15, r: 15, b: 40, l: 45 },
	}}
	className="w-full h-[220px]"
	config={{ responsive: true }}
	/>
	<div className="grid grid-cols-[7rem_minmax(0,1fr)] items-center gap-2">
	<span className="text-sm text-right text-gray-700">Optimizer</span>
	<div className="[&_label]:hidden [&_select]:py-1">
	<Dropdown
	label=""
	options={["sgd", "adam"]}
	activeOption={optimizerType}
	onChange={(newOptimizerType) => onOptimizerChange(newOptimizerType, optimizerParams)}
	/>
	</div>
	</div>
	{ optimizerType === 'sgd' && (
	<TrainingInputRow
	label="Learning Rate"
	value={optimizerParams.learningRate}
	onChange={(newLearningRate) => onOptimizerChange(optimizerType, {...optimizerParams, learningRate: newLearningRate})}
	/>
	)}
	{ optimizerType === 'adam' && (
	<>
	<TrainingInputRow
	label="Learning Rate"
	value={optimizerParams.learningRate}
	onChange={(newLearningRate) => onOptimizerChange(optimizerType, {...optimizerParams, learningRate: newLearningRate})}
	/>
	<TrainingInputRow
	label="Beta 1"
	value={optimizerParams.beta1}
	onChange={(newBeta1) => onOptimizerChange(optimizerType, {...optimizerParams, beta1: newBeta1})}
	/>
	<TrainingInputRow
	label="Beta 2"
	value={optimizerParams.beta2}
	onChange={(newBeta2) => onOptimizerChange(optimizerType, {...optimizerParams, beta2: newBeta2})}
	/>
	<TrainingInputRow
	label="Epsilon"
	value={optimizerParams.epsilon}
	onChange={(newEpsilon) => onOptimizerChange(optimizerType, {...optimizerParams, epsilon: newEpsilon})}
	/>
	</>
	)}
	<TrainingInputRow
	label="Batch Size"
	value={optimizerParams.batchSize}
	onChange={(newBatchSize) => onOptimizerChange(optimizerType, {...optimizerParams, batchSize: newBatchSize})}
	/>
	<TrainingInputRow
	label="Epochs"
	value={optimizerParams.epochs}
	onChange={(newEpochs) => onOptimizerChange(optimizerType, {...optimizerParams, epochs: newEpochs})}
	/>
	<div className="grid grid-cols-3 gap-2 mt-2 [&_button]:px-2 [&_button]:py-1 [&_button]:text-xs">
	<Button label="Train" onClick={onTrain} />
	<Button label="Pause" onClick={onPauseTraining} />
	<Button label="Reset" onClick={onResetTraining} />
	</div>
	</div>
	)}
	</div>
	</Card>
	)
	}