Add type definitions
This commit is contained in:
parent
49c8df369a
commit
69e7d9f60e
1 changed files with 57 additions and 0 deletions
57
llama_cpp/llama_types.py
Normal file
57
llama_cpp/llama_types.py
Normal file
|
@ -0,0 +1,57 @@
|
||||||
|
from typing import List, Optional, Dict, Literal
|
||||||
|
from typing_extensions import TypedDict
|
||||||
|
|
||||||
|
|
||||||
|
class EmbeddingUsage(TypedDict):
|
||||||
|
prompt_tokens: int
|
||||||
|
total_tokens: int
|
||||||
|
|
||||||
|
|
||||||
|
class EmbeddingData(TypedDict):
|
||||||
|
index: int
|
||||||
|
object: str
|
||||||
|
embedding: List[float]
|
||||||
|
|
||||||
|
|
||||||
|
class Embedding(TypedDict):
|
||||||
|
object: Literal["list"]
|
||||||
|
model: str
|
||||||
|
data: List[EmbeddingData]
|
||||||
|
usage: EmbeddingUsage
|
||||||
|
|
||||||
|
|
||||||
|
class CompletionLogprobs(TypedDict):
|
||||||
|
text_offset: List[int]
|
||||||
|
token_logprobs: List[float]
|
||||||
|
tokens: List[str]
|
||||||
|
top_logprobs: List[Dict[str, float]]
|
||||||
|
|
||||||
|
|
||||||
|
class CompletionChoice(TypedDict):
|
||||||
|
text: str
|
||||||
|
index: int
|
||||||
|
logprobs: Optional[CompletionLogprobs]
|
||||||
|
finish_reason: Optional[str]
|
||||||
|
|
||||||
|
|
||||||
|
class CompletionUsage(TypedDict):
|
||||||
|
prompt_tokens: int
|
||||||
|
completion_tokens: int
|
||||||
|
total_tokens: int
|
||||||
|
|
||||||
|
|
||||||
|
class CompletionChunk(TypedDict):
|
||||||
|
id: str
|
||||||
|
object: Literal["text_completion"]
|
||||||
|
created: int
|
||||||
|
model: str
|
||||||
|
choices: List[CompletionChoice]
|
||||||
|
|
||||||
|
|
||||||
|
class Completion(TypedDict):
|
||||||
|
id: str
|
||||||
|
object: Literal["text_completion"]
|
||||||
|
created: int
|
||||||
|
model: str
|
||||||
|
choices: List[CompletionChoice]
|
||||||
|
usage: CompletionUsage
|
Loading…
Reference in a new issue