{-# LANGUAGE OverloadedStrings #-}

module Circuit.LLM.BPE
  ( -- * Data Types
    BPEModel (..),
    BPEEncoding (..),
    BPEError (..),

    -- * Model Loading
    loadBPEModel,
    loadBPEModelWithPerf,

    -- * Encoding & Decoding
    encodeBPE,
    encodeBPEWithPerf,
    decodeBPE,
    decodeBPEWithPerf,

    -- * Display Functions
    prettifyBPEModel,
    prettifyEncoding,
  )
where

import Control.Exception (Exception, throwIO)
import Control.Monad (when)
import Data.ByteString (ByteString)
import Data.ByteString qualified as BS
import Data.Char (isAlpha, isDigit, isSpace)
import Data.IntMap.Strict (IntMap)
import Data.IntMap.Strict qualified as IntMap
import Data.List (minimumBy)
import Data.Map.Strict (Map)
import Data.Map.Strict qualified as Map
import Data.Ord (comparing)
import Data.Text (Text)
import Data.Text qualified as Text
import Data.Text.Encoding qualified as TE
import Data.Text.Encoding.Error qualified as TEE
import Data.Text.IO qualified as TIO
import Data.Vector.Unboxed (Vector)
import Data.Vector.Unboxed qualified as V
import Data.Word (Word32)

type Nanos = Integer

-- | BPE model loaded from .model file
-- Stores merge rules, vocabulary, and special token mappings for encoding/decoding
data BPEModel = BPEModel
  { -- | Version string (e.g., "simple-bpe v1")
    BPEModel -> Text
bpeVersion :: !Text,
    -- | Regex pattern for text splitting
    BPEModel -> ByteString
bpeRegex :: !ByteString,
    -- | Special tokens → IDs
    BPEModel -> Map Text Word32
bpeSpecialTokens :: !(Map Text Word32),
    -- | Special token IDs → tokens
    BPEModel -> IntMap Text
bpeReverseSpecial :: !(IntMap Text),
    -- | Pair → (new token, priority)
    BPEModel -> Map (Word32, Word32) (Word32, Int)
bpeMergeRules :: !(Map (Word32, Word32) (Word32, Int)),
    -- | Token ID → bytes (for decoding)
    BPEModel -> IntMap ByteString
bpeVocab :: !(IntMap ByteString),
    -- | Highest token ID in vocabulary
    BPEModel -> Word32
bpeMaxTokenId :: !Word32
  }
  deriving (BPEModel -> BPEModel -> Bool
(BPEModel -> BPEModel -> Bool)
-> (BPEModel -> BPEModel -> Bool) -> Eq BPEModel
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: BPEModel -> BPEModel -> Bool
== :: BPEModel -> BPEModel -> Bool
$c/= :: BPEModel -> BPEModel -> Bool
/= :: BPEModel -> BPEModel -> Bool
Eq, Int -> BPEModel -> ShowS
[BPEModel] -> ShowS
BPEModel -> [Char]
(Int -> BPEModel -> ShowS)
-> (BPEModel -> [Char]) -> ([BPEModel] -> ShowS) -> Show BPEModel
forall a.
(Int -> a -> ShowS) -> (a -> [Char]) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> BPEModel -> ShowS
showsPrec :: Int -> BPEModel -> ShowS
$cshow :: BPEModel -> [Char]
show :: BPEModel -> [Char]
$cshowList :: [BPEModel] -> ShowS
showList :: [BPEModel] -> ShowS
Show)

-- | Encoding result with metadata
data BPEEncoding = BPEEncoding
  { -- | Encoded token IDs
    BPEEncoding -> Vector Word32
encodedTokens :: !(Vector Word32),
    -- | Original input text
    BPEEncoding -> Text
originalText :: !Text,
    -- | Number of regex chunks processed
    BPEEncoding -> Int
numChunks :: !Int
  }
  deriving (BPEEncoding -> BPEEncoding -> Bool
(BPEEncoding -> BPEEncoding -> Bool)
-> (BPEEncoding -> BPEEncoding -> Bool) -> Eq BPEEncoding
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: BPEEncoding -> BPEEncoding -> Bool
== :: BPEEncoding -> BPEEncoding -> Bool
$c/= :: BPEEncoding -> BPEEncoding -> Bool
/= :: BPEEncoding -> BPEEncoding -> Bool
Eq, Int -> BPEEncoding -> ShowS
[BPEEncoding] -> ShowS
BPEEncoding -> [Char]
(Int -> BPEEncoding -> ShowS)
-> (BPEEncoding -> [Char])
-> ([BPEEncoding] -> ShowS)
-> Show BPEEncoding
forall a.
(Int -> a -> ShowS) -> (a -> [Char]) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> BPEEncoding -> ShowS
showsPrec :: Int -> BPEEncoding -> ShowS
$cshow :: BPEEncoding -> [Char]
show :: BPEEncoding -> [Char]
$cshowList :: [BPEEncoding] -> ShowS
showList :: [BPEEncoding] -> ShowS
Show)

-- | BPE operation errors
data BPEError
  = -- | File path and error message
    ModelParseError !FilePath !String
  | -- | Invalid token ID during decode
    InvalidTokenId !Word32
  | -- | Regex pattern compilation error
    RegexCompileError !String
  deriving (Int -> BPEError -> ShowS
[BPEError] -> ShowS
BPEError -> [Char]
(Int -> BPEError -> ShowS)
-> (BPEError -> [Char]) -> ([BPEError] -> ShowS) -> Show BPEError
forall a.
(Int -> a -> ShowS) -> (a -> [Char]) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> BPEError -> ShowS
showsPrec :: Int -> BPEError -> ShowS
$cshow :: BPEError -> [Char]
show :: BPEError -> [Char]
$cshowList :: [BPEError] -> ShowS
showList :: [BPEError] -> ShowS
Show, BPEError -> BPEError -> Bool
(BPEError -> BPEError -> Bool)
-> (BPEError -> BPEError -> Bool) -> Eq BPEError
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: BPEError -> BPEError -> Bool
== :: BPEError -> BPEError -> Bool
$c/= :: BPEError -> BPEError -> Bool
/= :: BPEError -> BPEError -> Bool
Eq)

instance Exception BPEError

-- | Load BPE model from .model file (Rust format)
--
-- File format:
-- Line 1: Version string ("simple-bpe v1")
-- Line 2: Regex pattern for text splitting
-- Line 3: Number of special tokens (integer)
-- Next N lines: Special token and its ID (e.g., "<|endoftext|> 256")
-- Remaining lines: Merge pairs - two token IDs per line (e.g., "65 66")
loadBPEModel :: FilePath -> IO BPEModel
loadBPEModel :: [Char] -> IO BPEModel
loadBPEModel [Char]
fp = do
  content <- [Char] -> IO Text
TIO.readFile [Char]
fp
  parseModelFile fp content

-- | Load BPE model with performance measurement
--
-- TODO: implement proper perf measurement with pure functions
loadBPEModelWithPerf :: FilePath -> IO (BPEModel, Map Text [Nanos])
loadBPEModelWithPerf :: [Char] -> IO (BPEModel, Map Text [Nanos])
loadBPEModelWithPerf [Char]
fp = do
  model <- [Char] -> IO BPEModel
loadBPEModel [Char]
fp
  let timings = Map k a
forall k a. Map k a
Map.empty
  pure (model, timings)

-- | Parse model file content into BPEModel structure
parseModelFile :: FilePath -> Text -> IO BPEModel
parseModelFile :: [Char] -> Text -> IO BPEModel
parseModelFile [Char]
fp Text
content = do
  let lns :: [Text]
lns = Text -> [Text]
Text.lines Text
content

  -- Validate minimum number of lines (version + regex + special token count)
  Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when ([Text] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Text]
lns Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
3) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$
    BPEError -> IO ()
