diff --git a/integration-tests/models/__snapshots__/test_idefics/test_idefics.json b/integration-tests/models/__snapshots__/test_idefics/test_idefics.json
index 1df5761d..0edd81b6 100644
--- a/integration-tests/models/__snapshots__/test_idefics/test_idefics.json
+++ b/integration-tests/models/__snapshots__/test_idefics/test_idefics.json
@@ -11,52 +11,52 @@
       },
       {
         "id": 4911,
-        "logprob": -5.3632812,
+        "logprob": -5.7773438,
         "text": "User"
       },
       {
         "id": 29901,
-        "logprob": -0.00762558,
+        "logprob": -0.0069999695,
         "text": ":"
       },
       {
         "id": 32000,
-        "logprob": -0.7739258,
+        "logprob": -0.8125,
         "text": "<fake_token_around_image>"
       },
       {
         "id": 32001,
-        "logprob": -9.775162e-05,
+        "logprob": -6.651878e-05,
         "text": "<image>"
       },
       {
         "id": 32000,
-        "logprob": -1.1920929e-07,
+        "logprob": -3.5762787e-07,
         "text": "<fake_token_around_image>"
       },
       {
         "id": 1815,
-        "logprob": -4.4140625,
+        "logprob": -4.2265625,
         "text": "Can"
       },
       {
         "id": 366,
-        "logprob": -0.01436615,
+        "logprob": -0.013977051,
         "text": "you"
       },
       {
         "id": 2649,
-        "logprob": -4.9414062,
+        "logprob": -4.4375,
         "text": "tell"
       },
       {
         "id": 592,
-        "logprob": -0.3005371,
+        "logprob": -0.29077148,
         "text": "me"
       },
       {
         "id": 263,
-        "logprob": -3.5703125,
+        "logprob": -4.2109375,
         "text": "a"
       },
       {
@@ -66,37 +66,37 @@
       },
       {
         "id": 3273,
-        "logprob": -1.9111328,
+        "logprob": -1.8671875,
         "text": "short"
       },
       {
         "id": 5828,
-        "logprob": -0.28881836,
+        "logprob": -0.26586914,
         "text": "story"
       },
       {
         "id": 2729,
-        "logprob": -3.4179688,
+        "logprob": -3.7460938,
         "text": "based"
       },
       {
         "id": 373,
-        "logprob": -0.00056886673,
+        "logprob": -0.0005350113,
         "text": "on"
       },
       {
         "id": 278,
-        "logprob": -0.14123535,
+        "logprob": -0.13867188,
         "text": "the"
       },
       {
         "id": 1967,
-        "logprob": -0.053985596,
+        "logprob": -0.06842041,
         "text": "image"
       },
       {
         "id": 29973,
-        "logprob": -0.15771484,
+        "logprob": -0.15319824,
         "text": "?"
       }
     ],
@@ -104,25 +104,25 @@
     "tokens": [
       {
         "id": 32002,
-        "logprob": -0.004295349,
+        "logprob": -0.0019445419,
         "special": true,
         "text": "<end_of_utterance>"
       },
       {
         "id": 29871,
-        "logprob": -7.43866e-05,
+        "logprob": -8.404255e-05,
         "special": false,
         "text": " "
       },
       {
         "id": 13,
-        "logprob": -2.3126602e-05,
+        "logprob": -1.7881393e-05,
         "special": false,
         "text": "\n"
       },
       {
         "id": 7900,
-        "logprob": -3.9339066e-06,
+        "logprob": -3.0994415e-06,
         "special": false,
         "text": "Ass"
       },
@@ -134,35 +134,36 @@
       },
       {
         "id": 29901,
-        "logprob": -2.6226044e-06,
+        "logprob": -3.2186508e-06,
         "special": false,
         "text": ":"
       },
       {
         "id": 319,
-        "logprob": -0.87841797,
+        "logprob": -0.9057617,
         "special": false,
         "text": " A"
       },
       {
-        "id": 521,
-        "logprob": -1.3837891,
+        "id": 696,
+        "logprob": -1.2314453,
         "special": false,
-        "text": " ch"
+        "text": " ro"
       },
       {
-        "id": 21475,
-        "logprob": -0.00051641464,
+        "id": 15664,
+        "logprob": -0.00024914742,
         "special": false,
-        "text": "icken"
+        "text": "oster"
       },
       {
-        "id": 338,
-        "logprob": -1.1435547,
+        "id": 15028,
+        "logprob": -1.1621094,
         "special": false,
-        "text": " is"
+        "text": " stands"
       }
