{-# LANGUAGE OverloadedStrings #-}

-- | GPT-2 architecture using hmatrix for linear algebra.
module Circuit.LLM.GPT
  ( -- * Normalisation and activation
    layerNorm,
    gelu,

    -- * Transformer block
    FeedForward (..),
    TransformerBlock (..),
    transformerBlock,

    -- * Full model
    GptConfig (..),
    Gpt (..),
    forward,
  )
where

import Numeric.LinearAlgebra
  ( Matrix,
    Vector,
    cmap,
    cols,
    fromLists,
    fromRows,
    rows,
    scale,
    subMatrix,
    sumElements,
    toList,
    toRows,
    tr,
    (|||),
  )
import Numeric.LinearAlgebra qualified as LA
import Numeric.LinearAlgebra.Data (maxElement)
import Prelude hiding (drop, sum, take)

-- | Layer normalisation along the last axis.
layerNorm :: Matrix Double -> Vector Double -> Vector Double -> Double -> Matrix Double
layerNorm :: Matrix Double
-> Vector Double -> Vector Double -> Double -> Matrix Double
layerNorm Matrix Double
x Vector Double
gamma Vector Double
beta Double
eps = [Vector Double] -> Matrix Double
forall t. Element t => [Vector t] -> Matrix t
fromRows ([Vector Double] -> Matrix Double)
-> [Vector Double] -> Matrix Double
forall a b. (a -> b) -> a -> b
$ (Vector Double -> Vector Double)
-> [Vector Double] -> [Vector Double]
forall a b. (a -> b) -> [a] -> [b]
map Vector Double -> Vector Double
normRow (Matrix Double -> [Vector Double]
forall t. Element t => Matrix t -> [Vector t]
toRows Matrix Double
x)
  where
    g :: [Double]
g = Vector Double -> [Double]
forall a. Storable a => Vector a -> [a]
toList Vector Double
gamma
    b :: [Double]
b = Vector Double -> [Double]
forall a. Storable a => Vector a -> [a]
toList Vector Double
beta
    d :: Double
d = Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Matrix Double -> Int
forall t. Matrix t -> Int
cols Matrix Double
x)
    normRow :: Vector Double -> Vector Double
normRow Vector Double
xi =
      let mu :: Double
mu = Vector Double -> Double
forall (c :: * -> *) e. Container c e => c e -> e
sumElements Vector Double
xi Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
d
          xi' :: Vector Double
xi' = (Double -> Double) -> Vector Double -> Vector Double
forall b (c :: * -> *) e.
(Element b, Container c e) =>
(e -> b) -> c e -> c b
cmap (Double -> Double -> Double
forall a. Num a => a -> a -> a
subtract Double
mu) Vector Double
xi
          var :: Double
