From 34174af8c8a3b4f32e82d18174593afe894834d1 Mon Sep 17 00:00:00 2001
From: drbh <david.richard.holtz@gmail.com>
Date: Wed, 18 Dec 2024 05:06:49 +0000
Subject: [PATCH] fix: bump idefics3 tests and snapshots

---
 .../test_flash_idefics3_next_simple_url.json  | 68 +++++++++----------
 integration-tests/models/test_idefics3.py     | 39 +----------
 2 files changed, 34 insertions(+), 73 deletions(-)

diff --git a/integration-tests/models/__snapshots__/test_idefics3/test_flash_idefics3_next_simple_url.json b/integration-tests/models/__snapshots__/test_idefics3/test_flash_idefics3_next_simple_url.json
index 052318df2c9..6bf2b93a2da 100644
--- a/integration-tests/models/__snapshots__/test_idefics3/test_flash_idefics3_next_simple_url.json
+++ b/integration-tests/models/__snapshots__/test_idefics3/test_flash_idefics3_next_simple_url.json
@@ -1,73 +1,67 @@
 {
   "details": {
     "best_of_sequences": null,
-    "finish_reason": "length",
-    "generated_tokens": 10,
+    "finish_reason": "eos_token",
+    "generated_tokens": 9,
     "prefill": [],
     "seed": null,
     "tokens": [
       {
-        "id": 578,
-        "logprob": -0.2475586,
+        "id": 2684,
+        "logprob": -0.24902344,
         "special": false,
-        "text": " The"
+        "text": " There"
       },
       {
-        "id": 2217,
-        "logprob": -0.017303467,
-        "special": false,
-        "text": " image"
-      },
-      {
-        "id": 62991,
-        "logprob": -0.7368164,
+        "id": 374,
+        "logprob": -0.0703125,
         "special": false,
-        "text": " depicts"
+        "text": " is"
       },
       {
-        "id": 279,
-        "logprob": -0.39990234,
+        "id": 264,
+        "logprob": -0.23535156,
         "special": false,
-        "text": " the"
+        "text": " a"
       },
       {
-        "id": 89675,
-        "logprob": -0.34350586,
+        "id": 35372,
+        "logprob": -0.125,
         "special": false,
-        "text": " Statue"
+        "text": " statue"
       },
       {
-        "id": 315,
-        "logprob": -0.0002901554,
+        "id": 304,
+        "logprob": -0.30273438,
         "special": false,
-        "text": " of"
+        "text": " in"
       },
       {
-        "id": 32492,
-        "logprob": -0.0009598732,
+        "id": 279,
+        "logprob": -0.20507812,
         "special": false,
-        "text": " Liberty"
+        "text": " the"
       },
       {
-        "id": 11,
-        "logprob": -0.2355957,
+        "id": 2217,
+        "logprob": -0.076171875,
         "special": false,
-        "text": ","
+        "text": " image"
       },
       {
-        "id": 264,
-        "logprob": -0.66503906,
+        "id": 13,
+        "logprob": -0.053710938,
         "special": false,
-        "text": " a"
+        "text": "."
       },
       {
-        "id": 97937,
-        "logprob": -0.9199219,
-        "special": false,
-        "text": " colossal"
+        "id": 128258,
+        "logprob": -0.011352539,
+        "special": true,
+        "text": "<end_of_utterance>"
       }
     ],
     "top_tokens": null
   },
-  "generated_text": " The image depicts the Statue of Liberty, a colossal"
+  "generated_text": " There is a statue in the image."
 }
diff --git a/integration-tests/models/test_idefics3.py b/integration-tests/models/test_idefics3.py
index 1f55872a066..80be2350fad 100644
--- a/integration-tests/models/test_idefics3.py
+++ b/integration-tests/models/test_idefics3.py
@@ -1,21 +1,9 @@
 import pytest
-import base64
-
-
-def get_chicken():
-    with open("integration-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.decode('utf-8')}"
 
 
 @pytest.fixture(scope="module")
 def flash_idefics3_next_handle(launcher):
-    with launcher(
-        "HuggingFaceM4/Idefics3-8B-Llama3",
-        max_total_tokens=3000,
-        max_batch_prefill_tokens=2501,
-        max_input_tokens=2500,
-    ) as handle:
+    with launcher("HuggingFaceM4/Idefics3-8B-Llama3") as handle:
         yield handle
 
 
@@ -25,26 +13,6 @@ async def flash_idefics3_next(flash_idefics3_next_handle):
     return flash_idefics3_next_handle.client
 
 
-# TODO: dont skip when token issue is resolved
-@pytest.mark.skip
-@pytest.mark.asyncio
-@pytest.mark.private
-async def test_flash_idefics3_next_simple_base64(
-    flash_idefics3_next, response_snapshot
-):
-    chicken = get_chicken()
-    query = "Write me a short story"
-    response = await flash_idefics3_next.generate(
-        f"<|begin_of_text|><|begin_of_text|>User:![]({chicken}){query}<end_of_utterance>\nAssistant:",
-        max_new_tokens=10,
-    )
-    assert (
-        response.generated_text == " A chicken is sitting on a pile of money."
-    ), f"{repr(response.generated_text)}"
-    # assert response.details.generated_tokens == 10
-    # assert response == response_snapshot
-
-
 @pytest.mark.asyncio
 @pytest.mark.private
 async def test_flash_idefics3_next_simple_url(flash_idefics3_next, response_snapshot):
@@ -57,8 +25,7 @@ async def test_flash_idefics3_next_simple_url(flash_idefics3_next, response_snap
     )
     print(response)
     assert (
-        response.generated_text
-        == " The image depicts the Statue of Liberty, a colossal"
+        response.generated_text == " There is a statue in the image."
     ), f"{repr(response.generated_text)}"
-    assert response.details.generated_tokens == 10
+    assert response.details.generated_tokens == 9
     assert response == response_snapshot