-    ]
+    ],
+    "top_tokens": null
   },
-  "generated_text": "\nAssistant: A chicken is"
+  "generated_text": "\nAssistant: A rooster stands"
 }
diff --git a/integration-tests/models/__snapshots__/test_idefics/test_idefics_load.json b/integration-tests/models/__snapshots__/test_idefics/test_idefics_load.json
index e0dfe96c..81cc1b19 100644
--- a/integration-tests/models/__snapshots__/test_idefics/test_idefics_load.json
+++ b/integration-tests/models/__snapshots__/test_idefics/test_idefics_load.json
@@ -12,556 +12,52 @@
         },
         {
           "id": 4911,
-          "logprob": -5.3476562,
+          "logprob": -5.7773438,
           "text": "User"
         },
         {
           "id": 29901,
-          "logprob": -0.0075531006,
+          "logprob": -0.0069999695,
           "text": ":"
         },
         {
           "id": 32000,
-          "logprob": -0.7729492,
+          "logprob": -0.8125,
           "text": "<fake_token_around_image>"
         },
         {
           "id": 32001,
-          "logprob": -9.787083e-05,
+          "logprob": -6.651878e-05,
           "text": "<image>"
         },
         {
           "id": 32000,
-          "logprob": -1.1920929e-07,
+          "logprob": -3.5762787e-07,
           "text": "<fake_token_around_image>"
         },
         {
           "id": 1815,
-          "logprob": -4.4296875,
+          "logprob": -4.2265625,
           "text": "Can"
         },
         {
           "id": 366,
-          "logprob": -0.01424408,
+          "logprob": -0.013977051,
           "text": "you"
         },
         {
           "id": 2649,
-          "logprob": -4.9335938,
+          "logprob": -4.4375,
           "text": "tell"
         },
         {
           "id": 592,
-          "logprob": -0.2993164,
+          "logprob": -0.29077148,
           "text": "me"
         },
         {
           "id": 263,
-          "logprob": -3.5664062,
-          "text": "a"
-        },
-        {
-          "id": 1407,
-          "logprob": -9.4453125,
-          "text": "very"
-        },
-        {
-          "id": 3273,
-          "logprob": -1.9306641,
-          "text": "short"
-        },
-        {
-          "id": 5828,
-          "logprob": -0.2836914,
-          "text": "story"
-        },
-        {
-          "id": 2729,
-          "logprob": -3.4179688,
-          "text": "based"
-        },
-        {
-          "id": 373,
-          "logprob": -0.00056934357,
-          "text": "on"
-        },
-        {
-          "id": 278,
-          "logprob": -0.13928223,
-          "text": "the"
-        },
-        {
-          "id": 1967,
-          "logprob": -0.05355835,
-          "text": "image"
-        },
-        {
-          "id": 29973,
-          "logprob": -0.15771484,
-          "text": "?"
-        }
-      ],
-      "seed": null,
-      "tokens": [
-        {
-          "id": 32002,
-          "logprob": -0.004333496,
-          "special": true,
-          "text": "<end_of_utterance>"
-        },
-        {
-          "id": 29871,
-          "logprob": -7.426739e-05,
-          "special": false,
-          "text": " "
-        },
-        {
-          "id": 13,
-          "logprob": -2.348423e-05,
-          "special": false,
-          "text": "\n"
-        },
-        {
-          "id": 7900,
-          "logprob": -3.9339066e-06,
-          "special": false,
-          "text": "Ass"
-        },
-        {
-          "id": 22137,
-          "logprob": 0.0,
-          "special": false,
-          "text": "istant"
-        },
-        {
-          "id": 29901,
-          "logprob": -2.861023e-06,
-          "special": false,
-          "text": ":"
-        },
-        {
-          "id": 319,
-          "logprob": -0.8828125,
-          "special": false,
-          "text": " A"
-        },
-        {
-          "id": 521,
-          "logprob": -1.3759766,
-          "special": false,
-          "text": " ch"
-        },
-        {
-          "id": 21475,
-          "logprob": -0.0005083084,
-          "special": false,
-          "text": "icken"
-        },
-        {
-          "id": 338,
-          "logprob": -1.1367188,
-          "special": false,
-          "text": " is"
-        }
-      ]
-    },
-    "generated_text": "\nAssistant: A chicken is"
-  },
-  {
-    "details": {
-      "best_of_sequences": null,
-      "finish_reason": "length",
-      "generated_tokens": 10,
-      "prefill": [
-        {
-          "id": 1,
-          "logprob": null,
-          "text": "<s>"
-        },
-        {
-          "id": 4911,
-          "logprob": -5.3476562,
-          "text": "User"
-        },
-        {
-          "id": 29901,
-          "logprob": -0.0075531006,
-          "text": ":"
-        },
-        {
-          "id": 32000,
-          "logprob": -0.7729492,
-          "text": "<fake_token_around_image>"
-        },
-        {
-          "id": 32001,
-          "logprob": -9.787083e-05,
-          "text": "<image>"
-        },
-        {
-          "id": 32000,
-          "logprob": -1.1920929e-07,
-          "text": "<fake_token_around_image>"
-        },
-        {
-          "id": 1815,
-          "logprob": -4.4296875,
-          "text": "Can"
-        },
-        {
-          "id": 366,
-          "logprob": -0.01423645,
-          "text": "you"
-        },
-        {
-          "id": 2649,
-          "logprob": -4.9335938,
-          "text": "tell"
-        },
-        {
-          "id": 592,
-          "logprob": -0.2993164,
-          "text": "me"
-        },
-        {
-          "id": 263,
-          "logprob": -3.5664062,
-          "text": "a"
-        },
-        {
-          "id": 1407,
-          "logprob": -9.4453125,
-          "text": "very"
-        },
-        {
-          "id": 3273,
-          "logprob": -1.9306641,
-          "text": "short"
-        },
-        {
-          "id": 5828,
-          "logprob": -0.2836914,
-          "text": "story"
-        },
-        {
-          "id": 2729,
-          "logprob": -3.4179688,
-          "text": "based"
-        },
-        {
-          "id": 373,
-          "logprob": -0.00056934357,
-          "text": "on"
-        },
-        {
-          "id": 278,
-          "logprob": -0.13928223,
-          "text": "the"
-        },
-        {
-          "id": 1967,
-          "logprob": -0.05355835,
-          "text": "image"
-        },
-        {
-          "id": 29973,
-          "logprob": -0.15771484,
-          "text": "?"
-        }
-      ],
-      "seed": null,
-      "tokens": [
-        {
-          "id": 32002,
-          "logprob": -0.004333496,
-          "special": true,
-          "text": "<end_of_utterance>"
-        },
-        {
-          "id": 29871,
-          "logprob": -7.4505806e-05,
-          "special": false,
-          "text": " "
-        },
-        {
-          "id": 13,
-          "logprob": -2.3722649e-05,
-          "special": false,
-          "text": "\n"
-        },
-        {
-          "id": 7900,
-          "logprob": -3.9339066e-06,
-          "special": false,
-          "text": "Ass"
-        },
-        {
-          "id": 22137,
-          "logprob": 0.0,
-          "special": false,
-          "text": "istant"
-        },
-        {
-          "id": 29901,
-          "logprob": -2.861023e-06,
-          "special": false,
-          "text": ":"
-        },
-        {
-          "id": 319,
-          "logprob": -0.8828125,
-          "special": false,
-          "text": " A"
-        },
-        {
-          "id": 521,
-          "logprob": -1.3759766,
-          "special": false,
-          "text": " ch"
-        },
-        {
-          "id": 21475,
-          "logprob": -0.00050878525,
-          "special": false,
-          "text": "icken"
-        },
-        {
-          "id": 338,
-          "logprob": -1.1367188,
-          "special": false,
-          "text": " is"
-        }
-      ]
-    },
-    "generated_text": "\nAssistant: A chicken is"
-  },
-  {
-    "details": {
-      "best_of_sequences": null,
-      "finish_reason": "length",
-      "generated_tokens": 10,
-      "prefill": [
-        {
-          "id": 1,
-          "logprob": null,
-          "text": "<s>"
-        },
-        {
-          "id": 4911,
-          "logprob": -5.3476562,
-          "text": "User"
-        },
-        {
-          "id": 29901,
-          "logprob": -0.0075531006,
-          "text": ":"
-        },
-        {
-          "id": 32000,
-          "logprob": -0.7729492,
-          "text": "<fake_token_around_image>"
-        },
-        {
-          "id": 32001,
-          "logprob": -9.775162e-05,
-          "text": "<image>"
-        },
-        {
-          "id": 32000,
-          "logprob": -1.1920929e-07,
-          "text": "<fake_token_around_image>"
-        },
-        {
-          "id": 1815,
-          "logprob": -4.4296875,
-          "text": "Can"
-        },
-        {
-          "id": 366,
-          "logprob": -0.01424408,
-          "text": "you"
-        },
-        {
-          "id": 2649,
-          "logprob": -4.9335938,
-          "text": "tell"
-        },
-        {
-          "id": 592,
-          "logprob": -0.2993164,
-          "text": "me"
-        },
-        {
-          "id": 263,
-          "logprob": -3.5664062,
-          "text": "a"
-        },
-        {
-          "id": 1407,
-          "logprob": -9.4453125,
-          "text": "very"
-        },
-        {
-          "id": 3273,
-          "logprob": -1.9306641,
-          "text": "short"
-        },
-        {
-          "id": 5828,
-          "logprob": -0.2836914,
-          "text": "story"
-        },
-        {
-          "id": 2729,
-          "logprob": -3.4179688,
-          "text": "based"
-        },
-        {
-          "id": 373,
-          "logprob": -0.00056934357,
-          "text": "on"
-        },
-        {
-          "id": 278,
-          "logprob": -0.13928223,
-          "text": "the"
-        },
-        {
-          "id": 1967,
-          "logprob": -0.05355835,
-          "text": "image"
-        },
-        {
-          "id": 29973,
-          "logprob": -0.15771484,
-          "text": "?"
-        }
-      ],
-      "seed": null,
-      "tokens": [
-        {
-          "id": 32002,
-          "logprob": -0.004333496,
-          "special": true,
-          "text": "<end_of_utterance>"
-        },
-        {
-          "id": 29871,
-          "logprob": -7.43866e-05,
-          "special": false,
-          "text": " "
-        },
-        {
-          "id": 13,
-          "logprob": -2.360344e-05,
-          "special": false,
-          "text": "\n"
-        },
-        {
-          "id": 7900,
-          "logprob": -3.9339066e-06,
-          "special": false,
-          "text": "Ass"
-        },
-        {
-          "id": 22137,
-          "logprob": 0.0,
-          "special": false,
-          "text": "istant"
-        },
-        {
-          "id": 29901,
-          "logprob": -2.7418137e-06,
-          "special": false,
-          "text": ":"
-        },
-        {
-          "id": 319,
-          "logprob": -0.8828125,
-          "special": false,
-          "text": " A"
-        },
-        {
-          "id": 521,
-          "logprob": -1.3759766,
-          "special": false,
-          "text": " ch"
-        },
-        {
-          "id": 21475,
-          "logprob": -0.00050878525,
-          "special": false,
-          "text": "icken"
-        },
-        {
-          "id": 338,
-          "logprob": -1.1367188,
-          "special": false,
-          "text": " is"
-        }
-      ]
-    },
-    "generated_text": "\nAssistant: A chicken is"
-  },
-  {
-    "details": {
-      "best_of_sequences": null,
-      "finish_reason": "length",
-      "generated_tokens": 10,
-      "prefill": [
-        {
-          "id": 1,
-          "logprob": null,
-          "text": "<s>"
-        },
-        {
-          "id": 4911,
-          "logprob": -5.3632812,
-          "text": "User"
-        },
-        {
-          "id": 29901,
-          "logprob": -0.00762558,
-          "text": ":"
-        },
-        {
-          "id": 32000,
-          "logprob": -0.7739258,
-          "text": "<fake_token_around_image>"
-        },
-        {
-          "id": 32001,
-          "logprob": -9.775162e-05,
-          "text": "<image>"
-        },
-        {
-          "id": 32000,
-          "logprob": -1.1920929e-07,
-          "text": "<fake_token_around_image>"
-        },
-        {
-          "id": 1815,
-          "logprob": -4.4140625,
-          "text": "Can"
-        },
-        {
-          "id": 366,
-          "logprob": -0.01436615,
-          "text": "you"
-        },
-        {
-          "id": 2649,
-          "logprob": -4.9414062,
-          "text": "tell"
-        },
-        {
-          "id": 592,
-          "logprob": -0.3005371,
-          "text": "me"
-        },
-        {
-          "id": 263,
-          "logprob": -3.5703125,
+          "logprob": -4.2109375,
           "text": "a"
         },
         {
@@ -571,37 +67,37 @@
         },
         {
           "id": 3273,
-          "logprob": -1.9111328,
+          "logprob": -1.8671875,
           "text": "short"
         },
         {
           "id": 5828,
-          "logprob": -0.28881836,
+          "logprob": -0.26586914,
           "text": "story"
         },
         {
           "id": 2729,
-          "logprob": -3.4179688,
+          "logprob": -3.7460938,
           "text": "based"
         },
         {
           "id": 373,
-          "logprob": -0.00056886673,
+          "logprob": -0.0005350113,
           "text": "on"
         },
         {
           "id": 278,
-          "logprob": -0.14123535,
+          "logprob": -0.13867188,
           "text": "the"
         },
         {
           "id": 1967,
-          "logprob": -0.053985596,
+          "logprob": -0.06842041,
           "text": "image"
         },
         {
           "id": 29973,
-          "logprob": -0.15771484,
+          "logprob": -0.15319824,
           "text": "?"
         }
       ],
@@ -609,25 +105,25 @@
       "tokens": [
         {
           "id": 32002,
-          "logprob": -0.004295349,
+          "logprob": -0.0019445419,
           "special": true,
           "text": "<end_of_utterance>"
         },
         {
           "id": 29871,
-          "logprob": -7.43866e-05,
+          "logprob": -8.416176e-05,
           "special": false,
           "text": " "
         },
         {
           "id": 13,
-          "logprob": -2.3126602e-05,
+          "logprob": -1.7881393e-05,
           "special": false,
           "text": "\n"
         },
         {
           "id": 7900,
-          "logprob": -3.9339066e-06,
+          "logprob": -3.0994415e-06,
           "special": false,
           "text": "Ass"
         },
@@ -639,36 +135,544 @@
         },
         {
           "id": 29901,
-          "logprob": -2.6226044e-06,
+          "logprob": -3.2186508e-06,
           "special": false,
           "text": ":"
         },
         {
           "id": 319,
-          "logprob": -0.87841797,
+          "logprob": -0.89941406,
           "special": false,
           "text": " A"
         },
         {
-          "id": 521,
-          "logprob": -1.3837891,
+          "id": 696,
+          "logprob": -1.234375,
           "special": false,
-          "text": " ch"
+          "text": " ro"
         },
         {
-          "id": 21475,
-          "logprob": -0.00051641464,
+          "id": 15664,
+          "logprob": -0.0002465248,
           "special": false,
-          "text": "icken"
+          "text": "oster"
         },
         {
-          "id": 338,
-          "logprob": -1.1435547,
+          "id": 15028,
+          "logprob": -1.1660156,
           "special": false,
-          "text": " is"
+          "text": " stands"
         }
-      ]
+      ],
+      "top_tokens": null
     },
