diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,11 @@ +# Changelog for `llama-cpp-hs` + +All notable changes to this project will be documented in this file. + +The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), +and this project adheres to the +[Haskell Package Versioning Policy](https://pvp.haskell.org/). + +## Unreleased + +## 0.1.0.0 - YYYY-MM-DD diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -0,0 +1,20 @@ +Copyright (c) 2025 Tushar Adhatrao + +Permission is hereby granted, free of charge, to any person obtaining +a copy of this software and associated documentation files (the +"Software"), to deal in the Software without restriction, including +without limitation the rights to use, copy, modify, merge, publish, +distribute, sublicense, and/or sell copies of the Software, and to +permit persons to whom the Software is furnished to do so, subject to +the following conditions: + +The above copyright notice and this permission notice shall be included +in all copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, +EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF +MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. +IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY +CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, +TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE +SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. diff --git a/README.md b/README.md new file mode 100644 --- /dev/null +++ b/README.md @@ -0,0 +1,86 @@ +# llama-cpp-hs + +Haskell bindings over [llama.cpp](https://github.com/ggml-org/llama.cpp) + +This package provides both low-level and high-level interfaces to interact with the LLaMA C++ inference engine via Haskell FFI. +It allows you to run LLMs locally in pure C/C++, with support for GPU acceleration and quantized models. + +## Features + +- Low-level access to the full LLaMA C API using Haskell FFI. +- Higher-level convenience functions for easier model interaction. +- Examples provided for quickly getting started. + +--- + +## Example Usage + +Check out the `/examples` directory to see how to load and query models directly from Haskell. + +--- + +## Setup + +### 1. Using Nix (Recommended) + +Ensure that [Nix](https://nixos.org/download.html) is installed on your system. + +Then, enter the development shell: + +```bash +nix-shell +``` + +Build the project using Stack: + +```bash +stack build +``` + +### 2. Using Stack (Manual Setup) + +If you prefer not to use Nix, follow these steps: + +1. Clone and install [`llama.cpp`](https://github.com/ggml-org/llama.cpp) manually. +2. Make sure `llama.h` is available at `/usr/local/include/` and compiled `libllama.a` or `libllama.so` at `/usr/local/lib/`. +3. Install Stack if you haven’t already: https://docs.haskellstack.org/en/stable/install_and_upgrade/ +4. Then proceed with: + +```bash +stack build +``` + +--- + +## Models + +To use this library, you'll need to download one of the many open-source GGUF models available on Hugging Face + +Search for compatible GGUF models: +- [Hugging Face GGUF Models](https://huggingface.co/models?search=gguf) + +--- + +## Current State + +The codebase is still under active development and may undergo breaking changes. Use it with caution in production environments. + +Pull requests, issues, and community contributions are highly encouraged! + +--- + +## Contributing + +Contributions are welcome! + +--- + +## License + +This project is licensed under [MIT](LICENSE). + +--- + +## Thank You + +Thanks to [ggml-org/llama.cpp](https://github.com/ggml-org/llama.cpp) for making local LLM inference fast, lightweight, and accessible! diff --git a/Setup.hs b/Setup.hs new file mode 100644 --- /dev/null +++ b/Setup.hs @@ -0,0 +1,2 @@ +import Distribution.Simple +main = defaultMain diff --git a/cbits/helper.c b/cbits/helper.c new file mode 100644 --- /dev/null +++ b/cbits/helper.c @@ -0,0 +1,92 @@ +#include +#include + +void _llama_backend_init() { + llama_backend_init(); +} + +void _llama_backend_free() { + llama_backend_free(); +} + +void llama_batch_init_into( + int32_t n_tokens, + int32_t embd, + int32_t n_seq_max, + struct llama_batch * out) { + + if (out == NULL) return; + + *out = (struct llama_batch){ + .n_tokens = n_tokens, + .token = NULL, // Set by Haskell + .embd = NULL, // Set by Haskell + .pos = NULL, // Set by Haskell + .n_seq_id = NULL, // Set by Haskell + .seq_id = NULL, // Set by Haskell + .logits = NULL, // Set by Haskell + }; + + // Call the original llama_batch_init + *out = llama_batch_init(n_tokens, embd, n_seq_max); +} + +int32_t llama_decode_wrap(struct llama_context * ctx, struct llama_batch *batch) { + llama_decode(ctx, *batch); +} + +void llama_sampler_chain_default_params_into(struct llama_sampler_chain_params * out) { + *out = llama_sampler_chain_default_params(); +} + +struct llama_model * llama_model_load_from_file_wrap (const char * path_model, const struct llama_model_params * params) { + return llama_model_load_from_file(path_model, *params); +} + +struct llama_context * llama_init_from_model_into(struct llama_model * model, const struct llama_context_params * params) { + return llama_init_from_model(model, *params); +} + +void llama_batch_free_wrap(struct llama_batch * batch) { + llama_batch_free(*batch); +} + +void llama_batch_get_one_into(llama_token * tokens + , int32_t n_tokens, struct llama_batch * out_batch) { + if (out_batch == NULL) { + return; // invalid output pointer + } + *out_batch = llama_batch_get_one(tokens, n_tokens); +} + +void llama_pooling_type_into(const struct llama_context * ctx, int * out) { + *out = llama_pooling_type(ctx); +} + +void llama_model_rope_type_into(const struct llama_model * model, int * out) { + *out = llama_model_rope_type(model); +} + +void llama_context_default_params_into(struct llama_context_params* out) { + if (out != NULL) { + *out = llama_context_default_params(); + } +} + +void llama_kv_cache_view_init_into(const struct llama_context * ctx, int32_t n_seq_max, struct llama_kv_cache_view* out) { + if (out != NULL) { + *out = llama_kv_cache_view_init(ctx, n_seq_max); + } +} + +void llama_perf_context_into(const struct llama_context * ctx, struct llama_perf_context_data* out ) { + if (out != NULL) { + *out = llama_perf_context(ctx); + } +} + +void llama_perf_sampler_into(const struct llama_sampler * chain, struct llama_perf_sampler_data* out) { + if (out != NULL) { + *out = llama_perf_sampler(chain); + } +} diff --git a/llama-cpp-hs.cabal b/llama-cpp-hs.cabal new file mode 100644 --- /dev/null +++ b/llama-cpp-hs.cabal @@ -0,0 +1,94 @@ +cabal-version: 1.12 + +-- This file has been generated from package.yaml by hpack version 0.38.0. +-- +-- see: https://github.com/sol/hpack + +name: llama-cpp-hs +version: 0.1.0.0 +synopsis: Haskell FFI bindings to the llama.cpp LLM inference library +description: Haskell bindings for [llama.cpp](https://github.com/ggerganov/llama.cpp), a performant, C++-based inference engine for running large language models (LLMs) like LLaMA, Mistral, Qwen, and others directly on local hardware. +category: ai, ffi, natural-language-processing +homepage: https://github.com/tusharad/llama-cpp-hs#readme +bug-reports: https://github.com/tusharad/llama-cpp-hs/issues +author: tushar +maintainer: tusharadhatrao@gmail.com +copyright: 2025 tushar +license: MIT +license-file: LICENSE +build-type: Simple +extra-source-files: + README.md + CHANGELOG.md + +source-repository head + type: git + location: https://github.com/tusharad/llama-cpp-hs + +library + exposed-modules: + Llama.Adapter + Llama.Backend + Llama.ChatTemplate + Llama.Context + Llama.Decode + Llama.Internal.Foreign + Llama.Internal.Foreign.Adapter + Llama.Internal.Foreign.Backend + Llama.Internal.Foreign.ChatTemplate + Llama.Internal.Foreign.Context + Llama.Internal.Foreign.Decode + Llama.Internal.Foreign.GGML + Llama.Internal.Foreign.KVCache + Llama.Internal.Foreign.Model + Llama.Internal.Foreign.Performance + Llama.Internal.Foreign.Sampler + Llama.Internal.Foreign.Split + Llama.Internal.Foreign.State + Llama.Internal.Foreign.Tokenize + Llama.Internal.Foreign.Vocab + Llama.Internal.Types + Llama.Internal.Types.Params + Llama.KVCache + Llama.Model + Llama.Performance + Llama.Sampler + Llama.Split + Llama.State + Llama.Tokenize + Llama.Vocab + other-modules: + Paths_llama_cpp_hs + hs-source-dirs: + src + ghc-options: -Wall -Wcompat -Widentities -Wincomplete-record-updates -Wincomplete-uni-patterns -Wmissing-home-modules -Wpartial-fields -Wredundant-constraints + include-dirs: + cbits + c-sources: + cbits/helper.c + extra-lib-dirs: + /usr/local/lib + extra-libraries: + llama + build-depends: + base >=4.7 && <5 + , bytestring >=0.9 && <0.13 + , derive-storable >=0.2 && <0.4 + default-language: Haskell2010 + +test-suite llama-cpp-hs-test + type: exitcode-stdio-1.0 + main-is: Spec.hs + other-modules: + Paths_llama_cpp_hs + hs-source-dirs: + test + ghc-options: -Wall -Wcompat -Widentities -Wincomplete-record-updates -Wincomplete-uni-patterns -Wmissing-home-modules -Wpartial-fields -Wredundant-constraints -threaded -rtsopts -with-rtsopts=-N + build-depends: + base >=4.7 && <5 + , bytestring >=0.9 && <0.13 + , derive-storable >=0.2 && <0.4 + , llama-cpp-hs + , tasty + , tasty-hunit + default-language: Haskell2010 diff --git a/src/Llama/Adapter.hs b/src/Llama/Adapter.hs new file mode 100644 --- /dev/null +++ b/src/Llama/Adapter.hs @@ -0,0 +1,92 @@ +{- | +Module : Llama.Adapter +Description : Adapter interface for llama-cpp +Copyright : (c) 2025 Tushar Adhatrao +License : MIT +Maintainer : Tushar Adhatrao +-} +module Llama.Adapter ( + initAdapterLora + , setAdapterLora + , rmAdapterLora + , clearAdapterLora + , applyAdapterCVec +) where + +import Foreign +import Foreign.C.String +import Llama.Internal.Foreign +import Llama.Internal.Types +import Foreign.C (CFloat) + +{- | Initialize a LoRA adapter from a file path. +This function wraps the C function 'llama_adapter_lora_init' +and returns a managed AdapterLora object. +-} +initAdapterLora :: Model -> FilePath -> IO (Either String AdapterLora) +initAdapterLora (Model modelFPtr) path = do + withForeignPtr modelFPtr $ \modelPtr -> do + withCString path $ \cPath -> do + adapter@(CLlamaAdapterLora adapterPtr) <- + c_llama_adapter_lora_init (CLlamaModel modelPtr) cPath + if adapter == CLlamaAdapterLora nullPtr + then return $ Left "Failed to initialize LoRA adapter" + else do + -- We directly use c_llama_adapter_lora_free as the finalizer + fp <- newForeignPtr p_llama_adapter_lora_free adapterPtr + return $ Right $ AdapterLora fp + +-- | Apply a LoRA adapter with a given scale. +setAdapterLora :: Context -> AdapterLora -> Float -> IO (Either String ()) +setAdapterLora (Context ctxFPtr) (AdapterLora adapterFPtr) scale = do + result <- withForeignPtr ctxFPtr $ \ctxPtr -> + withForeignPtr adapterFPtr $ \adapterPtr -> + c_llama_set_adapter_lora + (CLlamaContext ctxPtr) + (CLlamaAdapterLora adapterPtr) + (realToFrac scale) + if result == -1 + then return $ Left "Failed to set LoRA adapter" + else return $ Right () + +-- | Remove a previously applied LoRA adapter. +rmAdapterLora :: Context -> AdapterLora -> IO (Either String ()) +rmAdapterLora (Context ctxFPtr) (AdapterLora adapterFPtr) = do + result <- withForeignPtr ctxFPtr $ \ctxPtr -> + withForeignPtr adapterFPtr $ \adapterPtr -> + c_llama_rm_adapter_lora (CLlamaContext ctxPtr) (CLlamaAdapterLora adapterPtr) + if result == -1 + then return $ Left "Failed to remove LoRA adapter" + else return $ Right () + +-- | Clear all active LoRA adapters from the context. +clearAdapterLora :: Context -> IO () +clearAdapterLora (Context ctxFPtr) = + withForeignPtr ctxFPtr $ \ctxPtr -> + c_llama_clear_adapter_lora (CLlamaContext ctxPtr) + +-- | Apply a context vector (cvec). +applyAdapterCVec :: Context -> Maybe [Float] -> Int -> Int -> Int -> IO (Either String ()) +applyAdapterCVec (Context ctxFPtr) mValues n_embd il_start il_end = do + (ptr, len) <- case mValues of + Nothing -> pure (nullPtr, 0 :: Int) + Just vs -> do + arrayPtr <- floatArrayToPtr vs + pure (arrayPtr, length vs) + result <- withForeignPtr ctxFPtr $ \ctxPtr -> + c_llama_apply_adapter_cvec + (CLlamaContext ctxPtr) + ptr + (fromIntegral len) + (fromIntegral n_embd) + (fromIntegral il_start) + (fromIntegral il_end) + if result == -1 + then return $ Left "Failed to apply context vector" + else return $ Right () + +floatArrayToPtr :: [Float] -> IO (Ptr CFloat) +floatArrayToPtr xs = do + ptr <- mallocBytes (length xs * sizeOf (undefined :: CFloat)) + pokeArray ptr (map realToFrac xs) + return ptr diff --git a/src/Llama/Backend.hs b/src/Llama/Backend.hs new file mode 100644 --- /dev/null +++ b/src/Llama/Backend.hs @@ -0,0 +1,22 @@ +{- | +Module : Llama.Backend +Description : Backend related functions for llama-cpp +Copyright : (c) 2025 Tushar Adhatrao +License : MIT +Maintainer : Tushar Adhatrao +-} +module Llama.Backend ( + llamaBackendInit + , llamaBackendFree +) where +{- +Backend and intialization related functions for Llama.cpp +-} + +import Llama.Internal.Foreign.Backend + +llamaBackendInit :: IO () +llamaBackendInit = c_llama_backend_init + +llamaBackendFree :: IO () +llamaBackendFree = c_llama_backend_free diff --git a/src/Llama/ChatTemplate.hs b/src/Llama/ChatTemplate.hs new file mode 100644 --- /dev/null +++ b/src/Llama/ChatTemplate.hs @@ -0,0 +1,104 @@ +{- | +Module : Llama.ChatTemplate +Description : Chat template related functions for llama-cpp +Copyright : (c) 2025 Tushar Adhatrao +License : MIT +Maintainer : Tushar Adhatrao +-} +module Llama.ChatTemplate ( + ChatMessage (..) + , toCLlamaChatMessage + , chatApplyTemplate + , chatGetBuiltinTemplates + ) where + +import Data.Functor +import Foreign +import Foreign.C.String +import Llama.Internal.Foreign +import Llama.Internal.Types + +data ChatMessage = ChatMessage + { chatRole :: String + , chatContent :: String + } + deriving (Show, Eq) + +toCLlamaChatMessage :: ChatMessage -> IO LlamaChatMessage +toCLlamaChatMessage msg = do + rolePtr <- newCString (chatRole msg) + contentPtr <- newCString (chatContent msg) + return $ LlamaChatMessage rolePtr contentPtr + +-- | Apply a chat template to format a conversation. +chatApplyTemplate :: + -- | Optional custom template (uses built-in if Nothing) + Maybe String -> + -- | List of chat messages + [ChatMessage] -> + -- | Add assistant token at end? + Bool -> + -- | Buffer size (suggested: 4096) + Int -> + -- | Returns formatted string or error message + IO (Either String String) +chatApplyTemplate mTemplate messages addAssist bufferSize = do + let nMessages = length messages + bufSize = max bufferSize (nMessages * 64) -- Heuristic fallback + templateCString <- maybe (return nullPtr) newCString mTemplate + msgs <- mapM toCLlamaChatMessage messages + cMessages <- withArray msgs $ \ptr -> + return ptr + + allocaBytes bufSize $ \bufPtr -> do + result <- + c_llama_chat_apply_template + templateCString + cMessages + (fromIntegral nMessages) + (if addAssist then 1 else 0) + bufPtr + (fromIntegral bufSize) + + if result < 0 + then return $ Left "Failed to apply chat template" + else do + let actualSize = fromIntegral result + if actualSize >= bufSize + then -- Need larger buffer + allocaBytes actualSize $ \newBufPtr -> do + result' <- + c_llama_chat_apply_template + templateCString + cMessages + (fromIntegral nMessages) + (if addAssist then 1 else 0) + newBufPtr + (fromIntegral actualSize) + if result' < 0 + then return $ Left "Failed to apply chat template after resize" + else peekCString newBufPtr <&> Right + else do + str <- peekCString bufPtr + return $ Right str + +-- | Get list of available built-in chat templates. +chatGetBuiltinTemplates :: IO [String] +chatGetBuiltinTemplates = do + -- First get number of templates + numTemplates <- c_llama_chat_builtin_templates nullPtr 0 + let maxTemplates = fromIntegral numTemplates + + if numTemplates <= 0 + then return [] + else do + -- Allocate array of CString pointers + arrPtr <- mallocArray maxTemplates + result <- c_llama_chat_builtin_templates arrPtr (fromIntegral maxTemplates) + if result < 0 + then return [] + else do + cstrs <- peekArray maxTemplates arrPtr + strs <- mapM peekCString cstrs + free arrPtr + return strs diff --git a/src/Llama/Context.hs b/src/Llama/Context.hs new file mode 100644 --- /dev/null +++ b/src/Llama/Context.hs @@ -0,0 +1,98 @@ +{- | +Module : Llama.Context +Description : High level context interface for llama-cpp +Copyright : (c) 2025 Tushar Adhatrao +License : MIT +Maintainer : Tushar Adhatrao +-} +module Llama.Context ( + supportsRpc +, supportsGpuOffload +, supportsMLock +, supportsMMap +, getMaxDevices +, getTimeUs +, getContextSize +, getBatchSize +, getUnbatchedSize +, getMaxSeqCount +, getPoolingType +, detachThreadPool +, defaultContextParams +) where + +import Llama.Internal.Foreign +import Llama.Internal.Types +import Llama.Internal.Types.Params +import Foreign + +-- | Check if the backend supports remote procedure calls (RPC). +supportsRpc :: IO Bool +supportsRpc = toBool <$> llama_supports_rpc + +-- | Check if the backend supports GPU offloading. +supportsGpuOffload :: IO Bool +supportsGpuOffload = toBool <$> llama_supports_gpu_offload + +-- | Check if the backend supports locking model memory into RAM (no swapping). +supportsMLock :: IO Bool +supportsMLock = toBool <$> llama_supports_mlock + +-- | Check if the backend supports memory mapping models. +supportsMMap :: IO Bool +supportsMMap = toBool <$> llama_supports_mmap + +-- | Get maximum number of devices supported by the backend (e.g., GPUs). +getMaxDevices :: IO Int +getMaxDevices = fromIntegral <$> llama_max_devices + +-- | Get current time in microseconds since some unspecified epoch. +getTimeUs :: IO Int +getTimeUs = fromIntegral <$> llama_time_us + +-- | Get the maximum context size (n_ctx) of the model in the given context. +getContextSize :: Context -> IO Int +getContextSize (Context ctxFPtr) = + withForeignPtr ctxFPtr $ \ctxPtr -> + fromIntegral <$> llama_n_ctx (CLlamaContext ctxPtr) + +-- | Get the batch size (n_batch) used by the context. +getBatchSize :: Context -> IO Int +getBatchSize (Context ctxFPtr) = + withForeignPtr ctxFPtr $ \ctxPtr -> + fromIntegral <$> llama_n_batch (CLlamaContext ctxPtr) + +-- | Get the unbatched size (n_ubatch). +getUnbatchedSize :: Context -> IO Int +getUnbatchedSize (Context ctxFPtr) = + withForeignPtr ctxFPtr $ \ctxPtr -> + fromIntegral <$> llama_n_ubatch (CLlamaContext ctxPtr) + +-- | Get the maximum number of sequences supported. +getMaxSeqCount :: Context -> IO Int +getMaxSeqCount (Context ctxFPtr) = + withForeignPtr ctxFPtr $ \ctxPtr -> + fromIntegral <$> llama_n_seq_max (CLlamaContext ctxPtr) + +-- | Get the pooling type used by the context. +getPoolingType :: Context -> IO (Maybe LlamaPoolingType) +getPoolingType (Context ctxFPtr) = + withForeignPtr ctxFPtr $ \ctxPtr -> alloca $ \outPtr -> do + c_llama_pooling_type_into (CLlamaContext ctxPtr) outPtr + val <- peek outPtr + case val of + -1 -> return Nothing + _ -> return $ fromLlamaRopePoolingType val + +-- | Detach the internal threadpool from the context. +detachThreadPool :: Context -> IO () +detachThreadPool (Context ctxFPtr) = + withForeignPtr ctxFPtr $ \ctxPtr -> + c_llama_detach_threadpool (CLlamaContext ctxPtr) + +-- | Allocate and initialize a new 'LlamaContextParams' with defaults. +defaultContextParams :: IO LlamaContextParams +defaultContextParams = do + alloca $ \ptr -> do + c_llama_context_default_params_into (CLlamaContextParams ptr) + peek ptr diff --git a/src/Llama/Decode.hs b/src/Llama/Decode.hs new file mode 100644 --- /dev/null +++ b/src/Llama/Decode.hs @@ -0,0 +1,132 @@ +{- | +Module : Llama.Decode +Description : High level Decode interface for llama-cpp +Copyright : (c) 2025 Tushar Adhatrao +License : MIT +Maintainer : Tushar Adhatrao +-} +module Llama.Decode + ( batchInit + , batchGetOne + , freeBatch + , encodeBatch + , decodeBatch + , setThreadCount + , getThreadCount + , getBatchThreadCount + , setEmbeddingsEnabled + , areEmbeddingsEnabled + , setCausalAttention + , setThreadCounts + , setWarmupMode + , synchronizeContext + ) where + +import Foreign +import Llama.Internal.Foreign +import Llama.Internal.Types + +batchInit :: Int -> Int -> Int -> IO Batch +batchInit nTokens embd_ nSeqMax = do + let cTokens = fromIntegral nTokens + cEmb = fromIntegral embd_ + cSeqMax = fromIntegral nSeqMax + alloca $ \ptr -> do + c_llama_batch_init_into cTokens cEmb cSeqMax ptr + return $ Batch ptr + +-- | Create a batch from a list of tokens. +batchGetOne :: [LlamaToken] -> IO Batch +batchGetOne tokens = do + let nTokens = length tokens + alloca $ \ptr -> do + withArray tokens $ \tokensPtr -> do + c_llama_batch_get_one_into tokensPtr (fromIntegral nTokens) ptr + return $ Batch ptr + +-- | Free a batch of tokens allocated with initBatch +freeBatch :: Ptr LlamaBatch -> IO () +freeBatch = c_llama_batch_free_wrap + +-- | Encode tokens using the model context. +encodeBatch :: Context -> Batch -> IO (Either String ()) +encodeBatch (Context ctxFPtr) (Batch batchPtr) = do + result <- withForeignPtr ctxFPtr $ \ctxPtr -> + c_llama_encode (CLlamaContext ctxPtr) batchPtr + if result < 0 + then return $ Left "Encoding failed" + else return $ Right () + +-- | Decode tokens using the model context. +decodeBatch :: Context -> Batch -> IO (Either String ()) +decodeBatch (Context ctxFPtr) (Batch batchPtr) = do + result <- withForeignPtr ctxFPtr $ \ctxPtr -> + c_llama_decode_wrap (CLlamaContext ctxPtr) batchPtr + if result < 0 + then return $ Left "Decoding failed" + else return $ Right () + +-- | Set number of threads used for processing. +setThreadCount :: Context -> Int -> IO () +setThreadCount ctx_ n = setThreadCounts ctx_ n n + +-- | Set main and batch thread counts separately. +setThreadCounts :: Context -> Int -> Int -> IO () +setThreadCounts (Context ctxFPtr) nThreads nBatchThreads = + withForeignPtr ctxFPtr $ \ctxPtr -> + c_llama_set_n_threads + (CLlamaContext ctxPtr) + (fromIntegral nThreads) + (fromIntegral nBatchThreads) + +-- | Get current main thread count. +getThreadCount :: Context -> IO Int +getThreadCount (Context ctxFPtr) = + withForeignPtr ctxFPtr $ \ctxPtr -> + fromIntegral <$> c_llama_n_threads (CLlamaContext ctxPtr) + +-- | Get current batch thread count. +getBatchThreadCount :: Context -> IO Int +getBatchThreadCount (Context ctxFPtr) = + withForeignPtr ctxFPtr $ \ctxPtr -> + fromIntegral <$> c_llama_n_threads_batch (CLlamaContext ctxPtr) + +-- | Enable or disable embeddings output. +setEmbeddingsEnabled :: Context -> Bool -> IO () +setEmbeddingsEnabled (Context ctxFPtr) enabled = + withForeignPtr ctxFPtr $ \ctxPtr -> + c_llama_set_embeddings (CLlamaContext ctxPtr) (fromBool enabled) + +-- | Check if embeddings are enabled. +areEmbeddingsEnabled :: Context -> IO Bool +areEmbeddingsEnabled (Context ctxFPtr) = + withForeignPtr ctxFPtr $ \ctxPtr -> do + resPtr <- c_llama_get_embeddings (CLlamaContext ctxPtr) + toBool <$> peek resPtr + +-- | Set causal attention mode. +setCausalAttention :: Context -> Bool -> IO () +setCausalAttention (Context ctxFPtr) causal = + withForeignPtr ctxFPtr $ \ctxPtr -> + c_llama_set_causal_attn (CLlamaContext ctxPtr) (fromBool causal) + +-- | Set warmup mode (e.g. precompute KV cache). +setWarmupMode :: Context -> Bool -> IO () +setWarmupMode (Context ctxFPtr) warm = + withForeignPtr ctxFPtr $ \ctxPtr -> + c_llama_set_warmup (CLlamaContext ctxPtr) (fromBool warm) + +{- +type AbortCallback = Ptr () -> IO CInt +-- | Set an abort callback to cancel long-running ops from another thread. +setAbortCallback :: Context -> FunPtr AbortCallback -> Ptr () -> IO () +setAbortCallback (Context ctxFPtr) callback cbData = + withForeignPtr ctxFPtr $ \ctxPtr -> + c_llama_set_abort_callback (CLlamaContext ctxPtr) callback cbData + -} + +-- | Block until all async work is complete. +synchronizeContext :: Context -> IO () +synchronizeContext (Context ctxFPtr) = + withForeignPtr ctxFPtr $ \ctxPtr -> + c_llama_synchronize (CLlamaContext ctxPtr) diff --git a/src/Llama/Internal/Foreign.hs b/src/Llama/Internal/Foreign.hs new file mode 100644 --- /dev/null +++ b/src/Llama/Internal/Foreign.hs @@ -0,0 +1,29 @@ +module Llama.Internal.Foreign + ( module Llama.Internal.Foreign.Adapter + , module Llama.Internal.Foreign.ChatTemplate + , module Llama.Internal.Foreign.Context + , module Llama.Internal.Foreign.Decode + , module Llama.Internal.Foreign.KVCache + , module Llama.Internal.Foreign.Model + , module Llama.Internal.Foreign.Performance + , module Llama.Internal.Foreign.Sampler + , module Llama.Internal.Foreign.Split + , module Llama.Internal.Foreign.State + , module Llama.Internal.Foreign.Tokenize + , module Llama.Internal.Foreign.Vocab + , module Llama.Internal.Foreign.Backend + ) where + +import Llama.Internal.Foreign.Adapter +import Llama.Internal.Foreign.Backend +import Llama.Internal.Foreign.ChatTemplate +import Llama.Internal.Foreign.Context +import Llama.Internal.Foreign.Decode +import Llama.Internal.Foreign.KVCache +import Llama.Internal.Foreign.Model +import Llama.Internal.Foreign.Performance +import Llama.Internal.Foreign.Sampler +import Llama.Internal.Foreign.Split +import Llama.Internal.Foreign.State +import Llama.Internal.Foreign.Tokenize +import Llama.Internal.Foreign.Vocab diff --git a/src/Llama/Internal/Foreign/Adapter.hs b/src/Llama/Internal/Foreign/Adapter.hs new file mode 100644 --- /dev/null +++ b/src/Llama/Internal/Foreign/Adapter.hs @@ -0,0 +1,80 @@ +module Llama.Internal.Foreign.Adapter + ( p_llama_adapter_lora_free + , c_llama_adapter_lora_init + , c_llama_adapter_lora_free + , c_llama_set_adapter_lora + , c_llama_rm_adapter_lora + , c_llama_clear_adapter_lora + , c_llama_apply_adapter_cvec + ) where + +import Foreign +import Foreign.C +import Llama.Internal.Types + +{- | Foreign pointer to the C function for freeing LoRA adapters. +This allows Haskell's garbage collector to automatically free resources. +-} +foreign import ccall "&llama_adapter_lora_free" + p_llama_adapter_lora_free :: FinalizerPtr CLlamaAdapterLora + +{- | LLAMA_API struct llama_adapter_lora * llama_adapter_lora_init( +| struct llama_model * model, const char * path_lora); +-} +foreign import ccall "llama_adapter_lora_init" + c_llama_adapter_lora_init :: + CLlamaModel -> -- model + CString -> -- path_lora + IO CLlamaAdapterLora + +-- | LLAMA_API void llama_adapter_lora_free(struct llama_adapter_lora * adapter); +foreign import ccall "llama_adapter_lora_free" + c_llama_adapter_lora_free :: + CLlamaAdapterLora -> -- adapter + IO () + +{- | LLAMA_API int32_t llama_set_adapter_lora( +| struct llama_context * ctx, +| struct llama_adapter_lora * adapter, +| float scale); +-} +foreign import ccall "llama_set_adapter_lora" + c_llama_set_adapter_lora :: + CLlamaContext -> -- ctx + CLlamaAdapterLora -> -- adapter + CFloat -> -- scale + IO CInt + +{- | LLAMA_API int32_t llama_rm_adapter_lora( +| struct llama_context * ctx, +| struct llama_adapter_lora * adapter); +-} +foreign import ccall "llama_rm_adapter_lora" + c_llama_rm_adapter_lora :: + CLlamaContext -> -- ctx + CLlamaAdapterLora -> -- adapter + IO CInt + +-- | LLAMA_API void llama_clear_adapter_lora(struct llama_context * ctx); +foreign import ccall "llama_clear_adapter_lora" + c_llama_clear_adapter_lora :: + CLlamaContext -> -- ctx + IO () + +{- | LLAMA_API int32_t llama_apply_adapter_cvec( +| struct llama_context * ctx, +| const float * data, +| size_t len, +| int32_t n_embd, +| int32_t il_start, +| int32_t il_end); +-} +foreign import ccall "llama_apply_adapter_cvec" + c_llama_apply_adapter_cvec :: + CLlamaContext -> -- ctx + Ptr CFloat -> -- data (nullable) + CSize -> -- len + CInt -> -- n_embd + CInt -> -- il_start + CInt -> -- il_end + IO CInt diff --git a/src/Llama/Internal/Foreign/Backend.hs b/src/Llama/Internal/Foreign/Backend.hs new file mode 100644 --- /dev/null +++ b/src/Llama/Internal/Foreign/Backend.hs @@ -0,0 +1,11 @@ +module Llama.Internal.Foreign.Backend + ( -- * Backend + c_llama_backend_init + , c_llama_backend_free + ) where + +-- | Initialize the llama backend +foreign import ccall "_llama_backend_init" c_llama_backend_init :: IO () + +-- | Free the llama backend resources +foreign import ccall "_llama_backend_free" c_llama_backend_free :: IO () diff --git a/src/Llama/Internal/Foreign/ChatTemplate.hs b/src/Llama/Internal/Foreign/ChatTemplate.hs new file mode 100644 --- /dev/null +++ b/src/Llama/Internal/Foreign/ChatTemplate.hs @@ -0,0 +1,37 @@ +module Llama.Internal.Foreign.ChatTemplate + ( c_llama_chat_apply_template + , c_llama_chat_builtin_templates + ) where + +import Foreign +import Foreign.C +import Llama.Internal.Types + +{- | LLAMA_API int32_t llama_chat_apply_template( +| const char * tmpl, +| const struct llama_chat_message * chat, +| size_t n_msg, +| bool add_ass, +| char * buf, +| int32_t length); +-} +foreign import ccall "llama_chat_apply_template" + c_llama_chat_apply_template :: + CString -> -- tmpl (nullable) + Ptr LlamaChatMessage -> -- chat (array of messages) + CSize -> -- n_msg (number of messages) + CBool -> -- add_ass (add assistant token) + CString -> -- buf (output buffer) + CInt -> -- length (size of buffer) + + -- | Returns required buffer size (if > length) or bytes written + IO CInt + +-- | LLAMA_API int32_t llama_chat_builtin_templates(const char ** output, size_t len); +foreign import ccall "llama_chat_builtin_templates" + c_llama_chat_builtin_templates :: + Ptr CString -> -- output (array of template names) + CSize -> -- len (max number of names to write) + + -- | Returns total number of available templates + IO CInt diff --git a/src/Llama/Internal/Foreign/Context.hs b/src/Llama/Internal/Foreign/Context.hs new file mode 100644 --- /dev/null +++ b/src/Llama/Internal/Foreign/Context.hs @@ -0,0 +1,77 @@ +module Llama.Internal.Foreign.Context + ( llama_time_us + , llama_max_devices + , llama_supports_mmap + , llama_supports_mlock + , llama_supports_gpu_offload + , llama_supports_rpc + , llama_n_ctx + , llama_n_batch + , llama_n_ubatch + , llama_n_seq_max + , c_llama_pooling_type_into + , c_llama_get_kv_self + , c_llama_context_default_params_into + , c_llama_detach_threadpool + ) where + +import Foreign +import Foreign.C +import Llama.Internal.Types +import Llama.Internal.Types.Params + +-- LLAMA_API void llama_detach_threadpool(struct llama_context * ctx); +foreign import ccall "llama_detach_threadpool" + c_llama_detach_threadpool :: CLlamaContext -> IO () + +foreign import ccall "llama_context_default_params_into" + c_llama_context_default_params_into :: CLlamaContextParams -> IO () + +-- | Get current time in microseconds +foreign import ccall "llama_time_us" + llama_time_us :: IO Int64 + +-- | Maximum number of devices supported +foreign import ccall "llama_max_devices" + llama_max_devices :: IO CSize + +-- | Features supported by backend +foreign import ccall "llama_supports_mmap" + llama_supports_mmap :: IO CBool + +foreign import ccall "llama_supports_mlock" + llama_supports_mlock :: IO CBool + +foreign import ccall "llama_supports_gpu_offload" + llama_supports_gpu_offload :: IO CBool + +foreign import ccall "llama_supports_rpc" + llama_supports_rpc :: IO CBool + +-- | Context-related accessors +foreign import ccall "llama_n_ctx" + llama_n_ctx :: CLlamaContext -> IO CUInt + +foreign import ccall "llama_n_batch" + llama_n_batch :: CLlamaContext -> IO CUInt + +foreign import ccall "llama_n_ubatch" + llama_n_ubatch :: CLlamaContext -> IO CUInt + +foreign import ccall "llama_n_seq_max" + llama_n_seq_max :: CLlamaContext -> IO CUInt + +-- | Get KV cache from context +foreign import ccall "llama_get_kv_self" + c_llama_get_kv_self :: CLlamaContext -> IO CLlamaKVCache + +-- | Get pooling type from context +foreign import ccall safe "llama_pooling_type_into" + c_llama_pooling_type_into :: CLlamaContext -> Ptr CInt -> IO () + +-- llamaPoolingType :: CLlamaContext -> IO (Maybe LlamaPoolingType) +-- llamaPoolingType (CLlamaContext ctxPtr) = +-- alloca $ \outPtr -> do +-- c_llama_pooling_type_into ctxPtr outPtr +-- val <- peek outPtr +-- pure $ fromLlamaRopePoolingType val diff --git a/src/Llama/Internal/Foreign/Decode.hs b/src/Llama/Internal/Foreign/Decode.hs new file mode 100644 --- /dev/null +++ b/src/Llama/Internal/Foreign/Decode.hs @@ -0,0 +1,147 @@ +module Llama.Internal.Foreign.Decode ( + c_llama_batch_get_one_into + , c_llama_batch_init_into + , c_llama_batch_free_wrap + , c_llama_encode + , c_llama_decode_wrap + , c_llama_set_n_threads + , c_llama_n_threads + , c_llama_n_threads_batch + , c_llama_set_embeddings + , c_llama_set_causal_attn + , c_llama_set_warmup + , c_llama_set_abort_callback + , c_llama_synchronize + , c_llama_get_logits + , c_llama_get_logits_ith + , c_llama_get_embeddings + , c_llama_get_embeddings_ith + , c_llama_get_embeddings_seq +) where + +import Foreign +import Foreign.C +import Llama.Internal.Types + +-- | LLAMA_API struct llama_batch llama_batch_get_one(llama_token * tokens, int32_t n_tokens); +foreign import ccall "llama_batch_get_one_into" + c_llama_batch_get_one_into :: + Ptr LlamaToken -> -- tokens + CInt -> -- n_tokens + Ptr LlamaBatch -> -- out_batch + IO () + +-- | LLAMA_API struct llama_batch llama_batch_init(int32_t n_tokens, int32_t embd, int32_t n_seq_max); +foreign import ccall "llama_batch_init_into" + c_llama_batch_init_into :: + CInt -> -- n_tokens + CInt -> -- embd + CInt -> -- n_seq_max + Ptr LlamaBatch -> -- out + IO () + +-- | LLAMA_API void llama_batch_free(struct llama_batch batch); +foreign import ccall "llama_batch_free_wrap" + c_llama_batch_free_wrap :: Ptr LlamaBatch -> IO () + +-- | LLAMA_API int32_t llama_encode(struct llama_context * ctx, struct llama_batch batch); +foreign import ccall "llama_encode" + c_llama_encode :: + CLlamaContext -> -- ctx + Ptr LlamaBatch -> -- batch + IO CInt + +-- | LLAMA_API int32_t llama_decode(struct llama_context * ctx, struct llama_batch batch); +foreign import ccall "llama_decode_wrap" + c_llama_decode_wrap :: + CLlamaContext -> -- ctx + Ptr LlamaBatch -> -- batch + IO CInt + +-- | LLAMA_API void llama_set_n_threads(struct llama_context * ctx, int32_t n_threads, int32_t n_threads_batch); +foreign import ccall "llama_set_n_threads" + c_llama_set_n_threads :: + CLlamaContext -> -- ctx + CInt -> -- n_threads + CInt -> -- n_threads_batch + IO () + +-- | LLAMA_API int32_t llama_n_threads(struct llama_context * ctx); +foreign import ccall "llama_n_threads" + c_llama_n_threads :: + CLlamaContext -> -- ctx + IO CInt + +-- | LLAMA_API int32_t llama_n_threads_batch(struct llama_context * ctx); +foreign import ccall "llama_n_threads_batch" + c_llama_n_threads_batch :: + CLlamaContext -> -- ctx + IO CInt + +-- | LLAMA_API void llama_set_embeddings(struct llama_context * ctx, CBool embeddings); +foreign import ccall "llama_set_embeddings" + c_llama_set_embeddings :: + CLlamaContext -> -- ctx + CBool -> -- embeddings + IO () + +-- | LLAMA_API void llama_set_causal_attn(struct llama_context * ctx, CBool causal_attn); +foreign import ccall "llama_set_causal_attn" + c_llama_set_causal_attn :: + CLlamaContext -> -- ctx + CBool -> -- causal_attn + IO () + +-- | LLAMA_API void llama_set_warmup(struct llama_context * ctx, CBool warmup); +foreign import ccall "llama_set_warmup" + c_llama_set_warmup :: + CLlamaContext -> -- ctx + CBool -> -- warmup + IO () + +-- | LLAMA_API void llama_set_abort_callback(struct llama_context * ctx, ggml_abort_callback abort_callback, void * abort_callback_data); +foreign import ccall "llama_set_abort_callback" + c_llama_set_abort_callback :: + CLlamaContext -> -- ctx + FunPtr (Ptr () -> IO CInt) -> -- abort_callback + Ptr () -> -- abort_callback_data + IO () + +-- | LLAMA_API void llama_synchronize(struct llama_context * ctx); +foreign import ccall "llama_synchronize" + c_llama_synchronize :: + CLlamaContext -> -- ctx + IO () + +-- | LLAMA_API float * llama_get_logits(struct llama_context * ctx); +foreign import ccall "llama_get_logits" + c_llama_get_logits :: + CLlamaContext -> -- ctx + IO (Ptr CFloat) + +-- | LLAMA_API float * llama_get_logits_ith(struct llama_context * ctx, int32_t i); +foreign import ccall "llama_get_logits_ith" + c_llama_get_logits_ith :: + CLlamaContext -> -- ctx + CInt -> -- i + IO (Ptr CFloat) + +-- | LLAMA_API float * llama_get_embeddings(struct llama_context * ctx); +foreign import ccall "llama_get_embeddings" + c_llama_get_embeddings :: + CLlamaContext -> -- ctx + IO (Ptr CFloat) + +-- | LLAMA_API float * llama_get_embeddings_ith(struct llama_context * ctx, int32_t i); +foreign import ccall "llama_get_embeddings_ith" + c_llama_get_embeddings_ith :: + CLlamaContext -> -- ctx + CInt -> -- i + IO (Ptr CFloat) + +-- | LLAMA_API float * llama_get_embeddings_seq(struct llama_context * ctx, llama_seq_id seq_id); +foreign import ccall "llama_get_embeddings_seq" + c_llama_get_embeddings_seq :: + CLlamaContext -> -- ctx + LlamaSeqId -> -- seq_id + IO (Ptr CFloat) diff --git a/src/Llama/Internal/Foreign/GGML.hs b/src/Llama/Internal/Foreign/GGML.hs new file mode 100644 --- /dev/null +++ b/src/Llama/Internal/Foreign/GGML.hs @@ -0,0 +1,4 @@ +module Llama.Internal.Foreign.GGML where + +-- GGML_API void ggml_backend_load_all(void); +-- GGML_API int64_t ggml_time_us(void); diff --git a/src/Llama/Internal/Foreign/KVCache.hs b/src/Llama/Internal/Foreign/KVCache.hs new file mode 100644 --- /dev/null +++ b/src/Llama/Internal/Foreign/KVCache.hs @@ -0,0 +1,141 @@ +module Llama.Internal.Foreign.KVCache ( + + c_llama_kv_cache_view_init_into + , c_llama_kv_cache_view_free + , c_llama_kv_cache_view_update + , c_llama_kv_self_n_tokens + , c_llama_kv_self_used_cells + , c_llama_kv_self_clear + , c_llama_kv_self_seq_rm + , c_llama_kv_self_seq_cp + , c_llama_kv_self_seq_keep + , c_llama_kv_self_seq_add + , c_llama_kv_self_seq_div + , c_llama_kv_self_seq_pos_max + , c_llama_kv_self_defrag + , c_llama_kv_self_can_shift + , c_llama_kv_self_update + +) where + +import Foreign +import Foreign.C +import Llama.Internal.Types + +-- | LLAMA_API struct llama_kv_cache_view llama_kv_cache_view_init(const struct llama_context * ctx, int32_t n_seq_max); +foreign import ccall "llama_kv_cache_view_init_into" + c_llama_kv_cache_view_init_into :: + CLlamaContext -> -- ctx + CInt -> -- n_seq_max + Ptr LlamaKvCacheView -> + IO () + +-- | LLAMA_API void llama_kv_cache_view_free(struct llama_kv_cache_view * view); +foreign import ccall "llama_kv_cache_view_free" + c_llama_kv_cache_view_free :: + Ptr LlamaKvCacheView -> -- view + IO () + +-- |LLAMA_API void llama_kv_cache_view_update(const struct llama_context * ctx, struct llama_kv_cache_view * view); +foreign import ccall "llama_kv_cache_view_update" + c_llama_kv_cache_view_update :: + CLlamaContext -> -- ctx + Ptr LlamaKvCacheView -> -- view + IO () + +-- | LLAMA_API int32_t llama_kv_self_n_tokens(const struct llama_context * ctx); +foreign import ccall "llama_kv_self_n_tokens" + c_llama_kv_self_n_tokens :: + CLlamaContext -> -- ctx + IO CInt + +-- | LLAMA_API int32_t llama_kv_self_used_cells(const struct llama_context * ctx); +foreign import ccall "llama_kv_self_used_cells" + c_llama_kv_self_used_cells :: + CLlamaContext -> -- ctx + IO CInt + +-- | LLAMA_API void llama_kv_self_clear(struct llama_context * ctx); +foreign import ccall "llama_kv_self_clear" + c_llama_kv_self_clear :: + CLlamaContext -> -- ctx + IO () + +{- | LLAMA_API bool llama_kv_self_seq_rm( +| struct llama_context * ctx, llama_seq_id seq_id, llama_pos p0, llama_pos p1); +-} +foreign import ccall "llama_kv_self_seq_rm" + c_llama_kv_self_seq_rm :: + CLlamaContext -> -- ctx + LlamaSeqId -> -- seq_id + LlamaPos -> -- p0 + LlamaPos -> -- p1 + IO CBool + +{- | LLAMA_API void llama_kv_self_seq_cp( +| struct llama_context * ctx, llama_seq_id seq_id_src, llama_seq_id seq_id_dst, llama_pos p0, llama_pos p1); +-} +foreign import ccall "llama_kv_self_seq_cp" + c_llama_kv_self_seq_cp :: + CLlamaContext -> -- ctx + LlamaSeqId -> -- seq_id_src + LlamaSeqId -> -- seq_id_dst + LlamaPos -> -- p0 + LlamaPos -> -- p1 + IO () + +-- | LLAMA_API void llama_kv_self_seq_keep(struct llama_context * ctx, llama_seq_id seq_id); +foreign import ccall "llama_kv_self_seq_keep" + c_llama_kv_self_seq_keep :: + CLlamaContext -> -- ctx + LlamaSeqId -> -- seq_id + IO () + +{- | LLAMA_API void llama_kv_self_seq_add( +| struct llama_context * ctx, llama_seq_id seq_id, llama_pos p0, llama_pos p1, llama_pos delta); +-} +foreign import ccall "llama_kv_self_seq_add" + c_llama_kv_self_seq_add :: + CLlamaContext -> -- ctx + LlamaSeqId -> -- seq_id + LlamaPos -> -- p0 + LlamaPos -> -- p1 + LlamaPos -> -- delta + IO () + +{- | LLAMA_API void llama_kv_self_seq_div( +| struct llama_context * ctx, llama_seq_id seq_id, llama_pos p0, llama_pos p1, int d); +-} +foreign import ccall "llama_kv_self_seq_div" + c_llama_kv_self_seq_div :: + CLlamaContext -> -- ctx + LlamaSeqId -> -- seq_id + LlamaPos -> -- p0 + LlamaPos -> -- p1 + CInt -> -- d + IO () + +-- | LLAMA_API llama_pos llama_kv_self_seq_pos_max(struct llama_context * ctx, llama_seq_id seq_id); +foreign import ccall "llama_kv_self_seq_pos_max" + c_llama_kv_self_seq_pos_max :: + CLlamaContext -> -- ctx + LlamaSeqId -> -- seq_id + IO LlamaPos + +-- | LLAMA_API void llama_kv_self_defrag(struct llama_context * ctx); +foreign import ccall "llama_kv_self_defrag" + c_llama_kv_self_defrag :: + CLlamaContext -> -- ctx + IO () + +-- | LLAMA_API bool llama_kv_self_can_shift(const struct llama_context * ctx); +foreign import ccall "llama_kv_self_can_shift" + c_llama_kv_self_can_shift :: + Ptr CLlamaContext -> -- ctx + IO CBool + +-- | LLAMA_API void llama_kv_self_update(struct llama_context * ctx); +foreign import ccall "llama_kv_self_update" + c_llama_kv_self_update :: + CLlamaContext -> -- ctx + IO () diff --git a/src/Llama/Internal/Foreign/Model.hs b/src/Llama/Internal/Foreign/Model.hs new file mode 100644 --- /dev/null +++ b/src/Llama/Internal/Foreign/Model.hs @@ -0,0 +1,212 @@ +module Llama.Internal.Foreign.Model ( + + p_llama_free + , p_llama_model_free + , c_llama_free + , c_llama_init_from_model_wrap + , c_llama_model_default_params + , c_llama_model_load_from_file_wrap + , c_llama_model_free + , c_llama_model_get_vocab + , c_llama_model_load_from_splits + , c_llama_model_rope_type_into + , c_llama_model_n_ctx_train + , c_llama_model_n_embd + , c_llama_model_n_layer + , c_llama_model_n_head + , c_llama_model_n_head_kv + , c_llama_get_model + , c_llama_model_rope_freq_scale_train + , c_llama_vocab_type + , c_llama_model_meta_val_str + , c_llama_model_meta_count + , c_llama_model_meta_key_by_index + , c_llama_model_meta_val_str_by_index + , c_llama_model_desc + , c_llama_model_size + , c_llama_model_chat_template + , c_llama_model_n_params + , c_llama_model_has_encoder + , c_llama_model_has_decoder + , c_llama_model_decoder_start_token + , c_llama_model_is_recurrent + , c_llama_model_quantize + , c_llama_model_quantize_default_params + +) where + +import Foreign +import Foreign.C +import Llama.Internal.Types +import Llama.Internal.Types.Params + +-- | Foreign pointer to the C function for freeing the context. +foreign import ccall "&llama_free" + p_llama_free :: FinalizerPtr CLlamaContext + +-- | Foreign pointer to the C function for freeing the model. +foreign import ccall "&llama_model_free" + p_llama_model_free :: FinalizerPtr CLlamaModel + +foreign import ccall "llama_free" c_llama_free :: CLlamaContext -> IO () + +foreign import ccall "llama_init_from_model_into" + c_llama_init_from_model_wrap :: + CLlamaModel -> CLlamaContextParams -> IO CLlamaContext + +foreign import ccall "llama_model_default_params" + c_llama_model_default_params :: CLlamaModelParams -> IO () + +foreign import ccall "llama_model_load_from_file_wrap" + c_llama_model_load_from_file_wrap :: + CString -> CLlamaModelParams -> IO CLlamaModel + +foreign import ccall "llama_model_free" + c_llama_model_free :: CLlamaModel -> IO () + +foreign import ccall "llama_model_get_vocab" + c_llama_model_get_vocab :: CLlamaModel -> IO CLlamaVocab + +-- LLAMA_API struct llama_model * llama_model_load_from_splits( +-- const char ** paths, +-- size_t n_paths, +-- struct llama_model_params params); +foreign import ccall "llama_model_load_from_splits" + c_llama_model_load_from_splits :: + Ptr CString -> CSize -> CLlamaModelParams -> IO CLlamaModel + +-- | Get RoPE type from model +foreign import ccall safe "llama_model_rope_type_into" + c_llama_model_rope_type_into :: CLlamaModel -> Ptr CInt -> IO () + +-- | Get training context size +foreign import ccall "llama_model_n_ctx_train" + c_llama_model_n_ctx_train :: CLlamaModel -> IO Int32 + +-- | Get embedding dimension +foreign import ccall "llama_model_n_embd" + c_llama_model_n_embd :: CLlamaModel -> IO Int32 + +-- | Get number of layers +foreign import ccall "llama_model_n_layer" + c_llama_model_n_layer :: CLlamaModel -> IO Int32 + +-- | Get number of heads +foreign import ccall "llama_model_n_head" + c_llama_model_n_head :: CLlamaModel -> IO Int32 + +-- | Get number of key/value heads +foreign import ccall "llama_model_n_head_kv" + c_llama_model_n_head_kv :: CLlamaModel -> IO Int32 + +-- | Get model from context +foreign import ccall "llama_get_model" + c_llama_get_model :: CLlamaContext -> IO CLlamaModel + +-- For float llama_model_rope_freq_scale_train(const struct llama_model * model); +foreign import ccall "llama_model_rope_freq_scale_train" + c_llama_model_rope_freq_scale_train :: CLlamaModel -> IO CFloat + +-- For enum llama_vocab_type llama_vocab_type(const struct llama_vocab * vocab); +foreign import ccall "llama_vocab_type" + c_llama_vocab_type :: CLlamaVocab -> IO CInt + +-- | LLAMA_API int32_t llama_model_meta_val_str(const struct llama_model * model, const char * key, char * buf, size_t buf_size); +foreign import ccall "llama_model_meta_val_str" + c_llama_model_meta_val_str :: + CLlamaModel -> -- model + CString -> -- key + CString -> -- buffer + CSize -> -- buffer size + IO CInt + +-- | LLAMA_API int32_t llama_model_meta_count(const struct llama_model * model); +foreign import ccall "llama_model_meta_count" + c_llama_model_meta_count :: + CLlamaModel -> -- model + IO CInt + +-- | LLAMA_API int32_t llama_model_meta_key_by_index(const struct llama_model * model, int32_t i, char * buf, size_t buf_size); +foreign import ccall "llama_model_meta_key_by_index" + c_llama_model_meta_key_by_index :: + CLlamaModel -> -- model + CInt -> -- index + CString -> -- buffer + CSize -> -- buffer size + IO CInt + +-- | LLAMA_API int32_t llama_model_meta_val_str_by_index(const struct llama_model * model, int32_t i, char * buf, size_t buf_size); +foreign import ccall "llama_model_meta_val_str_by_index" + c_llama_model_meta_val_str_by_index :: + CLlamaModel -> -- model + CInt -> -- index + CString -> -- buffer + CSize -> -- buffer size + IO CInt + +-- | LLAMA_API int32_t llama_model_desc(const struct llama_model * model, char * buf, size_t buf_size); +foreign import ccall "llama_model_desc" + c_llama_model_desc :: + CLlamaModel -> -- model + CString -> -- buffer + CSize -> -- buffer size + IO CInt + +-- | LLAMA_API uint64_t llama_model_size(const struct llama_model * model); +foreign import ccall "llama_model_size" + c_llama_model_size :: + CLlamaModel -> -- model + IO Word64 + +-- | LLAMA_API const char * llama_model_chat_template(const struct llama_model * model, const char * name); +foreign import ccall "llama_model_chat_template" + c_llama_model_chat_template :: + CLlamaModel -> -- model + CString -> -- optional name (NULL = default) + IO CString + +-- | LLAMA_API uint64_t llama_model_n_params(const struct llama_model * model); +foreign import ccall "llama_model_n_params" + c_llama_model_n_params :: + CLlamaModel -> -- model + IO Word64 + +-- | LLAMA_API bool llama_model_has_encoder(const struct llama_model * model); +foreign import ccall "llama_model_has_encoder" + c_llama_model_has_encoder :: + CLlamaModel -> -- model + IO CBool + +-- | LLAMA_API bool llama_model_has_decoder(const struct llama_model * model); +foreign import ccall "llama_model_has_decoder" + c_llama_model_has_decoder :: + CLlamaModel -> -- model + IO CBool + +-- | LLAMA_API llama_token llama_model_decoder_start_token(const struct llama_model * model); +foreign import ccall "llama_model_decoder_start_token" + c_llama_model_decoder_start_token :: + CLlamaModel -> -- model + IO LlamaToken + +-- | LLAMA_API bool llama_model_is_recurrent(const struct llama_model * model); +foreign import ccall "llama_model_is_recurrent" + c_llama_model_is_recurrent :: + CLlamaModel -> -- model + IO CBool + +{- | LLAMA_API uint32_t llama_model_quantize( + const char * fname_inp, + const char * fname_out, + const llama_model_quantize_params * params); +-} +foreign import ccall "llama_model_quantize" + c_llama_model_quantize :: + CString -> -- fname_inp + CString -> -- fname_out + CLlamaModelQuantizeParams -> -- params + IO Word32 + +-- LLAMA_API struct llama_model_quantize_params llama_model_quantize_default_params(void); +foreign import ccall "llama_model_quantize_default_params" + c_llama_model_quantize_default_params :: IO CLlamaModelQuantizeParams diff --git a/src/Llama/Internal/Foreign/Performance.hs b/src/Llama/Internal/Foreign/Performance.hs new file mode 100644 --- /dev/null +++ b/src/Llama/Internal/Foreign/Performance.hs @@ -0,0 +1,51 @@ +module Llama.Internal.Foreign.Performance ( + + c_llama_perf_context + , c_llama_perf_context_print + , c_llama_perf_context_reset + , c_llama_perf_sampler + , c_llama_perf_sampler_print + , c_llama_perf_sampler_reset + +) where + +import Llama.Internal.Types +import Foreign + +-- | LLAMA_API struct llama_perf_context_data llama_perf_context(const struct llama_context * ctx); +foreign import ccall "llama_perf_context_into" + c_llama_perf_context :: + CLlamaContext -> -- ctx + Ptr LlamaPerfContextData -> -- output buffer + IO () + +-- | LLAMA_API void llama_perf_context_print(const struct llama_context * ctx); +foreign import ccall "llama_perf_context_print" + c_llama_perf_context_print :: + CLlamaContext -> -- ctx + IO () + +-- | LLAMA_API void llama_perf_context_reset(struct llama_context * ctx); +foreign import ccall "llama_perf_context_reset" + c_llama_perf_context_reset :: + CLlamaContext -> -- ctx + IO () + +-- | LLAMA_API struct llama_perf_sampler_data llama_perf_sampler(const struct llama_sampler * chain); +foreign import ccall "llama_perf_sampler_into" + c_llama_perf_sampler :: + (Ptr LlamaSampler) -> -- chain + Ptr LlamaPerfSamplerData -> -- output buffer + IO () + +-- | LLAMA_API void llama_perf_sampler_print(const struct llama_sampler * chain); +foreign import ccall "llama_perf_sampler_print" + c_llama_perf_sampler_print :: + (Ptr LlamaSampler) -> -- chain + IO () + +-- | LLAMA_API void llama_perf_sampler_reset(struct llama_sampler * chain); +foreign import ccall "llama_perf_sampler_reset" + c_llama_perf_sampler_reset :: + Ptr LlamaSampler -> -- chain + IO () diff --git a/src/Llama/Internal/Foreign/Sampler.hs b/src/Llama/Internal/Foreign/Sampler.hs new file mode 100644 --- /dev/null +++ b/src/Llama/Internal/Foreign/Sampler.hs @@ -0,0 +1,309 @@ +module Llama.Internal.Foreign.Sampler ( + c_llama_sampler_init_greedy + , c_llama_sampler_init_dist + , c_llama_sampler_init_top_k + , c_llama_sampler_init_top_p + , c_llama_sampler_init_min_p + , c_llama_sampler_init_typical + , c_llama_sampler_init_temp + , c_llama_sampler_init_temp_ext + , c_llama_sampler_init_xtc + , c_llama_sampler_init_top_n_sigma + , c_llama_sampler_init_mirostat + , c_llama_sampler_init_mirostat_v2 + , c_llama_sampler_init_grammar + , c_llama_sampler_init_grammar_lazy_patterns + , c_llama_sampler_init_penalties + , c_llama_sampler_init_dry + , c_llama_sampler_init_logit_bias + , c_llama_sampler_init_infill + , c_llama_sampler_sample + , c_llama_sampler_get_seed + , c_llama_sampler_chain_default_params_into + , c_llama_sampler_init + , c_llama_sampler_name + , c_llama_sampler_accept + , c_llama_sampler_apply + , c_llama_sampler_reset + , c_llama_sampler_clone + , p_llama_sampler_free + , c_llama_sampler_free + , c_llama_sampler_chain_init + , c_llama_sampler_chain_add + , c_llama_sampler_chain_get + , c_llama_sampler_chain_n + , c_llama_sampler_chain_remove +) where + +import Foreign +import Foreign.C +import Llama.Internal.Types +import Llama.Internal.Types.Params + +foreign import ccall "llama_sampler_chain_default_params_into" + c_llama_sampler_chain_default_params_into :: Ptr LlamaSamplerChainParams -> IO () + +-- | LLAMA_API struct llama_sampler * llama_sampler_init(const struct llama_sampler_i * iface, llama_sampler_context_t ctx); +foreign import ccall "llama_sampler_init" + c_llama_sampler_init :: + Ptr LlamaSamplerI -> -- iface (vtable) + LlamaSamplerContext -> -- ctx (opaque context) + IO (Ptr LlamaSampler) + +-- | LLAMA_API const char * llama_sampler_name(const struct llama_sampler * smpl); +foreign import ccall "llama_sampler_name" + c_llama_sampler_name :: + Ptr LlamaSampler -> -- smpl + IO CString + +-- | LLAMA_API void llama_sampler_accept(struct llama_sampler * smpl, llama_token token); +foreign import ccall "llama_sampler_accept" + c_llama_sampler_accept :: + Ptr LlamaSampler -> -- smpl + LlamaToken -> -- token + IO () + +-- | LLAMA_API void llama_sampler_apply(struct llama_sampler * smpl, llama_token_data_array * cur_p); +foreign import ccall "llama_sampler_apply" + c_llama_sampler_apply :: + Ptr LlamaSampler -> -- smpl + Ptr LlamaTokenDataArray -> -- cur_p + IO () + +-- | LLAMA_API void llama_sampler_reset(struct llama_sampler * smpl); +foreign import ccall "llama_sampler_reset" + c_llama_sampler_reset :: + Ptr LlamaSampler -> -- smpl + IO () + +-- | LLAMA_API struct llama_sampler * llama_sampler_clone(const struct llama_sampler * smpl); +foreign import ccall "llama_sampler_clone" + c_llama_sampler_clone :: + Ptr LlamaSampler -> -- smpl + IO (Ptr LlamaSampler) + +-- | Foreign pointer to the C function for freeing the SamplerChain. +foreign import ccall "&llama_sampler_free" + p_llama_sampler_free :: FinalizerPtr LlamaSampler + +-- | LLAMA_API void llama_sampler_free(struct llama_sampler * smpl); +foreign import ccall "llama_sampler_free" + c_llama_sampler_free :: + Ptr LlamaSampler -> -- smpl + IO () + +-- | LLAMA_API struct llama_sampler * llama_sampler_chain_init(struct llama_sampler_chain_params params); +foreign import ccall "llama_sampler_chain_init" + c_llama_sampler_chain_init :: + LlamaSamplerChainParams -> -- params + IO (Ptr LlamaSampler) + +-- | LLAMA_API void llama_sampler_chain_add(struct llama_sampler * chain, struct llama_sampler * smpl); +foreign import ccall "llama_sampler_chain_add" + c_llama_sampler_chain_add :: + Ptr LlamaSampler -> -- chain + Ptr LlamaSampler -> -- smpl + IO () + +-- | LLAMA_API struct llama_sampler * llama_sampler_chain_get(const struct llama_sampler * chain, int32_t i); +foreign import ccall "llama_sampler_chain_get" + c_llama_sampler_chain_get :: + Ptr LlamaSampler -> -- chain + CInt -> -- i + IO (Ptr LlamaSampler) + +-- | LLAMA_API int llama_sampler_chain_n(const struct llama_sampler * chain); +foreign import ccall "llama_sampler_chain_n" + c_llama_sampler_chain_n :: + Ptr LlamaSampler -> -- chain + IO CInt + +-- | LLAMA_API struct llama_sampler * llama_sampler_chain_remove(struct llama_sampler * chain, int32_t i); +foreign import ccall "llama_sampler_chain_remove" + c_llama_sampler_chain_remove :: + Ptr LlamaSampler -> -- chain + CInt -> -- i + IO (Ptr LlamaSampler) + +-- | LLAMA_API struct llama_sampler * llama_sampler_init_greedy(void); +foreign import ccall "llama_sampler_init_greedy" + c_llama_sampler_init_greedy :: + IO (Ptr LlamaSampler) + +-- | LLAMA_API struct llama_sampler * llama_sampler_init_dist(uint32_t seed); +foreign import ccall "llama_sampler_init_dist" + c_llama_sampler_init_dist :: + CUInt -> -- seed + IO (Ptr LlamaSampler) + +-- | LLAMA_API struct llama_sampler * llama_sampler_init_top_k(int32_t k); +foreign import ccall "llama_sampler_init_top_k" + c_llama_sampler_init_top_k :: + CInt -> -- k + IO (Ptr LlamaSampler) + +-- | LLAMA_API struct llama_sampler * llama_sampler_init_top_p(float p, size_t min_keep); +foreign import ccall "llama_sampler_init_top_p" + c_llama_sampler_init_top_p :: + CFloat -> -- p + CSize -> -- min_keep + IO (Ptr LlamaSampler) + +-- | LLAMA_API struct llama_sampler * llama_sampler_init_min_p(float p, size_t min_keep); +foreign import ccall "llama_sampler_init_min_p" + c_llama_sampler_init_min_p :: + CFloat -> -- p + CSize -> -- min_keep + IO (Ptr LlamaSampler) + +-- | LLAMA_API struct llama_sampler * llama_sampler_init_typical(float p, size_t min_keep); +foreign import ccall "llama_sampler_init_typical" + c_llama_sampler_init_typical :: + CFloat -> -- p + CSize -> -- min_keep + IO (Ptr LlamaSampler) + +-- | LLAMA_API struct llama_sampler * llama_sampler_init_temp(float t); +foreign import ccall "llama_sampler_init_temp" + c_llama_sampler_init_temp :: + CFloat -> -- t + IO (Ptr LlamaSampler) + +-- | LLAMA_API struct llama_sampler * llama_sampler_init_temp_ext(float t, float delta, float exponent); +foreign import ccall "llama_sampler_init_temp_ext" + c_llama_sampler_init_temp_ext :: + CFloat -> -- t + CFloat -> -- delta + CFloat -> -- exponent + IO (Ptr LlamaSampler) + +-- | LLAMA_API struct llama_sampler * llama_sampler_init_xtc(float p, float t, size_t min_keep, uint32_t seed); +foreign import ccall "llama_sampler_init_xtc" + c_llama_sampler_init_xtc :: + CFloat -> -- p + CFloat -> -- t + CSize -> -- min_keep + CUInt -> -- seed + IO (Ptr LlamaSampler) + +-- | LLAMA_API struct llama_sampler * llama_sampler_init_top_n_sigma(float n); +foreign import ccall "llama_sampler_init_top_n_sigma" + c_llama_sampler_init_top_n_sigma :: + CFloat -> -- n + IO (Ptr LlamaSampler) + +-- | LLAMA_API struct llama_sampler * llama_sampler_init_mirostat(int32_t n_vocab, uint32_t seed, float tau, float eta, int32_t m); +foreign import ccall "llama_sampler_init_mirostat" + c_llama_sampler_init_mirostat :: + CInt -> -- n_vocab + CUInt -> -- seed + CFloat -> -- tau + CFloat -> -- eta + CInt -> -- m + IO (Ptr LlamaSampler) + +-- | LLAMA_API struct llama_sampler * llama_sampler_init_mirostat_v2(uint32_t seed, float tau, float eta); +foreign import ccall "llama_sampler_init_mirostat_v2" + c_llama_sampler_init_mirostat_v2 :: + CUInt -> -- seed + CFloat -> -- tau + CFloat -> -- eta + IO (Ptr LlamaSampler) + +-- | LLAMA_API struct llama_sampler * llama_sampler_init_grammar(const struct llama_vocab * vocab, const char * grammar_str, const char * grammar_root); +foreign import ccall "llama_sampler_init_grammar" + c_llama_sampler_init_grammar :: + CLlamaVocab -> -- vocab + CString -> -- grammar_str + CString -> -- grammar_root + IO (Ptr LlamaSampler) + +{- | LLAMA_API struct llama_sampler * llama_sampler_init_grammar_lazy_patterns( +| const struct llama_vocab * vocab, +| const char * grammar_str, +| const char * grammar_root, +| const char ** trigger_patterns, +| size_t num_trigger_patterns, +| const llama_token * trigger_tokens, +| size_t num_trigger_tokens); +-} +foreign import ccall "llama_sampler_init_grammar_lazy_patterns" + c_llama_sampler_init_grammar_lazy_patterns :: + CLlamaVocab -> -- vocab + CString -> -- grammar_str + CString -> -- grammar_root + Ptr CString -> -- trigger_patterns (array of CStrings) + CSize -> -- num_trigger_patterns + Ptr LlamaToken -> -- trigger_tokens + CSize -> -- num_trigger_tokens + IO (Ptr LlamaSampler) + +{- | LLAMA_API struct llama_sampler * llama_sampler_init_penalties( +| int32_t penalty_last_n, +| float penalty_repeat, +| float penalty_freq, +| float penalty_present); +-} +foreign import ccall "llama_sampler_init_penalties" + c_llama_sampler_init_penalties :: + CInt -> -- penalty_last_n + CFloat -> -- penalty_repeat + CFloat -> -- penalty_freq + CFloat -> -- penalty_present + IO (Ptr LlamaSampler) + +{- | LLAMA_API struct llama_sampler * llama_sampler_init_dry( +| const struct llama_vocab * vocab, +| int32_t n_ctx_train, +| float dry_multiplier, +| float dry_base, +| int32_t dry_allowed_length, +| int32_t dry_penalty_last_n, +| const char ** seq_breakers, +| size_t num_breakers); +-} +foreign import ccall "llama_sampler_init_dry" + c_llama_sampler_init_dry :: + CLlamaVocab -> -- vocab + CInt -> -- n_ctx_train + CFloat -> -- dry_multiplier + CFloat -> -- dry_base + CInt -> -- dry_allowed_length + CInt -> -- dry_penalty_last_n + Ptr CString -> -- seq_breakers (array of CStrings) + CSize -> -- num_breakers + IO (Ptr LlamaSampler) + +{- | LLAMA_API struct llama_sampler * llama_sampler_init_logit_bias( +| int32_t n_vocab, +| int32_t n_logit_bias, +| const llama_logit_bias * logit_bias); +-} +foreign import ccall "llama_sampler_init_logit_bias" + c_llama_sampler_init_logit_bias :: + CInt -> -- n_vocab + CInt -> -- n_logit_bias + Ptr LlamaLogitBias -> -- logit_bias (opaque struct) + IO (Ptr LlamaSampler) + +-- | LLAMA_API struct llama_sampler * llama_sampler_init_infill(const struct llama_vocab * vocab); +foreign import ccall "llama_sampler_init_infill" + c_llama_sampler_init_infill :: + CLlamaVocab -> -- vocab + IO (Ptr LlamaSampler) + +-- | LLAMA_API uint32_t llama_sampler_get_seed(const struct llama_sampler * smpl); +foreign import ccall "llama_sampler_get_seed" + c_llama_sampler_get_seed :: + Ptr LlamaSampler -> -- smpl + + -- | Returns seed used by the sampler + IO CUInt + +-- | LLAMA_API llama_token llama_sampler_sample(struct llama_sampler * smpl, struct llama_context * ctx, int32_t idx); +foreign import ccall "llama_sampler_sample" + c_llama_sampler_sample :: + Ptr LlamaSampler -> -- smpl + CLlamaContext -> -- ctx + CInt -> -- idx + IO LlamaToken diff --git a/src/Llama/Internal/Foreign/Split.hs b/src/Llama/Internal/Foreign/Split.hs new file mode 100644 --- /dev/null +++ b/src/Llama/Internal/Foreign/Split.hs @@ -0,0 +1,46 @@ +module Llama.Internal.Foreign.Split ( + c_llama_split_path + , c_llama_split_prefix + , c_llama_print_system_info + , c_llama_log_set +) where + +import Foreign +import Foreign.C + +-- | LLAMA_API int llama_split_path(char * split_path, size_t maxlen, const char * path_prefix, int split_no, int split_count); +foreign import ccall "llama_split_path" + c_llama_split_path :: + CString -> -- split_path (output buffer) + CSize -> -- maxlen + CString -> -- path_prefix + CInt -> -- split_no + CInt -> -- split_count + + -- | Returns length of written string or required buffer size + IO CInt + +-- | LLAMA_API int llama_split_prefix(char * split_prefix, size_t maxlen, const char * split_path, int split_no, int split_count); +foreign import ccall "llama_split_prefix" + c_llama_split_prefix :: + CString -> -- split_prefix (output buffer) + CSize -> -- maxlen + CString -> -- split_path + CInt -> -- split_no + CInt -> -- split_count + + -- | Returns length of written string or required buffer size + IO CInt + +-- | LLAMA_API const char * llama_print_system_info(void); +foreign import ccall "llama_print_system_info" + c_llama_print_system_info :: + -- | Returns system info string + IO CString + +-- | LLAMA_API void llama_log_set(ggml_log_callback log_callback, void * user_data); +foreign import ccall "llama_log_set" + c_llama_log_set :: + FunPtr (CInt -> CString -> Ptr () -> IO ()) -> -- log_callback + Ptr () -> -- user_data + IO () diff --git a/src/Llama/Internal/Foreign/State.hs b/src/Llama/Internal/Foreign/State.hs new file mode 100644 --- /dev/null +++ b/src/Llama/Internal/Foreign/State.hs @@ -0,0 +1,92 @@ +module Llama.Internal.Foreign.State where + +import Foreign +import Foreign.C +import Llama.Internal.Types + +-- | LLAMA_API size_t llama_state_get_size(struct llama_context * ctx); +foreign import ccall "llama_state_get_size" + c_llama_state_get_size :: + CLlamaContext -> -- ctx + IO CSize + +-- | LLAMA_API size_t llama_state_get_data(struct llama_context * ctx, uint8_t * dst, size_t size); +foreign import ccall "llama_state_get_data" + c_llama_state_get_data :: + CLlamaContext -> -- ctx + Ptr Word8 -> -- dst (buffer) + CSize -> -- size + IO CSize + +-- | LLAMA_API size_t llama_state_set_data(struct llama_context * ctx, const uint8_t * src, size_t size); +foreign import ccall "llama_state_set_data" + c_llama_state_set_data :: + CLlamaContext -> -- ctx + Ptr Word8 -> -- src (buffer) + CSize -> -- size + IO CSize + +-- | LLAMA_API bool llama_state_load_file(struct llama_context * ctx, const char * path_session, llama_token * tokens_out, size_t n_token_capacity, size_t * n_token_count_out); +foreign import ccall "llama_state_load_file" + c_llama_state_load_file :: + CLlamaContext -> -- ctx + CString -> -- path_session + Ptr LlamaToken -> -- tokens_out + CSize -> -- n_token_capacity + Ptr CSize -> -- n_token_count_out + IO CBool + +-- | LLAMA_API bool llama_state_save_file(struct llama_context * ctx, const char * path_session, const llama_token * tokens, size_t n_token_count); +foreign import ccall "llama_state_save_file" + c_llama_state_save_file :: + CLlamaContext -> -- ctx + CString -> -- path_session + Ptr LlamaToken -> -- tokens + CSize -> -- n_token_count + IO CBool + +-- | LLAMA_API size_t llama_state_seq_get_size(struct llama_context * ctx, llama_seq_id seq_id); +foreign import ccall "llama_state_seq_get_size" + c_llama_state_seq_get_size :: + CLlamaContext -> -- ctx + LlamaSeqId -> -- seq_id + IO CSize + +-- | LLAMA_API size_t llama_state_seq_get_data(struct llama_context * ctx, uint8_t * dst, size_t size, llama_seq_id seq_id); +foreign import ccall "llama_state_seq_get_data" + c_llama_state_seq_get_data :: + CLlamaContext -> -- ctx + Ptr Word8 -> -- dst (buffer) + CSize -> -- size + LlamaSeqId -> -- seq_id + IO CSize + +-- | LLAMA_API size_t llama_state_seq_set_data(struct llama_context * ctx, const uint8_t * src, size_t size, llama_seq_id dest_seq_id); +foreign import ccall "llama_state_seq_set_data" + c_llama_state_seq_set_data :: + CLlamaContext -> -- ctx + Ptr Word8 -> -- src (buffer) + CSize -> -- size + LlamaSeqId -> -- dest_seq_id + IO CSize + +-- | LLAMA_API size_t llama_state_seq_save_file(struct llama_context * ctx, const char * filepath, llama_seq_id seq_id, const llama_token * tokens, size_t n_token_count); +foreign import ccall "llama_state_seq_save_file" + c_llama_state_seq_save_file :: + CLlamaContext -> -- ctx + CString -> -- filepath + LlamaSeqId -> -- seq_id + Ptr LlamaToken -> -- tokens + CSize -> -- n_token_count + IO CSize + +-- | LLAMA_API size_t llama_state_seq_load_file(struct llama_context * ctx, const char * filepath, llama_seq_id dest_seq_id, llama_token * tokens_out, size_t n_token_capacity, size_t * n_token_count_out); +foreign import ccall "llama_state_seq_load_file" + c_llama_state_seq_load_file :: + CLlamaContext -> -- ctx + CString -> -- filepath + LlamaSeqId -> -- dest_seq_id + Ptr LlamaToken -> -- tokens_out + CSize -> -- n_token_capacity + Ptr CSize -> -- n_token_count_out + IO CSize diff --git a/src/Llama/Internal/Foreign/Tokenize.hs b/src/Llama/Internal/Foreign/Tokenize.hs new file mode 100644 --- /dev/null +++ b/src/Llama/Internal/Foreign/Tokenize.hs @@ -0,0 +1,63 @@ +module Llama.Internal.Foreign.Tokenize where + +import Foreign +import Foreign.C +import Llama.Internal.Types + +{- | LLAMA_API int32_t llama_tokenize( +| const struct llama_vocab * vocab, +| const char * text, +| int32_t text_len, +| llama_token * tokens, +| int32_t n_tokens_max, +| bool add_special, +| bool parse_special); +-} +foreign import ccall "llama_tokenize" + c_llama_tokenize :: + CLlamaVocab -> -- vocab + CString -> -- text + CInt -> -- text_len + Ptr LlamaToken -> -- tokens + CInt -> -- n_tokens_max + CBool -> -- add_special + CBool -> -- parse_special + IO CInt + +{- | LLAMA_API int32_t llama_token_to_piece( +| const struct llama_vocab * vocab, +| llama_token token, +| char * buf, +| int32_t length, +| int32_t lstrip, +| bool special); +-} +foreign import ccall "llama_token_to_piece" + c_llama_token_to_piece :: + CLlamaVocab -> -- vocab + LlamaToken -> -- token + CString -> -- buf + CInt -> -- length + CInt -> -- lstrip + CBool -> -- special + IO CInt + +{- | LLAMA_API int32_t llama_detokenize( +| const struct llama_vocab * vocab, +| const llama_token * tokens, +| int32_t n_tokens, +| char * text, +| int32_t text_len_max, +| bool remove_special, +| bool unparse_special); +-} +foreign import ccall "llama_detokenize" + c_llama_detokenize :: + CLlamaVocab -> -- vocab + Ptr LlamaToken -> -- tokens + CInt -> -- n_tokens + CString -> -- text + CInt -> -- text_len_max + CBool -> -- remove_special + CBool -> -- unparse_special + IO CInt diff --git a/src/Llama/Internal/Foreign/Vocab.hs b/src/Llama/Internal/Foreign/Vocab.hs new file mode 100644 --- /dev/null +++ b/src/Llama/Internal/Foreign/Vocab.hs @@ -0,0 +1,133 @@ +module Llama.Internal.Foreign.Vocab where + +import Foreign.C +import Llama.Internal.Types + +--LLAMA_API int32_t llama_n_vocab (const struct llama_vocab * vocab) + +foreign import ccall "llama_n_vocab" c_llama_n_vocab :: CLlamaVocab -> IO CInt + +foreign import ccall "llama_vocab_n_tokens" + c_llama_vocab_n_tokens :: + CLlamaVocab -> IO CInt + +-- | LLAMA_API const char * llama_vocab_get_text(const struct llama_vocab * vocab, llama_token token); +foreign import ccall "llama_vocab_get_text" + c_llama_vocab_get_text :: + CLlamaVocab -> -- vocab + LlamaToken -> -- token + IO CString + +-- | LLAMA_API float llama_vocab_get_score(const struct llama_vocab * vocab, llama_token token); +foreign import ccall "llama_vocab_get_score" + c_llama_vocab_get_score :: + CLlamaVocab -> -- vocab + LlamaToken -> -- token + IO CFloat + +-- | LLAMA_API enum llama_token_attr llama_vocab_get_attr(const struct llama_vocab * vocab, llama_token token); +foreign import ccall "llama_vocab_get_attr" + c_llama_vocab_get_attr :: + CLlamaVocab -> -- vocab + LlamaToken -> -- token + + -- | Returns raw integer representing enum value + IO CInt + +-- | LLAMA_API bool llama_vocab_is_eog(const struct llama_vocab * vocab, llama_token token); +foreign import ccall "llama_vocab_is_eog" + c_llama_vocab_is_eog :: + CLlamaVocab -> -- vocab + LlamaToken -> -- token + IO CBool + +-- | LLAMA_API bool llama_vocab_is_control(const struct llama_vocab * vocab, llama_token token); +foreign import ccall "llama_vocab_is_control" + c_llama_vocab_is_control :: + CLlamaVocab -> -- vocab + LlamaToken -> -- token + IO CBool + +-- | LLAMA_API llama_token llama_vocab_bos(const struct llama_vocab * vocab); // beginning-of-sentence +foreign import ccall "llama_vocab_bos" + c_llama_vocab_bos :: + CLlamaVocab -> -- vocab + IO LlamaToken + +-- | LLAMA_API llama_token llama_vocab_eos(const struct llama_vocab * vocab); // end-of-sentence +foreign import ccall "llama_vocab_eos" + c_llama_vocab_eos :: + CLlamaVocab -> -- vocab + IO LlamaToken + +-- | LLAMA_API llama_token llama_vocab_eot(const struct llama_vocab * vocab); // end-of-turn +foreign import ccall "llama_vocab_eot" + c_llama_vocab_eot :: + CLlamaVocab -> -- vocab + IO LlamaToken + +-- | LLAMA_API llama_token llama_vocab_sep(const struct llama_vocab * vocab); // sentence separator +foreign import ccall "llama_vocab_sep" + c_llama_vocab_sep :: + CLlamaVocab -> -- vocab + IO LlamaToken + +-- | LLAMA_API llama_token llama_vocab_nl (const struct llama_vocab * vocab); // next-line +foreign import ccall "llama_vocab_nl" + c_llama_vocab_nl :: + CLlamaVocab -> -- vocab + IO LlamaToken + +-- | LLAMA_API llama_token llama_vocab_pad(const struct llama_vocab * vocab); // padding +foreign import ccall "llama_vocab_pad" + c_llama_vocab_pad :: + CLlamaVocab -> -- vocab + IO LlamaToken + +-- | LLAMA_API bool llama_vocab_get_add_bos(const struct llama_vocab * vocab); +foreign import ccall "llama_vocab_get_add_bos" + c_llama_vocab_get_add_bos :: + CLlamaVocab -> -- vocab + IO CBool + +-- | LLAMA_API bool llama_vocab_get_add_eos(const struct llama_vocab * vocab); +foreign import ccall "llama_vocab_get_add_eos" + c_llama_vocab_get_add_eos :: + CLlamaVocab -> -- vocab + IO CBool + +-- | LLAMA_API llama_token llama_vocab_fim_pre(const struct llama_vocab * vocab); +foreign import ccall "llama_vocab_fim_pre" + c_llama_vocab_fim_pre :: + CLlamaVocab -> -- vocab + IO LlamaToken + +-- | LLAMA_API llama_token llama_vocab_fim_suf(const struct llama_vocab * vocab); +foreign import ccall "llama_vocab_fim_suf" + c_llama_vocab_fim_suf :: + CLlamaVocab -> -- vocab + IO LlamaToken + +-- | LLAMA_API llama_token llama_vocab_fim_mid(const struct llama_vocab * vocab); +foreign import ccall "llama_vocab_fim_mid" + c_llama_vocab_fim_mid :: + CLlamaVocab -> -- vocab + IO LlamaToken + +-- | LLAMA_API llama_token llama_vocab_fim_pad(const struct llama_vocab * vocab); +foreign import ccall "llama_vocab_fim_pad" + c_llama_vocab_fim_pad :: + CLlamaVocab -> -- vocab + IO LlamaToken + +-- | LLAMA_API llama_token llama_vocab_fim_rep(const struct llama_vocab * vocab); +foreign import ccall "llama_vocab_fim_rep" + c_llama_vocab_fim_rep :: + CLlamaVocab -> -- vocab + IO LlamaToken + +-- | LLAMA_API llama_token llama_vocab_fim_sep(const struct llama_vocab * vocab); +foreign import ccall "llama_vocab_fim_sep" + c_llama_vocab_fim_sep :: + CLlamaVocab -> -- vocab + IO LlamaToken diff --git a/src/Llama/Internal/Types.hs b/src/Llama/Internal/Types.hs new file mode 100644 --- /dev/null +++ b/src/Llama/Internal/Types.hs @@ -0,0 +1,214 @@ +{-# LANGUAGE DeriveAnyClass #-} +{-# LANGUAGE DeriveGeneric #-} + +module Llama.Internal.Types + ( -- * Types + LlamaTokenData (..) + , LlamaTokenDataArray (..) + , LlamaBatch (..) + , AddBos (..) + , LlamaToken + , LlamaKvCacheView (..) + , LlamaSeqId + , LlamaPos + , LlamaChatMessage (..) + , LlamaSamplerI (..) + , LlamaSampler (..) + , LlamaSamplerContext + , LlamaPerfContextData (..) + , LlamaPerfSamplerData (..) + , LlamaLogitBias (..) + + -- * Raw pointers + , CLlamaVocab (..) + , CLlamaModel (..) + , CLlamaContext (..) + , CLlamaKVCache (..) + , CLlamaAdapterLora (..) + + -- * Managed handles + , Vocab (..) + , Batch (..) + , Model (..) + , Context (..) + , Sampler (..) + , KVCache (..) + , AdapterLora (..) + ) where + +import Foreign +import Foreign.C.String +import Foreign.C.Types +import Foreign.Storable.Generic +import GHC.Generics + +data AddBos = Always | Never + +data LlamaLogitBias = LlamaLogitBias + { tokenLogitBias :: LlamaToken + , bias :: Float + } + deriving (Show, Eq, Generic, GStorable) + +-- struct llama_adapter_lora + +newtype CLlamaAdapterLora = CLlamaAdapterLora (Ptr CLlamaAdapterLora) + deriving (Show, Eq) + +-- | Raw pointer to llama_vocab struct +newtype CLlamaVocab = CLlamaVocab (Ptr CLlamaVocab) + deriving (Show, Eq) + +-- | Raw pointer to llama_model struct +newtype CLlamaModel = CLlamaModel (Ptr CLlamaModel) + deriving (Show, Eq) + +-- | Raw pointer to llama_context struct +newtype CLlamaContext = CLlamaContext (Ptr CLlamaContext) + deriving (Show, Eq) + +newtype CLlamaKVCache = CLlamaKVCache (Ptr CLlamaKVCache) + deriving (Show, Eq) + +-- | Managed batch pointer with automatic cleanup +newtype Batch = Batch (Ptr LlamaBatch) + deriving (Show, Eq) + +-- | Managed vocabulary handle with automatic cleanup. +newtype Vocab = Vocab (ForeignPtr CLlamaVocab) + deriving (Show, Eq) + +-- | Managed model handle with automatic cleanup. +newtype Model = Model (ForeignPtr CLlamaModel) + deriving (Show, Eq) + +-- | Managed context handle with automatic cleanup. +newtype Context = Context (ForeignPtr CLlamaContext) + deriving (Show, Eq) + +-- | Managed sampler handle with automatic cleanup. +newtype Sampler = Sampler (ForeignPtr LlamaSampler) + deriving (Show, Eq) + +newtype KVCache = KVCache (ForeignPtr CLlamaKVCache) + deriving (Show, Eq) + +newtype AdapterLora = AdapterLora (ForeignPtr CLlamaAdapterLora) + deriving (Show, Eq) + +-- type PtrLlamaContext = Ptr LlamaContext +type LlamaToken = CInt +type LlamaPos = CInt +type LlamaSeqId = CInt + +data LlamaChatMessage = LlamaChatMessage + { role :: CString + -- ^ "system", "user", "assistant" + , content :: CString + -- ^ Message text + } + deriving (Show, Eq, Generic, GStorable) + +data LlamaTokenData = LlamaTokenData + { id :: LlamaToken + -- ^ token id + , logit :: CFloat + -- ^ log-odds of the token + , p :: CFloat + -- ^ probability of the token + } + deriving (Show, Eq, Generic, GStorable) + +data LlamaTokenDataArray = LlamaTokenDataArray + { data_ :: Ptr LlamaTokenData + , size_ :: CSize + , selected :: CLong + -- ^ this is the index in the data array (i.e. not the token id) + , sorted :: CBool + } + deriving (Eq, Show, Generic, GStorable) + +data LlamaBatch = LlamaBatch + { n_tokens :: CInt + , token :: Ptr CInt + , embd :: Ptr CFloat + , pos :: Ptr CInt + , n_seq_id :: Ptr CInt + , seq_id :: Ptr (Ptr CInt) + , logits :: Ptr CSChar + -- ^ TODO: rename this to "output" + } + deriving (Generic, GStorable, Eq, Show) + +-- | Corresponds to `struct llama_kv_cache_view_cell` +newtype LlamaKvCacheViewCell = LlamaKvCacheViewCell + { posKvCacheViewCell :: LlamaPos + -- ^ Position in the KV cache (may be negative if unpopulated) + } + deriving (Show, Eq, Generic, GStorable) + +-- | Corresponds to `struct llama_kv_cache_view` +data LlamaKvCacheView = LlamaKvCacheView + { n_cells :: CInt + -- ^ Total number of KV cache cells + , n_seq_max :: CInt + -- ^ Max sequences per cell (may not reflect actual) + , token_count :: CInt + -- ^ Total tokens in the cache + , used_cells :: CInt + -- ^ Number of populated cells + , max_contiguous :: CInt + -- ^ Max contiguous empty slots + , max_contiguous_idx :: CInt + -- ^ Index to start of max_contiguous (may be -1) + , cells :: Ptr LlamaKvCacheViewCell + -- ^ Array of view cells + , cells_sequences :: Ptr LlamaSeqId + -- ^ Sequence IDs for each cell (n_seq_max per cell) + } + deriving (Show, Eq, Generic, GStorable) + +-- | Corresponds to `llama_sampler_context_t` (void *) +type LlamaSamplerContext = Ptr () + +-- | Corresponds to `struct llama_sampler_i` +data LlamaSamplerI = LlamaSamplerI + { name :: FunPtr (Ptr LlamaSampler -> IO CString) + -- ^ Optional: const char * (*name)(const struct llama_sampler *) + , accept :: FunPtr (Ptr LlamaSampler -> LlamaToken -> IO ()) + -- ^ Optional: void (*accept)(struct llama_sampler *, llama_token) + , apply :: FunPtr (Ptr LlamaSampler -> Ptr LlamaTokenDataArray -> IO ()) + -- ^ Required: void (*apply)(struct llama_sampler *, llama_token_data_array *) + , reset :: FunPtr (Ptr LlamaSampler -> IO ()) + -- ^ Optional: void (*reset)(struct llama_sampler *) + , clone :: FunPtr (Ptr LlamaSampler -> IO (Ptr LlamaSampler)) + -- ^ Optional: struct llama_sampler * (*clone)(const struct llama_sampler *) + , free_ :: FunPtr (Ptr LlamaSampler -> IO ()) + -- ^ Optional: void (*free)(struct llama_sampler *) + } + deriving (Show, Eq, Generic, GStorable) + +-- | Corresponds to `struct llama_sampler` +data LlamaSampler = LlamaSampler + { iface :: Ptr LlamaSamplerI + -- ^ Pointer to interface vtable + , ctx :: LlamaSamplerContext + -- ^ Opaque context pointer + } + deriving (Generic, Show, Eq, GStorable) + +data LlamaPerfContextData = LlamaPerfContextData + { t_start_ms :: CDouble -- ^ Start time in milliseconds + , t_load_ms :: CDouble -- ^ Model load time in milliseconds + , t_p_eval_ms :: CDouble -- ^ Prefill evaluation time in milliseconds + , t_eval_ms :: CDouble -- ^ Generation evaluation time in milliseconds + , n_p_eval :: CInt -- ^ Number of tokens in prefill phase + , n_eval :: CInt -- ^ Number of tokens in generation phase + } + deriving (Show, Eq, Generic, GStorable) + +data LlamaPerfSamplerData = LlamaPerfSamplerData + { t_sample_ms :: CDouble -- ^ Sampling time in milliseconds + , n_sample :: CInt -- ^ Number of sampling operations + } + deriving (Show, Eq, Generic, GStorable) diff --git a/src/Llama/Internal/Types/Params.hs b/src/Llama/Internal/Types/Params.hs new file mode 100644 --- /dev/null +++ b/src/Llama/Internal/Types/Params.hs @@ -0,0 +1,566 @@ +{-# LANGUAGE DeriveAnyClass #-} +{-# LANGUAGE DeriveGeneric #-} + +module Llama.Internal.Types.Params + ( LlamaSamplerChainParams (..) + , LlamaContextParams (..) + , LlamaModelParams (..) + , LlamaModelQuantizeParams (..) + , GgmlType (..) + , LlamaRopeTypeScaling (..) + , LlamaPoolingType (..) + , LlamaAttentionType (..) + , LlamaSplitMode (..) + , CLlamaContextParams (..) + , CLlamaModelParams (..) + , CLlamaSamplerChainParams (..) + , CLlamaModelQuantizeParams (..) + , SamplerChainParamsPtr (..) + , ModelQuantizeParamsPtr (..) + , ContextParamsPtr (..) + , ModelParamsPtr (..) + , LlamaVocabType (..) + , fromLlamaRopePoolingType + , fromLlamaRopeTypeScaling + , fromLlamaRopeVocabType + ) where + +import Foreign +import Foreign.C.Types +import Foreign.Storable.Generic +import GHC.Generics + +-- | Raw pointer to llama_sampler_chain_params +newtype CLlamaSamplerChainParams = CLlamaSamplerChainParams (Ptr LlamaSamplerChainParams) + deriving (Show, Eq) + +-- | Raw pointer to llama_model_quantize_params +newtype CLlamaModelQuantizeParams = CLlamaModelQuantizeParams (Ptr LlamaModelQuantizeParams) + deriving (Show, Eq) + +-- | Raw pointer to llama_context_params +newtype CLlamaContextParams = CLlamaContextParams (Ptr LlamaContextParams) + deriving (Show, Eq) + +-- | Raw pointer to llama_model_params +newtype CLlamaModelParams = CLlamaModelParams (Ptr LlamaModelParams) + deriving (Show, Eq) + +-- | Safe wrapper for managed 'llama_sampler_chain_params' +newtype SamplerChainParamsPtr = SamplerChainParamsPtr (ForeignPtr CLlamaSamplerChainParams) + deriving (Show, Eq) + +-- | Safe wrapper for managed 'llama_model_quantize_params' +newtype ModelQuantizeParamsPtr = ModelQuantizeParamsPtr (ForeignPtr CLlamaModelQuantizeParams) + deriving (Show, Eq) + +-- | Safe wrapper for managed 'llama_context_params' +newtype ContextParamsPtr = ContextParamsPtr (ForeignPtr CLlamaContextParams) + deriving (Show, Eq) + +-- | Safe wrapper for managed 'llama_model_params' +newtype ModelParamsPtr = ModelParamsPtr (ForeignPtr CLlamaModelParams) + deriving (Show, Eq) + +-- Forward declarations for opaque C structs +data GgmlBackendDevT +data LlamaModelTensorBuftOverride +data LlamaModelKvOverride + +-- Enum for llama_split_mode (assumed to be a 32-bit integer) +newtype LlamaSplitMode = LlamaSplitMode Int32 + deriving (Eq, Show, Read, Generic, GStorable) + +-- Define the sampler chain parameters structure +newtype LlamaSamplerChainParams = LlamaSamplerChainParams + { noPerf :: CBool -- whether to measure performance timings + } + deriving (Show, Eq, Generic, GStorable) + +-- model quantization parameters +data LlamaModelQuantizeParams = LlamaModelQuantizeParams + { nthread :: Int32 + , ftype :: LlamaFtype + , outputTensorType :: GgmlType + , tokenEmbeddingType :: GgmlType + , allowRequantize :: Bool + , quantizeOutputTensor :: Bool + , onlyCopy :: Bool + , pure_ :: Bool + , keepSplit :: Bool + , imatrix :: Ptr () + , kvOverridesQuantizeParams :: Ptr () + , tensorTypes :: Ptr () + } + deriving (Show, Eq, Generic, GStorable) + +data LlamaContextParams = LlamaContextParams + { n_ctx :: CUInt + -- ^ text context, 0 = from model + , n_batch :: CUInt + -- ^ logical maximum batch size that can be submitted to llama_decode + , n_ubatch :: CUInt + -- ^ physical maximum batch size + , n_seq_max :: CUInt + -- ^ max number of sequences (i.e. distinct states for recurrent models) + , n_threads :: CInt + -- ^ number of threads to use for generation + , n_threads_batch :: CInt + -- ^ number of threads to use for batch processing + , rope_scaling_type :: LlamaRopeTypeScaling + -- ^ RoPE scaling type, from `enum llama_rope_scaling_type` + , pooling_type :: LlamaPoolingType + -- ^ whether to pool (sum) embedding results by sequence id + , attention_type :: LlamaAttentionType + -- ^ attention type to use for embeddings + , rope_freq_base :: CFloat + -- ^ RoPE base frequency, 0 = from model + , rope_freq_scale :: CFloat + -- ^ RoPE frequency scaling factor, 0 = from model + , yarn_ext_factor :: CFloat + -- ^ YaRN extrapolation mix factor, negative = from model + , yarn_attn_factor :: CFloat + -- ^ YaRN magnitude scaling factor + , yarn_beta_fast :: CFloat + -- ^ YaRN low correction dim + , yarn_beta_slow :: CFloat + -- ^ YaRN high correction dim + , yarn_orig_ctx :: CUInt + -- ^ YaRN original context size + , defrag_thold :: CFloat + -- ^ defragment the KV cache if holes/size > thold, < 0 disabled (default) + , cb_eval :: FunPtr (Ptr () -> IO ()) + , cb_eval_user_data :: Ptr () + , type_k :: GgmlType + -- ^ data type for K cache [EXPERIMENTAL] + , type_v :: GgmlType + -- ^ data type for V cache [EXPERIMENTAL] + , logits_all :: CBool + -- ^ the llama_decode() call computes all logits, not just the last one (DEPRECATED - set llama_batch.logits instead) + , embeddings :: CBool + -- ^ if true, extract embeddings (together with logits) + , offload_kqv :: CBool + -- ^ whether to offload the KQV ops (including the KV cache) to GPU + , flash_attn :: CBool + -- ^ whether to use flash attention [EXPERIMENTAL] + , no_perf :: CBool + -- ^ whether to measure performance timings + , abort_callback :: FunPtr (Ptr () -> IO CInt) + , abort_callback_data :: Ptr () + } + deriving (Show, Generic, GStorable) + +-- Haskell representation of the C struct llama_model_params +data LlamaModelParams = LlamaModelParams + { devices :: Ptr GgmlBackendDevT + -- ^ NULL-terminated list of devices for offloading + , tensorBuftOverrides :: Ptr LlamaModelTensorBuftOverride + -- ^ NULL-terminated list of buffer type overrides + , nGpuLayers :: Int32 + -- ^ Number of layers to store in VRAM + , splitMode :: LlamaSplitMode + -- ^ How to split the model across GPUs + , mainGpu :: Int32 + -- ^ GPU used when split_mode is LLAMA_SPLIT_MODE_NONE + , tensorSplit :: Ptr CFloat + -- ^ Proportion of model offloaded to each GPU + , progressCallback :: FunPtr (CFloat -> Ptr () -> IO CBool) + -- ^ Callback for progress updates (returns true to continue) + , progressCallbackUserData :: Ptr () + -- ^ User data for progress callback + , kvOverrides :: Ptr LlamaModelKvOverride + -- ^ Override for model metadata + , vocabOnly :: CBool + -- ^ Only load vocabulary, not weights + , useMmap :: CBool + -- ^ Use mmap if possible + , useMlock :: CBool + -- ^ Force model to stay in RAM + , checkTensors :: CBool + -- ^ Validate tensor data + } + deriving (Generic, GStorable, Eq, Show) + +-- Helpers +data LlamaRopeTypeScaling + = LLAMA_ROPE_SCALING_TYPE_UNSPECIFIED + | LLAMA_ROPE_SCALING_TYPE_NONE + | LLAMA_ROPE_SCALING_TYPE_LINEAR + | LLAMA_ROPE_SCALING_TYPE_YARN + | LLAMA_ROPE_SCALING_TYPE_LONGROPE + | LLAMA_ROPE_SCALING_TYPE_MAX_VALUE + deriving (Show, Eq) + +toLlamaRopeTypeScaling :: LlamaRopeTypeScaling -> CInt +toLlamaRopeTypeScaling LLAMA_ROPE_SCALING_TYPE_UNSPECIFIED = -1 +toLlamaRopeTypeScaling LLAMA_ROPE_SCALING_TYPE_NONE = 0 +toLlamaRopeTypeScaling LLAMA_ROPE_SCALING_TYPE_LINEAR = 1 +toLlamaRopeTypeScaling LLAMA_ROPE_SCALING_TYPE_YARN = 2 +toLlamaRopeTypeScaling LLAMA_ROPE_SCALING_TYPE_LONGROPE = 3 +toLlamaRopeTypeScaling LLAMA_ROPE_SCALING_TYPE_MAX_VALUE = 3 + +fromLlamaRopeTypeScaling :: CInt -> Maybe LlamaRopeTypeScaling +fromLlamaRopeTypeScaling (-1) = Just LLAMA_ROPE_SCALING_TYPE_UNSPECIFIED +fromLlamaRopeTypeScaling 0 = Just LLAMA_ROPE_SCALING_TYPE_NONE +fromLlamaRopeTypeScaling 1 = Just LLAMA_ROPE_SCALING_TYPE_LINEAR +fromLlamaRopeTypeScaling 2 = Just LLAMA_ROPE_SCALING_TYPE_YARN +fromLlamaRopeTypeScaling 3 = Just LLAMA_ROPE_SCALING_TYPE_LONGROPE +fromLlamaRopeTypeScaling _ = Nothing + +instance Storable LlamaRopeTypeScaling where + sizeOf _ = 4 + alignment _ = 4 + peek ptr = do + mbVal <- fromLlamaRopeTypeScaling <$> peek (castPtr ptr) + case mbVal of + Nothing -> error "Invalid LlamaRopeTypeScaling value" + Just val -> return val + poke ptr val = do + let val' = toLlamaRopeTypeScaling val + poke (castPtr ptr) val' + +data LlamaPoolingType + = LLAMA_POOLING_TYPE_UNSPECIFIED + | LLAMA_POOLING_TYPE_NONE + | LLAMA_POOLING_TYPE_MEAN + | LLAMA_POOLING_TYPE_CLS + | LLAMA_POOLING_TYPE_LAST + | LLAMA_POOLING_TYPE_RANK + deriving (Show, Eq, Ord, Generic) + +instance Storable LlamaPoolingType where + sizeOf _ = 4 + alignment _ = 4 + peek ptr = do + mbVal <- fromLlamaRopePoolingType <$> peek (castPtr ptr) + case mbVal of + Nothing -> error "Invalid LlamaPoolingType value" + Just val -> return val + poke ptr val = do + let val' = toLlamaRopePoolingType val + poke (castPtr ptr) val' + +toLlamaRopePoolingType :: LlamaPoolingType -> CInt +toLlamaRopePoolingType LLAMA_POOLING_TYPE_UNSPECIFIED = -1 +toLlamaRopePoolingType LLAMA_POOLING_TYPE_NONE = 0 +toLlamaRopePoolingType LLAMA_POOLING_TYPE_MEAN = 1 +toLlamaRopePoolingType LLAMA_POOLING_TYPE_CLS = 2 +toLlamaRopePoolingType LLAMA_POOLING_TYPE_LAST = 3 +toLlamaRopePoolingType LLAMA_POOLING_TYPE_RANK = 4 + +fromLlamaRopePoolingType :: CInt -> Maybe LlamaPoolingType +fromLlamaRopePoolingType (-1) = Just LLAMA_POOLING_TYPE_UNSPECIFIED +fromLlamaRopePoolingType 0 = Just LLAMA_POOLING_TYPE_NONE +fromLlamaRopePoolingType 1 = Just LLAMA_POOLING_TYPE_MEAN +fromLlamaRopePoolingType 2 = Just LLAMA_POOLING_TYPE_CLS +fromLlamaRopePoolingType 3 = Just LLAMA_POOLING_TYPE_LAST +fromLlamaRopePoolingType 4 = Just LLAMA_POOLING_TYPE_RANK +fromLlamaRopePoolingType _ = Nothing + +data LlamaAttentionType + = LLAMA_ATTENTION_TYPE_UNSPECIFIED + | LLAMA_ATTENTION_TYPE_CAUSAL + | LLAMA_ATTENTION_TYPE_NON_CAUSAL + deriving (Show, Eq, Ord, Generic) + +instance Storable LlamaAttentionType where + sizeOf _ = sizeOf (undefined :: CUInt) + + alignment _ = sizeOf (undefined :: CUInt) + + peek ptr = do + mbVal <- fromLlamaAttentionType <$> peek (castPtr ptr) + case mbVal of + Nothing -> error "Invalid LlamaAttentionType value" + Just val -> return val + poke ptr val = do + let val' = toLlamaAttentionType val + poke (castPtr ptr) val' + +toLlamaAttentionType :: LlamaAttentionType -> CInt +toLlamaAttentionType LLAMA_ATTENTION_TYPE_UNSPECIFIED = -1 +toLlamaAttentionType LLAMA_ATTENTION_TYPE_CAUSAL = 0 +toLlamaAttentionType LLAMA_ATTENTION_TYPE_NON_CAUSAL = 1 + +fromLlamaAttentionType :: CInt -> Maybe LlamaAttentionType +fromLlamaAttentionType (-1) = Just LLAMA_ATTENTION_TYPE_UNSPECIFIED +fromLlamaAttentionType 0 = Just LLAMA_ATTENTION_TYPE_CAUSAL +fromLlamaAttentionType 1 = Just LLAMA_ATTENTION_TYPE_NON_CAUSAL +fromLlamaAttentionType _ = Nothing + +data GgmlType + = GGML_TYPE_F32 + | GGML_TYPE_F16 + | GGML_TYPE_Q4_0 + | GGML_TYPE_Q4_1 + | GGML_TYPE_Q5_0 + | GGML_TYPE_Q5_1 + | GGML_TYPE_Q8_0 + | GGML_TYPE_Q8_1 + | GGML_TYPE_Q2_K + | GGML_TYPE_Q3_K + | GGML_TYPE_Q4_K + | GGML_TYPE_Q5_K + | GGML_TYPE_Q6_K + | GGML_TYPE_Q8_K + | GGML_TYPE_IQ2_XXS + | GGML_TYPE_IQ2_XS + | GGML_TYPE_IQ3_XXS + | GGML_TYPE_IQ1_S + | GGML_TYPE_IQ4_NL + | GGML_TYPE_IQ3_S + | GGML_TYPE_IQ2_S + | GGML_TYPE_IQ4_XS + | GGML_TYPE_I8 + | GGML_TYPE_I16 + | GGML_TYPE_I32 + | GGML_TYPE_I64 + | GGML_TYPE_F64 + | GGML_TYPE_IQ1_M + | GGML_TYPE_BF16 + | GGML_TYPE_TQ1_0 + | GGML_TYPE_TQ2_0 + | GGML_TYPE_COUNT + deriving (Show, Eq, Ord, Generic) + +toGgmlType :: GgmlType -> CInt +toGgmlType g = + case g of + GGML_TYPE_F32 -> 0 + GGML_TYPE_F16 -> 1 + GGML_TYPE_Q4_0 -> 2 + GGML_TYPE_Q4_1 -> 3 + GGML_TYPE_Q5_0 -> 6 + GGML_TYPE_Q5_1 -> 7 + GGML_TYPE_Q8_0 -> 8 + GGML_TYPE_Q8_1 -> 9 + GGML_TYPE_Q2_K -> 10 + GGML_TYPE_Q3_K -> 11 + GGML_TYPE_Q4_K -> 12 + GGML_TYPE_Q5_K -> 13 + GGML_TYPE_Q6_K -> 14 + GGML_TYPE_Q8_K -> 15 + GGML_TYPE_IQ2_XXS -> 16 + GGML_TYPE_IQ2_XS -> 17 + GGML_TYPE_IQ3_XXS -> 18 + GGML_TYPE_IQ1_S -> 19 + GGML_TYPE_IQ4_NL -> 20 + GGML_TYPE_IQ3_S -> 21 + GGML_TYPE_IQ2_S -> 22 + GGML_TYPE_IQ4_XS -> 23 + GGML_TYPE_I8 -> 24 + GGML_TYPE_I16 -> 25 + GGML_TYPE_I32 -> 26 + GGML_TYPE_I64 -> 27 + GGML_TYPE_F64 -> 28 + GGML_TYPE_IQ1_M -> 29 + GGML_TYPE_BF16 -> 30 + GGML_TYPE_TQ1_0 -> 34 + GGML_TYPE_TQ2_0 -> 35 + GGML_TYPE_COUNT -> 39 + +fromGgmlType :: CInt -> Maybe GgmlType +fromGgmlType num = + case num of + 0 -> Just GGML_TYPE_F32 + 1 -> Just GGML_TYPE_F16 + 2 -> Just GGML_TYPE_Q4_0 + 3 -> Just GGML_TYPE_Q4_1 + 6 -> Just GGML_TYPE_Q5_0 + 7 -> Just GGML_TYPE_Q5_1 + 8 -> Just GGML_TYPE_Q8_0 + 9 -> Just GGML_TYPE_Q8_1 + 10 -> Just GGML_TYPE_Q2_K + 11 -> Just GGML_TYPE_Q3_K + 12 -> Just GGML_TYPE_Q4_K + 13 -> Just GGML_TYPE_Q5_K + 14 -> Just GGML_TYPE_Q6_K + 15 -> Just GGML_TYPE_Q8_K + 16 -> Just GGML_TYPE_IQ2_XXS + 17 -> Just GGML_TYPE_IQ2_XS + 18 -> Just GGML_TYPE_IQ3_XXS + 19 -> Just GGML_TYPE_IQ1_S + 20 -> Just GGML_TYPE_IQ4_NL + 21 -> Just GGML_TYPE_IQ3_S + 22 -> Just GGML_TYPE_IQ2_S + 23 -> Just GGML_TYPE_IQ4_XS + 24 -> Just GGML_TYPE_I8 + 25 -> Just GGML_TYPE_I16 + 26 -> Just GGML_TYPE_I32 + 27 -> Just GGML_TYPE_I64 + 28 -> Just GGML_TYPE_F64 + 29 -> Just GGML_TYPE_IQ1_M + 30 -> Just GGML_TYPE_BF16 + 34 -> Just GGML_TYPE_TQ1_0 + 35 -> Just GGML_TYPE_TQ2_0 + 39 -> Just GGML_TYPE_COUNT + _ -> Nothing + +instance Storable GgmlType where + sizeOf _ = sizeOf (undefined :: CInt) + alignment _ = sizeOf (undefined :: CInt) + peek ptr = do + mbVal <- fromGgmlType <$> peek (castPtr ptr) + case mbVal of + Nothing -> error "Invalid GgmlType value" + Just val -> return val + poke ptr val = do + let val' = toGgmlType val + poke (castPtr ptr) val' + +data LlamaFtype + = LLAMA_FTYPE_ALL_F32 + | LLAMA_FTYPE_MOSTLY_F16 + | LLAMA_FTYPE_MOSTLY_Q4_0 + | LLAMA_FTYPE_MOSTLY_Q4_1 + | LLAMA_FTYPE_MOSTLY_Q8_0 + | LLAMA_FTYPE_MOSTLY_Q5_0 + | LLAMA_FTYPE_MOSTLY_Q5_1 + | LLAMA_FTYPE_MOSTLY_Q2_K + | LLAMA_FTYPE_MOSTLY_Q3_K_S + | LLAMA_FTYPE_MOSTLY_Q3_K_M + | LLAMA_FTYPE_MOSTLY_Q3_K_L + | LLAMA_FTYPE_MOSTLY_Q4_K_S + | LLAMA_FTYPE_MOSTLY_Q4_K_M + | LLAMA_FTYPE_MOSTLY_Q5_K_S + | LLAMA_FTYPE_MOSTLY_Q5_K_M + | LLAMA_FTYPE_MOSTLY_Q6_K + | LLAMA_FTYPE_MOSTLY_IQ2_XXS + | LLAMA_FTYPE_MOSTLY_IQ2_XS + | LLAMA_FTYPE_MOSTLY_Q2_K_S + | LLAMA_FTYPE_MOSTLY_IQ3_XS + | LLAMA_FTYPE_MOSTLY_IQ3_XXS + | LLAMA_FTYPE_MOSTLY_IQ1_S + | LLAMA_FTYPE_MOSTLY_IQ4_NL + | LLAMA_FTYPE_MOSTLY_IQ3_S + | LLAMA_FTYPE_MOSTLY_IQ3_M + | LLAMA_FTYPE_MOSTLY_IQ2_S + | LLAMA_FTYPE_MOSTLY_IQ2_M + | LLAMA_FTYPE_MOSTLY_IQ4_XS + | LLAMA_FTYPE_MOSTLY_IQ1_M + | LLAMA_FTYPE_MOSTLY_BF16 + | LLAMA_FTYPE_MOSTLY_TQ1_0 + | LLAMA_FTYPE_MOSTLY_TQ2_0 + | LLAMA_FTYPE_GUESSED + deriving (Show, Eq) + +toLlamaFtype :: LlamaFtype -> CInt +toLlamaFtype LLAMA_FTYPE_ALL_F32 = 0 +toLlamaFtype LLAMA_FTYPE_MOSTLY_F16 = 1 +toLlamaFtype LLAMA_FTYPE_MOSTLY_Q4_0 = 2 +toLlamaFtype LLAMA_FTYPE_MOSTLY_Q4_1 = 3 +toLlamaFtype LLAMA_FTYPE_MOSTLY_Q8_0 = 7 +toLlamaFtype LLAMA_FTYPE_MOSTLY_Q5_0 = 8 +toLlamaFtype LLAMA_FTYPE_MOSTLY_Q5_1 = 9 +toLlamaFtype LLAMA_FTYPE_MOSTLY_Q2_K = 10 +toLlamaFtype LLAMA_FTYPE_MOSTLY_Q3_K_S = 11 +toLlamaFtype LLAMA_FTYPE_MOSTLY_Q3_K_M = 12 +toLlamaFtype LLAMA_FTYPE_MOSTLY_Q3_K_L = 13 +toLlamaFtype LLAMA_FTYPE_MOSTLY_Q4_K_S = 14 +toLlamaFtype LLAMA_FTYPE_MOSTLY_Q4_K_M = 15 +toLlamaFtype LLAMA_FTYPE_MOSTLY_Q5_K_S = 16 +toLlamaFtype LLAMA_FTYPE_MOSTLY_Q5_K_M = 17 +toLlamaFtype LLAMA_FTYPE_MOSTLY_Q6_K = 18 +toLlamaFtype LLAMA_FTYPE_MOSTLY_IQ2_XXS = 19 +toLlamaFtype LLAMA_FTYPE_MOSTLY_IQ2_XS = 20 +toLlamaFtype LLAMA_FTYPE_MOSTLY_Q2_K_S = 21 +toLlamaFtype LLAMA_FTYPE_MOSTLY_IQ3_XS = 22 +toLlamaFtype LLAMA_FTYPE_MOSTLY_IQ3_XXS = 23 +toLlamaFtype LLAMA_FTYPE_MOSTLY_IQ1_S = 24 +toLlamaFtype LLAMA_FTYPE_MOSTLY_IQ4_NL = 25 +toLlamaFtype LLAMA_FTYPE_MOSTLY_IQ3_S = 26 +toLlamaFtype LLAMA_FTYPE_MOSTLY_IQ3_M = 27 +toLlamaFtype LLAMA_FTYPE_MOSTLY_IQ2_S = 28 +toLlamaFtype LLAMA_FTYPE_MOSTLY_IQ2_M = 29 +toLlamaFtype LLAMA_FTYPE_MOSTLY_IQ4_XS = 30 +toLlamaFtype LLAMA_FTYPE_MOSTLY_IQ1_M = 31 +toLlamaFtype LLAMA_FTYPE_MOSTLY_BF16 = 32 +toLlamaFtype LLAMA_FTYPE_MOSTLY_TQ1_0 = 36 +toLlamaFtype LLAMA_FTYPE_MOSTLY_TQ2_0 = 37 +toLlamaFtype LLAMA_FTYPE_GUESSED = 1024 + +fromLlamaFtype :: CInt -> Maybe LlamaFtype +fromLlamaFtype 0 = Just LLAMA_FTYPE_ALL_F32 +fromLlamaFtype 1 = Just LLAMA_FTYPE_MOSTLY_F16 +fromLlamaFtype 2 = Just LLAMA_FTYPE_MOSTLY_Q4_0 +fromLlamaFtype 3 = Just LLAMA_FTYPE_MOSTLY_Q4_1 +fromLlamaFtype 7 = Just LLAMA_FTYPE_MOSTLY_Q8_0 +fromLlamaFtype 8 = Just LLAMA_FTYPE_MOSTLY_Q5_0 +fromLlamaFtype 9 = Just LLAMA_FTYPE_MOSTLY_Q5_1 +fromLlamaFtype 10 = Just LLAMA_FTYPE_MOSTLY_Q2_K +fromLlamaFtype 11 = Just LLAMA_FTYPE_MOSTLY_Q3_K_S +fromLlamaFtype 12 = Just LLAMA_FTYPE_MOSTLY_Q3_K_M +fromLlamaFtype 13 = Just LLAMA_FTYPE_MOSTLY_Q3_K_L +fromLlamaFtype 14 = Just LLAMA_FTYPE_MOSTLY_Q4_K_S +fromLlamaFtype 15 = Just LLAMA_FTYPE_MOSTLY_Q4_K_M +fromLlamaFtype 16 = Just LLAMA_FTYPE_MOSTLY_Q5_K_S +fromLlamaFtype 17 = Just LLAMA_FTYPE_MOSTLY_Q5_K_M +fromLlamaFtype 18 = Just LLAMA_FTYPE_MOSTLY_Q6_K +fromLlamaFtype 19 = Just LLAMA_FTYPE_MOSTLY_IQ2_XXS +fromLlamaFtype 20 = Just LLAMA_FTYPE_MOSTLY_IQ2_XS +fromLlamaFtype 21 = Just LLAMA_FTYPE_MOSTLY_Q2_K_S +fromLlamaFtype 22 = Just LLAMA_FTYPE_MOSTLY_IQ3_XS +fromLlamaFtype 23 = Just LLAMA_FTYPE_MOSTLY_IQ3_XXS +fromLlamaFtype 24 = Just LLAMA_FTYPE_MOSTLY_IQ1_S +fromLlamaFtype 25 = Just LLAMA_FTYPE_MOSTLY_IQ4_NL +fromLlamaFtype 26 = Just LLAMA_FTYPE_MOSTLY_IQ3_S +fromLlamaFtype 27 = Just LLAMA_FTYPE_MOSTLY_IQ3_M +fromLlamaFtype 28 = Just LLAMA_FTYPE_MOSTLY_IQ2_S +fromLlamaFtype 29 = Just LLAMA_FTYPE_MOSTLY_IQ2_M +fromLlamaFtype 30 = Just LLAMA_FTYPE_MOSTLY_IQ4_XS +fromLlamaFtype 31 = Just LLAMA_FTYPE_MOSTLY_IQ1_M +fromLlamaFtype 32 = Just LLAMA_FTYPE_MOSTLY_BF16 +fromLlamaFtype 36 = Just LLAMA_FTYPE_MOSTLY_TQ1_0 +fromLlamaFtype 37 = Just LLAMA_FTYPE_MOSTLY_TQ2_0 +fromLlamaFtype 1024 = Just LLAMA_FTYPE_GUESSED +fromLlamaFtype _ = Nothing + +instance Storable LlamaFtype where + sizeOf _ = 4 + alignment _ = 4 + peek ptr = do + val <- peek (castPtr ptr) :: IO CInt + case fromLlamaFtype val of + Just v -> return v + Nothing -> error $ "Invalid LlamaFtype value: " ++ show val + poke ptr val = do + let cVal = toLlamaFtype val + poke (castPtr ptr) cVal + +data LlamaVocabType + = LLAMA_VOCAB_TYPE_NONE + | LLAMA_VOCAB_TYPE_SPM + | LLAMA_VOCAB_TYPE_BPE + | LLAMA_VOCAB_TYPE_WPM + | LLAMA_VOCAB_TYPE_UGM + | LLAMA_VOCAB_TYPE_RWKV + deriving (Show, Eq, Ord, Generic) + +instance Storable LlamaVocabType where + sizeOf _ = 4 + alignment _ = 4 + peek ptr = do + mbVal <- fromLlamaRopeVocabType <$> peek (castPtr ptr) + case mbVal of + Nothing -> error "Invalid LlamaVocabType value" + Just val -> return val + poke ptr val = do + let val' = toLlamaRopeVocabType val + poke (castPtr ptr) val' + +toLlamaRopeVocabType :: LlamaVocabType -> CInt +toLlamaRopeVocabType LLAMA_VOCAB_TYPE_NONE = 0 +toLlamaRopeVocabType LLAMA_VOCAB_TYPE_SPM = 1 +toLlamaRopeVocabType LLAMA_VOCAB_TYPE_BPE = 2 +toLlamaRopeVocabType LLAMA_VOCAB_TYPE_WPM = 3 +toLlamaRopeVocabType LLAMA_VOCAB_TYPE_UGM = 4 +toLlamaRopeVocabType LLAMA_VOCAB_TYPE_RWKV = 5 + +fromLlamaRopeVocabType :: CInt -> Maybe LlamaVocabType +fromLlamaRopeVocabType 0 = Just LLAMA_VOCAB_TYPE_NONE +fromLlamaRopeVocabType 1 = Just LLAMA_VOCAB_TYPE_SPM +fromLlamaRopeVocabType 2 = Just LLAMA_VOCAB_TYPE_BPE +fromLlamaRopeVocabType 3 = Just LLAMA_VOCAB_TYPE_WPM +fromLlamaRopeVocabType 4 = Just LLAMA_VOCAB_TYPE_UGM +fromLlamaRopeVocabType 5 = Just LLAMA_VOCAB_TYPE_RWKV +fromLlamaRopeVocabType _ = Nothing diff --git a/src/Llama/KVCache.hs b/src/Llama/KVCache.hs new file mode 100644 --- /dev/null +++ b/src/Llama/KVCache.hs @@ -0,0 +1,171 @@ +{- | +Module : Llama.KVCache +Description : High level KVCache interface for llama-cpp +Copyright : (c) 2025 Tushar Adhatrao +License : MIT +Maintainer : Tushar Adhatrao +-} +module Llama.KVCache ( + kvCacheViewInit +, kvSelfSeqAdd +, kvSelfSeqDiv +, kvSelfSeqPosMax +, kvSelfDefrag +, kvSelfCanShift +, kvSelfUpdate +, kvSelfSeqKeep +, kvSelfSeqCopy +, kvSelfSeqRemove +, kvSelfClear +, kvSelfUsedCells +, kvSelfNumTokens +, kvCacheViewUpdate +) where + +import Llama.Internal.Types +import Foreign +import Llama.Internal.Foreign + +{- +TODO: no free function for struct llama_kv_cache * +No one is using struct llama_kv_cache +-- | Get the KV cache associated with this context. +getKVCache :: Context -> IO KVCache +getKVCache (Context ctxFPtr) = + withForeignPtr ctxFPtr $ \ctxPtr -> do + ptr <- c_llama_get_kv_self (CLlamaContext ctxPtr) + -- We assume finalization is handled elsewhere or use a no-op finalizer + fp <- newForeignPtr_ ptr -- assumes Ptr () is actually Ptr CLlamaKVCache + return $ KVCache fp +-} + +-- | Convenience wrapper that allocates a LlamaKvCacheView and initializes it +kvCacheViewInit :: Context -> Int -> IO LlamaKvCacheView +kvCacheViewInit (Context fPtr) n_seq_max_ = do + withForeignPtr fPtr $ \contextPtr -> do + alloca $ \pView -> do + c_llama_kv_cache_view_init_into (CLlamaContext contextPtr) (fromIntegral n_seq_max_) pView + peek pView + +-- | Shift positions in a sequence by a delta. +kvSelfSeqAdd :: + Context -> + LlamaSeqId -> -- seq_id + LlamaPos -> -- p0 + LlamaPos -> -- p1 + LlamaPos -> -- delta + IO () +kvSelfSeqAdd (Context ctxFPtr) seqId p0 p1 delta = + withForeignPtr ctxFPtr $ \ctxPtr -> + c_llama_kv_self_seq_add (CLlamaContext ctxPtr) seqId p0 p1 delta + +-- | Divide positions in a sequence by a factor (used for attention windowing). +kvSelfSeqDiv :: + Context -> + LlamaSeqId -> -- seq_id + LlamaPos -> -- p0 + LlamaPos -> -- p1 + Int -> -- d + IO () +kvSelfSeqDiv (Context ctxFPtr) seqId p0 p1 d = + withForeignPtr ctxFPtr $ \ctxPtr -> + c_llama_kv_self_seq_div (CLlamaContext ctxPtr) seqId p0 p1 (fromIntegral d) + +-- | Get the maximum position stored in the KV cache for a given sequence. +kvSelfSeqPosMax :: + Context -> + LlamaSeqId -> -- seq_id + IO LlamaPos +kvSelfSeqPosMax (Context ctxFPtr) seqId = + withForeignPtr ctxFPtr $ \ctxPtr -> + c_llama_kv_self_seq_pos_max (CLlamaContext ctxPtr) seqId + +-- | Defragment the KV cache to optimize memory usage. +kvSelfDefrag :: + Context -> + IO () +kvSelfDefrag (Context ctxFPtr) = + withForeignPtr ctxFPtr $ \ctxPtr -> + c_llama_kv_self_defrag (CLlamaContext ctxPtr) + +-- | Check whether the KV cache can be shifted (e.g., due to full buffer). +kvSelfCanShift :: + Context -> + IO Bool +kvSelfCanShift (Context ctxFPtr) = + withForeignPtr ctxFPtr $ \ctxPtr -> + toBool <$> c_llama_kv_self_can_shift (castPtr ctxPtr) + +-- | Update the KV cache's internal state (e.g., after manual modifications). +kvSelfUpdate :: + Context -> + IO () +kvSelfUpdate (Context ctxFPtr) = + withForeignPtr ctxFPtr $ \ctxPtr -> + c_llama_kv_self_update (CLlamaContext ctxPtr) + +-- | Prevent a sequence from being removed during KV cache cleanup. +kvSelfSeqKeep :: + Context -> + LlamaSeqId -> -- seq_id + IO () +kvSelfSeqKeep (Context ctxFPtr) seqId = + withForeignPtr ctxFPtr $ \ctxPtr -> + c_llama_kv_self_seq_keep (CLlamaContext ctxPtr) seqId + +-- | Copy a range of positions from one sequence to another. +kvSelfSeqCopy :: + Context -> + LlamaSeqId -> -- seq_id_src + LlamaSeqId -> -- seq_id_dst + LlamaPos -> -- p0 + LlamaPos -> -- p1 + IO () +kvSelfSeqCopy (Context ctxFPtr) srcId dstId p0 p1 = + withForeignPtr ctxFPtr $ \ctxPtr -> + c_llama_kv_self_seq_cp (CLlamaContext ctxPtr) srcId dstId p0 p1 + +-- | Remove a range of positions from a sequence. +kvSelfSeqRemove :: + Context -> + LlamaSeqId -> -- seq_id + LlamaPos -> -- p0 + LlamaPos -> -- p1 + IO Bool +kvSelfSeqRemove (Context ctxFPtr) seqId p0 p1 = + withForeignPtr ctxFPtr $ \ctxPtr -> do + result <- c_llama_kv_self_seq_rm (CLlamaContext ctxPtr) seqId p0 p1 + return $ toBool result + +-- | Clear all sequences from the KV cache. +kvSelfClear :: + Context -> + IO () +kvSelfClear (Context ctxFPtr) = + withForeignPtr ctxFPtr $ \ctxPtr -> + c_llama_kv_self_clear (CLlamaContext ctxPtr) + +-- | Get number of used cells in the KV cache. +kvSelfUsedCells :: + Context -> + IO Int +kvSelfUsedCells (Context ctxFPtr) = + withForeignPtr ctxFPtr $ \ctxPtr -> + fromIntegral <$> c_llama_kv_self_used_cells (CLlamaContext ctxPtr) + +-- | Get total number of tokens currently stored in the KV cache. +kvSelfNumTokens :: + Context -> + IO Int +kvSelfNumTokens (Context ctxFPtr) = + withForeignPtr ctxFPtr $ \ctxPtr -> + fromIntegral <$> c_llama_kv_self_n_tokens (CLlamaContext ctxPtr) + +-- | Update the KV cache view to reflect current state. +kvCacheViewUpdate :: + Context -> + Ptr LlamaKvCacheView -> -- view + IO () +kvCacheViewUpdate (Context ctxFPtr) view = + withForeignPtr ctxFPtr $ \ctxPtr -> + c_llama_kv_cache_view_update (CLlamaContext ctxPtr) view diff --git a/src/Llama/Model.hs b/src/Llama/Model.hs new file mode 100644 --- /dev/null +++ b/src/Llama/Model.hs @@ -0,0 +1,338 @@ +{- | +Module : Llama.Model +Description : High level Model interface for llama-cpp +Copyright : (c) 2025 Tushar Adhatrao +License : MIT +Maintainer : Tushar Adhatrao +-} +module Llama.Model ( + defaultModelParams +, loadModelFromFile +, initContextFromModel +, getModelVocab +, getContextModel +, getVocabType +, getModelRoPEFreqScale +, getModelNumKVHeads +, getModelNumHeads +, getModelNumLayers +, getModelEmbeddingDim +, getModelTrainingContextSize +, getModelSize +, getModelChatTemplate +, getModelHasEncoder +, getModelNumParams +, getModelHasDecoder +, getModelDecoderStartToken +, getModelIsRecurrent +, quantizeModel +, quantizeModelDefault +, defaultQuantizeParams +, getModelMetaCount +, getModelMetaValue +, getModelMetaKeyByIndex +, getModelMetaValueByIndex +, getModelDescription +, loadModelFromSplits +, getModelRopeType +) where + +import Data.Functor +import Foreign +import Foreign.C.String +import Llama.Internal.Foreign +import Llama.Internal.Types +import Llama.Internal.Types.Params + +-- | Default model parameters +defaultModelParams :: IO LlamaModelParams +defaultModelParams = do + -- Convert to a pointer to pass to the C function + alloca $ \paramsPtr -> do + c_llama_model_default_params (CLlamaModelParams paramsPtr) + peek paramsPtr + +-- | Load a model from a file using the specified parameters +loadModelFromFile :: FilePath -> LlamaModelParams -> IO (Either String Model) +loadModelFromFile path params = do + withCString path $ \cPath -> do + withStorable params $ \paramsPtr -> do + model <- c_llama_model_load_from_file_wrap cPath (CLlamaModelParams paramsPtr) + if model == CLlamaModel nullPtr + then return $ Left "Failed to load model" + else do + let (CLlamaModel modelPtr) = model + fp <- newForeignPtr p_llama_model_free modelPtr + return $ Right $ Model fp + +-- | Create a context from a model using the specified parameters +initContextFromModel :: Model -> LlamaContextParams -> IO (Either String Context) +initContextFromModel (Model modelFPtr) params = do + withForeignPtr modelFPtr $ \modelPtr -> do + withStorable params $ \paramsPtr -> do + context <- + c_llama_init_from_model_wrap + (CLlamaModel modelPtr) + (CLlamaContextParams paramsPtr) + if context == CLlamaContext nullPtr + then return $ Left "Failed to initialize context" + else do + let (CLlamaContext contextPtr) = context + fp <- newForeignPtr p_llama_free contextPtr + return $ Right $ Context fp + +-- | Get the vocabulary from a model +getModelVocab :: Model -> IO (Either String Vocab) +getModelVocab (Model modelFPtr) = do + withForeignPtr modelFPtr $ \modelPtr -> do + vocab <- c_llama_model_get_vocab (CLlamaModel modelPtr) + if vocab == CLlamaVocab nullPtr + then return $ Left "Failed to get vocabulary" + else do + -- For now, assuming it's owned by the model and doesn't need separate freeing + let (CLlamaVocab vocabPtr) = vocab + -- Using a dummy finalizer since vocab is owned by the model + fp <- newForeignPtr_ vocabPtr + return $ Right $ Vocab fp + +-- | Convert Storable Haskell struct to pointer and run an action +withStorable :: Storable a => a -> (Ptr a -> IO b) -> IO b +withStorable x f = alloca $ \ptr -> do + poke ptr x + f ptr + +-- | Get the model associated with a context. +getContextModel :: Context -> IO Model +getContextModel (Context ctxFPtr) = + withForeignPtr ctxFPtr $ \ctxPtr -> do + CLlamaModel modelPtr <- c_llama_get_model (CLlamaContext ctxPtr) + fp <- newForeignPtr p_llama_model_free modelPtr + return $ Model fp + +-- | Get the vocabulary type. +getVocabType :: Vocab -> IO (Maybe LlamaVocabType) +getVocabType (Vocab vocabFPtr) = + withForeignPtr vocabFPtr $ \vocabPtr -> + fromLlamaRopeVocabType <$> c_llama_vocab_type (CLlamaVocab vocabPtr) + +-- | Get RoPE frequency scaling factor used during training. +getModelRoPEFreqScale :: Model -> IO Float +getModelRoPEFreqScale (Model modelFPtr) = + withForeignPtr modelFPtr $ \modelPtr -> + realToFrac <$> c_llama_model_rope_freq_scale_train (CLlamaModel modelPtr) + +-- | Get the number of key/value heads in the model. +getModelNumKVHeads :: Model -> IO Int +getModelNumKVHeads (Model modelFPtr) = + withForeignPtr modelFPtr $ \modelPtr -> + fromCInt <$> c_llama_model_n_head_kv (CLlamaModel modelPtr) + +-- | Get the number of attention heads in the model. +getModelNumHeads :: Model -> IO Int +getModelNumHeads (Model modelFPtr) = + withForeignPtr modelFPtr $ \modelPtr -> + fromCInt <$> c_llama_model_n_head (CLlamaModel modelPtr) + +-- | Get the number of transformer layers in the model. +getModelNumLayers :: Model -> IO Int +getModelNumLayers (Model modelFPtr) = + withForeignPtr modelFPtr $ \modelPtr -> + fromCInt <$> c_llama_model_n_layer (CLlamaModel modelPtr) + +-- | Get the embedding dimension of the model. +getModelEmbeddingDim :: Model -> IO Int +getModelEmbeddingDim (Model modelFPtr) = + withForeignPtr modelFPtr $ \modelPtr -> + fromCInt <$> c_llama_model_n_embd (CLlamaModel modelPtr) + +-- | Get the training context size of the model. +getModelTrainingContextSize :: Model -> IO Int +getModelTrainingContextSize (Model modelFPtr) = + withForeignPtr modelFPtr $ \modelPtr -> + fromCInt <$> c_llama_model_n_ctx_train (CLlamaModel modelPtr) + +fromCInt :: Int32 -> Int +fromCInt = fromIntegral + +-- | Get the size of the model in bytes +getModelSize :: Model -> IO Int64 +getModelSize (Model modelFPtr) = + withForeignPtr modelFPtr $ \modelPtr -> + fromIntegral <$> c_llama_model_size (CLlamaModel modelPtr) + +-- | Check if the model has an encoder +getModelHasEncoder :: Model -> IO Bool +getModelHasEncoder (Model modelFPtr) = do + withForeignPtr modelFPtr $ \modelPtr -> do + c_llama_model_has_encoder (CLlamaModel modelPtr) <&> (/= 0) + +-- | Get the chat template from a model, optionally by name +getModelChatTemplate :: Model -> Maybe String -> IO (Either String String) +getModelChatTemplate (Model modelFPtr) mName = do + withForeignPtr modelFPtr $ \modelPtr -> do + cName <- maybe (pure nullPtr) newCString mName + template <- c_llama_model_chat_template (CLlamaModel modelPtr) cName + if template == nullPtr + then return $ Left "Failed to get chat template" + else Right <$> peekCString template + +-- | Get the number of parameters in the model +getModelNumParams :: Model -> IO (Either String Int64) +getModelNumParams (Model modelFPtr) = do + withForeignPtr modelFPtr $ \modelPtr -> do + params <- c_llama_model_n_params (CLlamaModel modelPtr) + if params == 0 + then return $ Left "Failed to get number of parameters" + else return $ Right $ fromIntegral params + +-- | Check if the model has a decoder +getModelHasDecoder :: Model -> IO Bool +getModelHasDecoder (Model modelFPtr) = do + withForeignPtr modelFPtr $ \modelPtr -> do + c_llama_model_has_decoder (CLlamaModel modelPtr) <&> (/= 0) + +-- | Get the decoder start token from the model +getModelDecoderStartToken :: Model -> IO (Either String LlamaToken) +getModelDecoderStartToken (Model modelFPtr) = do + withForeignPtr modelFPtr $ \modelPtr -> do + token_ <- c_llama_model_decoder_start_token (CLlamaModel modelPtr) + if token_ == -1 + then return $ Left "Failed to get decoder start token" + else return $ Right token_ + +-- | Check if the model is recurrent +getModelIsRecurrent :: Model -> IO Bool +getModelIsRecurrent (Model modelFPtr) = do + withForeignPtr modelFPtr $ \modelPtr -> do + c_llama_model_is_recurrent (CLlamaModel modelPtr) <&> (/= 0) + +-- | Quantize a model from a file to another file using specified parameters +quantizeModel :: + FilePath -> + FilePath -> + LlamaModelQuantizeParams -> + IO (Either String Word32) +quantizeModel inpPath outPath params = do + withCString inpPath $ \cInpPath -> do + withCString outPath $ \cOutPath -> do + withStorable params $ \paramsPtr -> do + result <- + c_llama_model_quantize + cInpPath + cOutPath + (CLlamaModelQuantizeParams paramsPtr) + if result == 0 + then return $ Left "Failed to quantize model" + else return $ Right result + +-- | Quantize a model from a file to another file using default parameters +quantizeModelDefault :: FilePath -> FilePath -> IO (Either String Word32) +quantizeModelDefault inpPath outPath = do + params <- defaultQuantizeParams + quantizeModel inpPath outPath params + +-- | Get the default quantization parameters +defaultQuantizeParams :: IO LlamaModelQuantizeParams +defaultQuantizeParams = do + (CLlamaModelQuantizeParams paramsPtr) <- c_llama_model_quantize_default_params + peek paramsPtr + +-- | Get a metadata value as a string from a model +getModelMetaValue :: Model -> String -> IO (Either String String) +getModelMetaValue (Model modelFPtr) key = do + withForeignPtr modelFPtr $ \modelPtr -> do + withCString key $ \cKey -> do + allocaArray 256 $ \bufPtr -> do + result <- c_llama_model_meta_val_str (CLlamaModel modelPtr) cKey bufPtr 256 + if result == -1 + then return $ Left "Failed to get metadata value" + else do + str <- peekCString bufPtr + return $ Right str + +-- | Get the number of metadata entries in a model +getModelMetaCount :: Model -> IO Int +getModelMetaCount (Model modelFPtr) = do + withForeignPtr modelFPtr $ \modelPtr -> do + fromIntegral <$> c_llama_model_meta_count (CLlamaModel modelPtr) + +-- | Get a metadata key by index from a model +getModelMetaKeyByIndex :: Model -> Int -> IO (Either String String) +getModelMetaKeyByIndex (Model modelFPtr) index = do + withForeignPtr modelFPtr $ \modelPtr -> do + allocaArray 256 $ \bufPtr -> do + result <- + c_llama_model_meta_key_by_index + (CLlamaModel modelPtr) + (fromIntegral index) + bufPtr + 256 + if result == -1 + then return $ Left "Failed to get metadata key" + else do + str <- peekCString bufPtr + return $ Right str + +-- | Get a metadata value by index from a model +getModelMetaValueByIndex :: Model -> Int -> IO (Either String String) +getModelMetaValueByIndex (Model modelFPtr) index = do + withForeignPtr modelFPtr $ \modelPtr -> do + allocaArray 256 $ \bufPtr -> do + result <- + c_llama_model_meta_val_str_by_index + (CLlamaModel modelPtr) + (fromIntegral index) + bufPtr + 256 + if result == -1 + then return $ Left "Failed to get metadata value" + else do + str <- peekCString bufPtr + return $ Right str + +-- | Get a model description +getModelDescription :: Model -> IO String +getModelDescription (Model modelFPtr) = do + withForeignPtr modelFPtr $ \modelPtr -> do + allocaArray 256 $ \bufPtr -> do + _ <- c_llama_model_desc (CLlamaModel modelPtr) bufPtr 256 + peekCString bufPtr + +-- | Load a model from multiple file paths using specified parameters +loadModelFromSplits :: [FilePath] -> LlamaModelParams -> IO (Either String Model) +loadModelFromSplits paths params = do + withStorable params $ \paramsPtr -> do + pathsPtr <- newArrayOfPtrs paths + withForeignPtr pathsPtr $ \pathsPtr' -> do + model <- + c_llama_model_load_from_splits + pathsPtr' + (fromIntegral (length paths)) + (CLlamaModelParams paramsPtr) + if model == CLlamaModel nullPtr + then return $ Left "Failed to load model" + else do + let (CLlamaModel modelPtr) = model + fp <- newForeignPtr p_llama_model_free modelPtr + return $ Right $ Model fp + +-- | Get the RoPE type from a model +getModelRopeType :: Model -> IO (Maybe LlamaRopeTypeScaling) +getModelRopeType (Model modelFPtr) = do + withForeignPtr modelFPtr $ \modelPtr -> do + alloca $ \outPtr -> do + c_llama_model_rope_type_into (CLlamaModel modelPtr) outPtr + val <- peek outPtr + return $ fromLlamaRopeTypeScaling val + +newArrayOfPtrs :: [FilePath] -> IO (ForeignPtr CString) +newArrayOfPtrs xs = do + ptrs <- mallocBytes (length xs * sizeOf (undefined :: Ptr CString)) + mapM_ + ( \(i, x) -> withCString x $ \cstr -> + poke + (castPtr ptrs `plusPtr` (i * sizeOf (undefined :: Ptr CString))) + cstr + ) + (zip [0 ..] xs) + newForeignPtr_ ptrs diff --git a/src/Llama/Performance.hs b/src/Llama/Performance.hs new file mode 100644 --- /dev/null +++ b/src/Llama/Performance.hs @@ -0,0 +1,59 @@ +{- | +Module : Llama.Performance +Description : High level Performance interface for llama-cpp +Copyright : (c) 2025 Tushar Adhatrao +License : MIT +Maintainer : Tushar Adhatrao +-} +module Llama.Performance ( + printContextPerformance +, resetContextPerformance +, printSamplerPerformance +, resetSamplerPerformance +, getContextPerformance +, getSamplerPerformance + ) where + +import Llama.Internal.Types +import Foreign +import Llama.Internal.Foreign + +-- | Print performance information for a context +printContextPerformance :: Context -> IO () +printContextPerformance (Context ctxFPtr) = + withForeignPtr ctxFPtr $ \ctxPtr -> + c_llama_perf_context_print (CLlamaContext ctxPtr) + +-- | Reset performance information for a context +resetContextPerformance :: Context -> IO () +resetContextPerformance (Context ctxFPtr) = + withForeignPtr ctxFPtr $ \ctxPtr -> + c_llama_perf_context_reset (CLlamaContext ctxPtr) + +-- | Print performance information for a sampler chain +printSamplerPerformance :: Sampler -> IO () +printSamplerPerformance (Sampler samplerFPtr) = + withForeignPtr samplerFPtr $ \samplerPtr -> + c_llama_perf_sampler_print samplerPtr + +-- | Reset performance information for a sampler chain +resetSamplerPerformance :: Sampler -> IO () +resetSamplerPerformance (Sampler samplerFPtr) = + withForeignPtr samplerFPtr $ \samplerPtr -> + c_llama_perf_sampler_reset samplerPtr + +-- | Get performance data for a context +getContextPerformance :: Context -> IO LlamaPerfContextData +getContextPerformance (Context ctxFPtr) = do + alloca $ \perfDataPtr -> do + withForeignPtr ctxFPtr $ \ctxPtr -> do + c_llama_perf_context (CLlamaContext ctxPtr) perfDataPtr + peek perfDataPtr + +-- | Get performance data for a sampler chain +getSamplerPerformance :: Sampler -> IO LlamaPerfSamplerData +getSamplerPerformance (Sampler samplerFPtr) = do + alloca $ \perfDataPtr -> do + withForeignPtr samplerFPtr $ \samplerPtr -> do + c_llama_perf_sampler samplerPtr perfDataPtr + peek perfDataPtr diff --git a/src/Llama/Sampler.hs b/src/Llama/Sampler.hs new file mode 100644 --- /dev/null +++ b/src/Llama/Sampler.hs @@ -0,0 +1,392 @@ +{- | +Module : Llama.Sampler +Description : High level Sampler interface for llama-cpp +Copyright : (c) 2025 Tushar Adhatrao +License : MIT +Maintainer : Tushar Adhatrao +-} +module Llama.Sampler ( + defaultSamplerChainParams +, initSampler +, getSamplerName +, acceptTokenWithSampler +, applySampler +, resetSampler +, cloneSampler +, initSamplerChain +, addSamplerToChain +, getSamplerFromChain +, getSamplerChainLength +, removeSamplerFromChain +, initGreedySampler +, initDistributedSampler +, initTopKSampler +, initTopPSampler +, initMinPSampler +, initTypicalSampler +, initTempSampler +, initTempExtSampler +, initXTCSampler +, initTopNSigmaSampler +, initMirostatSampler +, initMirostatV2Sampler +, initGrammarSampler +, initGrammarLazyPatternsSampler +, initPenaltiesSampler +, initDrySampler +, initLogitBiasSampler +, initInfillSampler +, getSamplerSeed +, sampleWithSampler +) where + +import Llama.Internal.Types +import Llama.Internal.Types.Params +import Foreign +import Llama.Internal.Foreign +import Foreign.C.String + +-- | Convert Storable Haskell struct to pointer and run an action +withStorable :: Storable a => a -> (Ptr a -> IO b) -> IO b +withStorable x f = alloca $ \ptr -> do + poke ptr x + f ptr + +-- | Get the default parameters for a sampler chain +defaultSamplerChainParams :: IO LlamaSamplerChainParams +defaultSamplerChainParams = alloca $ \paramsPtr -> do + c_llama_sampler_chain_default_params_into paramsPtr + peek paramsPtr + +-- | Initialize a sampler +initSampler :: LlamaSamplerI -> LlamaSamplerContext -> IO (Either String Sampler) +initSampler iface_ ctx_ = do + ifacePtr <- mallocBytes (sizeOf (undefined :: LlamaSamplerI)) + poke ifacePtr iface_ + samplerPtr <- c_llama_sampler_init ifacePtr ctx_ + if samplerPtr == nullPtr + then return $ Left "Failed to initialize sampler" + else do + fp <- newForeignPtr p_llama_sampler_free samplerPtr + return $ Right (Sampler fp) + +-- | Get the name of a sampler +getSamplerName :: Sampler -> IO String +getSamplerName (Sampler samplerFPtr) = do + withForeignPtr samplerFPtr $ \samplerPtr -> do + name_ <- c_llama_sampler_name samplerPtr + peekCString name_ + +-- | Accept a token with a sampler +acceptTokenWithSampler :: Sampler -> LlamaToken -> IO () +acceptTokenWithSampler (Sampler samplerFPtr) token_ = do + withForeignPtr samplerFPtr $ \samplerPtr -> + c_llama_sampler_accept samplerPtr token_ + +-- | Apply a sampler to a token data array +applySampler :: Sampler -> LlamaTokenDataArray -> IO () +applySampler (Sampler samplerFPtr) tokenDataArray = do + withForeignPtr samplerFPtr $ \samplerPtr -> + withStorable tokenDataArray $ \tokenDataArrayPtr -> + c_llama_sampler_apply samplerPtr tokenDataArrayPtr + +-- | Reset a sampler +resetSampler :: Sampler -> IO () +resetSampler (Sampler samplerFPtr) = do + withForeignPtr samplerFPtr $ \samplerPtr -> + c_llama_sampler_reset (castPtr samplerPtr) + +-- | Clone a sampler +cloneSampler :: Sampler -> IO (Either String Sampler) +cloneSampler (Sampler samplerFPtr) = do + withForeignPtr samplerFPtr $ \samplerPtr -> do + clonedSamplerPtr <- c_llama_sampler_clone (castPtr samplerPtr) + if clonedSamplerPtr == nullPtr + then return $ Left "Failed to clone sampler" + else do + fp <- newForeignPtr p_llama_sampler_free clonedSamplerPtr + return $ Right $ Sampler fp + +-- | Initialize a sampler chain +initSamplerChain :: LlamaSamplerChainParams -> IO (Either String Sampler) +initSamplerChain params = do + chainPtr <- c_llama_sampler_chain_init params + if chainPtr == nullPtr + then return $ Left "Failed to initialize sampler chain" + else do + fp <- newForeignPtr p_llama_sampler_free chainPtr + return $ Right $ Sampler fp + +-- | Add a sampler to a sampler chain +addSamplerToChain :: Sampler -> Ptr LlamaSampler -> IO () +addSamplerToChain (Sampler chainFPtr) samplerPtr = do + withForeignPtr chainFPtr $ \chainPtr -> + c_llama_sampler_chain_add chainPtr samplerPtr + +getSamplerFromChain :: Sampler -> Int -> IO (Either String Sampler) +getSamplerFromChain (Sampler chainFPtr) index = do + withForeignPtr chainFPtr $ \chainPtr -> do + samplerPtr <- c_llama_sampler_chain_get chainPtr (fromIntegral index) + if samplerPtr == nullPtr + then return $ Left "Failed to get sampler from chain" + else do + fp <- newForeignPtr p_llama_sampler_free samplerPtr + return $ Right $ Sampler fp + +-- | Get the number of samplers in a sampler chain +getSamplerChainLength :: Sampler -> IO Int +getSamplerChainLength (Sampler chainFPtr) = do + withForeignPtr chainFPtr $ \chainPtr -> do + fromIntegral <$> c_llama_sampler_chain_n chainPtr + +-- | Remove a sampler from a sampler chain +removeSamplerFromChain :: Sampler -> Int -> IO (Either String Sampler) +removeSamplerFromChain (Sampler chainFPtr) index = do + withForeignPtr chainFPtr $ \chainPtr -> do + samplerPtr <- c_llama_sampler_chain_remove chainPtr (fromIntegral index) + if samplerPtr == nullPtr + then return $ Left "Failed to remove sampler from chain" + else do + fp <- newForeignPtr p_llama_sampler_free samplerPtr + return $ Right $ Sampler fp + +-- | Initialize a greedy sampler +initGreedySampler :: IO (Either String (Ptr LlamaSampler)) +initGreedySampler = do + samplerPtr <- c_llama_sampler_init_greedy + if samplerPtr == nullPtr + then return $ Left "Failed to initialize greedy sampler" + else do + return $ Right samplerPtr + +-- | Initialize a distributed sampler +initDistributedSampler :: Word32 -> IO (Either String Sampler) +initDistributedSampler seed = do + samplerPtr <- c_llama_sampler_init_dist (fromIntegral seed) + if samplerPtr == nullPtr + then return $ Left "Failed to initialize distributed sampler" + else do + fp <- newForeignPtr p_llama_sampler_free samplerPtr + return $ Right $ Sampler fp + +-- | Initialize a top-k sampler +initTopKSampler :: Int -> IO (Either String Sampler) +initTopKSampler k = do + samplerPtr <- c_llama_sampler_init_top_k (fromIntegral k) + if samplerPtr == nullPtr + then return $ Left "Failed to initialize top-k sampler" + else do + fp <- newForeignPtr p_llama_sampler_free samplerPtr + return $ Right $ Sampler fp + +-- | Initialize a top-p sampler +initTopPSampler :: Float -> Int -> IO (Either String Sampler) +initTopPSampler p1 minKeep = do + samplerPtr <- c_llama_sampler_init_top_p (realToFrac p1) (fromIntegral minKeep) + if samplerPtr == nullPtr + then return $ Left "Failed to initialize top-p sampler" + else do + fp <- newForeignPtr p_llama_sampler_free samplerPtr + return $ Right $ Sampler fp + +-- | Initialize a min-p sampler +initMinPSampler :: Float -> Int -> IO (Either String Sampler) +initMinPSampler p1 minKeep = do + samplerPtr <- c_llama_sampler_init_min_p (realToFrac p1) (fromIntegral minKeep) + if samplerPtr == nullPtr + then return $ Left "Failed to initialize min-p sampler" + else do + fp <- newForeignPtr p_llama_sampler_free samplerPtr + return $ Right $ Sampler fp + +-- | Initialize a typical sampler +initTypicalSampler :: Float -> Int -> IO (Either String Sampler) +initTypicalSampler p_ minKeep = do + samplerPtr <- c_llama_sampler_init_typical (realToFrac p_) (fromIntegral minKeep) + if samplerPtr == nullPtr + then return $ Left "Failed to initialize typical sampler" + else do + fp <- newForeignPtr p_llama_sampler_free samplerPtr + return $ Right $ Sampler fp + +-- | Initialize a temperature sampler +initTempSampler :: Float -> IO (Either String Sampler) +initTempSampler t = do + samplerPtr <- c_llama_sampler_init_temp (realToFrac t) + if samplerPtr == nullPtr + then return $ Left "Failed to initialize temperature sampler" + else do + fp <- newForeignPtr p_llama_sampler_free samplerPtr + return $ Right $ Sampler fp + +-- | Initialize an extended temperature sampler +initTempExtSampler :: Float -> Float -> Float -> IO (Either String Sampler) +initTempExtSampler t delta exponent_ = do + samplerPtr <- c_llama_sampler_init_temp_ext (realToFrac t) (realToFrac delta) (realToFrac exponent_) + if samplerPtr == nullPtr + then return $ Left "Failed to initialize extended temperature sampler" + else do + fp <- newForeignPtr p_llama_sampler_free samplerPtr + return $ Right $ Sampler fp + +-- | Initialize an XTC sampler +initXTCSampler :: Float -> Float -> Int -> Word32 -> IO (Either String Sampler) +initXTCSampler p1 t minKeep seed = do + samplerPtr <- c_llama_sampler_init_xtc (realToFrac p1) (realToFrac t) (fromIntegral minKeep) (fromIntegral seed) + if samplerPtr == nullPtr + then return $ Left "Failed to initialize XTC sampler" + else do + fp <- newForeignPtr p_llama_sampler_free samplerPtr + return $ Right $ Sampler fp + +-- | Initialize a top-N sigma sampler +initTopNSigmaSampler :: Float -> IO (Either String Sampler) +initTopNSigmaSampler n = do + samplerPtr <- c_llama_sampler_init_top_n_sigma (realToFrac n) + if samplerPtr == nullPtr + then return $ Left "Failed to initialize top-N sigma sampler" + else do + fp <- newForeignPtr p_llama_sampler_free samplerPtr + return $ Right $ Sampler fp + +-- | Initialize a Mirostat sampler +initMirostatSampler :: Int -> Word32 -> Float -> Float -> Int -> IO (Either String Sampler) +initMirostatSampler nVocab seed tau eta m = do + samplerPtr <- c_llama_sampler_init_mirostat (fromIntegral nVocab) (fromIntegral seed) (realToFrac tau) (realToFrac eta) (fromIntegral m) + if samplerPtr == nullPtr + then return $ Left "Failed to initialize Mirostat sampler" + else do + fp <- newForeignPtr p_llama_sampler_free samplerPtr + return $ Right $ Sampler fp + +-- | Initialize a Mirostat V2 sampler +initMirostatV2Sampler :: Word32 -> Float -> Float -> IO (Either String Sampler) +initMirostatV2Sampler seed tau eta = do + samplerPtr <- c_llama_sampler_init_mirostat_v2 (fromIntegral seed) (realToFrac tau) (realToFrac eta) + if samplerPtr == nullPtr + then return $ Left "Failed to initialize Mirostat V2 sampler" + else do + fp <- newForeignPtr p_llama_sampler_free samplerPtr + return $ Right $ Sampler fp + +-- | Initialize a grammar sampler +initGrammarSampler :: Vocab -> String -> String -> IO (Either String Sampler) +initGrammarSampler (Vocab vocab) grammarStr grammarRoot = do + withForeignPtr vocab $ \vocabPtr -> do + withCString grammarStr $ \grammarStrPtr -> do + withCString grammarRoot $ \grammarRootPtr -> do + samplerPtr <- c_llama_sampler_init_grammar (CLlamaVocab vocabPtr) grammarStrPtr grammarRootPtr + if samplerPtr == nullPtr + then return $ Left "Failed to initialize grammar sampler" + else do + fp <- newForeignPtr p_llama_sampler_free samplerPtr + return $ Right $ Sampler fp + +-- | Initialize a grammar sampler with lazy patterns +initGrammarLazyPatternsSampler :: Vocab -> String -> String -> [String] -> [LlamaToken] -> IO (Either String Sampler) +initGrammarLazyPatternsSampler (Vocab vocab) grammarStr grammarRoot triggerPatterns triggerTokens = do + withForeignPtr vocab $ \vocabPtr -> do + withCString grammarStr $ \grammarStrPtr -> do + withCString grammarRoot $ \grammarRootPtr -> do + triggerPatternsPtr <- newArrayOfPtrs triggerPatterns + withForeignPtr triggerPatternsPtr $ \triggerPatternsPtr' -> do + triggerTokensPtr <- mallocBytes (length triggerTokens * sizeOf (undefined :: LlamaToken)) + pokeArray triggerTokensPtr triggerTokens + samplerPtr <- c_llama_sampler_init_grammar_lazy_patterns + (CLlamaVocab vocabPtr) + grammarStrPtr + grammarRootPtr + triggerPatternsPtr' + (fromIntegral (length triggerPatterns)) + triggerTokensPtr + (fromIntegral (length triggerTokens)) + if samplerPtr == nullPtr + then return $ Left "Failed to initialize grammar lazy patterns sampler" + else do + fp <- newForeignPtr p_llama_sampler_free samplerPtr + return $ Right $ Sampler fp + +-- | Initialize a penalties sampler +initPenaltiesSampler :: Int -> Float -> Float -> Float -> IO (Either String Sampler) +initPenaltiesSampler penaltyLastN penaltyRepeat penaltyFreq penaltyPresent = do + samplerPtr <- c_llama_sampler_init_penalties + (fromIntegral penaltyLastN) + (realToFrac penaltyRepeat) + (realToFrac penaltyFreq) + (realToFrac penaltyPresent) + if samplerPtr == nullPtr + then return $ Left "Failed to initialize penalties sampler" + else do + fp <- newForeignPtr p_llama_sampler_free samplerPtr + return $ Right $ Sampler fp + +newArrayOfPtrs :: [String] -> IO (ForeignPtr CString) +newArrayOfPtrs xs = do + ptrs <- mallocBytes (length xs * sizeOf (undefined :: Ptr CString)) + mapM_ (\(i, x) -> withCString x $ \cstr -> poke (castPtr ptrs `plusPtr` (i * sizeOf (undefined :: Ptr CString))) cstr) (zip [0..] xs) + newForeignPtr_ ptrs + +-- Helper function +withCStringArray :: [String] -> (Ptr CString -> IO a) -> IO a +withCStringArray xs f = do + ptrs <- newArrayOfPtrs xs + withForeignPtr ptrs $ \ptrs' -> + f ptrs' + +-- | Initialize a dry sampler +initDrySampler :: Vocab -> Int -> Float -> Float -> Int -> Int -> [String] -> IO (Either String Sampler) +initDrySampler (Vocab vocab) nCtxTrain dryMultiplier dryBase dryAllowedLength dryPenaltyLastN seqBreakers = do + withForeignPtr vocab $ \vocabPtr -> do + withCStringArray seqBreakers $ \seqBreakersPtr -> do + samplerPtr <- c_llama_sampler_init_dry + (CLlamaVocab vocabPtr) + (fromIntegral nCtxTrain) + (realToFrac dryMultiplier) + (realToFrac dryBase) + (fromIntegral dryAllowedLength) + (fromIntegral dryPenaltyLastN) + seqBreakersPtr + (fromIntegral (length seqBreakers)) + if samplerPtr == nullPtr + then return $ Left "Failed to initialize dry sampler" + else do + fp <- newForeignPtr p_llama_sampler_free samplerPtr + return $ Right $ Sampler fp + +-- | Initialize a logit bias sampler +initLogitBiasSampler :: Int -> [LlamaLogitBias] -> IO (Either String Sampler) +initLogitBiasSampler nVocab logitBiases = do + allocaArray (length logitBiases) $ \logitBiasPtr -> do + pokeArray logitBiasPtr logitBiases + samplerPtr <- c_llama_sampler_init_logit_bias + (fromIntegral nVocab) + (fromIntegral (length logitBiases)) + logitBiasPtr + if samplerPtr == nullPtr + then return $ Left "Failed to initialize logit bias sampler" + else do + fp <- newForeignPtr p_llama_sampler_free samplerPtr + return $ Right $ Sampler fp + +-- | Initialize an infill sampler +initInfillSampler :: Vocab -> IO (Either String Sampler) +initInfillSampler (Vocab vocab) = do + withForeignPtr vocab $ \vocabPtr -> do + samplerPtr <- c_llama_sampler_init_infill (CLlamaVocab vocabPtr) + if samplerPtr == nullPtr + then return $ Left "Failed to initialize infill sampler" + else do + fp <- newForeignPtr p_llama_sampler_free samplerPtr + return $ Right $ Sampler fp + +-- | Get the seed used by a sampler +getSamplerSeed :: Sampler -> IO Word32 +getSamplerSeed (Sampler samplerFPtr) = do + withForeignPtr samplerFPtr (fmap fromIntegral . c_llama_sampler_get_seed) + +-- | Sample with a sampler +sampleWithSampler :: Sampler -> Context -> Int -> IO LlamaToken +sampleWithSampler (Sampler samplerFPtr) (Context ctxFPtr) idx = do + withForeignPtr samplerFPtr $ \samplerPtr -> do + withForeignPtr ctxFPtr $ \ctxPtr -> + c_llama_sampler_sample samplerPtr (CLlamaContext ctxPtr) (fromIntegral idx) diff --git a/src/Llama/Split.hs b/src/Llama/Split.hs new file mode 100644 --- /dev/null +++ b/src/Llama/Split.hs @@ -0,0 +1,46 @@ +{- | +Module : Llama.Split +Description : High level Split interface for llama-cpp +Copyright : (c) 2025 Tushar Adhatrao +License : MIT +Maintainer : Tushar Adhatrao +-} +module Llama.Split ( + splitPath + , splitPrefix + , printSystemInfo +) where + +import Llama.Internal.Foreign +import Foreign +import Foreign.C.String + +-- | Split a path into a prefix and a split number +splitPath :: FilePath -> Int -> Int -> IO (Either String String) +splitPath pathPrefix splitNo splitCount = do + withCString pathPrefix $ \cPathPrefix -> do + allocaArray 256 $ \splitPathPtr -> do + requiredLen <- c_llama_split_path splitPathPtr 256 cPathPrefix (fromIntegral splitNo) (fromIntegral splitCount) + if requiredLen < 0 + then return $ Left "Failed to split path" + else do + str <- peekCString splitPathPtr + return $ Right str + +-- | Get the prefix from a split path +splitPrefix :: FilePath -> Int -> Int -> IO (Either String String) +splitPrefix splitPath_ splitNo splitCount = do + withCString splitPath_ $ \cSplitPath -> do + allocaArray 256 $ \splitPrefixPtr -> do + requiredLen <- c_llama_split_prefix splitPrefixPtr 256 cSplitPath (fromIntegral splitNo) (fromIntegral splitCount) + if requiredLen < 0 + then return $ Left "Failed to get prefix from split path" + else do + str <- peekCString splitPrefixPtr + return $ Right str + +-- | Print system information +printSystemInfo :: IO String +printSystemInfo = do + systemInfo <- c_llama_print_system_info + peekCString systemInfo diff --git a/src/Llama/State.hs b/src/Llama/State.hs new file mode 100644 --- /dev/null +++ b/src/Llama/State.hs @@ -0,0 +1,119 @@ +{- | +Module : Llama.State +Description : High level State interface for llama-cpp +Copyright : (c) 2025 Tushar Adhatrao +License : MIT +Maintainer : Tushar Adhatrao +-} +module Llama.State ( + getStateSize +, getStateData +, setStateData +, loadStateFromFile +, saveStateToFile +, getSequenceStateSize +, setSequenceStateData +, saveSequenceStateToFile +, loadSequenceStateFromFile +) where + +import Llama.Internal.Foreign +import Llama.Internal.Types +import Foreign +import Foreign.C.String +import qualified Data.ByteString as BS +import Data.ByteString (ByteString) + +-- | Get the size of the state +getStateSize :: Context -> IO Word64 +getStateSize (Context ctxFPtr) = do + withForeignPtr ctxFPtr $ \ctxPtr -> + fromIntegral <$> c_llama_state_get_size (CLlamaContext ctxPtr) + +-- | Get the state data +getStateData :: Context -> IO ByteString +getStateData (Context ctxFPtr) = do + withForeignPtr ctxFPtr $ \ctxPtr -> do + size <- c_llama_state_get_size (CLlamaContext ctxPtr) + allocaBytes (fromIntegral size) $ \dstPtr -> do + _ <- c_llama_state_get_data (CLlamaContext ctxPtr) dstPtr size + BS.pack <$> peekArray (fromIntegral size) dstPtr + +-- | Set the state data +setStateData :: Context -> ByteString -> IO () +setStateData (Context ctxFPtr) bs = do + withForeignPtr ctxFPtr $ \ctxPtr -> do + withArray (BS.unpack bs) $ \srcPtr -> do + _ <- c_llama_state_set_data (CLlamaContext ctxPtr) srcPtr (fromIntegral (BS.length bs)) + return () + +-- | Load a state from a file +loadStateFromFile :: Context -> FilePath -> [LlamaToken] -> IO [LlamaToken] +loadStateFromFile (Context ctxFPtr) pathSession tokens = do + withForeignPtr ctxFPtr $ \ctxPtr -> do + withCString pathSession $ \cPathSession -> do + allocaArray (length tokens) $ \tokensOutPtr -> do + alloca $ \nTokenCountOutPtr -> do + success <- c_llama_state_load_file + (CLlamaContext ctxPtr) + cPathSession + tokensOutPtr + (fromIntegral (length tokens)) + nTokenCountOutPtr + if success == 0 + then return [] + else do + tokenCount <- peek nTokenCountOutPtr + peekArray (fromIntegral tokenCount) tokensOutPtr + +-- | Save a state to a file +saveStateToFile :: Context -> FilePath -> [LlamaToken] -> IO Bool +saveStateToFile (Context ctxFPtr) pathSession tokens = do + withForeignPtr ctxFPtr $ \ctxPtr -> do + withCString pathSession $ \cPathSession -> do + withArray tokens $ \tokensPtr -> do + toBool <$> c_llama_state_save_file + (CLlamaContext ctxPtr) + cPathSession + tokensPtr + (fromIntegral (length tokens)) + +-- | Get the size of a sequence in the state +getSequenceStateSize :: Context -> LlamaSeqId -> IO Word64 +getSequenceStateSize (Context ctxFPtr) seqId = do + withForeignPtr ctxFPtr $ \ctxPtr -> + fromIntegral <$> c_llama_state_seq_get_size (CLlamaContext ctxPtr) seqId + +-- | Set the state data for a sequence +setSequenceStateData :: Context -> ByteString -> LlamaSeqId -> IO Word64 +setSequenceStateData (Context ctxFPtr) bs seqId = do + withForeignPtr ctxFPtr $ \ctxPtr -> do + withArray (BS.unpack bs) $ \srcPtr -> do + fromIntegral <$> + c_llama_state_seq_set_data (CLlamaContext ctxPtr) srcPtr (fromIntegral (BS.length bs)) seqId + +-- | Save a sequence state to a file +saveSequenceStateToFile :: Context -> FilePath -> LlamaSeqId -> [LlamaToken] -> IO Word64 +saveSequenceStateToFile (Context ctxFPtr) filepath seqId tokens = do + withForeignPtr ctxFPtr $ \ctxPtr -> do + withArray tokens $ \tokenPtr -> do + withCString filepath $ \cfilepath -> do + fromIntegral <$> + c_llama_state_seq_save_file (CLlamaContext ctxPtr) cfilepath seqId tokenPtr (fromIntegral (length tokens)) + +-- | Load a sequence state from a file +loadSequenceStateFromFile :: Context -> FilePath -> LlamaSeqId -> [LlamaToken] -> IO [LlamaToken] +loadSequenceStateFromFile (Context ctxFPtr) filepath destSeqId tokens = do + withForeignPtr ctxFPtr $ \ctxPtr -> do + withCString filepath $ \cfilepath -> do + allocaArray (length tokens) $ \tokensOutPtr -> do + alloca $ \nTokenCountOutPtr -> do + _ <- c_llama_state_seq_load_file + (CLlamaContext ctxPtr) + cfilepath + destSeqId + tokensOutPtr + (fromIntegral (length tokens)) + nTokenCountOutPtr + tokenCount <- peek nTokenCountOutPtr + peekArray (fromIntegral tokenCount) tokensOutPtr diff --git a/src/Llama/Tokenize.hs b/src/Llama/Tokenize.hs new file mode 100644 --- /dev/null +++ b/src/Llama/Tokenize.hs @@ -0,0 +1,65 @@ +{- | +Module : Llama.Tokenize +Description : High level Tokenize interface for llama-cpp +Copyright : (c) 2025 Tushar Adhatrao +License : MIT +Maintainer : Tushar Adhatrao +-} +module Llama.Tokenize ( + tokenize + , tokenToPiece + , detokenize +) where + +import Llama.Internal.Foreign +import Llama.Internal.Types +import Foreign +import Foreign.C.String + +-- | Tokenize a string into tokens +tokenize :: Vocab -> String -> Bool -> Bool -> IO ([LlamaToken], Int) +tokenize (Vocab vocab) text addSpecial parseSpecial = do + withForeignPtr vocab $ \vocabPtr -> do + withCString text $ \cText -> do + let textLen = length text + allocaArray textLen $ \tokensPtr -> do + tokenCount <- c_llama_tokenize + (CLlamaVocab vocabPtr) + cText + (fromIntegral textLen) + tokensPtr + (fromIntegral textLen) + (fromBool addSpecial) + (fromBool parseSpecial) + tokens <- peekArray (fromIntegral tokenCount) tokensPtr + return (tokens, fromIntegral tokenCount) + +-- | Convert a token to a piece of text +tokenToPiece :: Vocab -> LlamaToken -> Bool -> IO String +tokenToPiece (Vocab vocab) token_ special = do + withForeignPtr vocab $ \vocabPtr -> do + allocaArray 256 $ \bufPtr -> do + _ <- c_llama_token_to_piece + (CLlamaVocab vocabPtr) + token_ + bufPtr + 256 + 0 + (fromBool special) + peekCString bufPtr + +-- | Detokenize tokens into a string +detokenize :: Vocab -> [LlamaToken] -> Bool -> Bool -> IO String +detokenize (Vocab vocab) tokens removeSpecial unparseSpecial = do + withForeignPtr vocab $ \vocabPtr -> do + allocaArray 256 $ \textPtr -> do + withArray tokens $ \tokensPtr -> do + _ <- c_llama_detokenize + (CLlamaVocab vocabPtr) + tokensPtr + (fromIntegral $ length tokens) + textPtr + 256 + (fromBool removeSpecial) + (fromBool unparseSpecial) + peekCString textPtr diff --git a/src/Llama/Vocab.hs b/src/Llama/Vocab.hs new file mode 100644 --- /dev/null +++ b/src/Llama/Vocab.hs @@ -0,0 +1,155 @@ +{- | +Module : Llama.Tokenize +Description : High level Tokenize interface for llama-cpp +Copyright : (c) 2025 Tushar Adhatrao +License : MIT +Maintainer : Tushar Adhatrao +-} +module Llama.Vocab ( + getVocabSize +, getVocabTokenCount +, getVocabTokenText +, getVocabTokenScore +, getVocabTokenAttr +, isVocabTokenEog +, isVocabTokenControl +, getVocabBosToken +, getVocabEosToken +, getVocabEotToken +, getVocabSepToken +, getVocabNlToken +, getVocabPadToken +, getVocabAddBOSToken +, getVocabAddEOSToken +, getVocabFIMPrefixToken +, getVocabFIMSuffixToken +, getVocabFIMMiddleToken +, getVocabFIMPADToken +, getVocabFIMSeparatorToken +) where + +import Llama.Internal.Foreign +import Llama.Internal.Types +import Foreign +import Foreign.C.String + +-- | Get the number of vocab entries +getVocabSize :: Vocab -> IO Int +getVocabSize (Vocab vocabFPtr) = do + withForeignPtr vocabFPtr $ \vocabPtr -> + fromIntegral <$> c_llama_n_vocab (CLlamaVocab vocabPtr) + +-- | Get the number of tokens in the vocab +getVocabTokenCount :: Vocab -> IO Int +getVocabTokenCount (Vocab vocabFPtr) = do + withForeignPtr vocabFPtr $ \vocabPtr -> + fromIntegral <$> c_llama_vocab_n_tokens (CLlamaVocab vocabPtr) + +-- | Get the text for a token +getVocabTokenText :: Vocab -> LlamaToken -> IO String +getVocabTokenText (Vocab vocabFPtr) token_ = do + withForeignPtr vocabFPtr $ \vocabPtr -> do + cText <- c_llama_vocab_get_text (CLlamaVocab vocabPtr) token_ + peekCString cText + +-- | Get the score for a token +getVocabTokenScore :: Vocab -> LlamaToken -> IO Float +getVocabTokenScore (Vocab vocabFPtr) token_ = do + withForeignPtr vocabFPtr $ \vocabPtr -> + realToFrac <$> c_llama_vocab_get_score (CLlamaVocab vocabPtr) token_ + +-- | Get the attribute for a token +getVocabTokenAttr :: Vocab -> LlamaToken -> IO Int +getVocabTokenAttr (Vocab vocabFPtr) token_ = do + withForeignPtr vocabFPtr $ \vocabPtr -> + fromIntegral <$> c_llama_vocab_get_attr (CLlamaVocab vocabPtr) token_ + +-- | Check if a token is end-of-grammar +isVocabTokenEog :: Vocab -> LlamaToken -> IO Bool +isVocabTokenEog (Vocab vocabFPtr) token_ = do + withForeignPtr vocabFPtr $ \vocabPtr -> + (/= 0) <$> c_llama_vocab_is_eog (CLlamaVocab vocabPtr) token_ + +-- | Check if a token is a control token +isVocabTokenControl :: Vocab -> LlamaToken -> IO Bool +isVocabTokenControl (Vocab vocabFPtr) token_ = do + withForeignPtr vocabFPtr $ \vocabPtr -> + (/= 0) <$> c_llama_vocab_is_control (CLlamaVocab vocabPtr) token_ + +-- | Get the beginning-of-sentence token +getVocabBosToken :: Vocab -> IO LlamaToken +getVocabBosToken (Vocab vocabFPtr) = do + withForeignPtr vocabFPtr $ \vocabPtr -> + c_llama_vocab_bos (CLlamaVocab vocabPtr) + +-- | Get the end-of-sentence token +getVocabEosToken :: Vocab -> IO LlamaToken +getVocabEosToken (Vocab vocabFPtr) = do + withForeignPtr vocabFPtr $ \vocabPtr -> + c_llama_vocab_eos (CLlamaVocab vocabPtr) + +-- | Get the end-of-turn token +getVocabEotToken :: Vocab -> IO LlamaToken +getVocabEotToken (Vocab vocabFPtr) = do + withForeignPtr vocabFPtr $ \vocabPtr -> + c_llama_vocab_eot (CLlamaVocab vocabPtr) + +-- | Get the sentence separator token +getVocabSepToken :: Vocab -> IO LlamaToken +getVocabSepToken (Vocab vocabFPtr) = do + withForeignPtr vocabFPtr $ \vocabPtr -> + c_llama_vocab_sep (CLlamaVocab vocabPtr) + +-- | Get the next-line token +getVocabNlToken :: Vocab -> IO LlamaToken +getVocabNlToken (Vocab vocabFPtr) = do + withForeignPtr vocabFPtr $ \vocabPtr -> + c_llama_vocab_nl (CLlamaVocab vocabPtr) + +-- | Get the padding token +getVocabPadToken :: Vocab -> IO LlamaToken +getVocabPadToken (Vocab vocabFPtr) = do + withForeignPtr vocabFPtr $ \vocabPtr -> + c_llama_vocab_pad (CLlamaVocab vocabPtr) + +-- | Get whether to add BOS token automatically +getVocabAddBOSToken :: Vocab -> IO Bool +getVocabAddBOSToken (Vocab vocabFPtr) = do + withForeignPtr vocabFPtr $ \vocabPtr -> + (/= 0) <$> c_llama_vocab_get_add_bos (CLlamaVocab vocabPtr) + +-- | Get whether to add EOS token automatically +getVocabAddEOSToken :: Vocab -> IO Bool +getVocabAddEOSToken (Vocab vocabFPtr) = do + withForeignPtr vocabFPtr $ \vocabPtr -> + (/= 0) <$> c_llama_vocab_get_add_eos (CLlamaVocab vocabPtr) + +-- | Get the FIM prefix token +getVocabFIMPrefixToken :: Vocab -> IO LlamaToken +getVocabFIMPrefixToken (Vocab vocabFPtr) = do + withForeignPtr vocabFPtr $ \vocabPtr -> + c_llama_vocab_fim_pre (CLlamaVocab vocabPtr) + +-- | Get the FIM suffix token +getVocabFIMSuffixToken :: Vocab -> IO LlamaToken +getVocabFIMSuffixToken (Vocab vocabFPtr) = do + withForeignPtr vocabFPtr $ \vocabPtr -> + c_llama_vocab_fim_suf (CLlamaVocab vocabPtr) + +-- | Get the FIM middle token +getVocabFIMMiddleToken :: Vocab -> IO LlamaToken +getVocabFIMMiddleToken (Vocab vocabFPtr) = do + withForeignPtr vocabFPtr $ \vocabPtr -> + c_llama_vocab_fim_mid (CLlamaVocab vocabPtr) + +-- | Get the FIM pad token +getVocabFIMPADToken :: Vocab -> IO LlamaToken +getVocabFIMPADToken (Vocab vocabFPtr) = do + withForeignPtr vocabFPtr $ \vocabPtr -> + c_llama_vocab_fim_pad (CLlamaVocab vocabPtr) + +-- | Get the FIM separator token +getVocabFIMSeparatorToken :: Vocab -> IO LlamaToken +getVocabFIMSeparatorToken (Vocab vocabFPtr) = do + withForeignPtr vocabFPtr $ \vocabPtr -> + c_llama_vocab_fim_sep (CLlamaVocab vocabPtr) diff --git a/test/Spec.hs b/test/Spec.hs new file mode 100644 --- /dev/null +++ b/test/Spec.hs @@ -0,0 +1,266 @@ +import Control.Exception (bracket) +import Data.Either +import qualified Llama.Backend as Backend +import qualified Llama.ChatTemplate as ChatTemplate +import Llama.Context +import Llama.Model +import Test.Tasty (TestTree, defaultMain, testGroup) +import Test.Tasty.HUnit + +main :: IO () +main = defaultMain tests + +tests :: TestTree +tests = do + testGroup + "tests" + [ testBackend + , testChatApplyTemplate + , checkContextFunctions + , checkIntFunctions + , checkBoolFunctions + , testModels + -- , synchronizeContextTests ctx + -- , warmupModeTests ctx + -- , batchInitTests + -- , batchGetOneTests + -- , freeBatchTests + -- , encodeDecodeTests ctx + -- , threadCountTests ctx + -- , embeddingsTests ctx + -- , causalAttentionTests ctx + ] + +testBackend :: TestTree +testBackend = + testGroup + "Llama.Backend tests" + [ testCase "llamaBackendInit and llamaBackendFree should not throw exceptions" $ do + Backend.llamaBackendInit + Backend.llamaBackendFree + return () + , testCase "llamaBackendInit and llamaBackendFree should be idempotent" $ do + Backend.llamaBackendInit + Backend.llamaBackendFree + Backend.llamaBackendFree + return () + , testCase "llamaBackendInit can be called multiple times" $ do + Backend.llamaBackendInit + Backend.llamaBackendInit + Backend.llamaBackendFree + Backend.llamaBackendFree + return () + , testCase "llamaBackendInit and llamaBackendFree should not throw exceptions" $ do + bracket Backend.llamaBackendInit (const Backend.llamaBackendFree) $ \_ -> return () + ] + +testChatApplyTemplate :: TestTree +testChatApplyTemplate = + testGroup + "chatApplyTemplate tests" + [ testCase "chatApplyTemplate with no template" $ do + messages <- + ChatTemplate.chatApplyTemplate + Nothing + [ChatTemplate.ChatMessage "role" "content"] + False + 4096 + assertBool "result is not Left" (isRight messages) + , -- testCase "chatApplyTemplate with custom template" $ do + -- messages <- ChatTemplate.chatApplyTemplate + -- (Just "template") [ChatTemplate.ChatMessage "role" "content"] False 4096 + -- assertBool "result is not Left" (isRight messages), -- throwing error as of now + testCase "chatApplyTemplate with empty messages" $ do + messages <- ChatTemplate.chatApplyTemplate Nothing [] False 4096 + assertBool "result is not Left" (isRight messages) + ] + +checkBoolFunctions :: TestTree +checkBoolFunctions = + testGroup + "Boolean Functions" + [ testCase "supportsRpc" $ do + result <- supportsRpc + assertBool "supportsRpc returns a boolean value" (not result) -- My env doesn't support + , testCase "supportsGpuOffload" $ do + result <- supportsGpuOffload + assertBool "supportsGpuOffload returns a boolean value" (not result) + , testCase "supportsMLock" $ do + result <- supportsMLock + assertBool "supportsMLock returns a boolean value" result + , testCase "supportsMMap" $ do + result <- supportsMMap + assertBool "supportsMMap returns a boolean value" result + ] + +checkIntFunctions :: TestTree +checkIntFunctions = + testGroup + "Integer Functions" + [ testCase "getMaxDevices" $ do + result <- getMaxDevices + assertBool "getMaxDevices returns a non-negative integer" (result >= 0) + , testCase "getTimeUs" $ do + result <- getTimeUs + assertBool "getTimeUs returns a non-negative integer" (result >= 0) + ] + +checkContextFunctions :: TestTree +checkContextFunctions = + testGroup + "Context Functions" + [ testCase "defaultContextParams" $ do + params <- defaultContextParams + assertBool + "defaultContextParams returns a valid LlamaContextParams" + (not $ null $ show params) + ] + +testModels :: TestTree +testModels = + testGroup + "Model functions" + [ testCase "loadModelFromFile" testLoadModelFromFile + , testCase "initContextFromModel" testInitContextFromModel + , testCase "getModelVocab" testGetModelVocab + ] + +testLoadModelFromFile :: Assertion +testLoadModelFromFile = do + params <- defaultModelParams + result <- loadModelFromFile "Qwen3-0.6B-Q4_K_M.gguf" params + case result of + Right _ -> assertBool "model loaded successfully" True + Left err -> assertFailure err + +testInitContextFromModel :: Assertion +testInitContextFromModel = do + params <- defaultModelParams + result <- loadModelFromFile "Qwen3-0.6B-Q4_K_M.gguf" params + case result of + Right model -> do + contextParams <- defaultContextParams + result1 <- initContextFromModel model contextParams + case result1 of + Right _ -> assertBool "context initialized successfully" True + Left err -> assertFailure err + Left err -> assertFailure err + +testGetModelVocab :: Assertion +testGetModelVocab = do + params <- defaultModelParams + result <- loadModelFromFile "Qwen3-0.6B-Q4_K_M.gguf" params + case result of + Right model -> do + result1 <- getModelVocab model + case result1 of + Right _ -> assertBool "vocab retrieved successfully" True + Left err -> assertFailure err + Left err -> assertFailure err + +{- +testGetContextModel :: Assertion +testGetContextModel = do + params <- defaultModelParams + result <- loadModelFromFile "Qwen3-0.6B-Q4_K_M.gguf" params + case result of + Right model -> do + contextParams <- defaultContextParams + result <- initContextFromModel model contextParams + case result of + Right context -> do + model' <- getContextModel context + assertEqual "models are equal" model model' + Left err -> assertFailure err + Left err -> assertFailure err + +modelParams <- liftIO defaultModelParams +eModel <- liftIO $ loadModelFromFile "Qwen3-0.6B-Q4_K_M.gguf" modelParams +case eModel of + Left err -> putStrLn $ "Something went wrong: " ++ err + Right model -> do + ctxParams <- defaultContextParams + eCtx <- initContextFromModel model ctxParams + case eCtx of + Left e -> putStrLn $ "Something went wrong" ++ e + Right ctx -> do + +batchInitTests :: TestTree +batchInitTests = testGroup "batchInit tests" + [ testCase "batchInit succeeds" $ do + batch <- batchInit 10 256 1 + freeBatch (getBatchPtr batch) + ] + +batchGetOneTests :: TestTree +batchGetOneTests = testGroup "batchGetOne tests" + [ testCase "batchGetOne succeeds" $ do + batch <- batchGetOne [1, 2, 3] + freeBatch (getBatchPtr batch) + ] + +freeBatchTests :: TestTree +freeBatchTests = testGroup "freeBatch tests" + [ testCase "freeBatch succeeds" $ do + batch <- batchInit 10 256 1 + freeBatch (getBatchPtr batch) + -- No assertion, just check if it doesn't crash + ] + +getBatchPtr :: Batch -> Ptr LlamaBatch +getBatchPtr (Batch ptr) = ptr + +encodeDecodeTests :: Context -> TestTree +encodeDecodeTests ctx = testGroup "encodeDecode tests" + [ testCase "encodeBatch succeeds" $ do + batch <- batchInit 10 256 1 + result <- encodeBatch ctx batch + assertBool "Encoding failed" (isRight result) + freeBatch (getBatchPtr batch) + -- decodeBatch test + result' <- decodeBatch ctx batch + assertBool "Decoding failed" (isRight result') + ] + +threadCountTests :: Context -> TestTree +threadCountTests ctx = testGroup "threadCount tests" + [ testCase "setThreadCount succeeds" $ do + setThreadCount ctx 4 + count <- getThreadCount ctx + assertBool "Thread count not set" (count == 4) + , testCase "getBatchThreadCount succeeds" $ do + setThreadCounts ctx 4 2 + count <- getBatchThreadCount ctx + assertBool "Batch thread count not set" (count == 2) + ] + +embeddingsTests :: Context -> TestTree +embeddingsTests ctx = testGroup "embeddings tests" + [ testCase "setEmbeddingsEnabled succeeds" $ do + setEmbeddingsEnabled ctx True + enabled <- areEmbeddingsEnabled ctx + assertBool "Embeddings not enabled" enabled + ] + +causalAttentionTests :: Context -> TestTree +causalAttentionTests ctx = testGroup "causalAttention tests" + [ testCase "setCausalAttention succeeds" $ do + setCausalAttention ctx True + -- No assertion, just check if it doesn't crash + ] + +warmupModeTests :: Context -> TestTree +warmupModeTests ctx = testGroup "warmupMode tests" + [ testCase "setWarmupMode succeeds" $ do + setWarmupMode ctx True + -- No assertion, just check if it doesn't crash + ] + +synchronizeContextTests :: Context -> TestTree +synchronizeContextTests ctx = testGroup "synchronizeContext tests" + [ testCase "synchronizeContext succeeds" $ do + synchronizeContext ctx + -- No assertion, just check if it doesn't crash + ] + +-}