-- | This module provides convenience functions built from the core of the RASP-L language.
--
-- It is based on Listing 3 of
-- "What Algorithms Can Transformers Learn", https://arxiv.org/abs/2310.16028,
-- by Zhou et al.
module RaskellLib
  ( -- * Logical Operations
    (?),
    shiftRight,
    toBool,
    mask,

    -- * Running Aggregations
    cumSum,
    maximum',
    minimum',
    argmax,
    argmin,

    -- * Aggregations with `Queries`
    numPrev,
    hasSeen,

    -- * Indexing with `Queries`
    firsts,
    lasts,
    indexSelect,

    -- * Token Comparisons
    leq,
    geq,
    lt,
    gt,

    -- * Sampling
    sample,

    -- * Compatibility
    where',
    sampleAutoregressive,
  )
where

import Data.Int (Int8)
import Data.Word (Word8)
import RaskellCore

-- | Use a boolean sequence to select between two sequences.
-- Also known in Python RASP-L as "where", see `where'`.
(?) :: [Bool] -> (Sequence, Sequence) -> Sequence
[Bool]
bs ? :: [Bool] -> (Sequence, Sequence) -> Sequence
? (Sequence
xs, Sequence
ys) = (Token -> Token -> Token) -> Sequence -> Sequence -> Sequence
seqMap (\Token
xm Token
ym -> if Token
xm Token -> Token -> Bool
forall a. Eq a => a -> a -> Bool
== Token
0 then Token
ym else Token
xm) Sequence
xms Sequence
yms
  where
    xms :: Sequence
xms = (Token -> Token -> Token) -> Sequence -> Sequence -> Sequence
seqMap (\Token
bt Token
x -> if Token
bt Token -> Token -> Bool
forall a. Eq a => a -> a -> Bool
== Token
1 then Token
x else Token
0) Sequence
bts Sequence
xs
    yms :: Sequence
yms = (Token -> Token -> Token) -> Sequence -> Sequence -> Sequence
seqMap (\Token
bt Token
y -> if Token
bt Token -> Token -> Bool
forall a. Eq a => a -> a -> Bool
== Token
0 then Token
y else Token
0) Sequence
bts Sequence
ys
    bts :: Sequence
bts = [Bool] -> Sequence
fromBoolSeq [Bool]
bs

-- | Use a boolean sequence to select between two sequences.
-- Provided for compatibility with Listing 3, but with
-- an apostrophe to avoid a name clash with the "where" keyword.
where' :: [Bool] -> Sequence -> Sequence -> Sequence
where' :: [Bool] -> Sequence -> Sequence -> Sequence
where' [Bool]
bs Sequence
xs Sequence
ys = [Bool]
bs [Bool] -> (Sequence, Sequence) -> Sequence
? (Sequence
xs, Sequence
ys)

-- | Shift a sequence to the right by a given number of elements,
-- filling the vacated positions with the provided `Token`.
shiftRight ::
  -- | Filler `Token`
  Token ->
  -- | Number of positions to shift
  Int8 ->
  -- | Input `Sequence`
  Sequence ->
  Sequence
shiftRight :: Token -> Token -> Sequence -> Sequence
shiftRight Token
filler Token
n Sequence
xs = Token
-> AggregationType
-> Sequence
-> Sequence
-> (Token -> Token -> Bool)
-> Sequence
-> Sequence
kqv Token
filler AggregationType
Mean Sequence
shiftedIdxs Sequence
idxs Token -> Token -> Bool
forall a. Eq a => a -> a -> Bool
(==) Sequence
xs
  where
    shiftedIdxs :: Sequence
shiftedIdxs = (Token -> Token) -> Sequence -> Sequence
forall a b. (a -> b) -> [a] -> [b]
map (Token -> Token -> Token
forall a. Num a => a -> a -> a
+ Token
n) Sequence
idxs
    idxs :: Sequence
idxs = Sequence -> Sequence
indices Sequence
xs

-- | Maps tokens onto bools using Python's "truthiness" rules.
toBool :: Token -> Bool
toBool :: Token -> Bool
toBool Token
x
  | Token
x Token -> Token -> Bool
forall a. Eq a => a -> a -> Bool
== Token
0 = Bool
False
  | Bool
otherwise = Bool
True

-- | Converts a list of bools to a sequence of tokens.
fromBoolSeq :: [Bool] -> Sequence
fromBoolSeq :: [Bool] -> Sequence
fromBoolSeq = (Bool -> Token) -> [Bool] -> Sequence
forall a b. (a -> b) -> [a] -> [b]
map Bool -> Token
fromBool

