mirror of
https://github.com/deepset-ai/haystack.git
synced 2025-07-14 04:20:38 +00:00
9e6a2e3cf9
* HuggingFaceTGIGenerator/HuggingFaceTGIChatGenerator check if model is deployed on free-tier
302 lines
12 KiB
Python
302 lines
12 KiB
Python
from unittest.mock import patch, MagicMock, Mock
|
|
|
|
import pytest
|
|
from huggingface_hub.inference._text_generation import TextGenerationStreamResponse, Token, StreamDetails, FinishReason
|
|
from huggingface_hub.utils import RepositoryNotFoundError
|
|
|
|
from haystack.components.generators import HuggingFaceTGIGenerator
|
|
from haystack.dataclasses import StreamingChunk
|
|
from haystack.utils.auth import Secret
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_list_inference_deployed_models():
|
|
with patch(
|
|
"haystack.components.generators.hugging_face_tgi.list_inference_deployed_models",
|
|
MagicMock(
|
|
return_value=["HuggingFaceH4/zephyr-7b-alpha", "HuggingFaceH4/zephyr-7b-alpha", "mistralai/Mistral-7B-v0.1"]
|
|
),
|
|
) as mock:
|
|
yield mock
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_check_valid_model():
|
|
with patch(
|
|
"haystack.components.generators.hugging_face_tgi.check_valid_model", MagicMock(return_value=None)
|
|
) as mock:
|
|
yield mock
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_text_generation():
|
|
with patch("huggingface_hub.InferenceClient.text_generation", autospec=True) as mock_text_generation:
|
|
mock_response = Mock()
|
|
mock_response.generated_text = "I'm fine, thanks."
|
|
details = Mock()
|
|
details.finish_reason = MagicMock(field1="value")
|
|
details.tokens = [1, 2, 3]
|
|
mock_response.details = details
|
|
mock_text_generation.return_value = mock_response
|
|
yield mock_text_generation
|
|
|
|
|
|
# used to test serialization of streaming_callback
|
|
def streaming_callback_handler(x):
|
|
return x
|
|
|
|
|
|
class TestHuggingFaceTGIGenerator:
|
|
def test_initialize_with_valid_model_and_generation_parameters(self, mock_check_valid_model):
|
|
model = "HuggingFaceH4/zephyr-7b-alpha"
|
|
generation_kwargs = {"n": 1}
|
|
stop_words = ["stop"]
|
|
streaming_callback = None
|
|
|
|
generator = HuggingFaceTGIGenerator(
|
|
model=model,
|
|
url=None,
|
|
token=None,
|
|
generation_kwargs=generation_kwargs,
|
|
stop_words=stop_words,
|
|
streaming_callback=streaming_callback,
|
|
)
|
|
|
|
assert generator.model == model
|
|
assert generator.generation_kwargs == {**generation_kwargs, **{"stop_sequences": ["stop"]}}
|
|
assert generator.tokenizer is None
|
|
assert generator.client is not None
|
|
assert generator.streaming_callback == streaming_callback
|
|
|
|
def test_to_dict(self, mock_check_valid_model):
|
|
# Initialize the HuggingFaceRemoteGenerator object with valid parameters
|
|
generator = HuggingFaceTGIGenerator(
|
|
token=Secret.from_env_var("ENV_VAR", strict=False),
|
|
generation_kwargs={"n": 5},
|
|
stop_words=["stop", "words"],
|
|
streaming_callback=lambda x: x,
|
|
)
|
|
|
|
# Call the to_dict method
|
|
result = generator.to_dict()
|
|
init_params = result["init_parameters"]
|
|
|
|
# Assert that the init_params dictionary contains the expected keys and values
|
|
assert init_params["model"] == "mistralai/Mistral-7B-v0.1"
|
|
assert init_params["token"] == {"env_vars": ["ENV_VAR"], "strict": False, "type": "env_var"}
|
|
assert init_params["generation_kwargs"] == {"n": 5, "stop_sequences": ["stop", "words"]}
|
|
|
|
def test_from_dict(self, mock_check_valid_model):
|
|
generator = HuggingFaceTGIGenerator(
|
|
model="mistralai/Mistral-7B-v0.1",
|
|
generation_kwargs={"n": 5},
|
|
stop_words=["stop", "words"],
|
|
streaming_callback=streaming_callback_handler,
|
|
)
|
|
# Call the to_dict method
|
|
result = generator.to_dict()
|
|
|
|
# now deserialize, call from_dict
|
|
generator_2 = HuggingFaceTGIGenerator.from_dict(result)
|
|
assert generator_2.model == "mistralai/Mistral-7B-v0.1"
|
|
assert generator_2.generation_kwargs == {"n": 5, "stop_sequences": ["stop", "words"]}
|
|
assert generator_2.streaming_callback is streaming_callback_handler
|
|
|
|
def test_initialize_with_invalid_url(self, mock_check_valid_model):
|
|
with pytest.raises(ValueError):
|
|
HuggingFaceTGIGenerator(model="mistralai/Mistral-7B-v0.1", url="invalid_url")
|
|
|
|
def test_initialize_with_url_but_invalid_model(self, mock_check_valid_model):
|
|
# When custom TGI endpoint is used via URL, model must be provided and valid HuggingFace Hub model id
|
|
mock_check_valid_model.side_effect = RepositoryNotFoundError("Invalid model id")
|
|
with pytest.raises(RepositoryNotFoundError):
|
|
HuggingFaceTGIGenerator(model="invalid_model_id", url="https://some_chat_model.com")
|
|
|
|
def test_generate_text_response_with_valid_prompt_and_generation_parameters(
|
|
self, mock_check_valid_model, mock_auto_tokenizer, mock_text_generation, mock_list_inference_deployed_models
|
|
):
|
|
model = "mistralai/Mistral-7B-v0.1"
|
|
|
|
generation_kwargs = {"n": 1}
|
|
stop_words = ["stop"]
|
|
streaming_callback = None
|
|
|
|
generator = HuggingFaceTGIGenerator(
|
|
model=model,
|
|
generation_kwargs=generation_kwargs,
|
|
stop_words=stop_words,
|
|
streaming_callback=streaming_callback,
|
|
)
|
|
generator.warm_up()
|
|
|
|
prompt = "Hello, how are you?"
|
|
response = generator.run(prompt)
|
|
|
|
# check kwargs passed to text_generation
|
|
# note how n was not passed to text_generation
|
|
_, kwargs = mock_text_generation.call_args
|
|
assert kwargs == {"details": True, "stop_sequences": ["stop"]}
|
|
|
|
assert isinstance(response, dict)
|
|
assert "replies" in response
|
|
assert "meta" in response
|
|
assert isinstance(response["replies"], list)
|
|
assert isinstance(response["meta"], list)
|
|
assert len(response["replies"]) == 1
|
|
assert len(response["meta"]) == 1
|
|
assert [isinstance(reply, str) for reply in response["replies"]]
|
|
|
|
def test_generate_multiple_text_responses_with_valid_prompt_and_generation_parameters(
|
|
self, mock_check_valid_model, mock_auto_tokenizer, mock_text_generation, mock_list_inference_deployed_models
|
|
):
|
|
model = "mistralai/Mistral-7B-v0.1"
|
|
generation_kwargs = {"n": 3}
|
|
stop_words = ["stop"]
|
|
streaming_callback = None
|
|
|
|
generator = HuggingFaceTGIGenerator(
|
|
model=model,
|
|
generation_kwargs=generation_kwargs,
|
|
stop_words=stop_words,
|
|
streaming_callback=streaming_callback,
|
|
)
|
|
generator.warm_up()
|
|
|
|
prompt = "Hello, how are you?"
|
|
response = generator.run(prompt)
|
|
|
|
# check kwargs passed to text_generation
|
|
# note how n was not passed to text_generation
|
|
_, kwargs = mock_text_generation.call_args
|
|
assert kwargs == {"details": True, "stop_sequences": ["stop"]}
|
|
|
|
assert isinstance(response, dict)
|
|
assert "replies" in response
|
|
assert "meta" in response
|
|
assert isinstance(response["replies"], list)
|
|
assert [isinstance(reply, str) for reply in response["replies"]]
|
|
|
|
assert isinstance(response["meta"], list)
|
|
assert len(response["replies"]) == 3
|
|
assert len(response["meta"]) == 3
|
|
assert [isinstance(reply, dict) for reply in response["meta"]]
|
|
|
|
def test_initialize_with_invalid_model(self, mock_check_valid_model):
|
|
model = "invalid_model"
|
|
generation_kwargs = {"n": 1}
|
|
stop_words = ["stop"]
|
|
streaming_callback = None
|
|
|
|
mock_check_valid_model.side_effect = ValueError("Invalid model path or url")
|
|
|
|
with pytest.raises(ValueError):
|
|
HuggingFaceTGIGenerator(
|
|
model=model,
|
|
generation_kwargs=generation_kwargs,
|
|
stop_words=stop_words,
|
|
streaming_callback=streaming_callback,
|
|
)
|
|
|
|
def test_generate_text_with_stop_words(
|
|
self, mock_check_valid_model, mock_auto_tokenizer, mock_text_generation, mock_list_inference_deployed_models
|
|
):
|
|
generator = HuggingFaceTGIGenerator()
|
|
generator.warm_up()
|
|
|
|
# Generate text response with stop words
|
|
response = generator.run("How are you?", generation_kwargs={"stop_words": ["stop", "words"]})
|
|
|
|
# check kwargs passed to text_generation
|
|
_, kwargs = mock_text_generation.call_args
|
|
assert kwargs == {"details": True, "stop_sequences": ["stop", "words"]}
|
|
|
|
# Assert that the response contains the generated replies
|
|
assert "replies" in response
|
|
assert isinstance(response["replies"], list)
|
|
assert len(response["replies"]) > 0
|
|
assert [isinstance(reply, str) for reply in response["replies"]]
|
|
|
|
# Assert that the response contains the metadata
|
|
assert "meta" in response
|
|
assert isinstance(response["meta"], list)
|
|
assert len(response["meta"]) > 0
|
|
assert [isinstance(reply, dict) for reply in response["replies"]]
|
|
|
|
def test_generate_text_with_custom_generation_parameters(
|
|
self, mock_check_valid_model, mock_auto_tokenizer, mock_text_generation, mock_list_inference_deployed_models
|
|
):
|
|
generator = HuggingFaceTGIGenerator()
|
|
generator.warm_up()
|
|
|
|
generation_kwargs = {"temperature": 0.8, "max_new_tokens": 100}
|
|
response = generator.run("How are you?", generation_kwargs=generation_kwargs)
|
|
|
|
# check kwargs passed to text_generation
|
|
_, kwargs = mock_text_generation.call_args
|
|
assert kwargs == {"details": True, "max_new_tokens": 100, "stop_sequences": [], "temperature": 0.8}
|
|
|
|
# Assert that the response contains the generated replies and the right response
|
|
assert "replies" in response
|
|
assert isinstance(response["replies"], list)
|
|
assert len(response["replies"]) > 0
|
|
assert [isinstance(reply, str) for reply in response["replies"]]
|
|
assert response["replies"][0] == "I'm fine, thanks."
|
|
|
|
# Assert that the response contains the metadata
|
|
assert "meta" in response
|
|
assert isinstance(response["meta"], list)
|
|
assert len(response["meta"]) > 0
|
|
assert [isinstance(reply, str) for reply in response["replies"]]
|
|
|
|
def test_generate_text_with_streaming_callback(
|
|
self, mock_check_valid_model, mock_auto_tokenizer, mock_text_generation, mock_list_inference_deployed_models
|
|
):
|
|
streaming_call_count = 0
|
|
|
|
# Define the streaming callback function
|
|
def streaming_callback_fn(chunk: StreamingChunk):
|
|
nonlocal streaming_call_count
|
|
streaming_call_count += 1
|
|
assert isinstance(chunk, StreamingChunk)
|
|
|
|
# Create an instance of HuggingFaceRemoteGenerator
|
|
generator = HuggingFaceTGIGenerator(streaming_callback=streaming_callback_fn)
|
|
generator.warm_up()
|
|
|
|
# Create a fake streamed response
|
|
# Don't remove self
|
|
def mock_iter(self):
|
|
yield TextGenerationStreamResponse(
|
|
generated_text=None, token=Token(id=1, text="I'm fine, thanks.", logprob=0.0, special=False)
|
|
)
|
|
yield TextGenerationStreamResponse(
|
|
generated_text=None,
|
|
token=Token(id=1, text="Ok bye", logprob=0.0, special=False),
|
|
details=StreamDetails(finish_reason=FinishReason.Length, generated_tokens=5),
|
|
)
|
|
|
|
mock_response = Mock(**{"__iter__": mock_iter})
|
|
mock_text_generation.return_value = mock_response
|
|
|
|
# Generate text response with streaming callback
|
|
response = generator.run("prompt")
|
|
|
|
# check kwargs passed to text_generation
|
|
_, kwargs = mock_text_generation.call_args
|
|
assert kwargs == {"details": True, "stop_sequences": [], "stream": True}
|
|
|
|
# Assert that the streaming callback was called twice
|
|
assert streaming_call_count == 2
|
|
|
|
# Assert that the response contains the generated replies
|
|
assert "replies" in response
|
|
assert isinstance(response["replies"], list)
|
|
assert len(response["replies"]) > 0
|
|
assert [isinstance(reply, str) for reply in response["replies"]]
|
|
|
|
# Assert that the response contains the metadata
|
|
assert "meta" in response
|
|
assert isinstance(response["meta"], list)
|
|
assert len(response["meta"]) > 0
|
|
assert [isinstance(reply, dict) for reply in response["replies"]]
|