forall e a. (HasCallStack, Exception e) => e -> IO a
throwIO (BPEError -> IO ()) -> BPEError -> IO ()
forall a b. (a -> b) -> a -> b
$
      [Char] -> [Char] -> BPEError
ModelParseError [Char]
fp [Char]
"File too short: missing header"

  case [Text]
lns of
    (Text
version : Text
regexLine : Text
numSpecialLine : [Text]
_) -> do
      let regexPattern :: ByteString
regexPattern = Text -> ByteString
TE.encodeUtf8 Text
regexLine
          numSpecialStr :: Text
numSpecialStr = Text
numSpecialLine

      -- Parse number of special tokens
      numSpecial <- case ReadS Int
forall a. Read a => ReadS a
reads (Text -> [Char]
Text.unpack Text
numSpecialStr) of
        [(Int
n, [Char]
"")] -> Int -> IO Int
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
n
        [(Int, [Char])]
_ -> BPEError -> IO Int
forall e a. (HasCallStack, Exception e) => e -> IO a
throwIO (BPEError -> IO Int) -> BPEError -> IO Int
forall a b. (a -> b) -> a -> b
$ [Char] -> [Char] -> BPEError
ModelParseError [Char]
fp ([Char]
"Invalid special token count: " [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ Text -> [Char]
Text.unpack Text
numSpecialStr)

      -- Parse special tokens
      let specialLines = Int -> [Text] -> [Text]
forall a. Int -> [a] -> [a]
take Int
numSpecial (Int -> [Text] -> [Text]
forall a. Int -> [a] -> [a]
drop Int
3 [Text]
lns)
      specialTokens <- parseSpecialTokens fp specialLines
      let reverseSpecial = [(Int, Text)] -> IntMap Text
forall a. [(Int, a)] -> IntMap a
IntMap.fromList [(Word32 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word32
tid, Text
tok) | (Text
tok, Word32
tid) <- Map Text Word32 -> [(Text, Word32)]
forall k a. Map k a -> [(k, a)]
Map.toList Map Text Word32
specialTokens]

      -- Parse merge rules
      let mergeLines = Int -> [Text] -> [Text]
forall a. Int -> [a] -> [a]
drop (Int
3 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
numSpecial) [Text]
lns
      (mergeRules, maxMergeId) <- parseMergeRules fp mergeLines numSpecial

      -- Build vocabulary: base bytes (0-255) + special tokens + merge results
      let vocab = Map Text Word32
-> Map (Word32, Word32) (Word32, Int)
-> Word32
-> IntMap ByteString
buildVocab Map Text Word32
specialTokens Map (Word32, Word32) (Word32, Int)
mergeRules Word32
maxMergeId

      pure $
        BPEModel
          { bpeVersion = version,
            bpeRegex = regexPattern,
            bpeSpecialTokens = specialTokens,
            bpeReverseSpecial = reverseSpecial,
            bpeMergeRules = mergeRules,
            bpeVocab = vocab,
            bpeMaxTokenId = maxMergeId
          }
    [Text]
_ -> BPEError -> IO BPEModel
forall e a. (HasCallStack, Exception e) => e -> IO a
throwIO (BPEError -> IO BPEModel) -> BPEError -> IO BPEModel
forall a b. (a -> b) -> a -> b
$ [Char] -> [Char] -> BPEError
ModelParseError [Char]
fp [Char]
"File too short: missing header"

-- | Parse special tokens from lines like "<|endoftext|> 256"
parseSpecialTokens :: FilePath -> [Text] -> IO (Map Text Word32)
parseSpecialTokens :: [Char] -> [Text] -> IO (Map Text Word32)
parseSpecialTokens [Char]
fp [Text]
lns = do
  let parseToken :: Text -> IO (Text, Word32)
parseToken Text
line = case Text -> [Text]
Text.words Text
line of
        [Text
tok, Text
idStr] -> case ReadS Word32
forall a. Read a => ReadS a
reads (Text -> [Char]
Text.unpack Text
idStr) of
          [(Word32
tid, [Char]
"")] -> (Text, Word32) -> IO (Text, Word32)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Text
tok, Word32
tid)
          [(Word32, [Char])]
