Fix array type signatures

abetlen · abetlen · commit 1545b22727f5 · 2023-03-31T02:08:20.000-04:00
diff --git a/llama_cpp/llama_cpp.py b/llama_cpp/llama_cpp.py
@@ -116,7 +116,7 @@ def llama_model_quantize(
 # Returns 0 on success
 def llama_eval(
     ctx: llama_context_p,
-    tokens: llama_token_p,
+    tokens: ctypes.Array[llama_token],
     n_tokens: c_int,
     n_past: c_int,
     n_threads: c_int,
@@ -136,7 +136,7 @@ def llama_eval(
 def llama_tokenize(
     ctx: llama_context_p,
     text: bytes,
-    tokens: llama_token_p,
+    tokens: ctypes.Array[llama_token],
     n_max_tokens: c_int,
     add_bos: c_bool,
 ) -> c_int:
@@ -176,7 +176,7 @@ def llama_n_embd(ctx: llama_context_p) -> c_int:
 # Can be mutated in order to change the probabilities of the next token
 # Rows: n_tokens
 # Cols: n_vocab
-def llama_get_logits(ctx: llama_context_p):
+def llama_get_logits(ctx: llama_context_p) -> ctypes.Array[c_float]:
     return _lib.llama_get_logits(ctx)
 
 
@@ -186,7 +186,7 @@ def llama_get_logits(ctx: llama_context_p):
 
 # Get the embeddings for the input
 # shape: [n_embd] (1-dimensional)
-def llama_get_embeddings(ctx: llama_context_p):
+def llama_get_embeddings(ctx: llama_context_p) -> ctypes.Array[c_float]:
     return _lib.llama_get_embeddings(ctx)
 
 
@@ -224,7 +224,7 @@ def llama_token_eos() -> llama_token:
 # TODO: improve the last_n_tokens interface ?
 def llama_sample_top_p_top_k(
     ctx: llama_context_p,
-    last_n_tokens_data: llama_token_p,
+    last_n_tokens_data: ctypes.Array[llama_token],
     last_n_tokens_size: c_int,
     top_k: c_int,
     top_p: c_float,
