{-# LANGUAGE OverloadedStrings #-}

-- | Load GPT-2 weights from a directory of raw float32 binary files.
--
-- Weight files follow a naming convention. Each .f32 file contains
-- raw little-endian IEEE 754 float32 values with no header.
--
-- Required files for a GPT-2 model:
--
-- @
--   wte.f32             [vocab_size, n_embd]
--   wpe.f32             [max_seq_len, n_embd]
--   hN.ln1.gamma.f32    [n_embd]             per block
--   hN.ln1.beta.f32     [n_embd]
--   hN.attn.qkv.w.f32   [n_embd, 3*n_embd]   fused Q,K,V projection
--   hN.attn.qkv.b.f32   [3*n_embd]
--   hN.attn.proj.w.f32  [n_embd, n_embd]
--   hN.attn.proj.b.f32  [n_embd]
--   hN.ln2.gamma.f32    [n_embd]
--   hN.ln2.beta.f32     [n_embd]
--   hN.mlp.fc.w.f32     [n_embd, 4*n_embd]
--   hN.mlp.fc.b.f32     [4*n_embd]
--   hN.mlp.proj.w.f32   [4*n_embd, n_embd]
--   hN.mlp.proj.b.f32   [n_embd]
--   lnf.gamma.f32       [n_embd]
--   lnf.beta.f32        [n_embd]
-- @
module Circuit.LLM.Weights
  ( -- * Loading
    loadGpt2,
    loadGpt2With,
    loadMatrix,
    loadVector,

    -- * GPT-2 architecture dimensions
    Gpt2Size (..),
    sizeConfig,
  )
where

import Circuit.LLM.GPT
  ( FeedForward (..),
    Gpt (..),
    GptConfig (..),
    TransformerBlock (..),
  )
import Data.ByteString (ByteString)
import Data.ByteString qualified as BS
import Data.Word (Word32)
import Foreign.Marshal.Alloc (alloca)
import Foreign.Ptr (Ptr, castPtr)
import Foreign.Storable (peek, poke)
import Numeric.LinearAlgebra
  ( Matrix,
    Vector,
    fromList,
    reshape,
    subMatrix,
    tr,
  )
import System.IO.Unsafe (unsafePerformIO)

----------------------------------------------------------------------
-- GPT-2 sizes
----------------------------------------------------------------------

-- | GPT-2 model size presets.
data Gpt2Size = Gpt2Small | Gpt2Medium | Gpt2Large | Gpt2Xl
  deriving (Int -> Gpt2Size -> ShowS
[Gpt2Size] -> ShowS
Gpt2Size -> [Char]
(Int -> Gpt2Size -> ShowS)
-> (Gpt2Size -> [Char]) -> ([Gpt2Size] -> ShowS) -> Show Gpt2Size
forall a.
(Int -> a -> ShowS) -> (a -> [Char]) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Gpt2Size -> ShowS
showsPrec :: Int -> Gpt2Size -> ShowS
$cshow :: Gpt2Size -> [Char]
show :: Gpt2Size -> [Char]
$cshowList :: [Gpt2Size] -> ShowS
showList :: [Gpt2Size] -> ShowS
Show, Gpt2Size -> Gpt2Size -> Bool
(Gpt2Size -> Gpt2Size -> Bool)
-> (Gpt2Size -> Gpt2Size -> Bool) -> Eq Gpt2Size
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Gpt2Size -> Gpt2Size -> Bool
== :: Gpt2Size -> Gpt2Size -> Bool
$c/= :: Gpt2Size -> Gpt2Size -> Bool
/= :: Gpt2Size -> Gpt2Size -> Bool
Eq)

-- | Model configuration for each preset.
sizeConfig :: Gpt2Size -> GptConfig
sizeConfig :: Gpt2Size -> GptConfig
sizeConfig Gpt2Size
Gpt2Small = Int -> Int -> Int -> Int -> GptConfig
GptConfig Int
50257 Int
768 Int
12 Int
12
sizeConfig Gpt2Size
Gpt2Medium = Int -> Int -> Int -> Int -> GptConfig
GptConfig Int
50257 Int
1024 Int
16 Int
24
sizeConfig Gpt2Size
Gpt2Large = Int -> Int -> Int -> Int -> GptConfig
GptConfig Int
50257 Int
1280 Int
20 Int
36
sizeConfig Gpt2Size
Gpt2Xl = Int -> Int -> Int -> Int -> GptConfig
GptConfig Int
50257 Int
1600 Int
25 Int
48

----------------------------------------------------------------------
-- Top-level loader
----------------------------------------------------------------------

-- | Load a GPT-2 model with a custom configuration (for non-standard dimensions).
loadGpt2With :: FilePath -> GptConfig -> IO Gpt
loadGpt2With :: [Char] -> GptConfig -> IO Gpt
loadGpt2With [Char]
dir GptConfig
cfg = do
  let nEmb :: Int
nEmb = GptConfig -> Int
gptNEmbd GptConfig
cfg
      nHead :: Int
nHead = GptConfig -> Int
gptNHead GptConfig
cfg
      nLayer :: Int
nLayer = GptConfig -> Int
gptNLayer GptConfig
cfg
      vocab :: Int
vocab = GptConfig -> Int
gptVocabSize GptConfig
cfg
      maxSeq :: Int
maxSeq = Int
1024

  wte <- [Char] -> Int -> Int -> IO (Matrix Double)
loadMatrix ([Char]
dir [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ [Char]
"/wte.f32") Int
vocab Int
nEmb
  wpe <- loadMatrix (dir ++ "/wpe.f32") maxSeq nEmb
  blocks <- mapM (loadBlock dir nEmb nHead) [0 .. nLayer - 1]
  lnG <- loadVector (dir ++ "/lnf.gamma.f32") nEmb
  lnB <- loadVector (dir ++ "/lnf.beta.f32") nEmb

  pure
    Gpt
      { gptWte = wte,
        gptWpe = wpe,
        gptBlocks = blocks,
        gptLnGamma = lnG,
        gptLnBeta = lnB,
        -- Tie output head to transposed token embeddings (weight tying)
        gptHead = tr wte, -- [n_embd, vocab]
        gptHeadB = zeroVector vocab
      }

-- | Load a full GPT-2 model from a weight directory.
loadGpt2 :: FilePath -> Gpt2Size -> IO Gpt
loadGpt2 :: [Char] -> Gpt2Size -> IO Gpt
loadGpt2 [Char]
dir Gpt2Size
sz = [Char] -> GptConfig -> IO Gpt
loadGpt2With [Char]
dir (Gpt2Size -> GptConfig
sizeConfig Gpt2Size
sz)

-- | Load one transformer block's weights.
loadBlock :: FilePath -> Int -> Int -> Int -> IO TransformerBlock
loadBlock :: [Char] -> Int -> Int -> Int -> IO TransformerBlock
loadBlock [Char]
dir Int
nEmb Int
_nHead Int
h = do
  let pfx :: [Char]
pfx = [Char]
dir [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ [Char]
"/h" [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ Int -> [Char]
forall a. Show a => a -> [Char]
show Int
h [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ [Char]
"."
      ffMul :: Int
ffMul = Int
4

  ln1G <- [Char] -> Int -> IO (Vector Double)
loadVector ([Char]
pfx [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ [Char]
"ln1.gamma.f32") Int
nEmb
  ln1B <- loadVector (pfx ++ "ln1.beta.f32") nEmb
  qkvW <- loadMatrix (pfx ++ "attn.qkv.w.f32") nEmb (3 * nEmb)
  _qkvB <- loadVector (pfx ++ "attn.qkv.b.f32") (3 * nEmb)
  projW <- loadMatrix (pfx ++ "attn.proj.w.f32") nEmb nEmb
  _projB <- loadVector (pfx ++ "attn.proj.b.f32") nEmb
  ln2G <- loadVector (pfx ++ "ln2.gamma.f32") nEmb
  ln2B <- loadVector (pfx ++ "ln2.beta.f32") nEmb
  fcW <- loadMatrix (pfx ++ "mlp.fc.w.f32") nEmb (ffMul * nEmb)
  fcB <- loadVector (pfx ++ "mlp.fc.b.f32") (ffMul * nEmb)
  proj2W <- loadMatrix (pfx ++ "mlp.proj.w.f32") (ffMul * nEmb) nEmb
  proj2B <- loadVector (pfx ++ "mlp.proj.b.f32") nEmb

  -- Split fused Q,K,V projection: columns [0:nEmb], [nEmb:2*nEmb], [2*nEmb:3*nEmb]
  let wQ = (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
nEmb, Int
nEmb) Matrix Double
qkvW
      wK = (Int, Int) -> (Int, Int) -> Matrix Double -> Matrix Double
forall a.
Element a =>
(Int, Int) -> (Int, Int) -> Matrix a -> Matrix a
subMatrix (Int
0, Int
nEmb) (Int
nEmb, Int
nEmb) Matrix Double
qkvW
      wV = (Int, Int) -> (Int, Int) -> Matrix Double -> Matrix Double
forall a.
Element a =>
(Int, Int) -> (Int, Int) -> Matrix a -> Matrix a
subMatrix (Int
0, Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
nEmb) (Int
nEmb, Int
nEmb) Matrix Double
qkvW

  pure
    TransformerBlock
      { tbAttnWq = wQ,
        tbAttnWk = wK,
        tbAttnWv = wV,
        tbAttnWo = projW,
        tbAttnLnGamma = ln1G,
        tbAttnLnBeta = ln1B,
        tbFfn =
          FeedForward
            { ffW1 = fcW,
              ffB1 = fcB,
              ffW2 = proj2W,
              ffB2 = proj2B
            },
        tbFfnLnGamma = ln2G,
        tbFfnLnBeta = ln2B
      }

----------------------------------------------------------------------
-- Binary file I/O
----------------------------------------------------------------------

-- | Load a matrix from a raw float32 (little-endian) file.
loadMatrix :: FilePath -> Int -> Int -> IO (Matrix Double)
loadMatrix :: [Char] -> Int -> Int -> IO (Matrix Double)
loadMatrix [Char]
path Int
rows_ Int
cols_ = do
  bs <- [Char] -> IO ByteString
BS.readFile [Char]
path
  let n = Int
rows_ Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
cols_
      expected = Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
4
      floats = Int -> [Double] -> [Double]
forall a. Int -> [a] -> [a]
take Int
n (ByteString -> [Double]
parseFloats ByteString
bs)
  if BS.length bs /= expected
    then error $ path ++ ": expected " ++ show expected ++ " bytes, got " ++ show (BS.length bs)
    else pure $ reshape cols_ $ fromList floats

-- | Load a vector from a raw float32 (little-endian) file.
loadVector :: FilePath -> Int -> IO (Vector Double)
loadVector :: [Char] -> Int -> IO (Vector Double)
loadVector [Char]
path Int
n = do
  bs <- [Char] -> IO ByteString
BS.readFile [Char]
path
  let expected = Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
4
      floats = Int -> [Double] -> [Double]
forall a. Int -> [a] -> [a]
take Int
n (ByteString -> [Double]
parseFloats ByteString
bs)
  if BS.length bs /= expected
    then error $ path ++ ": expected " ++ show expected ++ " bytes, got " ++ show (BS.length bs)
    else pure $ fromList floats

-- | Lazy parse of a ByteString into Double values (from float32 LE).
parseFloats :: ByteString -> [Double]
parseFloats :: ByteString -> [Double]
parseFloats ByteString
bs
  | ByteString -> Int
BS.length ByteString
bs Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
4 = []
  | Bool
otherwise =
      let (ByteString
chunk, ByteString
rest) = Int -> ByteString -> (ByteString, ByteString)
BS.splitAt Int
4 ByteString
bs
       in Float -> Double
forall a b. (Real a, Fractional b) => a -> b
realToFrac (Word32 -> Float
word32ToFloat (ByteString -> Word32
readWord32LE ByteString
chunk)) Double -> [Double] -> [Double]
forall a. a -> [a] -> [a]
: ByteString -> [Double]
parseFloats ByteString
rest

-- | Read a little-endian 32-bit word from 4 bytes.
readWord32LE :: ByteString -> Word32
readWord32LE :: ByteString -> Word32
readWord32LE ByteString
bs =
  Word8 -> Word32
forall a b. (Integral a, Num b) => a -> b
fromIntegral (HasCallStack => ByteString -> Int -> Word8
ByteString -> Int -> Word8
BS.index ByteString
bs Int
0)
    Word32 -> Word32 -> Word32
forall a. Num a => a -> a -> a
+ Word8 -> Word32
forall a b. (Integral a, Num b) => a -> b
fromIntegral (HasCallStack => ByteString -> Int -> Word8
ByteString -> Int -> Word8
BS.index ByteString
bs Int
1) Word32 -> Word32 -> Word32
forall a. Num a => a -> a -> a
* Word32
0x100
    Word32 -> Word32 -> Word32
forall a. Num a => a -> a -> a
+ Word8 -> Word32
forall a b. (Integral a, Num b) => a -> b
fromIntegral (HasCallStack => ByteString -> Int -> Word8
ByteString -> Int -> Word8
BS.index ByteString
bs Int
2) Word32 -> Word32 -> Word32
forall a. Num a => a -> a -> a
* Word32
0x10000
    Word32 -> Word32 -> Word32
forall a. Num a => a -> a -> a
+ Word8 -> Word32
forall a b. (Integral a, Num b) => a -> b
fromIntegral (HasCallStack => ByteString -> Int -> Word8
ByteString -> Int -> Word8
BS.index ByteString
bs Int
3) Word32 -> Word32 -> Word32
forall a. Num a => a -> a -> a
* Word32
0x1000000

-- | Reinterpret a Word32 as Float via pointer cast.
word32ToFloat :: Word32 -> Float
word32ToFloat :: Word32 -> Float
word32ToFloat Word32
w = IO Float -> Float
forall a. IO a -> a
unsafePerformIO (IO Float -> Float) -> IO Float -> Float
forall a b. (a -> b) -> a -> b
$
  (Ptr Word32 -> IO Float) -> IO Float
forall a b. Storable a => (Ptr a -> IO b) -> IO b
alloca ((Ptr Word32 -> IO Float) -> IO Float)
-> (Ptr Word32 -> IO Float) -> IO Float
forall a b. (a -> b) -> a -> b
$ \(Ptr Word32
p :: Ptr Word32) -> do
    Ptr Word32 -> Word32 -> IO ()
forall a. Storable a => Ptr a -> a -> IO ()
poke Ptr Word32
p Word32
w
    Ptr Float -> IO Float
forall a. Storable a => Ptr a -> IO a
peek (Ptr Word32 -> Ptr Float
forall a b. Ptr a -> Ptr b
castPtr Ptr Word32
p :: Ptr Float)

-- | Zero vector (placeholder).
zeroVector :: Int -> Vector Double
zeroVector :: Int -> Vector Double
zeroVector Int
n = [Double] -> Vector Double
forall a. Storable a => [a] -> Vector a
fromList (Int -> Double -> [Double]
forall a. Int -> a -> [a]
replicate Int
n Double
0)