Fix llama_cpp types

This commit is contained in:
Andrei Betlen 2023-05-05 13:54:22 -04:00
parent b6a9a0b6ba
commit 5e7ddfc3d6

View file

@ -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)