Add type definitions

abetlen · abetlen · commit 69e7d9f60ef9 · 2023-04-01T12:59:58.000-04:00
diff --git a/llama_cpp/llama_types.py b/llama_cpp/llama_types.py
@@ -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
