baalajimaestro/llama.cpp
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages 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
Show stats 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,
} }
], ],