diff --git a/llama_cpp/llama_cpp.py b/llama_cpp/llama_cpp.py index 0a35445..87d9249 100644 --- a/llama_cpp/llama_cpp.py +++ b/llama_cpp/llama_cpp.py @@ -8,6 +8,7 @@ from ctypes import ( c_void_p, c_bool, POINTER, + _Pointer, # type: ignore Structure, Array, c_uint8, @@ -252,9 +253,7 @@ _lib.llama_get_state_size.restype = c_size_t # Copies the state to the specified destination address. # Destination needs to have allocated enough memory. # Returns the number of bytes copied -def llama_copy_state_data( - ctx: llama_context_p, dest # type: Array[c_uint8] -) -> c_size_t: +def llama_copy_state_data(ctx: llama_context_p, dest: Array[c_uint8]) -> c_size_t: return _lib.llama_copy_state_data(ctx, dest) @@ -278,9 +277,9 @@ _lib.llama_set_state_data.restype = c_size_t def llama_load_session_file( ctx: llama_context_p, path_session: bytes, - tokens_out, # type: Array[llama_token] + tokens_out: Array[llama_token], n_token_capacity: c_size_t, - n_token_count_out, # type: Array[c_size_t] + n_token_count_out: _Pointer[c_size_t], ) -> c_size_t: return _lib.llama_load_session_file( ctx, path_session, tokens_out, n_token_capacity, n_token_count_out @@ -300,7 +299,7 @@ _lib.llama_load_session_file.restype = c_size_t def llama_save_session_file( ctx: llama_context_p, path_session: bytes, - tokens, # type: Array[llama_token] + tokens: Array[llama_token], n_token_count: c_size_t, ) -> c_size_t: return _lib.llama_save_session_file(ctx, path_session, tokens, n_token_count) @@ -321,7 +320,7 @@ _lib.llama_save_session_file.restype = c_size_t # Returns 0 on success def llama_eval( ctx: llama_context_p, - tokens, # type: Array[llama_token] + tokens: Array[llama_token], n_tokens: c_int, n_past: c_int, n_threads: c_int, @@ -440,8 +439,8 @@ _lib.llama_token_nl.restype = llama_token # @details Repetition penalty described in CTRL academic paper https://arxiv.org/abs/1909.05858, with negative logit fix. def llama_sample_repetition_penalty( ctx: llama_context_p, - candidates, # type: Array[llama_token_data] - last_tokens_data, # type: Array[llama_token] + candidates: _Pointer[llama_token_data], + last_tokens_data: Array[llama_token], last_tokens_size: c_int, penalty: c_float, ): @@ -463,8 +462,8 @@ _lib.llama_sample_repetition_penalty.restype = None # @details Frequency and presence penalties described in OpenAI API https://platform.openai.com/docs/api-reference/parameter-details. def llama_sample_frequency_and_presence_penalties( ctx: llama_context_p, - candidates, # type: Array[llama_token_data] - last_tokens_data, # type: Array[llama_token] + candidates: _Pointer[llama_token_data], + last_tokens_data: Array[llama_token], last_tokens_size: c_int, alpha_frequency: c_float, alpha_presence: c_float, @@ -491,10 +490,7 @@ _lib.llama_sample_frequency_and_presence_penalties.restype = None # @details Sorts candidate tokens by their logits in descending order and calculate probabilities based on logits. -def llama_sample_softmax( - ctx: llama_context_p, - candidates # type: Array[llama_token_data] -): +def llama_sample_softmax(ctx: llama_context_p, candidates: _Pointer[llama_token_data]): return _lib.llama_sample_softmax(ctx, candidates) @@ -507,10 +503,10 @@ _lib.llama_sample_softmax.restype = None # @details Top-K sampling described in academic paper "The Curious Case of Neural Text Degeneration" https://arxiv.org/abs/1904.09751 def llama_sample_top_k( - ctx: llama_context_p, - candidates, # type: Array[llama_token_data] - k: c_int, - min_keep: c_size_t = c_size_t(1) + ctx: llama_context_p, + candidates: _Pointer[llama_token_data], + k: c_int, + min_keep: c_size_t = c_size_t(1), ): return _lib.llama_sample_top_k(ctx, candidates, k, min_keep) @@ -526,10 +522,10 @@ _lib.llama_sample_top_k.restype = None # @details Nucleus sampling described in academic paper "The Curious Case of Neural Text Degeneration" https://arxiv.org/abs/1904.09751 def llama_sample_top_p( - ctx: llama_context_p, - candidates, # type: Array[llama_token_data] - p: c_float, - min_keep: c_size_t = c_size_t(1) + ctx: llama_context_p, + candidates: _Pointer[llama_token_data], + p: c_float, + min_keep: c_size_t = c_size_t(1), ): return _lib.llama_sample_top_p(ctx, candidates, p, min_keep) @@ -546,9 +542,9 @@ _lib.llama_sample_top_p.restype = None # @details Tail Free Sampling described in https://www.trentonbricken.com/Tail-Free-Sampling/. def llama_sample_tail_free( ctx: llama_context_p, - candidates, # type: Array[llama_token_data] + candidates: _Pointer[llama_token_data], z: c_float, - min_keep: c_size_t = c_size_t(1) + min_keep: c_size_t = c_size_t(1), ): return _lib.llama_sample_tail_free(ctx, candidates, z, min_keep) @@ -565,9 +561,9 @@ _lib.llama_sample_tail_free.restype = None # @details Locally Typical Sampling implementation described in the paper https://arxiv.org/abs/2202.00666. def llama_sample_typical( ctx: llama_context_p, - candidates, # type: Array[llama_token_data] - p: c_float, - min_keep: c_size_t = c_size_t(1) + candidates: _Pointer[llama_token_data], + p: c_float, + min_keep: c_size_t = c_size_t(1), ): return _lib.llama_sample_typical(ctx, candidates, p, min_keep) @@ -582,9 +578,7 @@ _lib.llama_sample_typical.restype = None def llama_sample_temperature( - ctx: llama_context_p, - candidates, # type: Array[llama_token_data] - temp: c_float + ctx: llama_context_p, candidates: _Pointer[llama_token_data], temp: c_float ): return _lib.llama_sample_temperature(ctx, candidates, temp) @@ -605,11 +599,11 @@ _lib.llama_sample_temperature.restype = None # @param mu Maximum cross-entropy. This value is initialized to be twice the target cross-entropy (`2 * tau`) and is updated in the algorithm based on the error between the target and observed surprisal. def llama_sample_token_mirostat( ctx: llama_context_p, - candidates, # type: Array[llama_token_data] + candidates: _Pointer[llama_token_data], tau: c_float, - eta: c_float, + eta: c_float, m: c_int, - mu # type: Array[c_float] + mu: _Pointer[c_float], ) -> llama_token: return _lib.llama_sample_token_mirostat(ctx, candidates, tau, eta, m, mu) @@ -632,10 +626,10 @@ _lib.llama_sample_token_mirostat.restype = llama_token # @param mu Maximum cross-entropy. This value is initialized to be twice the target cross-entropy (`2 * tau`) and is updated in the algorithm based on the error between the target and observed surprisal. def llama_sample_token_mirostat_v2( ctx: llama_context_p, - candidates, # type: Array[llama_token_data] - tau: c_float, + candidates: _Pointer[llama_token_data], + tau: c_float, eta: c_float, - mu # type: Array[c_float] + mu: _Pointer[c_float], ) -> llama_token: return _lib.llama_sample_token_mirostat_v2(ctx, candidates, tau, eta, mu) @@ -652,8 +646,7 @@ _lib.llama_sample_token_mirostat_v2.restype = llama_token # @details Selects the token with the highest probability. def llama_sample_token_greedy( - ctx: llama_context_p, - candidates # type: Array[llama_token_data] + ctx: llama_context_p, candidates: _Pointer[llama_token_data] ) -> llama_token: return _lib.llama_sample_token_greedy(ctx, candidates) @@ -667,8 +660,7 @@ _lib.llama_sample_token_greedy.restype = llama_token # @details Randomly selects a token from the candidates based on their probabilities. def llama_sample_token( - ctx: llama_context_p, - candidates # type: Array[llama_token_data] + ctx: llama_context_p, candidates: _Pointer[llama_token_data] ) -> llama_token: return _lib.llama_sample_token(ctx, candidates)