-- | Computes the cumulative sum of a boolean sequence.
cumSum :: [Bool] -> Sequence
cumSum :: [Bool] -> Sequence
cumSum [Bool]
bs = Selector -> Sequence
selWidth (Sequence -> Sequence -> (Token -> Token -> Bool) -> Selector
selectCausal Sequence
bTokens Sequence
bTokens Token -> Token -> Bool
forall {p}. Token -> p -> Bool
first)
  where
    bTokens :: Sequence
bTokens = [Bool] -> Sequence
fromBoolSeq [Bool]
bs
    first :: Token -> p -> Bool
first Token
x p
_ = Token -> Bool
toBool Token
x

-- | Masks a `Sequence` with a boolean sequence, using the provided `Token` as the mask.
mask :: Token -> [Bool] -> Sequence -> Sequence
mask :: Token -> [Bool] -> Sequence -> Sequence
mask Token
maskT [Bool]
bs Sequence
xs = [Bool]
bs [Bool] -> (Sequence, Sequence) -> Sequence
? (Sequence
xs, Sequence
xs Sequence -> Token -> Sequence
`filledWith` Token
maskT)

-- | Computes the running maximum of a `Sequence`.
maximum' :: Sequence -> Sequence
maximum' :: Sequence -> Sequence
maximum' Sequence
xs = Sequence
-> Sequence -> (Token -> Token -> Bool) -> Sequence -> Sequence
maxKQV Sequence
xs Sequence
xs Token -> Token -> Bool
forall {p} {p}. p -> p -> Bool
always Sequence
xs
  where
    always :: p -> p -> Bool
always p
_ p
_ = Bool
True

-- | Computes the running minimum of a `Sequence`.
minimum' :: Sequence -> Sequence
minimum' :: Sequence -> Sequence
minimum' Sequence
xs = Sequence
-> Sequence -> (Token -> Token -> Bool) -> Sequence -> Sequence
minKQV Sequence
xs Sequence
xs Token -> Token -> Bool
forall {p} {p}. p -> p -> Bool
always Sequence
xs
  where
    always :: p -> p -> Bool
always p
_ p
_ = Bool
True

-- | Computes the indices of the running maximum values in a `Sequence`.
argmax :: Sequence -> Sequence
argmax :: Sequence -> Sequence
argmax Sequence
xs = Sequence
-> Sequence -> (Token -> Token -> Bool) -> Sequence -> Sequence
maxKQV Sequence
xs Sequence
maxs Token -> Token -> Bool
forall a. Eq a => a -> a -> Bool
(==) (Sequence -> Sequence
indicesOf Sequence
xs)
  where
    maxs :: Sequence
maxs = Sequence -> Sequence
maximum' Sequence
xs

-- | Computes the indices of the running minimum values in a `Sequence`.
argmin :: Sequence -> Sequence
argmin :: Sequence -> Sequence
argmin Sequence
xs = Sequence
-> Sequence -> (Token -> Token -> Bool) -> Sequence -> Sequence
maxKQV Sequence
xs Sequence
mins Token -> Token -> Bool
forall a. Eq a => a -> a -> Bool
(==) (Sequence -> Sequence
indicesOf Sequence
xs)
  where
    mins :: Sequence
mins = Sequence -> Sequence
minimum' Sequence
xs

-- | Computes the number of previous tokens in a `Sequence` that are equal to each `Token` from `Queries`.
numPrev :: Sequence -> Queries -> Sequence
numPrev :: Sequence -> Sequence -> Sequence
numPrev Sequence
xs Sequence
queries = Selector -> Sequence
selWidth (Sequence -> Sequence -> (Token -> Token -> Bool) -> Selector
selectCausal Sequence
xs Sequence
queries Token -> Token -> Bool
forall a. Eq a => a -> a -> Bool
(==))

-- | Returns 1s where the `Token` from the `Queries` has been seen before in the `Sequence`.
hasSeen :: Sequence -> Queries -> Sequence
hasSeen :: Sequence -> Sequence -> Sequence
hasSeen Sequence
xs Sequence
queries = Token
-> AggregationType
-> Sequence
-> Sequence
-> (Token -> Token -> Bool)
-> Sequence
-> Sequence
kqv Token
0 AggregationType
Max Sequence
xs Sequence
queries Token -> Token -> Bool
forall a. Eq a => a -> a -> Bool
(==) (Sequence
queries Sequence -> Token -> Sequence
`filledWith` Token
1)

