phoenix-oss/llama-stack-mirror
0
0
Fork 1
mirror of https://github.com/meta-llama/llama-stack.git synced 2025-07-29 15:23:51 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Use InferenceClient.text_generation for TGI inference

Browse source
This commit is contained in:
Celina Hanouti 2024-09-06 17:56:27 +02:00
parent 7aa50934bf
commit 031dbc0e45
1 changed files with 50 additions and 77 deletions
Download patch file Download diff file Expand all files Collapse all files

127
llama_toolchain/inference/adapters/tgi/tgi.py
View file

@ -9,9 +9,8 @@ from typing import AsyncGenerator
from huggingface_hub import InferenceClient from huggingface_hub import InferenceClient
from llama_models.llama3.api.chat_format import ChatFormat from llama_models.llama3.api.chat_format import ChatFormat
from llama_models.llama3.api.datatypes import Message, StopReason from llama_models.llama3.api.datatypes import StopReason
from llama_models.llama3.api.tokenizer import Tokenizer from llama_models.llama3.api.tokenizer import Tokenizer
from llama_models.sku_list import resolve_model
from llama_toolchain.inference.api import * from llama_toolchain.inference.api import *
from llama_toolchain.inference.api.api import ( # noqa: F403 from llama_toolchain.inference.api.api import ( # noqa: F403
@ -19,6 +18,7 @@ from llama_toolchain.inference.api.api import ( # noqa: F403
ChatCompletionResponse, ChatCompletionResponse,
ChatCompletionResponseStreamChunk, ChatCompletionResponseStreamChunk,
) )
from llama_toolchain.inference.prepare_messages import prepare_messages
from .config import TGIImplConfig from .config import TGIImplConfig
@ -49,26 +49,6 @@ class TGIAdapter(Inference):
async def completion(self, request: CompletionRequest) -> AsyncGenerator: async def completion(self, request: CompletionRequest) -> AsyncGenerator:
raise NotImplementedError() raise NotImplementedError()
def _convert_messages(self, messages: list[Message]) -> List[Message]: # type: ignore
tgi_messages = []
for message in messages:
if message.role == "ipython":
role = "tool"
else:
role = message.role
tgi_messages.append({"role": role, "content": message.content})
return tgi_messages
def resolve_hf_model(self, model_name: str) -> str:
model = resolve_model(model_name)
assert (
model is not None
and model.descriptor(shorten_default_variant=True) in HF_SUPPORTED_MODELS
), f"Unsupported model: {model_name}, use one of the supported models: {','.join(HF_SUPPORTED_MODELS.keys())}"
return HF_SUPPORTED_MODELS.get(model.descriptor(shorten_default_variant=True))
def get_chat_options(self, request: ChatCompletionRequest) -> dict: def get_chat_options(self, request: ChatCompletionRequest) -> dict:
options = {} options = {}
if request.sampling_params is not None: if request.sampling_params is not None:
@ -79,27 +59,36 @@ class TGIAdapter(Inference):
return options return options
async def chat_completion(self, request: ChatCompletionRequest) -> AsyncGenerator: async def chat_completion(self, request: ChatCompletionRequest) -> AsyncGenerator:
messages = prepare_messages(request)
model_input = self.formatter.encode_dialog_prompt(messages)
prompt = self.tokenizer.decode(model_input.tokens)
model_info = self.client.get_endpoint_info(model=self.config.url)
max_new_tokens = min(
request.sampling_params.max_tokens or model_info["max_total_tokens"],
model_info["max_total_tokens"] - len(model_input.tokens) - 1,
)
options = self.get_chat_options(request) options = self.get_chat_options(request)
messages = self._convert_messages(request.messages)
if not request.stream: if not request.stream:
response = self.client.chat_completion( response = self.client.text_generation(
messages=messages, prompt=prompt,
stream=False, stream=False,
details=True,
max_new_tokens=max_new_tokens,
stop_sequences=["<|eom_id|>", "<|eot_id|>"],
**options, **options,
) )
stop_reason = None stop_reason = None
if response.choices[0].finish_reason: if response.details.finish_reason:
if ( if response.details.finish_reason == "stop":
response.choices[0].finish_reason == "stop_sequence"
or response.choices[0].finish_reason == "eos_token"
):
stop_reason = StopReason.end_of_turn stop_reason = StopReason.end_of_turn
elif response.choices[0].finish_reason == "length": elif response.details.finish_reason == "length":
stop_reason = StopReason.out_of_tokens stop_reason = StopReason.out_of_tokens
completion_message = self.formatter.decode_assistant_message_from_content( completion_message = self.formatter.decode_assistant_message_from_content(
response.choices[0].message.content, response.generated_text,
stop_reason, stop_reason,
) )
yield ChatCompletionResponse( yield ChatCompletionResponse(
@ -117,32 +106,22 @@ class TGIAdapter(Inference):
buffer = "" buffer = ""
ipython = False ipython = False
stop_reason = None stop_reason = None
tokens = []
for chunk in self.client.chat_completion( for response in self.client.text_generation(
messages=messages, stream=True, **options prompt=prompt,
stream=True,
details=True,
max_new_tokens=max_new_tokens,
stop_sequences=["<|eom_id|>", "<|eot_id|>"],
**options,
): ):
if chunk.choices[0].finish_reason: token_result = response.token
if (
stop_reason is None
and chunk.choices[0].finish_reason == "stop_sequence"
) or (
stop_reason is None
and chunk.choices[0].finish_reason == "eos_token"
):
stop_reason = StopReason.end_of_turn
elif (
stop_reason is None
and chunk.choices[0].finish_reason == "length"
):
stop_reason = StopReason.out_of_tokens
break
text = chunk.choices[0].delta.content buffer += token_result.text
if text is None: tokens.append(token_result.id)
continue
# check if its a tool call ( aka starts with <|python_tag|> ) if not ipython and buffer.startswith("<|python_tag|>"):
if not ipython and text.startswith("<|python_tag|>"):
ipython = True ipython = True
yield ChatCompletionResponseStreamChunk( yield ChatCompletionResponseStreamChunk(
event=ChatCompletionResponseEvent( event=ChatCompletionResponseEvent(
@ -153,25 +132,27 @@ class TGIAdapter(Inference):
), ),
) )
) )
buffer += text buffer = buffer[len("<|python_tag|>") :]
continue continue
if ipython: if token_result.text == "<|eot_id|>":
if text == "<|eot_id|>": stop_reason = StopReason.end_of_turn
stop_reason = StopReason.end_of_turn text = ""
text = "" elif token_result.text == "<|eom_id|>":
continue stop_reason = StopReason.end_of_message
elif text == "<|eom_id|>": text = ""
stop_reason = StopReason.end_of_message else:
text = "" text = token_result.text
continue
buffer += text if ipython:
delta = ToolCallDelta( delta = ToolCallDelta(
content=text, content=text,
parse_status=ToolCallParseStatus.in_progress, parse_status=ToolCallParseStatus.in_progress,
) )
else:
delta = text
if stop_reason is None:
yield ChatCompletionResponseStreamChunk( yield ChatCompletionResponseStreamChunk(
event=ChatCompletionResponseEvent( event=ChatCompletionResponseEvent(
event_type=ChatCompletionResponseEventType.progress, event_type=ChatCompletionResponseEventType.progress,
@ -179,20 +160,12 @@ class TGIAdapter(Inference):
stop_reason=stop_reason, stop_reason=stop_reason,
) )
) )
else:
buffer += text if stop_reason is None:
yield ChatCompletionResponseStreamChunk( stop_reason = StopReason.out_of_tokens
event=ChatCompletionResponseEvent(
event_type=ChatCompletionResponseEventType.progress,
delta=text,
stop_reason=stop_reason,
)
)
# parse tool calls and report errors # parse tool calls and report errors
message = self.formatter.decode_assistant_message_from_content( message = self.formatter.decode_assistant_message(tokens, stop_reason)
buffer, stop_reason
)
parsed_tool_calls = len(message.tool_calls) > 0 parsed_tool_calls = len(message.tool_calls) > 0
if ipython and not parsed_tool_calls: if ipython and not parsed_tool_calls:
yield ChatCompletionResponseStreamChunk( yield ChatCompletionResponseStreamChunk(