wehub-resource-sync
314 lines
10 KiB
Python
314 lines
10 KiB
Python
import uuid
|
|
from typing import Any, NamedTuple
|
|
from unittest import mock
|
|
|
|
import pandas as pd
|
|
import pytest
|
|
import torch
|
|
|
|
from mlflow.exceptions import MlflowException
|
|
from mlflow.models import infer_signature
|
|
from mlflow.transformers.llm_inference_utils import (
|
|
_get_default_task_for_llm_inference_task,
|
|
_get_finish_reason,
|
|
_get_output_and_usage_from_tensor,
|
|
_get_stopping_criteria,
|
|
_get_token_usage,
|
|
convert_messages_to_prompt,
|
|
infer_signature_from_llm_inference_task,
|
|
preprocess_llm_inference_input,
|
|
)
|
|
from mlflow.types.llm import (
|
|
CHAT_MODEL_INPUT_SCHEMA,
|
|
CHAT_MODEL_OUTPUT_SCHEMA,
|
|
COMPLETIONS_MODEL_INPUT_SCHEMA,
|
|
COMPLETIONS_MODEL_OUTPUT_SCHEMA,
|
|
)
|
|
|
|
|
|
def test_infer_signature_from_llm_inference_task():
|
|
signature = infer_signature_from_llm_inference_task("llm/v1/completions")
|
|
assert signature.inputs == COMPLETIONS_MODEL_INPUT_SCHEMA
|
|
assert signature.outputs == COMPLETIONS_MODEL_OUTPUT_SCHEMA
|
|
|
|
signature = infer_signature_from_llm_inference_task("llm/v1/chat")
|
|
assert signature.inputs == CHAT_MODEL_INPUT_SCHEMA
|
|
assert signature.outputs == CHAT_MODEL_OUTPUT_SCHEMA
|
|
|
|
signature = infer_signature("hello", "world")
|
|
with pytest.raises(MlflowException, match=r".*llm/v1/completions.*signature"):
|
|
infer_signature_from_llm_inference_task("llm/v1/completions", signature)
|
|
|
|
|
|
class DummyTokenizer:
|
|
def __call__(self, text: str, **kwargs):
|
|
input_ids = list(map(int, text.split(" ")))
|
|
return {"input_ids": torch.tensor([input_ids])}
|
|
|
|
def decode(self, tensor, **kwargs):
|
|
if isinstance(tensor, torch.Tensor):
|
|
tensor = tensor.tolist()
|
|
return " ".join([str(x) for x in tensor])
|
|
|
|
def convert_tokens_to_ids(self, tokens: list[str]):
|
|
return [int(x) for x in tokens]
|
|
|
|
def tokenize(self, text: str):
|
|
return [x for x in text.split(" ") if x]
|
|
|
|
def apply_chat_template(self, messages: list[dict[str, str]], **kwargs):
|
|
return " ".join(message["content"] for message in messages)
|
|
|
|
|
|
def test_apply_chat_template():
|
|
data1 = [{"role": "A", "content": "one"}, {"role": "B", "content": "two"}]
|
|
# Test that the function modifies the data in place for Chat task
|
|
prompt = convert_messages_to_prompt(data1, DummyTokenizer())
|
|
assert prompt == "one two"
|
|
|
|
with pytest.raises(MlflowException, match=r"Input messages should be list of"):
|
|
convert_messages_to_prompt([["one", "two"]], DummyTokenizer())
|
|
|
|
|
|
class _TestCase(NamedTuple):
|
|
data: Any
|
|
params: Any
|
|
expected_data: Any
|
|
expected_params: Any
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"case",
|
|
[
|
|
# Case 0: Data only includes prompt
|
|
_TestCase(
|
|
data=pd.DataFrame({"prompt": ["Hello world!"]}),
|
|
params={},
|
|
expected_data=["Hello world!"],
|
|
expected_params={},
|
|
),
|
|
# Case 1: Data includes prompt and params
|
|
_TestCase(
|
|
data=pd.DataFrame({
|
|
"prompt": ["Hello world!"],
|
|
"temperature": [0.7],
|
|
"max_tokens": [100],
|
|
"stop": [None],
|
|
}),
|
|
params={},
|
|
expected_data=["Hello world!"],
|
|
expected_params={
|
|
"temperature": 0.7,
|
|
# max_tokens is replaced with max_new_tokens
|
|
"max_new_tokens": 100,
|
|
# do not pass `stop` to params as it is None
|
|
},
|
|
),
|
|
# Case 2: Params are passed if not specified in data
|
|
_TestCase(
|
|
data=pd.DataFrame({
|
|
"prompt": ["Hello world!"],
|
|
}),
|
|
params={
|
|
"temperature": 0.7,
|
|
"max_tokens": 100,
|
|
"stop": ["foo", "bar"],
|
|
},
|
|
expected_data=["Hello world!"],
|
|
expected_params={
|
|
"temperature": 0.7,
|
|
"max_new_tokens": 100,
|
|
# Stopping criteria is _StopSequenceMatchCriteria instance
|
|
# "stop": ...
|
|
},
|
|
),
|
|
# Case 3: Data overrides params
|
|
_TestCase(
|
|
data=pd.DataFrame({
|
|
"messages": [
|
|
[
|
|
{"role": "user", "content": "Hello!"},
|
|
{"role": "assistant", "content": "Hi!"},
|
|
]
|
|
],
|
|
"temperature": [0.1],
|
|
"max_tokens": [100],
|
|
"stop": [["foo", "bar"]],
|
|
}),
|
|
params={
|
|
"temperature": [0.2],
|
|
"max_tokens": [200],
|
|
"stop": ["foo", "bar", "baz"],
|
|
},
|
|
expected_data=[
|
|
[
|
|
{"role": "user", "content": "Hello!"},
|
|
{"role": "assistant", "content": "Hi!"},
|
|
]
|
|
],
|
|
expected_params={
|
|
"temperature": 0.1,
|
|
"max_new_tokens": 100,
|
|
},
|
|
),
|
|
# Case 4: Batch input
|
|
_TestCase(
|
|
data=pd.DataFrame({
|
|
"prompt": ["Hello!", "Hi", "Hola"],
|
|
"temperature": [0.1, 0.2, 0.3],
|
|
"max_tokens": [None, 200, 300],
|
|
}),
|
|
params={
|
|
"temperature": 0.4,
|
|
"max_tokens": 400,
|
|
},
|
|
expected_data=["Hello!", "Hi", "Hola"],
|
|
# The values in the first data is used, otherwise params
|
|
expected_params={
|
|
"temperature": 0.1,
|
|
"max_new_tokens": 400,
|
|
},
|
|
),
|
|
# Case 5: Raw dict input
|
|
_TestCase(
|
|
data={
|
|
"messages": [
|
|
{"role": "user", "content": "Hello!"},
|
|
{"role": "assistant", "content": "Hi!"},
|
|
],
|
|
"temperature": 0.1,
|
|
"max_tokens": 100,
|
|
"stop": ["foo", "bar"],
|
|
},
|
|
params={},
|
|
expected_data=[
|
|
[
|
|
{"role": "user", "content": "Hello!"},
|
|
{"role": "assistant", "content": "Hi!"},
|
|
]
|
|
],
|
|
expected_params={
|
|
"temperature": 0.1,
|
|
"max_new_tokens": 100,
|
|
},
|
|
),
|
|
],
|
|
)
|
|
def test_preprocess_llm_inference_input(case):
|
|
task = "llm/v1/completions" if "prompt" in case.data else "llm/v1/chat"
|
|
flavor_config = {"inference_task": task, "source_model_name": "test"}
|
|
|
|
with mock.patch(
|
|
"mlflow.transformers.llm_inference_utils._get_stopping_criteria"
|
|
) as mock_get_stopping_criteria:
|
|
data, params = preprocess_llm_inference_input(case.data, case.params, flavor_config)
|
|
|
|
# Test that OpenAI params are separated from data and replaced with Hugging Face params
|
|
assert data == case.expected_data
|
|
if "stopping_criteria" in params:
|
|
assert params.pop("stopping_criteria") is not None
|
|
mock_get_stopping_criteria.assert_called_once_with(["foo", "bar"], "test")
|
|
assert params == case.expected_params
|
|
|
|
|
|
def test_preprocess_llm_inference_input_raise_if_key_invalid():
|
|
# Missing input key
|
|
with pytest.raises(MlflowException, match=r"Transformer model saved with"):
|
|
preprocess_llm_inference_input(
|
|
pd.DataFrame({"invalid_key": [1, 2, 3]}),
|
|
flavor_config={"inference_task": "llm/v1/completions"},
|
|
)
|
|
|
|
# Unmatched key (should be "messages" for chat task)
|
|
with pytest.raises(MlflowException, match=r"Transformer model saved with"):
|
|
preprocess_llm_inference_input(
|
|
pd.DataFrame({"prompt": ["Hi"]}), flavor_config={"inference_task": "llm/v1/chat"}
|
|
)
|
|
|
|
|
|
def test_stopping_criteria():
|
|
with mock.patch("transformers.AutoTokenizer.from_pretrained") as mock_from_pretrained:
|
|
mock_from_pretrained.return_value = DummyTokenizer()
|
|
|
|
stopping_criteria = _get_stopping_criteria(stop=None, model_name=None)
|
|
assert stopping_criteria is None
|
|
|
|
input_ids = torch.tensor([[1, 2, 3, 4, 5]])
|
|
scores = torch.ones(1, 5)
|
|
|
|
stopping_criteria = _get_stopping_criteria(stop="5", model_name="my/model")
|
|
stopping_criteria_matches = [f(input_ids, scores) for f in stopping_criteria]
|
|
assert stopping_criteria_matches == [True, True]
|
|
|
|
stopping_criteria = _get_stopping_criteria(stop=["100", "5"], model_name="my/model")
|
|
stopping_criteria_matches = [f(input_ids, scores) for f in stopping_criteria]
|
|
assert stopping_criteria_matches == [False, False, True, True]
|
|
|
|
|
|
def test_output_dict_for_completions():
|
|
prompt = "1 2 3"
|
|
output_tensor = [1, 2, 3, 4, 5]
|
|
flavor_config = {"source_model_name": "gpt2"}
|
|
model_config = {"max_new_tokens": 2}
|
|
inference_task = "llm/v1/completions"
|
|
|
|
pipeline = mock.MagicMock()
|
|
pipeline.tokenizer = DummyTokenizer()
|
|
|
|
output_dict = _get_output_and_usage_from_tensor(
|
|
prompt, output_tensor, pipeline, flavor_config, model_config, inference_task
|
|
)
|
|
|
|
# Test UUID validity
|
|
uuid.UUID(output_dict["id"])
|
|
|
|
assert output_dict["object"] == "text_completion"
|
|
assert output_dict["model"] == "gpt2"
|
|
|
|
assert output_dict["choices"][0]["text"] == "4 5"
|
|
assert output_dict["choices"][0]["finish_reason"] == "length"
|
|
|
|
usage = output_dict["usage"]
|
|
assert usage["prompt_tokens"] + usage["completion_tokens"] == usage["total_tokens"]
|
|
|
|
|
|
def test_token_usage():
|
|
prompt = "1 2 3"
|
|
output_tensor = [1, 2, 3, 4, 5]
|
|
|
|
pipeline = mock.MagicMock()
|
|
pipeline.tokenizer = DummyTokenizer()
|
|
|
|
usage = _get_token_usage(prompt, output_tensor, pipeline, {})
|
|
assert usage["prompt_tokens"] == 3
|
|
assert usage["completion_tokens"] == 2
|
|
assert usage["total_tokens"] == 5
|
|
|
|
|
|
def test_finish_reason():
|
|
assert _get_finish_reason(total_tokens=20, completion_tokens=10, model_config={}) == "stop"
|
|
|
|
assert (
|
|
_get_finish_reason(
|
|
total_tokens=20, completion_tokens=10, model_config={"max_new_tokens": 10}
|
|
)
|
|
== "length"
|
|
)
|
|
|
|
assert (
|
|
_get_finish_reason(total_tokens=20, completion_tokens=10, model_config={"max_length": 15})
|
|
== "length"
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("inference_task", "expected_task"),
|
|
[
|
|
("llm/v1/completions", "text-generation"),
|
|
("llm/v1/chat", "text-generation"),
|
|
(None, None),
|
|
],
|
|
)
|
|
def test_default_task_for_llm_inference_task(inference_task, expected_task):
|
|
assert _get_default_task_for_llm_inference_task(inference_task) == expected_task
|