Models.hs

{- | Provides interface for creating and using neural network models.
-}


-- 'TypeFamilies' are needed to instantiate 'Synapse.Tensors.DType'.
{-# LANGUAGE TypeFamilies #-}


module Synapse.NN.Models
    ( -- * Common for models
      InputSize (InputSize)

      -- * 'SequentialModel' datatype
      
    , SequentialModel (SequentialModel, unSequentialModel)
    , buildSequentialModel

    , layerPrefix
    ) where


import Synapse.Tensors (DType, SingletonOps(singleton))

import Synapse.Autograd (SymbolIdentifier (SymbolIdentifier))

import Synapse.NN.Layers.Layer(AbstractLayer(..), Layer, LayerConfiguration)

import Data.Maybe (fromMaybe)
import Data.Foldable (foldl')


-- | 'InputSize' newtype wraps 'Int' - amount of features of input that the model should support (@InputSize 3@ means that model supports any matrix with size (x, 3)).
newtype InputSize = InputSize Int


-- | 'SequentialModel' datatype represents any model grouping layers linearly.
newtype SequentialModel a = SequentialModel
    { unSequentialModel :: [Layer a]  -- ^ Returns layers of 'SequentialModel'.
    }

instance Show a => Show (SequentialModel a) where
    show (SequentialModel layers) = go 1 layers
      where
        go _ [] = ""
        go i (x:xs) = "Layer " ++ show i ++ " parameters: " ++ show (getParameters (layerPrefix mempty i) x) ++ ";\n" ++ go (i + 1) xs

-- | Builds sequential model using input size and layer configurations to ensure that layers are compatible with each other.
buildSequentialModel :: InputSize -> [LayerConfiguration (Layer a)] -> SequentialModel a
buildSequentialModel (InputSize i) layerConfigs = SequentialModel $ go i layerConfigs
  where
    go _ [] = []
    go prevSize (l:ls) = let layer = l prevSize
                             outputMaybe = outputSize layer
                             output = fromMaybe prevSize outputMaybe
                         in layer : go output ls

type instance DType (SequentialModel a) = a


-- | Forms prefix for layers according to 'Synapse.NN.Layers.Layer.AbstractLayer' requirements.
layerPrefix :: SymbolIdentifier -> Int -> SymbolIdentifier
layerPrefix prefix i = mconcat [prefix, SymbolIdentifier "l", SymbolIdentifier (show i), SymbolIdentifier "w"]

instance AbstractLayer SequentialModel where
    inputSize = inputSize . head . unSequentialModel
    outputSize = outputSize . head . unSequentialModel

    nParameters = foldl' (\parameters layer -> parameters + nParameters layer) 0 . unSequentialModel
    getParameters prefix =
        snd . foldl' (\(i, acc) layer -> (i + 1, acc ++ getParameters (layerPrefix prefix i) layer)) (1, []) . unSequentialModel
    updateParameters (SequentialModel model) = SequentialModel . go model
      where
        go [] _ = []
        go (layer:layers) parameters = let (x, parameters') = splitAt (nParameters layer) parameters
                                       in updateParameters layer x : go layers parameters'

    symbolicForward prefix input =
        snd . foldl' (\(i, (mat, loss)) layer -> let (mat', newLoss) = symbolicForward (layerPrefix prefix i) mat layer
                                                 in (i + 1, (mat', loss + newLoss)))
              (1, (input, singleton 0)) . unSequentialModel