diff --git a/COMP4027.cabal b/COMP4027.cabal
index 72612e1..3aae927 100644
--- a/COMP4027.cabal
+++ b/COMP4027.cabal
@@ -30,6 +30,7 @@ library
   build-depends:
       ad ==4.5.2
     , base ^>=4.14.1.0
+    , extra ==1.7.11
     , random ==1.2.0
     , JuicyPixels ==3.3.5
     , JuicyPixels-util ==0.2
diff --git a/src/MLP_testing_MNIST.hs b/src/MLP_testing_MNIST.hs
new file mode 100644
index 0000000..8d4d6e7
--- /dev/null
+++ b/src/MLP_testing_MNIST.hs
@@ -0,0 +1,52 @@
+module MLP_testing_MNIST where
+
+{- |
+  In this module are a number of functions that have been used to test models
+  that have been trained on the MNIST dataset.
+
+  The functions below have been defined to test models that have been trained
+  on the MNIST dataset. The testMNIST function obtains the MNIST test dataset
+  and uses this to test the model. In testModel, each sample is tested to see
+  if the one hot encoded prediction is the same as the target, if so then 1
+  is added to a sum, otherwise 0 is added. The final sum tells us how many
+  samples (out of 10000) the model predicted correctly.
+-}
+
+import Data.Bool (bool)
+import Data.Tuple.Extra (both)
+import Data.Time.Clock -- clock to measure convergence time
+import Types -- all custom types/classes for parallel multi-layer perceptrons
+import MNIST -- MNIST dataset for target function. Source: https://git.thm.de/akwl20/neural-network-haskell
+import Data.Matrix as Matrix -- toList operation for matrices
+import MLP_utils -- for feedFrwd and toLoss
+import qualified Data.Text.IO as TXTIO
+import qualified Data.Text as TXT
+
+-- Returns bool indicating whether the sample was predicted correctly.
+testSample :: [UnactivatedLayer] -> ([Double],[Double]) -> Bool
+testSample model (input,target) =
+  target == (predictClass . outputs . head $ feedFrwd model input)
+  where
+  predictClass :: [Double] -> [Double]
+  predictClass xs = [if x == maximum xs then 1.0 else 0.0 | x <- xs]
+  -- ^ Converts a softmax output to a one hot vector prediction.
+
+-- | Returns the number of samples predicted correctly.
+testModel :: [UnactivatedLayer] -> [([Double], [Double])] -> Int
+testModel model testData = sum $ map (bool 0 1 . testSample model) testData
+
+-- | Measures the time taken to test a model that has been trained on the MNIST
+-- dataset. Also prints how many samples out of the 10000-sample MNIST test-set
+-- the model predicted correctly.
+timedRun :: IO ()
+timedRun = do
+  let preprocess :: (Real a, Fractional b) => [(Matrix a, Matrix a)] -> [([b], [b])]
+      preprocess = map $ both $ map realToFrac . Matrix.toList
+  model <- read . TXT.unpack <$> TXTIO.readFile "trained_model" :: IO [UnactivatedLayer]
+  testData <- preprocess <$> getTestSamples
+  t0 <- getCurrentTime
+  let numTruePositives = testModel model testData
+  -- Force evaluation by printing the # of true positives before taking the time:
+  putStrLn ("Accuracy: " ++ show numTruePositives ++ " out of " ++ show (length testData))
+  t1 <- getCurrentTime
+  putStrLn ("Time to test: " ++ show (diffUTCTime t1 t0))
diff --git a/src/MLP_testing_MNIST.lhs b/src/MLP_testing_MNIST.lhs
deleted file mode 100644
index f1ee4d3..0000000
--- a/src/MLP_testing_MNIST.lhs
+++ /dev/null
@@ -1,56 +0,0 @@
-> module MLP_testing_MNIST where
-
-In this module are a number of functions that have been used to test models that have been trained on the MNIST dataset.
-
-> import Data.Time.Clock -- clock to measure convergence time
-> import Types -- all custom types/classes for parallel multi-layer perceptrons
-> import MNIST -- MNIST dataset for target function. Source: https://git.thm.de/akwl20/neural-network-haskell
-> import Data.Matrix -- toList operation for matrices
-> import MLP_utils -- for feedFrwd and toLoss
-
-> import qualified Data.Text.IO as TXTIO
-> import qualified Data.Text as TXT
-
-The functions below have been defined to test models that have been trained on the MNIST dataset. The testMNIST function obtains the MNIST test dataset and uses this to test the model.
-In testModel, each sample is tested to see if the one hot encoded prediction is the same as the target, if so then 1 is added to a sum, otherwise 0 is added. The final sum tells us
-how many samples (out of 10000) the model predicted correctly. 
-
-> testSample :: [UnactivatedLayer] -> ([Double],[Double]) -> Bool
-> testSample model (input,target) = true_pos
->     where 
->     all_outs = feedFrwd model input
->     true_pos = target == (toPred . outputs . head $ all_outs)
-
-> testModel :: [UnactivatedLayer] -> [([Double], [Double])] -> Int
-> testModel model testData = sum $ map (toInt . testSample model) testData
->     where toInt x = case x of
->                          False -> 0
->                          True -> 1
-
-> testMNIST :: [UnactivatedLayer] -> IO Int
-> testMNIST model = do 
->     testData <- getTestSamples
->     let format (x,l) = (realToFrac <$> toList x, realToFrac <$> toList l)
->     let testData' = format <$> testData
->     return $ testModel model testData'
-
-The toPred function converts a softmax output to a one hot vector prediction.
-
-> toPred :: [Double] -> [Double]
-> toPred xs = [if x == maximum xs then 1.0 else 0.0 | x <- xs]
-
-The timedRun function below measures testing time for the model and will print how many samples out of 10000 the model predicted correctly.
-
-> timedRun = do 
->     t0 <- getCurrentTime
->     model <- read . TXT.unpack <$> TXTIO.readFile "trained_model" :: IO [UnactivatedLayer]
->     err <- testMNIST model
->     testData <- getTestSamples
->     putStrLn ("Accuracy: " ++ show err ++ " out of " ++ show (length testData))
->     t1 <- getCurrentTime
->     putStrLn ("Time to test: " ++ show (diffUTCTime t1 t0))
-
-The main function below is configured to test a model trained on the MNIST dataset.
-
-> main :: IO ()
-> main = timedRun