-- | Finds the first occurrence of each query token in a `Sequence`.
firsts :: Token -> Sequence -> Queries -> Sequence
firsts :: Token -> Sequence -> Sequence -> Sequence
firsts Token
filler Sequence
xs Sequence
queries = Token
-> AggregationType
-> Sequence
-> Sequence
-> (Token -> Token -> Bool)
-> Sequence
-> Sequence
kqv Token
filler AggregationType
Min Sequence
xs Sequence
queries Token -> Token -> Bool
forall a. Eq a => a -> a -> Bool
(==) (Sequence -> Sequence
indicesOf Sequence
xs)

-- | Finds the last occurrence of each query token in a `Sequence`.
lasts :: Token -> Sequence -> Queries -> Sequence
lasts :: Token -> Sequence -> Sequence -> Sequence
lasts Token
filler Sequence
xs Sequence
queries = Token
-> AggregationType
-> Sequence
-> Sequence
-> (Token -> Token -> Bool)
-> Sequence
-> Sequence
kqv Token
filler AggregationType
Max Sequence
xs Sequence
queries Token -> Token -> Bool
forall a. Eq a => a -> a -> Bool
(==) (Sequence -> Sequence
indicesOf Sequence
xs)

-- | Selects the tokens from a `Sequence` at the indices provided by another sequence.
indexSelect :: Token -> Sequence -> Sequence -> Sequence
indexSelect :: Token -> Sequence -> Sequence -> Sequence
indexSelect Token
filler Sequence
xs Sequence
idxs = Token
-> AggregationType
-> Sequence
-> Sequence
-> (Token -> Token -> Bool)
-> Sequence
-> Sequence
kqv Token
filler AggregationType
Max (Sequence -> Sequence
indicesOf Sequence
xs) Sequence
idxs Token -> Token -> Bool
forall a. Eq a => a -> a -> Bool
(==) Sequence
xs

leq :: Token -> Token -> Bool
leq :: Token -> Token -> Bool
leq = Token -> Token -> Bool
forall a. Ord a => a -> a -> Bool
(<=)

geq :: Token -> Token -> Bool
geq :: Token -> Token -> Bool
geq = Token -> Token -> Bool
forall a. Ord a => a -> a -> Bool
(>=)

lt :: Token -> Token -> Bool
lt :: Token -> Token -> Bool
lt = Token -> Token -> Bool
forall a. Ord a => a -> a -> Bool
(<)

gt :: Token -> Token -> Bool
gt :: Token -> Token -> Bool
gt = Token -> Token -> Bool
forall a. Ord a => a -> a -> Bool
(>)

-- | Greedily and autoregressively sample the output of a RASP-L program on a sequence.
sample ::
  -- | End of sequence token
  Token ->
  -- | RASP-L program to extend the sequence
  (Sequence -> Sequence) ->
  -- | Initial/prompt sequence
  Sequence ->
  -- | Number of steps to decode
  Word8 ->
  -- | Output (including prompt)
  Sequence
sample :: Token -> (Sequence -> Sequence) -> Sequence -> Word8 -> Sequence
sample Token
_ Sequence -> Sequence
_ Sequence
xs Word8
0 = Sequence
xs
sample Token
endOfSequence Sequence -> Sequence
prog Sequence
xs Word8
n
  | Sequence -> Token
forall a. HasCallStack => [a] -> a
last Sequence
xs Token -> Token -> Bool
forall a. Eq a => a -> a -> Bool
== Token
endOfSequence = Sequence
xs
  | Bool
otherwise = Token -> (Sequence -> Sequence) -> Sequence -> Word8 -> Sequence
sample Token
endOfSequence Sequence -> Sequence
prog (Sequence
xs Sequence -> Sequence -> Sequence
forall a. [a] -> [a] -> [a]
++ [Sequence -> Token
forall a. HasCallStack => [a] -> a
last (Sequence -> Token) -> Sequence -> Token
forall a b. (a -> b) -> a -> b
$ Sequence -> Sequence
prog Sequence
xs]) (Word8
n Word8 -> Word8 -> Word8
forall a. Num a => a -> a -> a
- Word8
1)

-- | Greedily and autoregressively sample the output of a RASP-L program on a sequence.
--
-- Provided for compatibility with Listing 3.
sampleAutoregressive :: Token -> (Sequence -> Sequence) -> Sequence -> Word8 -> Sequence
sampleAutoregressive :: Token -> (Sequence -> Sequence) -> Sequence -> Word8 -> Sequence
sampleAutoregressive = Token -> (Sequence -> Sequence) -> Sequence -> Word8 -> Sequence
sample