_ -> BPEError -> IO (Text, Word32)
forall e a. (HasCallStack, Exception e) => e -> IO a
throwIO (BPEError -> IO (Text, Word32)) -> BPEError -> IO (Text, Word32)
forall a b. (a -> b) -> a -> b
$ [Char] -> [Char] -> BPEError
ModelParseError [Char]
fp ([Char]
"Invalid token ID: " [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ Text -> [Char]
Text.unpack Text
line)
        [Text]
_ -> BPEError -> IO (Text, Word32)
forall e a. (HasCallStack, Exception e) => e -> IO a
throwIO (BPEError -> IO (Text, Word32)) -> BPEError -> IO (Text, Word32)
forall a b. (a -> b) -> a -> b
$ [Char] -> [Char] -> BPEError
ModelParseError [Char]
fp ([Char]
"Invalid special token line: " [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ Text -> [Char]
Text.unpack Text
line)
  tokens <- (Text -> IO (Text, Word32)) -> [Text] -> IO [(Text, Word32)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM Text -> IO (Text, Word32)
parseToken [Text]
lns
  pure $ Map.fromList tokens

-- | Parse merge rules from lines like "65 66"
-- Returns (merge rules map, max token ID).
-- Python: idx starts at 256, increments per merge line.
parseMergeRules :: FilePath -> [Text] -> Int -> IO (Map (Word32, Word32) (Word32, Int), Word32)
parseMergeRules :: [Char]
-> [Text] -> Int -> IO (Map (Word32, Word32) (Word32, Int), Word32)
parseMergeRules [Char]
fp [Text]
lns Int
_specialCount = do
  let startIdx :: Word32
startIdx = Word32
256 :: Word32 -- Python starts merge tokens at 256
      parseRule :: (Int, Text) -> IO ((Word32, Word32), (Word32, Int))
parseRule (Int
idx, Text
line) = case Text -> [Text]
Text.words Text
line of
        [Text
id1Str, Text
id2Str] -> case (ReadS Word32
forall a. Read a => ReadS a
reads (Text -> [Char]
Text.unpack Text
id1Str), ReadS Word32
forall a. Read a => ReadS a
reads (Text -> [Char]
Text.unpack Text
id2Str)) of
          ([(Word32
id1, [Char]
"")], [(Word32
id2, [Char]
"")]) ->
            let newTokenId :: Word32
newTokenId = Word32
startIdx Word32 -> Word32 -> Word32
forall a. Num a => a -> a -> a
+ Int -> Word32
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
idx
             in ((Word32, Word32), (Word32, Int))
-> IO ((Word32, Word32), (Word32, Int))
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((Word32
id1, Word32
id2), (Word32
newTokenId, Int
idx))
          ([(Word32, [Char])], [(Word32, [Char])])
_ -> BPEError -> IO ((Word32, Word32), (Word32, Int))
forall e a. (HasCallStack, Exception e) => e -> IO a
throwIO (BPEError -> IO ((Word32, Word32), (Word32, Int)))
-> BPEError -> IO ((Word32, Word32), (Word32, Int))
forall a b. (a -> b) -> a -> b
$ [Char] -> [Char] -> BPEError
ModelParseError [Char]
fp ([Char]
"Invalid merge rule: " [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ Text -> [Char]
Text.unpack Text
line)
        [Text]
_ -> BPEError -> IO ((Word32, Word32), (Word32, Int))
forall e a. (HasCallStack, Exception e) => e -> IO a
throwIO (BPEError -> IO ((Word32, Word32), (Word32, Int)))
-> BPEError -> IO ((Word32, Word32), (Word32, Int))
forall a b. (a -> b) -> a -> b
$ [Char] -> [Char] -> BPEError
ModelParseError [Char]
fp ([Char]
"Invalid merge rule format: " [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ Text -> [Char]
Text.unpack Text
line)

  rules <- ((Int, Text) -> IO ((Word32, Word32), (Word32, Int)))
-> [(Int, Text)] -> IO [((Word32, Word32), (Word32, Int))]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (Int, Text) -> IO ((Word32, Word32), (Word32, Int))
parseRule ([Int] -> [Text] -> [(Int, Text)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0 ..] [Text]
lns)
  let mergeMap = [((Word32, Word32), (Word32, Int))]
-> Map (Word32, Word32) (Word32, Int)
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList [((Word32, Word32), (Word32, Int))]
rules
      maxId =
        if [((Word32, Word32), (Word32, Int))] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [((Word32, Word32), (Word32, Int))]
rules
          then Word32
255
          else [Word32] -> Word32
forall a. Ord a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
maximum [Word32
tid | ((Word32, Word32)
_, (Word32
tid, Int
_)) <- [((Word32, Word32), (Word32, Int))]
rules]
  pure (mergeMap, maxId)

-- | Build vocabulary from base bytes, special tokens, and merge rules
-- Uses lazy approach: only base bytes (0-255) + special tokens initially
-- Merged token bytes computed on-demand during decode
buildVocab :: Map Text Word32 -> Map (Word32, Word32) (Word32, Int) -> Word32 -> IntMap ByteString
buildVocab :: Map Text Word32
-> Map (Word32, Word32) (Word32, Int)
-> Word32
-> IntMap ByteString
buildVocab Map Text Word32
specialTokens Map (Word32, Word32) (Word32, Int)
_mergeRules Word32
_maxId =
  let -- Base vocabulary: bytes 0-255
      baseVocab :: IntMap ByteString
baseVocab = [(Int, ByteString)] -> IntMap ByteString
forall a. [(Int, a)] -> IntMap a
IntMap.fromList [(Int
i, Word8 -> ByteString
BS.singleton (Int -> Word8
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
i)) | Int
i <- [Int
0 .. Int
255]]

      -- Special tokens
      specialVocab :: IntMap ByteString
specialVocab =
        [(Int, ByteString)] -> IntMap ByteString
forall a. [(Int, a)] -> IntMap a
IntMap.fromList
          [(Word32 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word32
tid, Text -> ByteString
TE.encodeUtf8 Text
tok) | (Text
tok, Word32
tid) <- Map Text Word32 -> [(Text, Word32)]
forall k a. Map k a -> [(k, a)]
Map.toList Map Text Word32
specialTokens]
   in -- Combined: base bytes + special tokens
      -- Merged tokens will be computed lazily during decode
      IntMap ByteString -> IntMap ByteString -> IntMap ByteString
forall a. IntMap a -> IntMap a -> IntMap a
IntMap.union IntMap ByteString
specialVocab IntMap ByteString
baseVocab

-- | Encode text using BPE model
--
-- Algorithm:
-- 1. Split text by regex pattern into chunks
-- 2. For each chunk:
--    - Check if it's a special token → encode directly
--    - Otherwise: convert to bytes → apply BPE merges
-- 3. Concatenate all results
encodeBPE :: BPEModel -> Text -> BPEEncoding
encodeBPE :: BPEModel -> Text -> BPEEncoding
encodeBPE BPEModel
model Text
text =
  let chunks :: [Text]
chunks = ByteString -> Text -> [Text]
splitByRegex (BPEModel -> ByteString
bpeRegex BPEModel
model) Text
text
      encodedChunks :: [Vector Word32]
encodedChunks = (Text -> Vector Word32) -> [Text] -> [Vector Word32]
forall a b. (a -> b) -> [a] -> [b]
map (BPEModel -> Text -> Vector Word32
encodeChunk BPEModel
model) [Text]
chunks
      allTokens :: Vector Word32
allTokens = [Vector Word32] -> Vector Word32
forall a. Unbox a => [Vector a] -> Vector a
V.concat [Vector Word32]
encodedChunks
   in BPEEncoding
        { encodedTokens :: Vector Word32
encodedTokens = Vector Word32
allTokens,
          originalText :: Text
originalText = Text
text,
          numChunks :: Int
numChunks = [Text] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Text]
chunks
        }

-- | Encode with performance measurement
--
-- TODO: implement proper perf measurement with pure functions
encodeBPEWithPerf :: BPEModel -> Text -> IO (BPEEncoding, Map Text [Nanos])
encodeBPEWithPerf :: BPEModel -> Text -> IO (BPEEncoding, Map Text [Nanos])
encodeBPEWithPerf BPEModel
model Text
text = do
  let result :: BPEEncoding
result = BPEModel -> Text -> BPEEncoding
encodeBPE BPEModel
model Text
text
      timings :: Map k a
timings = Map k a
forall k a. Map k a
Map.empty
  (BPEEncoding, Map Text [Nanos])
-> IO (BPEEncoding, Map Text [Nanos])
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (BPEEncoding
result, Map Text [Nanos]
forall k a. Map k a
timings)

-- | Split text using the GPT-2 tokenizer regex pattern.
--
-- Pattern (from openai/tiktoken):
--   '(?:[sdmt]|ll|ve|re) | ?\p{L}+ | ?\p{N}+ | ?[^\s\p{L}\p{N}]+ | \s+(?!\S) | \s+
--
-- Implemented manually since regex-tdfa doesn't support \p{}.
splitByRegex :: ByteString -> Text -> [Text]
splitByRegex :: ByteString -> Text -> [Text]
splitByRegex ByteString
_pattern = Text -> [Text]
gpt2Split
  where
    gpt2Split :: Text -> [Text]
    gpt2Split :: Text -> [Text]
gpt2Split Text
t
      | Text -> Bool
Text.null Text
t = []
      | Bool
otherwise =
          case Text -> Maybe (Text, Text)
matchGpt2Token Text
t of
            Maybe (Text, Text)
Nothing -> Text -> [Text]
gpt2Split (Int -> Text -> Text
Text.drop Int
1 Text
t) -- skip unrecognized char
            Just (Text
tok, Text
rest) -> Text
tok Text -> [Text] -> [Text]
forall a. a -> [a] -> [a]
: Text -> [Text]
gpt2Split Text
rest

    -- Try each alternative in order, return first match
    matchGpt2Token :: Text -> Maybe (Text, Text)
    matchGpt2Token :: Text -> Maybe (Text, Text)
matchGpt2Token Text
t =
      Text -> Maybe (Text, Text)
matchContraction Text
t
        Maybe (Text, Text) -> Maybe (Text, Text) -> Maybe (Text, Text)
forall a. Maybe a -> Maybe a -> Maybe a
<|> Text -> Maybe (Text, Text)
matchOptSpaceLetters Text
t
        Maybe (Text, Text) -> Maybe (Text, Text) -> Maybe (Text, Text)
forall a. Maybe a -> Maybe a -> Maybe a
<|> Text -> Maybe (Text, Text)
matchOptSpaceDigits Text
t
        Maybe (Text, Text) -> Maybe (Text, Text) -> Maybe (Text, Text)
forall a. Maybe a -> Maybe a -> Maybe a
<|> Text -> Maybe (Text, Text)
matchOptSpacePunct Text
t
        Maybe (Text, Text) -> Maybe (Text, Text) -> Maybe (Text, Text)
forall a. Maybe a -> Maybe a -> Maybe a
<|> Text -> Maybe (Text, Text)
matchTrailingSpace Text
t
        Maybe (Text, Text) -> Maybe (Text, Text) -> Maybe (Text, Text)
forall a. Maybe a -> Maybe a -> Maybe a
<|> Text -> Maybe (Text, Text)
matchSpace Text
t

    -- Contractions: '(?:[sdmt]|ll|ve|re)
    matchContraction :: Text -> Maybe (Text, Text)
matchContraction Text
t
      | Int -> Text -> Text
Text.take Int
3 Text
t Text -> Text -> Bool
forall a. Eq a => a -> a -> Bool
== Text
"'ll" = (Text, Text) -> Maybe (Text, Text)
forall a. a -> Maybe a
Just (Int -> Text -> Text
Text.take Int
3 Text
t, Int -> Text -> Text
Text.drop Int
3 Text
t)
      | Int -> Text -> Text
Text.take Int
3 Text
t Text -> Text -> Bool
forall a. Eq a => a -> a -> Bool
== Text
"'ve" = (Text, Text) -> Maybe (Text, Text)
forall a. a -> Maybe a
Just (Int -> Text -> Text
Text.take Int
3 Text
t, Int -> Text -> Text
Text.drop Int
3 Text
t)
      | Int -> Text -> Text
Text.take Int
3 Text
t Text -> Text -> Bool
forall a. Eq a => a -> a -> Bool
== Text
"'re" = (Text, Text) -> Maybe (Text, Text)
forall a. a -> Maybe a
Just (Int -> Text -> Text
Text.take Int
3 Text
t, Int -> Text -> Text
Text.drop Int
3 Text
t)
      | Text -> Int
Text.length Text
t Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
2,
        HasCallStack => Text -> Int -> Char
Text -> Int -> Char
Text.index Text
t Int
0 Char -> Char -> Bool
forall a. Eq a => a -> a -> Bool
== Char
'\'',
        HasCallStack => Text -> Int -> Char
Text -> Int -> Char
Text.index Text
t Int
1 Char -> [Char] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [Char
's', Char
'd', Char
'm', Char
't'] =
          (Text, Text) -> Maybe (Text, Text)
forall a. a -> Maybe a
Just (Int -> Text -> Text
Text.take Int
2 Text
t, Int -> Text -> Text
Text.drop Int
2 Text
t)
      | Bool
otherwise = Maybe (Text, Text)
forall a. Maybe a
Nothing

    -- Optional space followed by letters:  ?\p{L}+
    matchOptSpaceLetters :: Text -> Maybe (Text, Text)
matchOptSpaceLetters Text
t =
      let (Text
sp, Text
rest1) = Text -> (Text, Text)
matchOptSpace Text
t
          (Text
letters, Text
rest2) = (Char -> Bool) -> Text -> (Text, Text)
Text.span Char -> Bool
isAlpha Text
rest1
       in if Text -> Bool
Text.null Text
letters
            then Maybe (Text, Text)
forall a. Maybe a
Nothing
            else (Text, Text) -> Maybe (Text, Text)
forall a. a -> Maybe a
Just (Text
sp Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
letters, Text
rest2)

    -- Optional space followed by digits:  ?\p{N}+
    matchOptSpaceDigits :: Text -> Maybe (Text, Text)
matchOptSpaceDigits Text
t =
      let (Text
sp, Text
rest1) = Text -> (Text, Text)
matchOptSpace Text
t
          (Text
digits, Text
rest2) = (Char -> Bool) -> Text -> (Text, Text)
Text.span Char -> Bool
isDigit Text
rest1
       in if Text -> Bool
Text.null Text
digits
            then Maybe (Text, Text)
forall a. Maybe a
Nothing
            else (Text, Text) -> Maybe (Text, Text)
forall a. a -> Maybe a
Just (Text
sp Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
digits, Text
rest2)

    -- Optional space followed by punctuation:  ?[^\s\p{L}\p{N}]+
    matchOptSpacePunct :: Text -> Maybe (Text, Text)
matchOptSpacePunct Text
t =
      let (Text
sp, Text
rest1) = Text -> (Text, Text)
matchOptSpace Text
t
          (Text
punct, Text
rest2) = (Char -> Bool) -> Text -> (Text, Text)
Text.span (\Char
c -> Bool -> Bool
not (Char -> Bool
isSpace Char
c) Bool -> Bool -> Bool
&& Bool -> Bool
not (Char -> Bool
isAlpha Char
c) Bool -> Bool -> Bool
&& Bool -> Bool
not (Char -> Bool
isDigit Char
c)) Text
rest1
       in if Text -> Bool
Text.null Text
punct
            then Maybe (Text, Text)
forall a. Maybe a
Nothing
            else (Text, Text) -> Maybe (Text, Text)
forall a. a -> Maybe a
Just (Text
sp Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
punct, Text
rest2)

    -- Trailing whitespace: \s+(?!\S)  = whitespace at end of string
    matchTrailingSpace :: Text -> Maybe (Text, Text)
matchTrailingSpace Text
t =
      let (Text
sp, Text
rest) = (Char -> Bool) -> Text -> (Text, Text)
Text.span Char -> Bool
isSpace Text
t
       in if Text -> Bool
Text.null Text
sp Bool -> Bool -> Bool
|| Bool -> Bool
not (Text -> Bool
Text.null Text
rest)
            then Maybe (Text, Text)
forall a. Maybe a
Nothing
            else (Text, Text) -> Maybe (Text, Text)
forall a. a -> Maybe a
Just (Text
sp, Text
rest)

    -- Other whitespace: \s+
    matchSpace :: Text -> Maybe (Text, Text)
matchSpace Text
t =
      let (Text
sp, Text
rest) = (Char -> Bool) -> Text -> (Text, Text)
Text.span Char -> Bool
isSpace Text
t
       in if Text -> Bool
Text.null Text
sp
            then Maybe (Text, Text)
forall a. Maybe a
Nothing
            else (Text, Text) -> Maybe (Text, Text)
forall a. a -> Maybe a
Just (Text
sp, Text
rest)

    -- Match optional single space
    matchOptSpace :: Text -> (Text, Text)
matchOptSpace Text
t
      | Bool -> Bool
not (Text -> Bool
Text.null Text
t), Char -> Bool
isSpace (HasCallStack => Text -> Char
Text -> Char
Text.head Text
t) = (Int -> Text -> Text
Text.take Int
1 Text
t, Int -> Text -> Text
Text.drop Int
1 Text
t)
      | Bool
otherwise = (Text
"", Text
t)

    (<|>) :: Maybe a -> Maybe a -> Maybe a
    Maybe a
Nothing <|> :: forall a. Maybe a -> Maybe a -> Maybe a
<|> Maybe a
y = Maybe a
y
    Maybe a
x <|> Maybe a
_ = Maybe a
x

-- | Encode a single text chunk
encodeChunk :: BPEModel -> Text -> Vector Word32
encodeChunk :: BPEModel -> Text -> Vector Word32
encodeChunk BPEModel
model Text
chunk =
  case Text -> Map Text Word32 -> Maybe Word32
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup Text
chunk (BPEModel -> Map Text Word32
bpeSpecialTokens BPEModel
model) of
    Just Word32
tokenId -> Word32 -> Vector Word32
forall a. Unbox a => a -> Vector a
V.singleton Word32
tokenId
    Maybe Word32
Nothing -> BPEModel -> ByteString -> Vector Word32
encodeBytes BPEModel
model (Text -> ByteString
TE.encodeUtf8 Text
chunk)

-- | Encode bytes using BPE merge rules
encodeBytes :: BPEModel -> ByteString -> Vector Word32
encodeBytes :: BPEModel -> ByteString -> Vector Word32
encodeBytes BPEModel
model ByteString
bs =
  let initialTokens :: Vector Word32
initialTokens = [Word32] -> Vector Word32
forall a. Unbox a => [a] -> Vector a
V.fromList [Word8 -> Word32
forall a b. (Integral a, Num b) => a -> b
fromIntegral (HasCallStack => ByteString -> Int -> Word8
ByteString -> Int -> Word8
BS.index ByteString
bs Int
i) | Int
i <- [Int
0 .. ByteString -> Int
BS.length ByteString
bs Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]]
      finalTokens :: Vector Word32
finalTokens = Map (Word32, Word32) (Word32, Int)
-> Vector Word32 -> Vector Word32
applyMerges (BPEModel -> Map (Word32, Word32) (Word32, Int)
bpeMergeRules BPEModel
model) Vector Word32
initialTokens
   in Vector Word32
finalTokens

-- | Apply BPE merges iteratively until no more mergeable pairs
applyMerges :: Map (Word32, Word32) (Word32, Int) -> Vector Word32 -> Vector Word32
applyMerges :: Map (Word32, Word32) (Word32, Int)
-> Vector Word32 -> Vector Word32
applyMerges Map (Word32, Word32) (Word32, Int)
rules Vector Word32
tokens =
  case Map (Word32, Word32) (Word32, Int)
-> Vector Word32 -> Maybe ((Word32, Word32), Word32, Int)
findBestPair Map (Word32, Word32) (Word32, Int)
rules Vector Word32
tokens of
    Maybe ((Word32, Word32), Word32, Int)
Nothing -> Vector Word32
tokens
    Just ((Word32, Word32)
pair, Word32
newToken, Int
_priority) ->
      let merged :: Vector Word32
merged = Vector Word32 -> (Word32, Word32) -> Word32 -> Vector Word32
mergePair Vector Word32
tokens (Word32, Word32)
pair Word32
newToken
       in Map (Word32, Word32) (Word32, Int)
-> Vector Word32 -> Vector Word32
applyMerges Map (Word32, Word32) (Word32, Int)
rules Vector Word32
merged

-- | Find the best pair to merge (lowest priority = earliest in training)
findBestPair :: Map (Word32, Word32) (Word32, Int) -> Vector Word32 -> Maybe ((Word32, Word32), Word32, Int)
findBestPair :: Map (Word32, Word32) (Word32, Int)
-> Vector Word32 -> Maybe ((Word32, Word32), Word32, Int)
findBestPair Map (Word32, Word32) (Word32, Int)
rules Vector Word32
tokens =
  let len :: Int
len = Vector Word32 -> Int
forall a. Unbox a => Vector a -> Int
V.length Vector Word32
tokens
      pairs :: [(Word32, Word32)]
pairs = [(Vector Word32 -> Int -> Word32
forall a. Unbox a => Vector a -> Int -> a
V.unsafeIndex Vector Word32
tokens Int
i, Vector Word32 -> Int -> Word32
forall a. Unbox a => Vector a -> Int -> a
V.unsafeIndex Vector Word32
tokens (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)) | Int
i <- [Int
0 .. Int
len Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
2]]
      validPairs :: [((Word32, Word32), Word32, Int)]
validPairs =
        [ ((Word32, Word32)
pair, Word32
newToken, Int
priority)
        | (Word32, Word32)
pair <- [(Word32, Word32)]
pairs,
          Just (Word32
newToken, Int
priority) <- [(Word32, Word32)
-> Map (Word32, Word32) (Word32, Int) -> Maybe (Word32, Int)
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup (Word32, Word32)
pair Map (Word32, Word32) (Word32, Int)
rules]
        ]
   in if [((Word32, Word32), Word32, Int)] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [((Word32, Word32), Word32, Int)]
validPairs
        then Maybe ((Word32, Word32), Word32, Int)
forall a. Maybe a
Nothing
        else ((Word32, Word32), Word32, Int)
-> Maybe ((Word32, Word32), Word32, Int)
forall a. a -> Maybe a
Just (((Word32, Word32), Word32, Int)
 -> Maybe ((Word32, Word32), Word32, Int))
-> ((Word32, Word32), Word32, Int)
-> Maybe ((Word32, Word32), Word32, Int)
forall a b. (a -> b) -> a -> b
$ (((Word32, Word32), Word32, Int)
 -> ((Word32, Word32), Word32, Int) -> Ordering)
-> [((Word32, Word32), Word32, Int)]
-> ((Word32, Word32), Word32, Int)
forall (t :: * -> *) a.
Foldable t =>
(a -> a -> Ordering) -> t a -> a
minimumBy ((((Word32, Word32), Word32, Int) -> Int)
-> ((Word32, Word32), Word32, Int)
-> ((Word32, Word32), Word32, Int)
-> Ordering
forall a b. Ord a => (b -> a) -> b -> b -> Ordering
comparing (\((Word32, Word32)
_, Word32
_, Int
p) -> Int
p)) [((Word32, Word32), Word32, Int)]
validPairs

-- | Merge all occurrences of a pair into a single token
mergePair :: Vector Word32 -> (Word32, Word32) -> Word32 -> Vector Word32
mergePair :: Vector Word32 -> (Word32, Word32) -> Word32 -> Vector Word32
mergePair Vector Word32
tokens (Word32
tok1, Word32
tok2) Word32
newToken =
  let len :: Int
len = Vector Word32 -> Int
forall a. Unbox a => Vector a -> Int
V.length Vector Word32
tokens
      go :: Int -> [Word32] -> [Word32]
go Int
i [Word32]
acc
        | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
len = [Word32] -> [Word32]
forall a. [a] -> [a]
reverse [Word32]
acc
        | Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
len Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1 = [Word32] -> [Word32]
forall a. [a] -> [a]
reverse (Vector Word32 -> Int -> Word32
forall a. Unbox a => Vector a -> Int -> a
V.unsafeIndex Vector Word32
tokens Int
i Word32 -> [Word32] -> [Word32]
forall a. a -> [a] -> [a]
: [Word32]
acc)
        | Bool
otherwise =
            let current :: Word32
current = Vector Word32 -> Int -> Word32
forall a. Unbox a => Vector a -> Int -> a
V.unsafeIndex Vector Word32
tokens Int
i
                next :: Word32
next = Vector Word32 -> Int -> Word32
forall a. Unbox a => Vector a -> Int -> a
V.unsafeIndex Vector Word32
tokens (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
             in if Word32
current Word32 -> Word32 -> Bool
forall a. Eq a => a -> a -> Bool
== Word32
tok1 Bool -> Bool -> Bool
&& Word32
next Word32 -> Word32 -> Bool
forall a. Eq a => a -> a -> Bool
== Word32
tok2
                  then Int -> [Word32] -> [Word32]
go (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
2) (Word32
newToken Word32 -> [Word32] -> [Word32]
forall a. a -> [a] -> [a]
: [Word32]
acc)
                  else Int -> [Word32] -> [Word32]
go (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Word32
current Word32 -> [Word32] -> [Word32]
forall a. a -> [a] -> [a]
: [Word32]
acc)
      merged :: [Word32]
merged = Int -> [Word32] -> [Word32]
go Int
0 []
   in [Word32] -> Vector Word32
forall a. Unbox a => [a] -> Vector a
V.fromList [Word32]
merged

-- | Decode token IDs back to text
--
-- Algorithm:
-- 1. For each token ID, lookup bytes in vocabulary
-- 2. If not found in initial vocab, compute recursively from merge rules
-- 3. Concatenate all bytes
-- 4. Decode as UTF-8 (lossy)
decodeBPE :: BPEModel -> Vector Word32 -> Text
decodeBPE :: BPEModel -> Vector Word32 -> Text
decodeBPE BPEModel
model Vector Word32
tokens =
  let bytesList :: [ByteString]
bytesList = [BPEModel -> Word32 -> ByteString
lookupTokenBytes BPEModel
model Word32
tid | Word32
tid <- Vector Word32 -> [Word32]
forall a. Unbox a => Vector a -> [a]
V.toList Vector Word32
tokens]
      allBytes :: ByteString
allBytes = [ByteString] -> ByteString
BS.concat [ByteString]
bytesList
   in OnDecodeError -> ByteString -> Text
TE.decodeUtf8With OnDecodeError
TEE.lenientDecode ByteString
allBytes

-- | Decode with performance measurement
--
-- TODO: implement proper perf measurement with pure functions
decodeBPEWithPerf :: BPEModel -> Vector Word32 -> IO (Text, Map Text [Nanos])
decodeBPEWithPerf :: BPEModel -> Vector Word32 -> IO (Text, Map Text [Nanos])
decodeBPEWithPerf BPEModel
model Vector Word32
tokens = do
  let result :: Text
result = BPEModel -> Vector Word32 -> Text
decodeBPE BPEModel
model Vector Word32
tokens
      timings :: Map k a
timings = Map k a
forall k a. Map k a
Map.empty -- Placeholder
  (Text, Map Text [Nanos]) -> IO (Text, Map Text [Nanos])
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Text
result, Map Text [Nanos]
forall k a. Map k a
timings)

-- | Lookup bytes for a token ID (with lazy vocabulary building)
-- First tries direct lookup in vocab, then computes recursively if needed
lookupTokenBytes :: BPEModel -> Word32 -> ByteString
lookupTokenBytes :: BPEModel -> Word32 -> ByteString
lookupTokenBytes BPEModel
model Word32
tokenId =
  case Int -> IntMap ByteString -> Maybe ByteString
forall a. Int -> IntMap a -> Maybe a
IntMap.lookup (Word32 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word32
tokenId) (BPEModel -> IntMap ByteString
bpeVocab BPEModel
model) of
    Just ByteString
bytes -> ByteString
bytes
    Maybe ByteString
Nothing ->
      -- Token not in vocab - need to compute it from merge rules
      -- This happens for merged tokens not pre-computed
      BPEModel -> Word32 -> ByteString
computeTokenBytes BPEModel
model Word32
tokenId

-- | Compute bytes for a merged token ID (recursive)
--
-- Finds the merge rule that created this token and concatenates component bytes
computeTokenBytes :: BPEModel -> Word32 -> ByteString
computeTokenBytes :: BPEModel -> Word32 -> ByteString
computeTokenBytes BPEModel
model Word32
tokenId =
  case Map (Word32, Word32) (Word32, Int)
-> Word32 -> Maybe (Word32, Word32)
findMergeForToken (BPEModel -> Map (Word32, Word32) (Word32, Int)
bpeMergeRules BPEModel
model) Word32
tokenId of
    Maybe (Word32, Word32)
Nothing -> ByteString
BS.empty
    Just (Word32
tok1, Word32
tok2) ->
      let bytes1 :: ByteString
bytes1 = BPEModel -> Word32 -> ByteString
lookupTokenBytes BPEModel
model Word32
tok1
          bytes2 :: ByteString
bytes2 = BPEModel -> Word32 -> ByteString
lookupTokenBytes BPEModel
model Word32
tok2
       in ByteString -> ByteString -> ByteString
BS.append ByteString
bytes1 ByteString
bytes2

-- | Find the merge rule that produced a given token ID
--
-- Searches through merge rules to find which pair merged into this token
findMergeForToken :: Map (Word32, Word32) (Word32, Int) -> Word32 -> Maybe (Word32, Word32)
findMergeForToken :: Map (Word32, Word32) (Word32, Int)
-> Word32 -> Maybe (Word32, Word32)
findMergeForToken Map (Word32, Word32) (Word32, Int)
rules Word32
targetId =
  let matches :: [((Word32, Word32), Word32)]
matches = [((Word32, Word32)
pair, Word32
tid) | ((Word32, Word32)
pair, (Word32
tid, Int
_)) <- Map (Word32, Word32) (Word32, Int)
-> [((Word32, Word32), (Word32, Int))]
forall k a. Map k a -> [(k, a)]
Map.toList Map (Word32, Word32) (Word32, Int)
rules, Word32
tid Word32 -> Word32 -> Bool
forall a. Eq a => a -> a -> Bool
== Word32
targetId]
   in case [((Word32, Word32), Word32)]
matches of
        (((Word32, Word32)
pair, Word32
_) : [((Word32, Word32), Word32)]
_) -> (Word32, Word32) -> Maybe (Word32, Word32)
forall a. a -> Maybe a
Just (Word32, Word32)
pair
        [] -> Maybe (Word32, Word32)
forall a. Maybe a
Nothing

-- | Pretty-print BPE model information
prettifyBPEModel :: BPEModel -> String
prettifyBPEModel :: BPEModel -> [Char]
prettifyBPEModel BPEModel
model =
  [[Char]] -> [Char]
unlines ([[Char]] -> [Char]) -> [[Char]] -> [Char]
forall a b. (a -> b) -> a -> b
$
    [ [Char]
"BPE Model:",
      [Char]
"  Version: " [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ Text -> [Char]
Text.unpack (BPEModel -> Text
bpeVersion BPEModel
model),
      [Char]
"  Regex: " [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ ByteString -> [Char]
forall a. Show a => a -> [Char]
show (BPEModel -> ByteString
bpeRegex BPEModel
model),
      [Char]
"  Special tokens: " [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ Int -> [Char]
forall a. Show a => a -> [Char]
show (Map Text Word32 -> Int
forall k a. Map k a -> Int
Map.size (BPEModel -> Map Text Word32
bpeSpecialTokens BPEModel
model)),
      [Char]
"  Merge rules: " [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ Int -> [Char]
forall a. Show a => a -> [Char]
show (Map (Word32, Word32) (Word32, Int) -> Int
forall k a. Map k a -> Int
Map.size (BPEModel -> Map (Word32, Word32) (Word32, Int)
bpeMergeRules BPEModel
model)),
      [Char]
"  Vocab size: " [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ Int -> [Char]
forall a. Show a => a -> [Char]
show (IntMap ByteString -> Int
forall a. IntMap a -> Int
IntMap.size (BPEModel -> IntMap ByteString
bpeVocab BPEModel
model)),
      [Char]
"  Max token ID: " [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ Word32 -> [Char]
forall a. Show a => a -> [Char]
show (BPEModel -> Word32
bpeMaxTokenId BPEModel
model),
      [Char]
"",
      [Char]
"Special tokens:"
    ]
      [[Char]] -> [[Char]] -> [[Char]]
forall a. [a] -> [a] -> [a]
++ [[Char]]
specialTokenList
  where
    specialTokenList :: [[Char]]
specialTokenList =
      [ [Char]
"  " [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ Text -> [Char]
Text.unpack Text
tok [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ [Char]
" -> " [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ Word32 -> [Char]
forall a. Show a => a -> [Char]
show Word32
tid
      | (Text
tok, Word32
tid) <- Map Text Word32 -> [(Text, Word32)]
forall k a. Map k a -> [(k, a)]
Map.toList (BPEModel -> Map Text Word32
bpeSpecialTokens BPEModel
model)
      ]

-- | Pretty-print encoding result
prettifyEncoding :: BPEEncoding -> String
prettifyEncoding :: BPEEncoding -> [Char]
prettifyEncoding BPEEncoding
enc =
  [[Char]] -> [Char]
unlines
    [ [Char]
"BPE Encoding:",
      [Char]
"  Original text: " [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ Text -> [Char]
forall a. Show a => a -> [Char]
show (BPEEncoding -> Text
originalText BPEEncoding
enc),
      [Char]
"  Token count: " [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ Int -> [Char]
forall a. Show a => a -> [Char]
show (Vector Word32 -> Int
forall a. Unbox a => Vector a -> Int
V.length (BPEEncoding -> Vector Word32
encodedTokens BPEEncoding
enc)),
      [Char]
"  Chunks processed: " [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ Int -> [Char]
forall a. Show a => a -> [Char]
show (BPEEncoding -> Int
numChunks BPEEncoding
enc),
      [Char]
"",
      [Char]
"Token IDs:",
      [Char]
"  " [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ [Word32] -> [Char]
forall a. Show a => a -> [Char]
show (Vector Word32 -> [Word32]
forall a. Unbox a => Vector a -> [a]
V.toList (BPEEncoding -> Vector Word32
encodedTokens BPEEncoding
enc))
    ]