{-# LANGUAGE DataKinds #-}
{-# LANGUAGE KindSignatures #-}
module Torch.GraduallyTyped.NN.Transformer.T5.Base where
import Data.Singletons (SingI (sing))
import GHC.TypeLits (Nat)
import Torch.GraduallyTyped.Device (Device (..), DeviceType (..), SDevice)
import Torch.GraduallyTyped.NN.Class (ModelSpec)
import Torch.GraduallyTyped.NN.Transformer.T5.Common (T5ModelF, t5ModelSpec)
import Torch.GraduallyTyped.NN.Transformer.Type (STransformerHead, STransformerStyle (ST5), TransformerHead, TransformerStyle (T5))
import Torch.GraduallyTyped.NN.Type (HasDropout, SHasDropout)
import Torch.GraduallyTyped.Prelude.TypeLits (SNat)
import Torch.GraduallyTyped.RequiresGradient (Gradient, RequiresGradient, SGradient)
import Torch.GraduallyTyped.Shape.Type (Dim (..), Name (..), Size (..))
type T5BaseNumLayers = 12
t5BaseNumLayers :: SNat T5BaseNumLayers
t5BaseNumLayers :: SNat T5BaseNumLayers
t5BaseNumLayers = forall {k} (a :: k). SingI a => Sing a
sing
type T5BaseHeadDim = 'Dim ('Name "*") ('Size 12)
type T5BaseHeadEmbedDim = 'Dim ('Name "*") ('Size 64)
type T5BaseEmbedDim = 'Dim ('Name "*") ('Size 768)
type T5BaseInputEmbedDim = 'Dim ('Name "*") ('Size 768)
type T5BaseFFNDim = 'Dim ('Name "*") ('Size 3072)
type T5BaseVocabDim = 'Dim ('Name "*") ('Size 32128)
type T5Base
(transformerHead :: TransformerHead)
(gradient :: Gradient RequiresGradient)
(device :: Device (DeviceType Nat))
(hasDropout :: HasDropout) =
T5ModelF 'T5 transformerHead T5BaseNumLayers T5BaseNumLayers gradient device T5BaseHeadDim T5BaseHeadEmbedDim T5BaseEmbedDim T5BaseInputEmbedDim T5BaseFFNDim T5BaseVocabDim hasDropout
t5BaseSpec ::
STransformerHead transformerHead ->
SGradient gradient ->
SDevice device ->
SHasDropout hasDropout ->
ModelSpec (T5Base transformerHead gradient device hasDropout)
t5BaseSpec :: forall (transformerHead :: TransformerHead)
(gradient :: Gradient RequiresGradient)
(device :: Device (DeviceType Natural)) (hasDropout :: HasDropout).
STransformerHead transformerHead
-> SGradient gradient
-> SDevice device
-> SHasDropout hasDropout
-> ModelSpec (T5Base transformerHead gradient device hasDropout)
t5BaseSpec STransformerHead transformerHead
transformerHead = forall (style :: TransformerStyle)
(transformerHead :: TransformerHead) (numEncoderLayers :: Natural)
(numDecoderLayers :: Natural)
(gradient :: Gradient RequiresGradient)
(device :: Device (DeviceType Natural))
(headDim :: Dim (Name Symbol) (Size Natural))
(headEmbedDim :: Dim (Name Symbol) (Size Natural))
(embedDim :: Dim (Name Symbol) (Size Natural))
(inputEmbedDim :: Dim (Name Symbol) (Size Natural))
(ffnDim :: Dim (Name Symbol) (Size Natural))
(vocabDim :: Dim (Name Symbol) (Size Natural))
(hasDropout :: HasDropout).
(SingI headDim, SingI headEmbedDim, SingI embedDim,
SingI inputEmbedDim, SingI ffnDim, SingI vocabDim) =>
STransformerStyle style
-> STransformerHead transformerHead
-> SNat numEncoderLayers
-> SNat numDecoderLayers
-> SGradient gradient
-> SDevice device
-> SHasDropout hasDropout
-> ModelSpec
(T5ModelF
style
transformerHead
numEncoderLayers
numDecoderLayers
gradient
device
headDim
headEmbedDim
embedDim
inputEmbedDim
ffnDim
vocabDim
hasDropout)
t5ModelSpec STransformerStyle 'T5
ST5 STransformerHead transformerHead
transformerHead SNat T5BaseNumLayers
t5BaseNumLayers SNat T5BaseNumLayers
t5BaseNumLayers