-- | Token sampling and auto-regressive text generation.
module Circuit.LLM.Inference
  ( -- * Sampling
    greedySample,
    temperatureSample,
    topKSample,

    -- * Generation
    generate,

    -- * Utilities
    argmax,
    softmaxV,
    lastRow,
  )
where

import Circuit.LLM.BPE (BPEEncoding (..), BPEModel (..), decodeBPE, encodeBPE)
import Circuit.LLM.GPT (Gpt, GptConfig, forward)
import Data.List (sortOn)
import Data.Ord (Down (..))
import Data.Text (Text)
import Data.Vector.Unboxed qualified as VU
import Numeric.LinearAlgebra
  ( Matrix,
    Vector,
    cmap,
    fromList,
    maxElement,
    maxIndex,
    sumElements,
    toList,
    toRows,
  )
import System.Random (RandomGen, randomR)

----------------------------------------------------------------------
-- Sampling strategies
----------------------------------------------------------------------

-- | Greedy: pick the token with the highest logit.
greedySample :: Vector Double -> Int
greedySample :: Vector Double -> Int
greedySample = Vector Double -> Int
Vector Double -> IndexOf Vector
forall (c :: * -> *) e. Container c e => c e -> IndexOf c
maxIndex

-- | Temperature sampling.
temperatureSample :: (RandomGen g) => Double -> Vector Double -> g -> (Int, g)
temperatureSample :: forall g. RandomGen g => Double -> Vector Double -> g -> (Int, g)
temperatureSample Double
temp Vector Double
logits g
g =
  let scaled :: Vector Double
scaled = (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. Fractional a => a -> a -> a
/ Double
temp) Vector Double
logits
      probs :: Vector Double