var = Vector Double -> Double
forall (c :: * -> *) e. Container c e => c e -> e
sumElements (Vector Double
xi' Vector Double -> Vector Double -> Vector Double
forall a. Num a => a -> a -> a
* Vector Double
xi') Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
d Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
eps
          invStd :: Double
invStd = Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double -> Double
forall a. Floating a => a -> a
sqrt Double
var
          scaled :: Vector Double
scaled = Double -> Vector Double -> Vector Double
forall t (c :: * -> *). Linear t c => t -> c t -> c t
scale Double
invStd Vector Double
xi'
       in [Double] -> Vector Double
forall a. Storable a => [a] -> Vector a
LA.fromList ([Double] -> Vector Double) -> [Double] -> Vector Double
forall a b. (a -> b) -> a -> b
$ (Double -> Double -> Double -> Double)
-> [Double] -> [Double] -> [Double] -> [Double]
forall a b c d. (a -> b -> c -> d) -> [a] -> [b] -> [c] -> [d]
zipWith3 (\Double
v Double
gi Double
bi -> Double
v Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
gi Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
bi) (Vector Double -> [Double]
forall a. Storable a => Vector a -> [a]
toList Vector Double
scaled) [Double]
g [Double]
b

-- | GELU activation (approximation).
gelu :: Matrix Double -> Matrix Double
gelu :: Matrix Double -> Matrix Double
gelu = (Double -> Double) -> Matrix Double -> Matrix Double
forall b (c :: * -> *) e.
(Element b, Container c e) =>
(e -> b) -> c e -> c b
cmap Double -> Double
forall a. Floating a => a -> a
f
  where
    f :: a -> a
f a
v =
      let x' :: a
x' = a
1.59577 a -> a -> a
forall a. Num a => a -> a -> a
* a
v a -> a -> a
forall a. Num a => a -> a -> a
* (a
1 a -> a -> a
forall a. Num a => a -> a -> a
+ a
0.044715 a -> a -> a
forall a. Num a => a -> a -> a
* a
v a -> a -> a
forall a. Num a => a -> a -> a
* a
v)
       in a
v a -> a -> a
forall a. Fractional a => a -> a -> a
/ (a
1 a -> a -> a
forall a. Num a => a -> a -> a
+ a -> a
forall a. Floating a => a -> a
exp (-a
x'))

-- | Feed-forward network.
data FeedForward = FeedForward
  { FeedForward -> Matrix Double
ffW1 :: Matrix Double,
    FeedForward -> Vector Double
ffB1 :: Vector Double,
    FeedForward -> Matrix Double
ffW2 :: Matrix Double,
    FeedForward -> Vector Double
ffB2 :: Vector Double
  }

feedForward :: FeedForward -> Matrix Double -> Matrix Double
feedForward :: FeedForward -> Matrix Double -> Matrix Double
feedForward FeedForward
ff Matrix Double
x =
  let h :: Matrix Double
h = Matrix Double -> Matrix Double
gelu (Matrix Double
x Matrix Double -> Matrix Double -> Matrix Double
forall t. Numeric t => Matrix t -> Matrix t -> Matrix t
LA.<> FeedForward -> Matrix Double
ffW1 FeedForward
ff Matrix Double -> Matrix Double -> Matrix Double
forall a. Num a => a -> a -> a
+ Vector Double -> Matrix Double
forall a. Storable a => Vector a -> Matrix a
LA.asRow (FeedForward -> Vector Double
ffB1 FeedForward
ff))
   in Matrix Double
h Matrix Double -> Matrix Double -> Matrix Double
forall t. Numeric t => Matrix t -> Matrix t -> Matrix t
LA.<> FeedForward -> Matrix Double
ffW2 FeedForward
ff Matrix Double -> Matrix Double -> Matrix Double
forall a. Num a => a -> a -> a
+ Vector Double -> Matrix Double
forall a. Storable a => Vector a -> Matrix a
LA.asRow (FeedForward -> Vector Double
ffB2 FeedForward
ff)

-- | Causal mask.
causalMask :: Int -> Matrix Double
causalMask :: Int -> Matrix Double
causalMask Int
n = [[Double]] -> Matrix Double
forall t. Element t => [[t]] -> Matrix t
fromLists [[if Int
j Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
i then Double
0 else -(Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
0) | Int
j <- [Int
0 .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]] | Int
i <- [Int
0 .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]]

-- | Row-wise softmax.
softmax :: Matrix Double -> Matrix Double
softmax :: Matrix Double -> Matrix Double
softmax Matrix Double
x = [Vector Double] -> Matrix Double
forall t. Element t => [Vector t] -> Matrix t
fromRows [Vector Double -> Vector Double
forall {c :: * -> *} {t}.
(Container c t, Floating t, Linear t c) =>
c t -> c t
softmaxRow Vector Double
row | Vector Double
row <- Matrix Double -> [Vector Double]
forall t. Element t => Matrix t -> [Vector t]
toRows Matrix Double
x]
  where
    softmaxRow :: c t -> c t
softmaxRow c t
v =
      let mx :: t
mx = c t -> t
forall (c :: * -> *) e. Container c e => c e -> e
maxElement c t
v
          shifted :: c t
shifted = (t -> t) -> c t -> c t
forall b (c :: * -> *) e.
(Element b, Container c e) =>
(e -> b) -> c e -> c b
cmap (\t
xi -> t -> t
forall a. Floating a => a -> a
exp (t
xi t -> t -> t
forall a. Num a => a -> a -> a
- t
mx)) c t
v
          s :: t
s = c t -> t
forall (c :: * -> *) e. Container c e => c e -> e
sumElements c t
shifted
       in t -> c t -> c t
forall t (c :: * -> *). Linear t c => t -> c t -> c t
scale (t
1 t -> t -> t
forall a. Fractional a => a -> a -> a
/ t
s) c t
shifted

-- | Scaled dot-product attention.
scaledDotProductAttention :: Matrix Double -> Matrix Double -> Matrix Double -> Matrix Double -> Matrix Double
scaledDotProductAttention :: Matrix Double
-> Matrix Double -> Matrix Double -> Matrix Double -> Matrix Double
scaledDotProductAttention Matrix Double
q Matrix Double
k Matrix Double
v Matrix Double
mask =
  let dk :: Double
dk = Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Matrix Double -> Int
forall t. Matrix t -> Int
cols Matrix Double
k) :: Double
      scores :: Matrix Double
scores = Double -> Matrix Double -> Matrix Double
forall t (c :: * -> *). Linear t c => t -> c t -> c t
scale (Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double -> Double
forall a. Floating a => a -> a
sqrt Double
dk) (Matrix Double
q Matrix Double -> Matrix Double -> Matrix Double
forall t. Numeric t => Matrix t -> Matrix t -> Matrix t
LA.<> Matrix Double -> Matrix Double
forall m mt. Transposable m mt => m -> mt
tr Matrix Double
k)
      masked :: Matrix Double
masked = Matrix Double
scores Matrix Double -> Matrix Double -> Matrix Double
forall a. Num a => a -> a -> a
+ Matrix Double
mask
      attn :: Matrix Double
attn = Matrix Double -> Matrix Double
softmax Matrix Double
masked
   in Matrix Double
attn Matrix Double -> Matrix Double -> Matrix Double
forall t. Numeric t => Matrix t -> Matrix t -> Matrix t
LA.<> Matrix Double
v

-- | Multi-head self-attention.
multiHeadAttention :: Int -> Matrix Double -> Matrix Double -> Matrix Double -> Matrix Double -> Matrix Double -> Matrix Double -> Matrix Double
multiHeadAttention :: Int
-> Matrix Double
-> Matrix Double
-> Matrix Double
-> Matrix Double
-> Matrix Double
-> Matrix Double
-> Matrix Double
multiHeadAttention Int
nHead Matrix Double
x Matrix Double
wQ Matrix Double
wK Matrix Double
wV Matrix Double
wO Matrix Double
mask =
  let headDim :: Int
headDim = Matrix Double -> Int
forall t. Matrix t -> Int
cols Matrix Double
wQ Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
nHead
      wSlice :: Matrix Double -> Int -> Matrix Double
wSlice Matrix Double
w Int
h = (Int, Int) -> (Int, Int) -> Matrix Double -> Matrix Double
forall a.
Element a =>
(Int, Int) -> (Int, Int) -> Matrix a -> Matrix a
subMatrix (Int
0, Int
h Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
headDim) (Matrix Double -> Int
forall t. Matrix t -> Int
rows Matrix Double
w, Int
headDim) Matrix Double
w
      heads :: [Matrix Double]
heads =
        [ Matrix Double
-> Matrix Double -> Matrix Double -> Matrix Double -> Matrix Double
scaledDotProductAttention (Matrix Double
x Matrix Double -> Matrix Double -> Matrix Double
forall t. Numeric t => Matrix t -> Matrix t -> Matrix t
LA.<> Matrix Double -> Int -> Matrix Double
wSlice Matrix Double
wQ Int
h) (Matrix Double
x Matrix Double -> Matrix Double -> Matrix Double
forall t. Numeric t => Matrix t -> Matrix t -> Matrix t
LA.<> Matrix Double -> Int -> Matrix Double
wSlice Matrix Double
wK Int
h) (Matrix Double
x Matrix Double -> Matrix Double -> Matrix Double
forall t. Numeric t => Matrix t -> Matrix t -> Matrix t
LA.<> Matrix Double -> Int -> Matrix Double
wSlice Matrix Double
wV Int
h) Matrix Double
mask
        | Int
h <- [Int
0 .. Int
nHead Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
        ]
   in [Matrix Double] -> Matrix Double
forall {t}. Element t => [Matrix t] -> Matrix t
concatCols [Matrix Double]
heads Matrix Double -> Matrix Double -> Matrix Double
forall t. Numeric t => Matrix t -> Matrix t -> Matrix t
LA.<> Matrix Double
wO
  where
    concatCols :: [Matrix t] -> Matrix t
concatCols [Matrix t
a] = Matrix t
a
    concatCols (Matrix t
a : [Matrix t]
as) = Matrix t
a Matrix t -> Matrix t -> Matrix t
forall t. Element t => Matrix t -> Matrix t -> Matrix t
||| [Matrix t] -> Matrix t
concatCols [Matrix t]
as
    concatCols [] = [Char] -> Matrix t
forall a. HasCallStack => [Char] -> a
error [Char]
"no heads"

-- | One transformer block.
data TransformerBlock = TransformerBlock
  { TransformerBlock -> Matrix Double
tbAttnWq :: Matrix Double,
    TransformerBlock -> Matrix Double
tbAttnWk :: Matrix Double,
    TransformerBlock -> Matrix Double
tbAttnWv :: Matrix Double,
    TransformerBlock -> Matrix Double
tbAttnWo :: Matrix Double,
    TransformerBlock -> Vector Double
tbAttnLnGamma :: Vector Double,
    TransformerBlock -> Vector Double
tbAttnLnBeta :: Vector Double,
    TransformerBlock -> FeedForward
tbFfn :: FeedForward,
    TransformerBlock -> Vector Double
tbFfnLnGamma :: Vector Double,
    TransformerBlock -> Vector Double
tbFfnLnBeta :: Vector Double
  }

transformerBlock :: Int -> TransformerBlock -> Matrix Double -> Matrix Double
transformerBlock :: Int -> TransformerBlock -> Matrix Double -> Matrix Double
transformerBlock Int
nHead TransformerBlock
tb Matrix Double
x =
  let attnOut :: Matrix Double
attnOut =
        Int
-> Matrix Double
-> Matrix Double
-> Matrix Double
-> Matrix Double
-> Matrix Double
-> Matrix Double
-> Matrix Double
multiHeadAttention
          Int
nHead
          Matrix Double
x
          (TransformerBlock -> Matrix Double
tbAttnWq TransformerBlock
tb)
          (TransformerBlock -> Matrix Double
tbAttnWk TransformerBlock
tb)
          (TransformerBlock -> Matrix Double
tbAttnWv TransformerBlock
tb)
          (TransformerBlock -> Matrix Double
tbAttnWo TransformerBlock
tb)
          (Int -> Matrix Double
causalMask (Matrix Double -> Int
forall t. Matrix t -> Int
rows Matrix Double
x))
      attnRes :: Matrix Double
attnRes = Matrix Double
x Matrix Double -> Matrix Double -> Matrix Double
forall a. Num a => a -> a -> a
+ Matrix Double
attnOut
      attnNorm :: Matrix Double
attnNorm = Matrix Double
-> Vector Double -> Vector Double -> Double -> Matrix Double
layerNorm Matrix Double
attnRes (TransformerBlock -> Vector Double
tbAttnLnGamma TransformerBlock
tb) (TransformerBlock -> Vector Double
tbAttnLnBeta TransformerBlock
tb) Double
1e-5
      ffnOut :: Matrix Double
ffnOut = FeedForward -> Matrix Double -> Matrix Double
feedForward (TransformerBlock -> FeedForward
tbFfn TransformerBlock
tb) Matrix Double
attnNorm
      ffnRes :: Matrix Double
ffnRes = Matrix Double
attnNorm Matrix Double -> Matrix Double -> Matrix Double
forall a. Num a => a -> a -> a
+ Matrix Double
ffnOut
   in Matrix Double
-> Vector Double -> Vector Double -> Double -> Matrix Double
layerNorm Matrix Double
ffnRes (TransformerBlock -> Vector Double
tbFfnLnGamma TransformerBlock
tb) (TransformerBlock -> Vector Double
tbFfnLnBeta TransformerBlock
tb) Double
1e-5

-- | GPT-2 model configuration.
data GptConfig = GptConfig
  { GptConfig -> Int
gptVocabSize :: Int,
    GptConfig -> Int
gptNEmbd :: Int,
    GptConfig -> Int
gptNHead :: Int,
    GptConfig -> Int
gptNLayer :: Int
  }

-- | GPT-2 model parameters.
data Gpt = Gpt
  { Gpt -> Matrix Double
gptWte :: Matrix Double,
    Gpt -> Matrix Double
gptWpe :: Matrix Double,
    Gpt -> [TransformerBlock]
gptBlocks :: [TransformerBlock],
    Gpt -> Vector Double
gptLnGamma :: Vector Double,
    Gpt -> Vector Double
gptLnBeta :: Vector Double,
    Gpt -> Matrix Double
gptHead :: Matrix Double,
    Gpt -> Vector Double
gptHeadB :: Vector Double
  }

-- | Full forward pass.
forward :: GptConfig -> Gpt -> [Int] -> Matrix Double
forward :: GptConfig -> Gpt -> [Int] -> Matrix Double
forward GptConfig
cfg Gpt
m [Int]
inputIds =
  let seqLen :: Int
seqLen = [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
inputIds
      tokEmb :: Matrix Double
tokEmb = [Vector Double] -> Matrix Double
forall t. Element t => [Vector t] -> Matrix t
fromRows [Matrix Double -> [Vector Double]
forall t. Element t => Matrix t -> [Vector t]
toRows (Gpt -> Matrix Double
gptWte Gpt
m) [Vector Double] -> Int -> Vector Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
i | Int
i <- [Int]
inputIds]
      posEmb :: Matrix Double
posEmb = (Int, Int) -> (Int, Int) -> Matrix Double -> Matrix Double
forall a.
Element a =>
(Int, Int) -> (Int, Int) -> Matrix a -> Matrix a
subMatrix (Int
0, Int
0) (Int
seqLen, Matrix Double -> Int
forall t. Matrix t -> Int
cols (Gpt -> Matrix Double
gptWpe Gpt
m)) (Gpt -> Matrix Double
gptWpe Gpt
m)
      x :: Matrix Double
x = Matrix Double
tokEmb Matrix Double -> Matrix Double -> Matrix Double
forall a. Num a => a -> a -> a
+ Matrix Double
posEmb
      x' :: Matrix Double
x' = (Matrix Double -> TransformerBlock -> Matrix Double)
-> Matrix Double -> [TransformerBlock] -> Matrix Double
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl ((TransformerBlock -> Matrix Double -> Matrix Double)
-> Matrix Double -> TransformerBlock -> Matrix Double
forall a b c. (a -> b -> c) -> b -> a -> c
flip (Int -> TransformerBlock -> Matrix Double -> Matrix Double
transformerBlock (GptConfig -> Int
gptNHead GptConfig
cfg))) Matrix Double
x (Gpt -> [TransformerBlock]
gptBlocks Gpt
m)
      x'' :: Matrix Double
x'' = Matrix Double
-> Vector Double -> Vector Double -> Double -> Matrix Double
layerNorm Matrix Double
x' (Gpt -> Vector Double
gptLnGamma Gpt
m) (Gpt -> Vector Double
gptLnBeta Gpt
m) Double
1e-5
   in Matrix Double
x'' Matrix Double -> Matrix Double -> Matrix Double
forall t. Numeric t => Matrix t -> Matrix t -> Matrix t
LA.<> Gpt -> Matrix Double
gptHead Gpt
m Matrix Double -> Matrix Double -> Matrix Double
forall a. Num a => a -> a -> a
+ Vector Double -> Matrix Double
forall a. Storable a => Vector a -> Matrix a
LA.asRow (Gpt -> Vector Double
gptHeadB Gpt
m)