feat: Implement streaming for Functionary v2 + Bug fixes (#1419)
* set up streaming for v2 * assert v2 streaming, fix tool_call vs function_call * fix streaming with tool_choice/function_call * make functions return 1 function call only when 'auto' * fix --------- Co-authored-by: Andrei <abetlen@gmail.com>
This commit is contained in:
Jeffrey Fong
GitHub
parent
f9b7221c8f
commit
1f56c648c3
1 changed files with 443 additions and 133 deletions
|
|
@ -1894,6 +1894,8 @@ def functionary_v1_v2_chat_handler(
|
||||||
function_call = (
|
function_call = (
|
||||||
tool_choice if isinstance(tool_choice, str) else tool_choice["function"]
|
tool_choice if isinstance(tool_choice, str) else tool_choice["function"]
|
||||||
)
|
)
|
||||||
|
elif function_call is not None:
|
||||||
|
pass
|
||||||
else:
|
else:
|
||||||
function_call = "auto"
|
function_call = "auto"
|
||||||
|
|
||||||
|
|
@ -1930,11 +1932,10 @@ def functionary_v1_v2_chat_handler(
|
||||||
logits_processor=logits_processor,
|
logits_processor=logits_processor,
|
||||||
grammar=grammar,
|
grammar=grammar,
|
||||||
)
|
)
|
||||||
|
if stream is False:
|
||||||
completion_or_completion_chunks["choices"][0]["text"] = completion_or_completion_chunks["choices"][0]["text"].lstrip()
|
completion_or_completion_chunks["choices"][0]["text"] = completion_or_completion_chunks["choices"][0]["text"].lstrip()
|
||||||
return _convert_completion_to_chat(completion_or_completion_chunks, stream=stream) # type: ignore
|
return _convert_completion_to_chat(completion_or_completion_chunks, stream=stream) # type: ignore
|
||||||
|
|
||||||
assert stream is False # TODO: support stream mode
|
|
||||||
|
|
||||||
def get_grammar(function_call):
|
def get_grammar(function_call):
|
||||||
function_body = None
|
function_body = None
|
||||||
for function in functions or []:
|
for function in functions or []:
|
||||||
|
|
@ -1968,7 +1969,7 @@ def functionary_v1_v2_chat_handler(
|
||||||
|
|
||||||
return grammar
|
return grammar
|
||||||
|
|
||||||
def create_completion(stop):
|
def create_completion(prompt, stop, grammar):
|
||||||
completion = cast(llama_types.Completion, llama.create_completion(
|
completion = cast(llama_types.Completion, llama.create_completion(
|
||||||
prompt=prompt,
|
prompt=prompt,
|
||||||
temperature=temperature,
|
temperature=temperature,
|
||||||
|
|
@ -1976,7 +1977,7 @@ def functionary_v1_v2_chat_handler(
|
||||||
top_k=top_k,
|
top_k=top_k,
|
||||||
min_p=min_p,
|
min_p=min_p,
|
||||||
typical_p=typical_p,
|
typical_p=typical_p,
|
||||||
stream=False,
|
stream=stream,
|
||||||
stop=stop,
|
stop=stop,
|
||||||
max_tokens=max_tokens,
|
max_tokens=max_tokens,
|
||||||
presence_penalty=presence_penalty,
|
presence_penalty=presence_penalty,
|
||||||
|
|
@ -1997,6 +1998,315 @@ def functionary_v1_v2_chat_handler(
|
||||||
function_calls, function_bodies = [], []
|
function_calls, function_bodies = [], []
|
||||||
completion_tokens = 0
|
completion_tokens = 0
|
||||||
|
|
||||||
|
def generate_streaming(tools, functions, function_call, prompt):
|
||||||
|
assert version == "v2", "Streaming for v1 is not supported"
|
||||||
|
|
||||||
|
chunk_id, chunk_created = None, None
|
||||||
|
|
||||||
|
# If tool_choice/function_call is provided
|
||||||
|
if isinstance(function_call, dict):
|
||||||
|
prompt += f"{function_call['name']}\n{CONTENT_TOKEN}"
|
||||||
|
grammar = get_grammar(function_call["name"])
|
||||||
|
stops = [STOP_TOKEN, FROM_TOKEN]
|
||||||
|
tool_id = "".join([random.choice(string.ascii_letters + string.digits) for _ in range(24)])
|
||||||
|
completion = create_completion(prompt=prompt, stop=stops, grammar=grammar)
|
||||||
|
completion_text = ""
|
||||||
|
first = True
|
||||||
|
for chunk in completion:
|
||||||
|
# Yield the tool/function name first
|
||||||
|
if first:
|
||||||
|
if tools is not None:
|
||||||
|
func_call_dict = {
|
||||||
|
"tool_calls": [
|
||||||
|
{
|
||||||
|
"index": 0,
|
||||||
|
"id": "call_" + tool_id,
|
||||||
|
"type": "function",
|
||||||
|
"function": {"name": function_call["name"], "arguments": ""},
|
||||||
|
}
|
||||||
|
]
|
||||||
|
}
|
||||||
|
else:
|
||||||
|
func_call_dict = {"function_call": {"name": function_call["name"], "arguments": ""}}
|
||||||
|
yield llama_types.CreateChatCompletionStreamResponse(
|
||||||
|
id="chat" + chunk["id"],
|
||||||
|
object="chat.completion.chunk",
|
||||||
|
created=chunk["created"],
|
||||||
|
model=chunk["model"],
|
||||||
|
choices=[
|
||||||
|
{"index": 0, "logprobs": None, "delta": {"role": None, "content": None, **func_call_dict}}
|
||||||
|
],
|
||||||
|
)
|
||||||
|
first = False
|
||||||
|
if tools is not None:
|
||||||
|
func_call_dict = {
|
||||||
|
"tool_calls": [
|
||||||
|
{
|
||||||
|
"index": 0,
|
||||||
|
"id": "call_" + tool_id,
|
||||||
|
"type": "function",
|
||||||
|
"function": {
|
||||||
|
"name": None,
|
||||||
|
"arguments": chunk["choices"][0]["text"].rstrip(),
|
||||||
|
},
|
||||||
|
}
|
||||||
|
]
|
||||||
|
}
|
||||||
|
else:
|
||||||
|
func_call_dict = {"function_call": {"name": None, "arguments": chunk["choices"][0]["text"].rstrip()}}
|
||||||
|
if len(chunk["choices"][0]["text"].rstrip()) > 0:
|
||||||
|
yield llama_types.CreateChatCompletionStreamResponse(
|
||||||
|
id="chat" + chunk["id"],
|
||||||
|
object="chat.completion.chunk",
|
||||||
|
created=chunk["created"],
|
||||||
|
model=chunk["model"],
|
||||||
|
choices=[
|
||||||
|
{
|
||||||
|
"index": 0,
|
||||||
|
"logprobs": chunk["choices"][0]["logprobs"],
|
||||||
|
"delta": {
|
||||||
|
"role": None,
|
||||||
|
"content": None,
|
||||||
|
**func_call_dict,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
],
|
||||||
|
)
|
||||||
|
# Yield tool_call/function_call stop message
|
||||||
|
yield llama_types.CreateChatCompletionStreamResponse(
|
||||||
|
id="chat" + chunk["id"],
|
||||||
|
object="chat.completion.chunk",
|
||||||
|
created=chunk["created"],
|
||||||
|
model=chunk["model"],
|
||||||
|
choices=[
|
||||||
|
{
|
||||||
|
"index": 0,
|
||||||
|
"finish_reason": "tool_calls" if tools is not None else "function_call",
|
||||||
|
"logprobs": None,
|
||||||
|
"delta": {
|
||||||
|
"role": None, "content": None, "function_call": None, "tool_calls": None
|
||||||
|
},
|
||||||
|
}
|
||||||
|
],
|
||||||
|
)
|
||||||
|
# If "auto" or no tool_choice/function_call
|
||||||
|
elif isinstance(function_call, str) and function_call == "auto":
|
||||||
|
tool_index = 0
|
||||||
|
while True:
|
||||||
|
# Generate function name first
|
||||||
|
grammar = None
|
||||||
|
stops = CONTENT_TOKEN
|
||||||
|
completion = create_completion(prompt=prompt, stop=stops, grammar=grammar)
|
||||||
|
completion_text = ""
|
||||||
|
for chunk in completion:
|
||||||
|
completion_text += chunk["choices"][0]["text"]
|
||||||
|
if chunk_id is None:
|
||||||
|
chunk_id = chunk["id"]
|
||||||
|
if chunk_created is None:
|
||||||
|
chunk_created = chunk["created"]
|
||||||
|
function_name = completion_text.strip()
|
||||||
|
if function_name == "all":
|
||||||
|
prompt += "all\n<|content|>"
|
||||||
|
# Yield the first empty message for content
|
||||||
|
yield llama_types.CreateChatCompletionStreamResponse(
|
||||||
|
id="chat" + chunk_id,
|
||||||
|
model=chunk["model"],
|
||||||
|
created=chunk_created,
|
||||||
|
object="chat.completion.chunk",
|
||||||
|
choices=[
|
||||||
|
{
|
||||||
|
"index": 0,
|
||||||
|
"delta": {"role": "assistant", "content": ""},
|
||||||
|
"logprobs": None,
|
||||||
|
"finish_reason": None,
|
||||||
|
}
|
||||||
|
],
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
prompt += f"{function_name}\n<|content|>"
|
||||||
|
grammar = get_grammar(function_name)
|
||||||
|
tool_id = "".join([random.choice(string.ascii_letters + string.digits) for _ in range(24)])
|
||||||
|
if tools is not None:
|
||||||
|
func_call_dict = {
|
||||||
|
"tool_calls": [
|
||||||
|
{
|
||||||
|
"index": tool_index,
|
||||||
|
"id": "call_" + tool_id,
|
||||||
|
"type": "function",
|
||||||
|
"function": {"name": function_name, "arguments": ""},
|
||||||
|
}
|
||||||
|
]
|
||||||
|
}
|
||||||
|
else:
|
||||||
|
func_call_dict = {"function_call": {"name": function_name, "arguments": ""}}
|
||||||
|
# Stream function name
|
||||||
|
yield llama_types.CreateChatCompletionStreamResponse(
|
||||||
|
id="chat" + chunk_id,
|
||||||
|
object="chat.completion.chunk",
|
||||||
|
created=chunk_created,
|
||||||
|
model=chunk["model"],
|
||||||
|
choices=[
|
||||||
|
{
|
||||||
|
"index": 0,
|
||||||
|
"logprobs": chunk["choices"][0]["logprobs"],
|
||||||
|
"delta": {
|
||||||
|
"role": "assistant",
|
||||||
|
"content": None,
|
||||||
|
**func_call_dict,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
],
|
||||||
|
)
|
||||||
|
# Generate content
|
||||||
|
stops = [RECIPIENT_TOKEN, STOP_TOKEN]
|
||||||
|
completion = create_completion(prompt=prompt, stop=stops, grammar=grammar)
|
||||||
|
if function_name == "all":
|
||||||
|
completion_text = ""
|
||||||
|
stop_sequence, buffer, is_end = "\n<|from|>assistant\n<|recipient|>", [], False
|
||||||
|
for i, chunk in enumerate(completion):
|
||||||
|
completion_text += chunk["choices"][0]["text"]
|
||||||
|
if is_end:
|
||||||
|
buffer.append(chunk["choices"][0]["text"].strip(" "))
|
||||||
|
if stop_sequence.startswith("".join(buffer)):
|
||||||
|
continue
|
||||||
|
else:
|
||||||
|
buffer.pop()
|
||||||
|
while len(buffer) > 0:
|
||||||
|
yield llama_types.CreateChatCompletionStreamResponse(
|
||||||
|
id="chat" + chunk_id,
|
||||||
|
object="chat.completion.chunk",
|
||||||
|
created=chunk_created,
|
||||||
|
model=chunk["model"],
|
||||||
|
choices=[
|
||||||
|
{
|
||||||
|
"index": 0,
|
||||||
|
"logprobs": chunk["choices"][0]["logprobs"],
|
||||||
|
"delta": {
|
||||||
|
"role": "assistant", "content": buffer.pop(0)
|
||||||
|
},
|
||||||
|
}
|
||||||
|
],
|
||||||
|
)
|
||||||
|
is_end = False
|
||||||
|
elif chunk["choices"][0]["text"] == "\n":
|
||||||
|
is_end = True
|
||||||
|
buffer.append(chunk["choices"][0]["text"].strip(" "))
|
||||||
|
continue
|
||||||
|
|
||||||
|
if len(buffer) == 0 and len(chunk["choices"][0]["text"]) > 0:
|
||||||
|
yield llama_types.CreateChatCompletionStreamResponse(
|
||||||
|
id="chat" + chunk_id,
|
||||||
|
object="chat.completion.chunk",
|
||||||
|
created=chunk_created,
|
||||||
|
model=chunk["model"],
|
||||||
|
choices=[
|
||||||
|
{
|
||||||
|
"index": 0,
|
||||||
|
"logprobs": chunk["choices"][0]["logprobs"],
|
||||||
|
"delta": {
|
||||||
|
"role": "assistant",
|
||||||
|
"content": chunk["choices"][0]["text"] if i > 0 else chunk["choices"][0]["text"].lstrip()
|
||||||
|
},
|
||||||
|
}
|
||||||
|
],
|
||||||
|
)
|
||||||
|
# Check whether the model wants to generate another turn
|
||||||
|
if "<|from|> assistant" in completion_text or "<|from|>assistant" in completion_text:
|
||||||
|
if completion_text.endswith("\n<|from|>assistant\n"):
|
||||||
|
cleaned_completion_text = completion_text[:-len("\n<|from|>assistant\n")].strip()
|
||||||
|
elif completion_text.endswith("\n<|from|> assistant\n"):
|
||||||
|
cleaned_completion_text = completion_text[:-len("\n<|from|> assistant\n")].strip()
|
||||||
|
else:
|
||||||
|
cleaned_completion_text = completion_text.strip()
|
||||||
|
prompt += f"{cleaned_completion_text}\n<|from|>assistant\n<|recipient|>"
|
||||||
|
else:
|
||||||
|
# Yield stop message
|
||||||
|
yield llama_types.CreateChatCompletionStreamResponse(
|
||||||
|
id="chat" + chunk_id,
|
||||||
|
model=chunk["model"],
|
||||||
|
created=chunk_created,
|
||||||
|
object="chat.completion.chunk",
|
||||||
|
choices=[
|
||||||
|
{
|
||||||
|
"index": 0,
|
||||||
|
"delta": {},
|
||||||
|
"logprobs": None,
|
||||||
|
"finish_reason": "stop",
|
||||||
|
}
|
||||||
|
],
|
||||||
|
)
|
||||||
|
break
|
||||||
|
else:
|
||||||
|
# Check whether the model wants to generate another turn
|
||||||
|
completion_text = ""
|
||||||
|
for chunk in completion:
|
||||||
|
completion_text += chunk["choices"][0]["text"]
|
||||||
|
if len(chunk["choices"][0]["text"].rstrip()) > 0:
|
||||||
|
if tools is not None:
|
||||||
|
func_call_dict = {
|
||||||
|
"tool_calls": [
|
||||||
|
{
|
||||||
|
"index": tool_index,
|
||||||
|
"id": "call_" + tool_id,
|
||||||
|
"type": "function",
|
||||||
|
"function": {
|
||||||
|
"name": None,
|
||||||
|
"arguments": chunk["choices"][0]["text"].rstrip(),
|
||||||
|
},
|
||||||
|
}
|
||||||
|
]
|
||||||
|
}
|
||||||
|
else:
|
||||||
|
func_call_dict = {"function_call": {"name": None, "arguments": chunk["choices"][0]["text"].rstrip()}}
|
||||||
|
yield llama_types.CreateChatCompletionStreamResponse(
|
||||||
|
id="chat" + chunk_id,
|
||||||
|
object="chat.completion.chunk",
|
||||||
|
created=chunk_created,
|
||||||
|
model=chunk["model"],
|
||||||
|
choices=[
|
||||||
|
{
|
||||||
|
"index": 0,
|
||||||
|
"logprobs": chunk["choices"][0]["logprobs"],
|
||||||
|
"delta": {
|
||||||
|
"role": None,
|
||||||
|
"content": None,
|
||||||
|
**func_call_dict,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
],
|
||||||
|
)
|
||||||
|
prompt += completion_text.strip()
|
||||||
|
grammar = None
|
||||||
|
completion = create_completion(prompt=prompt, stop=stops, grammar=grammar)
|
||||||
|
completion_text += "".join([chunk["choices"][0]["text"] for chunk in completion])
|
||||||
|
if ("<|from|> assistant" in completion_text or "<|from|>assistant" in completion_text) and tools is not None:
|
||||||
|
prompt += "\n<|from|>assistant\n<|recipient|>"
|
||||||
|
tool_index += 1
|
||||||
|
else:
|
||||||
|
# Yield tool_call/function_call stop message
|
||||||
|
yield llama_types.CreateChatCompletionStreamResponse(
|
||||||
|
id="chat" + chunk_id,
|
||||||
|
object="chat.completion.chunk",
|
||||||
|
created=chunk_created,
|
||||||
|
model=chunk["model"],
|
||||||
|
choices=[
|
||||||
|
{
|
||||||
|
"index": 0,
|
||||||
|
"finish_reason": "tool_calls" if tools is not None else "function_call",
|
||||||
|
"logprobs": None,
|
||||||
|
"delta": {
|
||||||
|
"role": None, "content": None, "function_call": None, "tool_calls": None
|
||||||
|
},
|
||||||
|
}
|
||||||
|
],
|
||||||
|
)
|
||||||
|
break
|
||||||
|
|
||||||
|
if stream is not False:
|
||||||
|
return generate_streaming(
|
||||||
|
tools=tools, functions=functions, function_call=function_call, prompt=prompt
|
||||||
|
)
|
||||||
|
else:
|
||||||
if version == "v1":
|
if version == "v1":
|
||||||
# If no or "auto" tool_choice/function_call
|
# If no or "auto" tool_choice/function_call
|
||||||
if isinstance(function_call, str) and function_call == "auto":
|
if isinstance(function_call, str) and function_call == "auto":
|
||||||
|
|
|
||||||
Loading…
Add table
Reference in a new issue