Skip to content

Commit

Permalink
Adding openai completions as an endpoint
Browse files Browse the repository at this point in the history
  • Loading branch information
@nv-braf
nv-braf committed Mar 6, 2024
1 parent e701b84 commit 1d35b33
Show file tree
Hide file tree
Showing 2 changed files with 124 additions and 12 deletions.
110 changes: 102 additions & 8 deletions src/c++/perf_analyzer/genai-pa/genai_pa/llm_inputs/llm_inputs.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -36,7 +36,8 @@ class InputFormat(Enum):


class OutputFormat(Enum):
OPENAI = auto()
OPENAI_CHAT_COMPLETIONS = auto()
OPENAI_COMPLETIONS = auto()
TRTLLM = auto()
VLLM = auto()

Expand Down Expand Up @@ -232,8 +233,14 @@ def _convert_generic_json_to_output_format(
add_stream: bool,
model_name: str = "",
) -> Dict:
if output_format == OutputFormat.OPENAI:
output_json = LlmInputs._convert_generic_json_to_openai_format(
if output_format == OutputFormat.OPENAI_CHAT_COMPLETIONS:
output_json = (
LlmInputs._convert_generic_json_to_openai_chat_completions_format(
generic_dataset, add_model_name, add_stream, model_name
)
)
elif output_format == OutputFormat.OPENAI_COMPLETIONS:
output_json = LlmInputs._convert_generic_json_to_openai_completions_format(
generic_dataset, add_model_name, add_stream, model_name
)
elif output_format == OutputFormat.VLLM:
Expand All @@ -248,7 +255,7 @@ def _convert_generic_json_to_output_format(
return output_json

@classmethod
def _convert_generic_json_to_openai_format(
def _convert_generic_json_to_openai_chat_completions_format(
cls,
dataset_json: Dict,
add_model_name: bool,
Expand All @@ -261,10 +268,35 @@ def _convert_generic_json_to_openai_format(
user_role_headers,
_,
) = LlmInputs._determine_json_feature_roles(dataset_json)
pa_json = LlmInputs._populate_openai_output_json(
pa_json = LlmInputs._populate_openai_chat_completions_output_json(
dataset_json,
system_role_headers,
user_role_headers,
add_model_name,
add_stream,
model_name,
)

return pa_json

@classmethod
def _convert_generic_json_to_openai_completions_format(
cls,
dataset_json: Dict,
add_model_name: bool,
add_stream: bool,
model_name: str = "",
) -> Dict:
(
system_role_headers,
user_role_headers,
text_input_headers,
) = LlmInputs._determine_json_feature_roles(dataset_json)
pa_json = LlmInputs._populate_openai_completions_output_json(
dataset_json,
system_role_headers,
user_role_headers,
text_input_headers,
add_model_name,
add_stream,
model_name,
Expand Down Expand Up @@ -333,7 +365,7 @@ def _determine_json_feature_roles(
return system_role_headers, user_role_headers, text_input_headers

@classmethod
def _populate_openai_output_json(
def _populate_openai_chat_completions_output_json(
cls,
dataset_json: Dict,
system_role_headers: List[str],
Expand All @@ -348,7 +380,7 @@ def _populate_openai_output_json(
pa_json["data"][0]["payload"].append({"messages": []})

for header, content in entry.items():
new_message = LlmInputs._create_new_message(
new_message = LlmInputs._create_new_openai_chat_completions_message(
header, system_role_headers, user_role_headers, content
)

Expand All @@ -362,6 +394,39 @@ def _populate_openai_output_json(

return pa_json

@classmethod
def _populate_openai_completions_output_json(
cls,
dataset_json: Dict,
system_role_headers: List[str],
user_role_headers: List[str],
text_input_headers: List[str],
add_model_name: bool,
add_stream: bool,
model_name: str = "",
) -> Dict:
pa_json = LlmInputs._create_empty_openai_pa_json()

for index, entry in enumerate(dataset_json["rows"]):
pa_json["data"][0]["payload"].append({"prompt": []})

for header, content in entry.items():
new_prompt = LlmInputs._create_new_prompt(
header,
system_role_headers,
user_role_headers,
text_input_headers,
content,
)

pa_json = LlmInputs._add_new_prompt_to_json(pa_json, index, new_prompt)

pa_json = LlmInputs._add_optional_tags_to_openai_json(
pa_json, index, add_model_name, add_stream, model_name
)

return pa_json

@classmethod
def _populate_vllm_output_json(
cls,
Expand Down Expand Up @@ -410,7 +475,7 @@ def _create_empty_vllm_pa_json(cls) -> Dict:
return empty_pa_json

@classmethod
def _create_new_message(
def _create_new_openai_chat_completions_message(
cls,
header: str,
system_role_headers: List[str],
Expand All @@ -432,6 +497,26 @@ def _create_new_message(

return new_message

@classmethod
def _create_new_prompt(
cls,
header: str,
system_role_headers: List[str],
user_role_headers: List[str],
text_input_headers: List[str],
content: str,
) -> Optional[str]:
new_prompt = ""

if (
header in system_role_headers
or header in user_role_headers
or header in text_input_headers
):
new_prompt = content

return new_prompt

@classmethod
def _create_new_text_input(
cls,
Expand Down Expand Up @@ -470,6 +555,15 @@ def _add_new_text_input_to_json(

return pa_json

@classmethod
def _add_new_prompt_to_json(
cls, pa_json: Dict, index: int, new_prompt: str
) -> Dict:
if new_prompt:
pa_json["data"][0]["payload"][index]["prompt"].append(new_prompt)

return pa_json

@classmethod
def _add_optional_tags_to_openai_json(
cls,
Expand Down
26 changes: 22 additions & 4 deletions src/c++/perf_analyzer/genai-pa/tests/test_llm_inputs.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -104,7 +104,7 @@ def test_llm_inputs_error_in_server_response(self):
_ = LlmInputs.create_llm_inputs(
input_type=InputType.URL,
input_format=InputFormat.OPENAI,
output_format=OutputFormat.OPENAI,
output_format=OutputFormat.OPENAI_CHAT_COMPLETIONS,
model_name=OPEN_ORCA,
starting_index=LlmInputs.DEFAULT_STARTING_INDEX,
length=int(LlmInputs.DEFAULT_LENGTH * 100),
Expand Down Expand Up @@ -160,7 +160,7 @@ def test_convert_default_json_to_pa_format(self, default_configured_url):
input_format=InputFormat.OPENAI, dataset=dataset
)
pa_json = LlmInputs._convert_generic_json_to_output_format(
output_format=OutputFormat.OPENAI,
output_format=OutputFormat.OPENAI_CHAT_COMPLETIONS,
generic_dataset=dataset_json,
add_model_name=False,
add_stream=False,
Expand All @@ -176,7 +176,7 @@ def test_create_openai_llm_inputs_cnn_dailymail(self):
pa_json = LlmInputs.create_llm_inputs(
input_type=InputType.URL,
input_format=InputFormat.OPENAI,
output_format=OutputFormat.OPENAI,
output_format=OutputFormat.OPENAI_CHAT_COMPLETIONS,
model_name=CNN_DAILY_MAIL,
)

Expand All @@ -192,7 +192,7 @@ def test_write_to_file(self):
pa_json = LlmInputs.create_llm_inputs(
input_type=InputType.URL,
input_format=InputFormat.OPENAI,
output_format=OutputFormat.OPENAI,
output_format=OutputFormat.OPENAI_CHAT_COMPLETIONS,
model_name=OPEN_ORCA,
add_model_name=True,
add_stream=True,
Expand Down Expand Up @@ -223,3 +223,21 @@ def test_create_openai_to_vllm(self):

assert pa_json is not None
assert len(pa_json["data"]) == LlmInputs.DEFAULT_LENGTH

def test_create_openai_to_completions(self):
"""
Test conversion of openai to completions
"""
pa_json = LlmInputs.create_llm_inputs(
input_type=InputType.URL,
input_format=InputFormat.OPENAI,
output_format=OutputFormat.OPENAI_COMPLETIONS,
model_name=OPEN_ORCA,
add_model_name=False,
add_stream=True,
)

os.remove(DEFAULT_INPUT_DATA_JSON)

assert pa_json is not None
assert len(pa_json["data"][0]["payload"]) == LlmInputs.DEFAULT_LENGTH

0 comments on commit 1d35b33

Please sign in to comment.