Repair idefics integration tests.

This commit is contained in:
Nicolas Patry 2024-03-21 18:06:06 +01:00
parent de6cb15fa5
commit 12fac2e03f

View File

@ -1,4 +1,5 @@
import pytest
import base64
@pytest.fixture(scope="module")
@ -14,11 +15,17 @@ async def idefics(idefics_handle):
await idefics_handle.health(300)
return idefics_handle.client
def get_chicken():
with open("integrations-tests/images/chicken_on_money.png", "rb") as image_file:
encoded_string = base64.b64encode(image_file.read())
return f"data:image/png;base64,{encoded_string}"
@pytest.mark.asyncio
async def test_idefics(idefics, response_snapshot):
chicken = get_chicken()
response = await idefics.generate(
"User:![](https://huggingface.co/spaces/HuggingFaceM4/idefics_playground/resolve/6af279110440b7ba0c43f6a0be93b9926f062c42/example_images/chicken_on_money.png)Can you tell me a very short story based on the image?",
f"User:![]({chicken})Can you tell me a very short story based on the image?",
max_new_tokens=10,
decoder_input_details=True,
)
@ -31,7 +38,7 @@ async def test_idefics(idefics, response_snapshot):
async def test_idefics_load(idefics, generate_load, response_snapshot):
responses = await generate_load(
idefics,
"User:![](https://huggingface.co/spaces/HuggingFaceM4/idefics_playground/resolve/6af279110440b7ba0c43f6a0be93b9926f062c42/example_images/chicken_on_money.png)Can you tell me a very short story based on the image?",
f"User:![]({chicken})Can you tell me a very short story based on the image?",
max_new_tokens=10,
n=4,
)