Add return type annotations for embeddings and logits

abetlen · abetlen · commit b5f3e746275b · 2023-05-05T14:22:55.000-04:00
diff --git a/llama_cpp/llama_cpp.py b/llama_cpp/llama_cpp.py
@@ -387,7 +387,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):  # type: (...) -> Array[float] # type: ignore
     return _lib.llama_get_logits(ctx)
 
 
@@ -397,7 +397,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):  # type: (...) -> Array[float] # type: ignore
     return _lib.llama_get_embeddings(ctx)
 
 