probs = Vector Double -> Vector Double
softmaxV Vector Double
scaled
      (Double
r, g
g') = (Double, Double) -> g -> (Double, g)
forall g. RandomGen g => (Double, Double) -> g -> (Double, g)
forall a g. (Random a, RandomGen g) => (a, a) -> g -> (a, g)
randomR (Double
0, Double
1) g
g
   in ([Double] -> Double -> Int
sampleCategorical (Vector Double -> [Double]
forall a. Storable a => Vector a -> [a]
toList Vector Double
probs) Double
r, g
g')

-- | Top-K sampling.
topKSample :: (RandomGen g) => Int -> Double -> Vector Double -> g -> (Int, g)
topKSample :: forall g.
RandomGen g =>
Int -> Double -> Vector Double -> g -> (Int, g)
topKSample Int
k Double
temp Vector Double
logits g
g =
  let vals :: [Double]
vals = Vector Double -> [Double]
forall a. Storable a => Vector a -> [a]
toList Vector Double
logits
      indexed :: [(Int, Double)]
indexed = [Int] -> [Double] -> [(Int, Double)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0 :: Int ..] [Double]
vals
      topk :: [(Int, Double)]
topk = Int -> [(Int, Double)] -> [(Int, Double)]
forall a. Int -> [a] -> [a]
take Int
k (((Int, Double) -> Down Double)
-> [(Int, Double)] -> [(Int, Double)]
forall b a. Ord b => (a -> b) -> [a] -> [a]
sortOn (Double -> Down Double
forall a. a -> Down a
Down (Double -> Down Double)
-> ((Int, Double) -> Double) -> (Int, Double) -> Down Double
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int, Double) -> Double
forall a b. (a, b) -> b
snd) [(Int, Double)]
indexed)
      topIds :: [Int]
topIds = ((Int, Double) -> Int) -> [(Int, Double)] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map (Int, Double) -> Int
forall a b. (a, b) -> a
fst [(Int, Double)]
topk
      topVals :: [Double]
topVals = ((Int, Double) -> Double) -> [(Int, Double)] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map (Int, Double) -> Double
forall a b. (a, b) -> b
snd [(Int, Double)]
topk
      scaled :: [Double]
scaled = (Double -> Double) -> [Double] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map (Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
temp) [Double]
topVals
      probs :: Vector Double
probs = Vector Double -> Vector Double
softmaxV ([Double] -> Vector Double
forall a. Storable a => [a] -> Vector a
fromList [Double]
scaled)
      (Double
r, g
g') = (Double, Double) -> g -> (Double, g)
forall g. RandomGen g => (Double, Double) -> g -> (Double, g)
forall a g. (Random a, RandomGen g) => (a, a) -> g -> (a, g)
randomR (Double
0, Double
1) g
g
   in ([Int]
topIds [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!! [Double] -> Double -> Int
sampleCategorical (Vector Double -> [Double]
forall a. Storable a => Vector a -> [a]
toList Vector Double
probs) Double
r, g
g')

----------------------------------------------------------------------
-- Auto-regressive generation
----------------------------------------------------------------------

-- | Auto-regressive text generation using greedy sampling.
generate ::
  GptConfig -> Gpt -> BPEModel -> Text -> Int -> IO Text
generate :: GptConfig -> Gpt -> BPEModel -> Text -> Int -> IO Text
generate GptConfig
cfg Gpt
model BPEModel
bpe Text
prompt Int
maxNewTokens = do
  let enc :: BPEEncoding
enc = BPEModel -> Text -> BPEEncoding
encodeBPE BPEModel
bpe Text
prompt
      tokens :: [Word32]
tokens = Vector Word32 -> [Word32]
forall a. Unbox a => Vector a -> [a]
VU.toList (BPEEncoding -> Vector Word32
encodedTokens BPEEncoding
enc)
      initialIds :: [Int]
initialIds = (Word32 -> Int) -> [Word32] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map Word32 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral [Word32]
tokens -- Word32 -> Int
  resultIds <- [Int] -> Int -> IO [Int]
go [Int]
initialIds Int
maxNewTokens
  let resultEnc = BPEEncoding
enc {encodedTokens = VU.fromList (map fromIntegral resultIds)}
  pure $ decodeBPE bpe (encodedTokens resultEnc)
  where
    go :: [Int] -> Int -> IO [Int]
go [Int]
toks Int
0 = [Int] -> IO [Int]
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [Int]
toks
    go [Int]
toks Int
n = do
      let logits :: Matrix Double
logits = GptConfig -> Gpt -> [Int] -> Matrix Double
forward GptConfig
cfg Gpt
model [Int]
toks
          nextLogits :: Vector Double
nextLogits = Matrix Double -> Vector Double
lastRow Matrix Double
logits
          nextToken :: Int
nextToken = Vector Double -> Int
greedySample Vector Double
nextLogits
      [Int] -> Int -> IO [Int]
go ([Int]
toks [Int] -> [Int] -> [Int]
forall a. [a] -> [a] -> [a]
++ [Int
nextToken]) (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)

----------------------------------------------------------------------
-- Utilities
----------------------------------------------------------------------

-- | Softmax over a vector (numerically stable).
softmaxV :: Vector Double -> Vector Double
softmaxV :: Vector Double -> Vector Double
softmaxV Vector Double
v =
  let mx :: Double
mx = Vector Double -> Double
forall (c :: * -> *) e. Container c e => c e -> e
maxElement Vector Double
v
      shifted :: Vector Double
shifted = (Double -> Double) -> Vector Double -> Vector Double
forall b (c :: * -> *) e.
(Element b, Container c e) =>
(e -> b) -> c e -> c b
cmap (\Double
x -> Double -> Double
forall a. Floating a => a -> a
exp (Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
mx)) Vector Double
v
      s :: Double
s = Vector Double -> Double
forall (c :: * -> *) e. Container c e => c e -> e
sumElements Vector Double
shifted
   in (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. Fractional a => a -> a -> a
/ Double
s) Vector Double
shifted

-- | Argmax of a vector.
argmax :: Vector Double -> Int
argmax :: Vector Double -> Int
argmax = Vector Double -> Int
Vector Double -> IndexOf Vector
forall (c :: * -> *) e. Container c e => c e -> IndexOf c
maxIndex

-- | Extract the last row of a matrix.
lastRow :: Matrix Double -> Vector Double
lastRow :: Matrix Double -> Vector Double
lastRow = [Vector Double] -> Vector Double
forall a. HasCallStack => [a] -> a
last ([Vector Double] -> Vector Double)
-> (Matrix Double -> [Vector Double])
-> Matrix Double
-> Vector Double
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Matrix Double -> [Vector Double]
forall t. Element t => Matrix t -> [Vector t]
toRows

-- | Sample from a categorical distribution.
sampleCategorical :: [Double] -> Double -> Int
sampleCategorical :: [Double] -> Double -> Int
sampleCategorical [Double]
probs Double
r = Int -> Double -> [Double] -> Int
forall {t} {t}. (Ord t, Num t, Num t) => t -> t -> [t] -> t
go Int
0 Double
r [Double]
probs
  where
    go :: t -> t -> [t] -> t
go t
_ t
_ [] = t
0
    go t
i t
acc (t
p : [t]
ps)
      | t
acc t -> t -> Bool
forall a. Ord a => a -> a -> Bool
< t
p = t
i
      | Bool
otherwise = t -> t -> [t] -> t
go (t
i t -> t -> t
forall a. Num a => a -> a -> a
+ t
1) (t
acc t -> t -> t
forall a. Num a => a -> a -> a
- t
p) [t]
ps