theroyallab/tabbyAPI
1
0
Fork 0
mirror of https://github.com/theroyallab/tabbyAPI.git synced 2026-03-15 00:07:28 +00:00
Code Issues Packages Projects Releases Wiki Activity

API: Fix response creation

Browse Source 
Change chat completion and text completion responses to be more
flexible.

Signed-off-by: kingbri <bdashore3@proton.me>
This commit is contained in:
kingbri
2024-02-07 23:36:08 -05:00
committed by Brian Dashore
parent 0af6a38af3
commit c02fe4d1db
4 changed files with 50 additions and 37 deletions
Download Patch File Download Diff File Expand all files Collapse all files

53
OAI/utils/completion.py
View File

@@ -17,32 +17,35 @@ from OAI.types.completion import (
from OAI.types.common import UsageStats
def create_completion_response(**kwargs):
def create_completion_response(generation: dict, model_name: Optional[str]):
"""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 = None
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],
)
token_probs = unwrap(generation.get("token_probs"), {})
if token_probs:
logprobs = unwrap(generation.get("logprobs"), [])
offset = unwrap(generation.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"), ""),
text=unwrap(generation.get("text"), ""),
logprobs=logprob_response,
)
prompt_tokens = unwrap(kwargs.get("prompt_tokens"), 0)
completion_tokens = unwrap(kwargs.get("completion_tokens"), 0)
prompt_tokens = unwrap(generation.get("prompt_tokens"), 0)
completion_tokens = unwrap(generation.get("completion_tokens"), 0)
response = CompletionResponse(
choices=[choice],
model=unwrap(kwargs.get("model_name"), ""),
model=unwrap(model_name, ""),
usage=UsageStats(
prompt_tokens=prompt_tokens,
completion_tokens=completion_tokens,
@@ -53,17 +56,18 @@ def create_completion_response(**kwargs):
return response
def create_chat_completion_response(
text: str,
prompt_tokens: Optional[int],
completion_tokens: Optional[int],
model_name: Optional[str],
):
def create_chat_completion_response(generation: dict, model_name: Optional[str]):
"""Create a chat completion response from the provided text."""
message = ChatCompletionMessage(role="assistant", content=unwrap(text, ""))
message = ChatCompletionMessage(
role="assistant", content=unwrap(generation.get("text"), "")
)
choice = ChatCompletionRespChoice(finish_reason="Generated", message=message)
prompt_tokens = unwrap(generation.get("prompt_tokens"), 0)
completion_tokens = unwrap(generation.get("completion_tokens"), 0)
response = ChatCompletionResponse(
choices=[choice],
model=unwrap(model_name, ""),
@@ -79,15 +83,18 @@ def create_chat_completion_response(
def create_chat_completion_stream_chunk(
const_id: str,
text: Optional[str] = None,
generation: Optional[dict] = 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)
message = ChatCompletionMessage(
role="assistant", content=unwrap(generation.get("text"), "")
)
# The finish reason can be None
choice = ChatCompletionStreamChoice(finish_reason=finish_reason, delta=message)