-    "generated_text": "\nAssistant: A chicken is"
+    "generated_text": "\nAssistant: A rooster stands"
+  },
+  {
+    "details": {
+      "best_of_sequences": null,
+      "finish_reason": "length",
+      "generated_tokens": 10,
+      "prefill": [
+        {
+          "id": 1,
+          "logprob": null,
+          "text": "<s>"
+        },
+        {
+          "id": 4911,
+          "logprob": -5.7890625,
+          "text": "User"
+        },
+        {
+          "id": 29901,
+          "logprob": -0.0070152283,
+          "text": ":"
+        },
+        {
+          "id": 32000,
+          "logprob": -0.8125,
+          "text": "<fake_token_around_image>"
+        },
+        {
+          "id": 32001,
+          "logprob": -6.651878e-05,
+          "text": "<image>"
+        },
+        {
+          "id": 32000,
+          "logprob": -3.5762787e-07,
+          "text": "<fake_token_around_image>"
+        },
+        {
+          "id": 1815,
+          "logprob": -4.2265625,
+          "text": "Can"
+        },
+        {
+          "id": 366,
+          "logprob": -0.014190674,
+          "text": "you"
+        },
+        {
+          "id": 2649,
+          "logprob": -4.4140625,
+          "text": "tell"
+        },
+        {
+          "id": 592,
+          "logprob": -0.2919922,
+          "text": "me"
+        },
+        {
+          "id": 263,
+          "logprob": -4.2109375,
+          "text": "a"
+        },
+        {
+          "id": 1407,
+          "logprob": -9.4375,
+          "text": "very"
+        },
+        {
+          "id": 3273,
+          "logprob": -1.8720703,
+          "text": "short"
+        },
+        {
+          "id": 5828,
+          "logprob": -0.26904297,
+          "text": "story"
+        },
+        {
+          "id": 2729,
+          "logprob": -3.7675781,
+          "text": "based"
+        },
+        {
+          "id": 373,
+          "logprob": -0.0005402565,
+          "text": "on"
+        },
+        {
+          "id": 278,
+          "logprob": -0.13867188,
+          "text": "the"
+        },
+        {
+          "id": 1967,
+          "logprob": -0.068359375,
+          "text": "image"
+        },
+        {
+          "id": 29973,
+          "logprob": -0.15539551,
+          "text": "?"
+        }
+      ],
+      "seed": null,
+      "tokens": [
+        {
+          "id": 32002,
+          "logprob": -0.0019168854,
+          "special": true,
+          "text": "<end_of_utterance>"
+        },
+        {
+          "id": 29871,
+          "logprob": -8.392334e-05,
+          "special": false,
+          "text": " "
+        },
+        {
+          "id": 13,
+          "logprob": -1.7642975e-05,
+          "special": false,
+          "text": "\n"
+        },
+        {
+          "id": 7900,
+          "logprob": -3.0994415e-06,
+          "special": false,
+          "text": "Ass"
+        },
+        {
+          "id": 22137,
+          "logprob": 0.0,
+          "special": false,
+          "text": "istant"
+        },
+        {
+          "id": 29901,
+          "logprob": -3.0994415e-06,
+          "special": false,
+          "text": ":"
+        },
+        {
+          "id": 319,
+          "logprob": -0.90722656,
+          "special": false,
+          "text": " A"
+        },
+        {
+          "id": 696,
+          "logprob": -1.2373047,
+          "special": false,
+          "text": " ro"
+        },
+        {
+          "id": 15664,
+          "logprob": -0.00024938583,
+          "special": false,
+          "text": "oster"
+        },
+        {
+          "id": 15028,
+          "logprob": -1.1708984,
+          "special": false,
+          "text": " stands"
+        }
+      ],
+      "top_tokens": null
+    },
+    "generated_text": "\nAssistant: A rooster stands"
+  },
+  {
+    "details": {
+      "best_of_sequences": null,
+      "finish_reason": "length",
+      "generated_tokens": 10,
+      "prefill": [
+        {
+          "id": 1,
+          "logprob": null,
+          "text": "<s>"
+        },
+        {
+          "id": 4911,
+          "logprob": -5.7890625,
+          "text": "User"
+        },
+        {
+          "id": 29901,
+          "logprob": -0.0070152283,
+          "text": ":"
+        },
+        {
+          "id": 32000,
+          "logprob": -0.8125,
+          "text": "<fake_token_around_image>"
+        },
+        {
+          "id": 32001,
+          "logprob": -6.663799e-05,
+          "text": "<image>"
+        },
+        {
+          "id": 32000,
+          "logprob": -3.5762787e-07,
+          "text": "<fake_token_around_image>"
+        },
+        {
+          "id": 1815,
+          "logprob": -4.2265625,
+          "text": "Can"
+        },
+        {
+          "id": 366,
+          "logprob": -0.014190674,
+          "text": "you"
+        },
+        {
+          "id": 2649,
+          "logprob": -4.4140625,
+          "text": "tell"
+        },
+        {
+          "id": 592,
+          "logprob": -0.2919922,
+          "text": "me"
+        },
+        {
+          "id": 263,
+          "logprob": -4.2109375,
+          "text": "a"
+        },
+        {
+          "id": 1407,
+          "logprob": -9.4375,
+          "text": "very"
+        },
+        {
+          "id": 3273,
+          "logprob": -1.8720703,
+          "text": "short"
+        },
+        {
+          "id": 5828,
+          "logprob": -0.26904297,
+          "text": "story"
+        },
+        {
+          "id": 2729,
+          "logprob": -3.7675781,
+          "text": "based"
+        },
+        {
+          "id": 373,
+          "logprob": -0.0005402565,
+          "text": "on"
+        },
+        {
+          "id": 278,
+          "logprob": -0.13867188,
+          "text": "the"
+        },
+        {
+          "id": 1967,
+          "logprob": -0.068359375,
+          "text": "image"
+        },
+        {
+          "id": 29973,
+          "logprob": -0.15539551,
+          "text": "?"
+        }
+      ],
+      "seed": null,
+      "tokens": [
+        {
+          "id": 32002,
+          "logprob": -0.0019168854,
+          "special": true,
+          "text": "<end_of_utterance>"
+        },
+        {
+          "id": 29871,
+          "logprob": -8.404255e-05,
+          "special": false,
+          "text": " "
+        },
+        {
+          "id": 13,
+          "logprob": -1.7642975e-05,
+          "special": false,
+          "text": "\n"
+        },
+        {
+          "id": 7900,
+          "logprob": -3.0994415e-06,
+          "special": false,
+          "text": "Ass"
+        },
+        {
+          "id": 22137,
+          "logprob": 0.0,
+          "special": false,
+          "text": "istant"
+        },
+        {
+          "id": 29901,
+          "logprob": -3.0994415e-06,
+          "special": false,
+          "text": ":"
+        },
+        {
+          "id": 319,
+          "logprob": -0.90722656,
+          "special": false,
+          "text": " A"
+        },
+        {
+          "id": 696,
+          "logprob": -1.2373047,
+          "special": false,
+          "text": " ro"
+        },
+        {
+          "id": 15664,
+          "logprob": -0.00024938583,
+          "special": false,
+          "text": "oster"
+        },
+        {
+          "id": 15028,
+          "logprob": -1.1708984,
+          "special": false,
+          "text": " stands"
+        }
+      ],
+      "top_tokens": null
+    },
+    "generated_text": "\nAssistant: A rooster stands"
+  },
+  {
+    "details": {
+      "best_of_sequences": null,
+      "finish_reason": "length",
+      "generated_tokens": 10,
+      "prefill": [
+        {
+          "id": 1,
+          "logprob": null,
+          "text": "<s>"
+        },
+        {
+          "id": 4911,
+          "logprob": -5.7890625,
+          "text": "User"
+        },
+        {
+          "id": 29901,
+          "logprob": -0.0070152283,
+          "text": ":"
+        },
+        {
+          "id": 32000,
+          "logprob": -0.8125,
+          "text": "<fake_token_around_image>"
+        },
+        {
+          "id": 32001,
+          "logprob": -6.663799e-05,
+          "text": "<image>"
+        },
+        {
+          "id": 32000,
+          "logprob": -3.5762787e-07,
+          "text": "<fake_token_around_image>"
+        },
+        {
+          "id": 1815,
+          "logprob": -4.2265625,
+          "text": "Can"
+        },
+        {
+          "id": 366,
+          "logprob": -0.014190674,
+          "text": "you"
+        },
+        {
+          "id": 2649,
+          "logprob": -4.4140625,
+          "text": "tell"
+        },
+        {
+          "id": 592,
+          "logprob": -0.2919922,
+          "text": "me"
+        },
+        {
+          "id": 263,
+          "logprob": -4.2109375,
+          "text": "a"
+        },
+        {
+          "id": 1407,
+          "logprob": -9.4375,
+          "text": "very"
+        },
+        {
+          "id": 3273,
+          "logprob": -1.8720703,
+          "text": "short"
+        },
+        {
+          "id": 5828,
+          "logprob": -0.26904297,
+          "text": "story"
+        },
+        {
+          "id": 2729,
+          "logprob": -3.7675781,
+          "text": "based"
+        },
+        {
+          "id": 373,
+          "logprob": -0.0005402565,
+          "text": "on"
+        },
+        {
+          "id": 278,
+          "logprob": -0.13867188,
+          "text": "the"
+        },
+        {
+          "id": 1967,
+          "logprob": -0.068359375,
+          "text": "image"
+        },
+        {
+          "id": 29973,
+          "logprob": -0.15539551,
+          "text": "?"
+        }
+      ],
+      "seed": null,
+      "tokens": [
+        {
+          "id": 32002,
+          "logprob": -0.0019159317,
+          "special": true,
+          "text": "<end_of_utterance>"
+        },
+        {
+          "id": 29871,
+          "logprob": -8.404255e-05,
+          "special": false,
+          "text": " "
+        },
+        {
+          "id": 13,
+          "logprob": -1.7642975e-05,
+          "special": false,
+          "text": "\n"
+        },
+        {
+          "id": 7900,
+          "logprob": -3.0994415e-06,
+          "special": false,
+          "text": "Ass"
+        },
+        {
+          "id": 22137,
+          "logprob": 0.0,
+          "special": false,
+          "text": "istant"
+        },
+        {
+          "id": 29901,
+          "logprob": -3.0994415e-06,
+          "special": false,
+          "text": ":"
+        },
+        {
+          "id": 319,
+          "logprob": -0.90722656,
+          "special": false,
+          "text": " A"
+        },
+        {
+          "id": 696,
+          "logprob": -1.2373047,
+          "special": false,
+          "text": " ro"
+        },
+        {
+          "id": 15664,
+          "logprob": -0.00024938583,
+          "special": false,
+          "text": "oster"
+        },
+        {
+          "id": 15028,
+          "logprob": -1.1708984,
+          "special": false,
+          "text": " stands"
+        }
+      ],
+      "top_tokens": null
+    },
+    "generated_text": "\nAssistant: A rooster stands"
   }
 ]
diff --git a/server/pyproject.toml b/server/pyproject.toml
index 6e83d0e3..1babf749 100644
--- a/server/pyproject.toml
+++ b/server/pyproject.toml
@@ -25,8 +25,8 @@ opentelemetry-instrumentation-grpc = "^0.36b0"
 hf-transfer = "^0.1.2"
 sentencepiece = "^0.1.97"
 tokenizers = "^0.13.3"
-huggingface-hub = "^0.14.1"
-transformers = "^4.31.0"
+huggingface-hub = "^0.16.4"
+transformers = "^4.32.1"
 einops = "^0.6.1"
 texttable = { version = "^1.6.7", optional = true }
 datasets = { version = "^2.14.0", optional = true }