mirror of
https://github.com/theroyallab/tabbyAPI.git
synced 2026-03-14 15:57:27 +00:00
0af6a38af3
Returns token offsets, selected tokens, probabilities of tokens post-sampling, and normalized probability of selecting a token pre-sampling (for efficiency purposes). Only for text completions. Chat completions in a later commit. Signed-off-by: kingbri <bdashore3@proton.me>
100 lines
2.9 KiB
Python
100 lines
2.9 KiB
Python
""" Utility functions for the OpenAI server. """
|
|
from typing import Optional
|
|
|
|
from common.utils import unwrap
|
|
from OAI.types.chat_completion import (
|
|
ChatCompletionMessage,
|
|
ChatCompletionRespChoice,
|
|
ChatCompletionStreamChunk,
|
|
ChatCompletionResponse,
|
|
ChatCompletionStreamChoice,
|
|
)
|
|
from OAI.types.completion import (
|
|
CompletionResponse,
|
|
CompletionRespChoice,
|
|
CompletionLogProbs,
|
|
)
|
|
from OAI.types.common import UsageStats
|
|
|
|
|
|
def create_completion_response(**kwargs):
|
|
"""Create a completion response from the provided text."""
|
|
|
|
token_probs = unwrap(kwargs.get("token_probs"), {})
|
|
logprobs = unwrap(kwargs.get("logprobs"), [])
|
|
offset = unwrap(kwargs.get("offset"), [])
|
|
|
|
logprob_response = CompletionLogProbs(
|
|
text_offset=offset if isinstance(offset, list) else [offset],
|
|
token_logprobs=token_probs.values(),
|
|
tokens=token_probs.keys(),
|
|
top_logprobs=logprobs if isinstance(logprobs, list) else [logprobs],
|
|
)
|
|
|
|
choice = CompletionRespChoice(
|
|
finish_reason="Generated",
|
|
text=unwrap(kwargs.get("text"), ""),
|
|
logprobs=logprob_response,
|
|
)
|
|
|
|
prompt_tokens = unwrap(kwargs.get("prompt_tokens"), 0)
|
|
completion_tokens = unwrap(kwargs.get("completion_tokens"), 0)
|
|
|
|
response = CompletionResponse(
|
|
choices=[choice],
|
|
model=unwrap(kwargs.get("model_name"), ""),
|
|
usage=UsageStats(
|
|
prompt_tokens=prompt_tokens,
|
|
completion_tokens=completion_tokens,
|
|
total_tokens=prompt_tokens + completion_tokens,
|
|
),
|
|
)
|
|
|
|
return response
|
|
|
|
|
|
def create_chat_completion_response(
|
|
text: str,
|
|
prompt_tokens: Optional[int],
|
|
completion_tokens: Optional[int],
|
|
model_name: Optional[str],
|
|
):
|
|
"""Create a chat completion response from the provided text."""
|
|
message = ChatCompletionMessage(role="assistant", content=unwrap(text, ""))
|
|
|
|
choice = ChatCompletionRespChoice(finish_reason="Generated", message=message)
|
|
|
|
response = ChatCompletionResponse(
|
|
choices=[choice],
|
|
model=unwrap(model_name, ""),
|
|
usage=UsageStats(
|
|
prompt_tokens=prompt_tokens,
|
|
completion_tokens=completion_tokens,
|
|
total_tokens=prompt_tokens + completion_tokens,
|
|
),
|
|
)
|
|
|
|
return response
|
|
|
|
|
|
def create_chat_completion_stream_chunk(
|
|
const_id: str,
|
|
text: Optional[str] = None,
|
|
model_name: Optional[str] = None,
|
|
finish_reason: Optional[str] = None,
|
|
):
|
|
"""Create a chat completion stream chunk from the provided text."""
|
|
if finish_reason:
|
|
message = {}
|
|
else:
|
|
message = ChatCompletionMessage(role="assistant", content=text)
|
|
|
|
# The finish reason can be None
|
|
choice = ChatCompletionStreamChoice(finish_reason=finish_reason, delta=message)
|
|
|
|
chunk = ChatCompletionStreamChunk(
|
|
id=const_id, choices=[choice], model=unwrap(model_name, "")
|
|
)
|
|
|
|
return chunk
|