alihan/llama-cpp-python
1
0
Fork 0
mirror of https://github.com/abetlen/llama-cpp-python.git synced 2023-09-07 17:34:22 +03:00
Code Issues Packages Projects Releases Wiki Activity

Add support for logprobs

Browse Source
This commit is contained in:
Andrei Betlen
2023-03-23 15:51:05 -04:00
parent 3d6eb32c76
commit e63ea4dbbc
1 changed files with 7 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
llama_cpp/llama.py
View File

@@ -44,6 +44,7 @@ class Llama:
max_tokens: int = 16, max_tokens: int = 16,
temperature: float = 0.8, temperature: float = 0.8,
top_p: float = 0.95, top_p: float = 0.95,
logprobs: Optional[int] = None,
echo: bool = False, echo: bool = False,
stop: List[str] = [], stop: List[str] = [],
repeat_penalty: float = 1.1, repeat_penalty: float = 1.1,
@@ -105,6 +106,11 @@ class Llama:
if suffix is not None: if suffix is not None:
text = text + suffix text = text + suffix
if logprobs is not None:
logprobs = llama_cpp.llama_get_logits(
self.ctx,
)[:logprobs]
return { return {
"id": f"cmpl-{str(uuid.uuid4())}", # Likely to change "id": f"cmpl-{str(uuid.uuid4())}", # Likely to change
"object": "text_completion", "object": "text_completion",
@@ -114,7 +120,7 @@ class Llama:
{ {
"text": text, "text": text,
"index": 0, "index": 0,
"logprobs": None, "logprobs": logprobs,
"finish_reason": finish_reason, "finish_reason": finish_reason,
} }
], ],