From e808222dbf8d19875ba4ca00f4fb7d75c3c9717c Mon Sep 17 00:00:00 2001
From: Nicolas Patry <patry.nicolas@protonmail.com>
Date: Tue, 5 Dec 2023 09:40:29 +0000
Subject: [PATCH] Working around falcon tests.

---
 .../test_flash_falcon_load.json               | 1544 +++--------------
 integration-tests/models/test_flash_falcon.py |    2 +-
 2 files changed, 213 insertions(+), 1333 deletions(-)

diff --git a/integration-tests/models/__snapshots__/test_flash_falcon/test_flash_falcon_load.json b/integration-tests/models/__snapshots__/test_flash_falcon/test_flash_falcon_load.json
index 2d6c45fe..84b16f0b 100644
--- a/integration-tests/models/__snapshots__/test_flash_falcon/test_flash_falcon_load.json
+++ b/integration-tests/models/__snapshots__/test_flash_falcon/test_flash_falcon_load.json
@@ -6,376 +6,96 @@
       "generated_tokens": 10,
       "prefill": [
         {
-          "id": 50,
+          "id": 1562,
           "logprob": null,
-          "text": "G"
-        },
-        {
-          "id": 330,
-          "logprob": -5.96875,
-          "text": "ir"
-        },
-        {
-          "id": 1622,
-          "logprob": -5.6132812,
-          "text": "af"
-        },
-        {
-          "id": 249,
-          "logprob": -6.5,
-          "text": "at"
-        },
-        {
-          "id": 1480,
-          "logprob": -8.0703125,
-          "text": "ron"
+          "text": "What"
         },
         {
           "id": 304,
-          "logprob": -2.328125,
+          "logprob": -1.5371094,
           "text": " is"
         },
         {
-          "id": 23866,
-          "logprob": -9.59375,
-          "text": " obsessed"
+          "id": 12604,
+          "logprob": -8.7109375,
+          "text": " Deep"
         },
         {
-          "id": 335,
-          "logprob": -0.04840088,
-          "text": " with"
+          "id": 9078,
+          "logprob": -1.9384766,
+          "text": " Learning"
         },
         {
-          "id": 26680,
-          "logprob": -4.0039062,
-          "text": " gir"
-        },
-        {
-          "id": 1903,
-          "logprob": -0.07525635,
-          "text": "aff"
-        },
-        {
-          "id": 255,
-          "logprob": -0.006778717,
-          "text": "es"
-        },
-        {
-          "id": 23,
-          "logprob": -1.546875,
-          "text": ","
-        },
-        {
-          "id": 248,
-          "logprob": -4.328125,
-          "text": " the"
-        },
-        {
-          "id": 758,
-          "logprob": -3.734375,
-          "text": " most"
-        },
-        {
-          "id": 21735,
-          "logprob": -5.109375,
-          "text": " glorious"
-        },
-        {
-          "id": 5985,
-          "logprob": -2.09375,
-          "text": " animal"
-        },
-        {
-          "id": 313,
-          "logprob": -1.1845703,
-          "text": " on"
-        },
-        {
-          "id": 248,
-          "logprob": -0.77783203,
-          "text": " the"
-        },
-        {
-          "id": 1936,
-          "logprob": -2.3789062,
-          "text": " face"
-        },
-        {
-          "id": 275,
-          "logprob": -0.004436493,
-          "text": " of"
-        },
-        {
-          "id": 414,
-          "logprob": -1.9667969,
-          "text": " this"
-        },
-        {
-          "id": 6490,
-          "logprob": -2.0449219,
-          "text": " Earth"
-        },
-        {
-          "id": 25,
-          "logprob": -0.28198242,
-          "text": "."
-        },
-        {
-          "id": 401,
-          "logprob": -7.9140625,
-          "text": " G"
-        },
-        {
-          "id": 6013,
-          "logprob": -2.2753906,
-          "text": "ira"
-        },
-        {
-          "id": 694,
-          "logprob": -0.6230469,
-          "text": "ft"
-        },
-        {
-          "id": 1480,
-          "logprob": -0.20922852,
-          "text": "ron"
-        },
-        {
-          "id": 9369,
-          "logprob": -4.5507812,
-          "text": " believes"
-        },
-        {
-          "id": 455,
-          "logprob": -4.5664062,
-          "text": " all"
-        },
-        {
-          "id": 599,
-          "logprob": -2.7382812,
-          "text": " other"
-        },
-        {
-          "id": 5632,
-          "logprob": -0.21887207,
-          "text": " animals"
-        },
-        {
-          "id": 362,
-          "logprob": -0.76660156,
-          "text": " are"
-        },
-        {
-          "id": 23981,
-          "logprob": -4.9960938,
-          "text": " irrelevant"
-        },
-        {
-          "id": 635,
-          "logprob": -4.2304688,
-          "text": " when"
-        },
-        {
-          "id": 4354,
-          "logprob": -0.51416016,
-          "text": " compared"
-        },
-        {
-          "id": 271,
-          "logprob": -0.103637695,
-          "text": " to"
-        },
-        {
-          "id": 248,
-          "logprob": -0.58496094,
-          "text": " the"
-        },
-        {
-          "id": 21735,
-          "logprob": -3.6816406,
-          "text": " glorious"
-        },
-        {
-          "id": 64398,
-          "logprob": -1.8222656,
-          "text": " majesty"
-        },
-        {
-          "id": 275,
-          "logprob": -0.23596191,
-          "text": " of"
-        },
-        {
-          "id": 248,
-          "logprob": -0.35424805,
-          "text": " the"
-        },
-        {
-          "id": 26680,
-          "logprob": -0.24621582,
-          "text": " gir"
-        },
-        {
-          "id": 23226,
-          "logprob": -0.029632568,
-          "text": "affe"
-        },
-        {
-          "id": 25,
-          "logprob": -0.17382812,
-          "text": "."
-        },
-        {
-          "id": 193,
-          "logprob": -1.3955078,
-          "text": "\n"
-        },
-        {
-          "id": 23626,
-          "logprob": -10.0625,
-          "text": "Daniel"
-        },
-        {
-          "id": 37,
-          "logprob": -4.59375,
-          "text": ":"
-        },
-        {
-          "id": 23090,
-          "logprob": -6.9414062,
-          "text": " Hello"
-        },
-        {
-          "id": 23,
-          "logprob": -0.9926758,
-          "text": ","
-        },
-        {
-          "id": 29033,
-          "logprob": -2.2324219,
-          "text": " Gir"
-        },
-        {
-          "id": 1622,
-          "logprob": -0.10760498,
-          "text": "af"
-        },
-        {
-          "id": 249,
-          "logprob": -0.04260254,
-          "text": "at"
-        },
-        {
-          "id": 1480,
-          "logprob": -0.0024738312,
-          "text": "ron"
-        },
-        {
-          "id": 12,
-          "logprob": -1.4267578,
-          "text": "!"
-        },
-        {
-          "id": 193,
-          "logprob": -1.0996094,
-          "text": "\n"
-        },
-        {
-          "id": 50,
-          "logprob": -0.057159424,
-          "text": "G"
-        },
-        {
-          "id": 330,
-          "logprob": -0.13146973,
-          "text": "ir"
-        },
-        {
-          "id": 1622,
-          "logprob": -0.007095337,
-          "text": "af"
-        },
-        {
-          "id": 249,
-          "logprob": -0.008460999,
-          "text": "at"
-        },
-        {
-          "id": 1480,
-          "logprob": -0.0006928444,
-          "text": "ron"
-        },
-        {
-          "id": 37,
-          "logprob": -0.0074768066,
-          "text": ":"
+          "id": 42,
+          "logprob": -0.32958984,
+          "text": "?"
         }
       ],
       "seed": null,
       "tokens": [
-        {
-          "id": 23090,
-          "logprob": -1.8251953,
-          "special": false,
-          "text": " Hello"
-        },
-        {
-          "id": 23,
-          "logprob": -0.31713867,
-          "special": false,
-          "text": ","
-        },
-        {
-          "id": 8156,
-          "logprob": -0.23913574,
-          "special": false,
-          "text": " Daniel"
-        },
-        {
-          "id": 12,
-          "logprob": -0.5698242,
-          "special": false,
-          "text": "!"
-        },
         {
           "id": 193,
-          "logprob": -0.6152344,
+          "logprob": -0.27563477,
           "special": false,
           "text": "\n"
         },
         {
-          "id": 23626,
-          "logprob": -0.41992188,
+          "id": 27319,
+          "logprob": -1.1015625,
           "special": false,
-          "text": "Daniel"
+          "text": "Deep"
         },
         {
-          "id": 37,
-          "logprob": -0.00233078,
+          "id": 2860,
+          "logprob": -0.5292969,
           "special": false,
-          "text": ":"
+          "text": " learning"
         },
         {
-          "id": 1634,
-          "logprob": -2.0625,
+          "id": 304,
+          "logprob": -0.26538086,
           "special": false,
-          "text": " What"
+          "text": " is"
         },
         {
-          "id": 18,
-          "logprob": -1.4707031,
+          "id": 241,
+          "logprob": -0.41430664,
           "special": false,
-          "text": "'"
+          "text": " a"
         },
         {
-          "id": 94,
-          "logprob": -0.026260376,
+          "id": 30095,
+          "logprob": -1.2724609,
           "special": false,
-          "text": "s"
+          "text": " subset"
+        },
+        {
+          "id": 275,
+          "logprob": -0.0115737915,
+          "special": false,
+          "text": " of"
+        },
+        {
+          "id": 3414,
+          "logprob": -0.2409668,
+          "special": false,
+          "text": " machine"
+        },
+        {
+          "id": 2860,
+          "logprob": -0.008323669,
+          "special": false,
+          "text": " learning"
+        },
+        {
+          "id": 325,
+          "logprob": -1.1533203,
+          "special": false,
+          "text": " that"
         }
       ]
     },
-    "generated_text": " Hello, Daniel!\nDaniel: What's"
+    "generated_text": "\nDeep learning is a subset of machine learning that"
   },
   {
     "details": {
@@ -384,376 +104,96 @@
       "generated_tokens": 10,
       "prefill": [
         {
-          "id": 50,
+          "id": 1562,
           "logprob": null,
-          "text": "G"
-        },
-        {
-          "id": 330,
-          "logprob": -5.96875,
-          "text": "ir"
-        },
-        {
-          "id": 1622,
-          "logprob": -5.6132812,
-          "text": "af"
-        },
-        {
-          "id": 249,
-          "logprob": -6.5,
-          "text": "at"
-        },
-        {
-          "id": 1480,
-          "logprob": -8.078125,
-          "text": "ron"
+          "text": "What"
         },
         {
           "id": 304,
-          "logprob": -2.3261719,
+          "logprob": -1.5371094,
           "text": " is"
         },
         {
-          "id": 23866,
-          "logprob": -9.59375,
-          "text": " obsessed"
+          "id": 12604,
+          "logprob": -8.7109375,
+          "text": " Deep"
         },
         {
-          "id": 335,
-          "logprob": -0.048431396,
-          "text": " with"
+          "id": 9078,
+          "logprob": -1.9384766,
+          "text": " Learning"
         },
         {
-          "id": 26680,
-          "logprob": -3.9980469,
-          "text": " gir"
-        },
-        {
-          "id": 1903,
-          "logprob": -0.07531738,
-          "text": "aff"
-        },
-        {
-          "id": 255,
-          "logprob": -0.006793976,
-          "text": "es"
-        },
-        {
-          "id": 23,
-          "logprob": -1.5478516,
-          "text": ","
-        },
-        {
-          "id": 248,
-          "logprob": -4.3320312,
-          "text": " the"
-        },
-        {
-          "id": 758,
-          "logprob": -3.734375,
-          "text": " most"
-        },
-        {
-          "id": 21735,
-          "logprob": -5.109375,
-          "text": " glorious"
-        },
-        {
-          "id": 5985,
-          "logprob": -2.0917969,
-          "text": " animal"
-        },
-        {
-          "id": 313,
-          "logprob": -1.1826172,
-          "text": " on"
-        },
-        {
-          "id": 248,
-          "logprob": -0.7758789,
-          "text": " the"
-        },
-        {
-          "id": 1936,
-          "logprob": -2.3808594,
-          "text": " face"
-        },
-        {
-          "id": 275,
-          "logprob": -0.0044403076,
-          "text": " of"
-        },
-        {
-          "id": 414,
-          "logprob": -1.9658203,
-          "text": " this"
-        },
-        {
-          "id": 6490,
-          "logprob": -2.0449219,
-          "text": " Earth"
-        },
-        {
-          "id": 25,
-          "logprob": -0.28198242,
-          "text": "."
-        },
-        {
-          "id": 401,
-          "logprob": -7.921875,
-          "text": " G"
-        },
-        {
-          "id": 6013,
-          "logprob": -2.2734375,
-          "text": "ira"
-        },
-        {
-          "id": 694,
-          "logprob": -0.6245117,
-          "text": "ft"
-        },
-        {
-          "id": 1480,
-          "logprob": -0.21008301,
-          "text": "ron"
-        },
-        {
-          "id": 9369,
-          "logprob": -4.5507812,
-          "text": " believes"
-        },
-        {
-          "id": 455,
-          "logprob": -4.5664062,
-          "text": " all"
-        },
-        {
-          "id": 599,
-          "logprob": -2.7382812,
-          "text": " other"
-        },
-        {
-          "id": 5632,
-          "logprob": -0.21850586,
-          "text": " animals"
-        },
-        {
-          "id": 362,
-          "logprob": -0.7675781,
-          "text": " are"
-        },
-        {
-          "id": 23981,
-          "logprob": -4.9960938,
-          "text": " irrelevant"
-        },
-        {
-          "id": 635,
-          "logprob": -4.2304688,
-          "text": " when"
-        },
-        {
-          "id": 4354,
-          "logprob": -0.51220703,
-          "text": " compared"
-        },
-        {
-          "id": 271,
-          "logprob": -0.103637695,
-          "text": " to"
-        },
-        {
-          "id": 248,
-          "logprob": -0.5839844,
-          "text": " the"
-        },
-        {
-          "id": 21735,
-          "logprob": -3.6816406,
-          "text": " glorious"
-        },
-        {
-          "id": 64398,
-          "logprob": -1.8203125,
-          "text": " majesty"
-        },
-        {
-          "id": 275,
-          "logprob": -0.23596191,
-          "text": " of"
-        },
-        {
-          "id": 248,
-          "logprob": -0.35375977,
-          "text": " the"
-        },
-        {
-          "id": 26680,
-          "logprob": -0.24707031,
-          "text": " gir"
-        },
-        {
-          "id": 23226,
-          "logprob": -0.02960205,
-          "text": "affe"
-        },
-        {
-          "id": 25,
-          "logprob": -0.17346191,
-          "text": "."
-        },
-        {
-          "id": 193,
-          "logprob": -1.3935547,
-          "text": "\n"
-        },
-        {
-          "id": 23626,
-          "logprob": -10.0625,
-          "text": "Daniel"
-        },
-        {
-          "id": 37,
-          "logprob": -4.59375,
-          "text": ":"
-        },
-        {
-          "id": 23090,
-          "logprob": -6.9375,
-          "text": " Hello"
-        },
-        {
-          "id": 23,
-          "logprob": -0.9941406,
-          "text": ","
-        },
-        {
-          "id": 29033,
-          "logprob": -2.2304688,
-          "text": " Gir"
-        },
-        {
-          "id": 1622,
-          "logprob": -0.10748291,
-          "text": "af"
-        },
-        {
-          "id": 249,
-          "logprob": -0.04260254,
-          "text": "at"
-        },
-        {
-          "id": 1480,
-          "logprob": -0.0024719238,
-          "text": "ron"
-        },
-        {
-          "id": 12,
-          "logprob": -1.4287109,
-          "text": "!"
-        },
-        {
-          "id": 193,
-          "logprob": -1.1005859,
-          "text": "\n"
-        },
-        {
-          "id": 50,
-          "logprob": -0.057037354,
-          "text": "G"
-        },
-        {
-          "id": 330,
-          "logprob": -0.13049316,
-          "text": "ir"
-        },
-        {
-          "id": 1622,
-          "logprob": -0.0071029663,
-          "text": "af"
-        },
-        {
-          "id": 249,
-          "logprob": -0.008460999,
-          "text": "at"
-        },
-        {
-          "id": 1480,
-          "logprob": -0.00068998337,
-          "text": "ron"
-        },
-        {
-          "id": 37,
-          "logprob": -0.007484436,
-          "text": ":"
+          "id": 42,
+          "logprob": -0.32958984,
+          "text": "?"
         }
       ],
       "seed": null,
       "tokens": [
-        {
-          "id": 23090,
-          "logprob": -1.8261719,
-          "special": false,
-          "text": " Hello"
-        },
-        {
-          "id": 23,
-          "logprob": -0.3173828,
-          "special": false,
-          "text": ","
-        },
-        {
-          "id": 8156,
-          "logprob": -0.23840332,
-          "special": false,
-          "text": " Daniel"
-        },
-        {
-          "id": 12,
-          "logprob": -0.5703125,
-          "special": false,
-          "text": "!"
-        },
         {
           "id": 193,
-          "logprob": -0.61376953,
+          "logprob": -0.27563477,
           "special": false,
           "text": "\n"
         },
         {
-          "id": 23626,
-          "logprob": -0.4194336,
+          "id": 27319,
+          "logprob": -1.1015625,
           "special": false,
-          "text": "Daniel"
+          "text": "Deep"
         },
         {
-          "id": 37,
-          "logprob": -0.0023326874,
+          "id": 2860,
+          "logprob": -0.5292969,
           "special": false,
-          "text": ":"
+          "text": " learning"
         },
         {
-          "id": 1634,
-          "logprob": -2.0625,
+          "id": 304,
+          "logprob": -0.26538086,
           "special": false,
-          "text": " What"
+          "text": " is"
         },
         {
-          "id": 18,
-          "logprob": -1.5273438,
+          "id": 241,
+          "logprob": -0.41430664,
           "special": false,
-          "text": "'"
+          "text": " a"
         },
         {
-          "id": 94,
-          "logprob": -0.007896423,
+          "id": 30095,
+          "logprob": -1.2724609,
           "special": false,
-          "text": "s"
+          "text": " subset"
+        },
+        {
+          "id": 275,
+          "logprob": -0.0115737915,
+          "special": false,
+          "text": " of"
+        },
+        {
+          "id": 3414,
+          "logprob": -0.2409668,
+          "special": false,
+          "text": " machine"
+        },
+        {
+          "id": 2860,
+          "logprob": -0.008323669,
+          "special": false,
+          "text": " learning"
+        },
+        {
+          "id": 325,
+          "logprob": -1.1533203,
+          "special": false,
+          "text": " that"
         }
       ]
     },
-    "generated_text": " Hello, Daniel!\nDaniel: What's"
+    "generated_text": "\nDeep learning is a subset of machine learning that"
   },
   {
     "details": {
@@ -762,376 +202,96 @@
       "generated_tokens": 10,
       "prefill": [
         {
-          "id": 50,
+          "id": 1562,
           "logprob": null,
-          "text": "G"
-        },
-        {
-          "id": 330,
-          "logprob": -5.96875,
-          "text": "ir"
-        },
-        {
-          "id": 1622,
-          "logprob": -5.6132812,
-          "text": "af"
-        },
-        {
-          "id": 249,
-          "logprob": -6.5,
-          "text": "at"
-        },
-        {
-          "id": 1480,
-          "logprob": -8.0703125,
-          "text": "ron"
+          "text": "What"
         },
         {
           "id": 304,
-          "logprob": -2.328125,
+          "logprob": -1.5371094,
           "text": " is"
         },
         {
-          "id": 23866,
-          "logprob": -9.59375,
-          "text": " obsessed"
+          "id": 12604,
+          "logprob": -8.7109375,
+          "text": " Deep"
         },
         {
-          "id": 335,
-          "logprob": -0.04840088,
-          "text": " with"
+          "id": 9078,
+          "logprob": -1.9384766,
+          "text": " Learning"
         },
         {
-          "id": 26680,
-          "logprob": -4.0039062,
-          "text": " gir"
-        },
-        {
-          "id": 1903,
-          "logprob": -0.07525635,
-          "text": "aff"
-        },
-        {
-          "id": 255,
-          "logprob": -0.006778717,
-          "text": "es"
-        },
-        {
-          "id": 23,
-          "logprob": -1.546875,
-          "text": ","
-        },
-        {
-          "id": 248,
-          "logprob": -4.328125,
-          "text": " the"
-        },
-        {
-          "id": 758,
-          "logprob": -3.734375,
-          "text": " most"
-        },
-        {
-          "id": 21735,
-          "logprob": -5.109375,
-          "text": " glorious"
-        },
-        {
-          "id": 5985,
-          "logprob": -2.09375,
-          "text": " animal"
-        },
-        {
-          "id": 313,
-          "logprob": -1.1845703,
-          "text": " on"
-        },
-        {
-          "id": 248,
-          "logprob": -0.77783203,
-          "text": " the"
-        },
-        {
-          "id": 1936,
-          "logprob": -2.3789062,
-          "text": " face"
-        },
-        {
-          "id": 275,
-          "logprob": -0.004436493,
-          "text": " of"
-        },
-        {
-          "id": 414,
-          "logprob": -1.9667969,
-          "text": " this"
-        },
-        {
-          "id": 6490,
-          "logprob": -2.0449219,
-          "text": " Earth"
-        },
-        {
-          "id": 25,
-          "logprob": -0.28198242,
-          "text": "."
-        },
-        {
-          "id": 401,
-          "logprob": -7.9140625,
-          "text": " G"
-        },
-        {
-          "id": 6013,
-          "logprob": -2.2753906,
-          "text": "ira"
-        },
-        {
-          "id": 694,
-          "logprob": -0.6230469,
-          "text": "ft"
-        },
-        {
-          "id": 1480,
-          "logprob": -0.20922852,
-          "text": "ron"
-        },
-        {
-          "id": 9369,
-          "logprob": -4.5507812,
-          "text": " believes"
-        },
-        {
-          "id": 455,
-          "logprob": -4.5664062,
-          "text": " all"
-        },
-        {
-          "id": 599,
-          "logprob": -2.7382812,
-          "text": " other"
-        },
-        {
-          "id": 5632,
-          "logprob": -0.21887207,
-          "text": " animals"
-        },
-        {
-          "id": 362,
-          "logprob": -0.76660156,
-          "text": " are"
-        },
-        {
-          "id": 23981,
-          "logprob": -4.9960938,
-          "text": " irrelevant"
-        },
-        {
-          "id": 635,
-          "logprob": -4.2304688,
-          "text": " when"
-        },
-        {
-          "id": 4354,
-          "logprob": -0.51416016,
-          "text": " compared"
-        },
-        {
-          "id": 271,
-          "logprob": -0.103637695,
-          "text": " to"
-        },
-        {
-          "id": 248,
-          "logprob": -0.58496094,
-          "text": " the"
-        },
-        {
-          "id": 21735,
-          "logprob": -3.6816406,
-          "text": " glorious"
-        },
-        {
-          "id": 64398,
-          "logprob": -1.8222656,
-          "text": " majesty"
-        },
-        {
-          "id": 275,
-          "logprob": -0.23596191,
-          "text": " of"
-        },
-        {
-          "id": 248,
-          "logprob": -0.35424805,
-          "text": " the"
-        },
-        {
-          "id": 26680,
-          "logprob": -0.24621582,
-          "text": " gir"
-        },
-        {
-          "id": 23226,
-          "logprob": -0.029632568,
-          "text": "affe"
-        },
-        {
-          "id": 25,
-          "logprob": -0.17382812,
-          "text": "."
-        },
-        {
-          "id": 193,
-          "logprob": -1.3955078,
-          "text": "\n"
-        },
-        {
-          "id": 23626,
-          "logprob": -10.0625,
-          "text": "Daniel"
-        },
-        {
-          "id": 37,
-          "logprob": -4.59375,
-          "text": ":"
-        },
-        {
-          "id": 23090,
-          "logprob": -6.9414062,
-          "text": " Hello"
-        },
-        {
-          "id": 23,
-          "logprob": -0.9926758,
-          "text": ","
-        },
-        {
-          "id": 29033,
-          "logprob": -2.2324219,
-          "text": " Gir"
-        },
-        {
-          "id": 1622,
-          "logprob": -0.10760498,
-          "text": "af"
-        },
-        {
-          "id": 249,
-          "logprob": -0.04260254,
-          "text": "at"
-        },
-        {
-          "id": 1480,
-          "logprob": -0.0024738312,
-          "text": "ron"
-        },
-        {
-          "id": 12,
-          "logprob": -1.4267578,
-          "text": "!"
-        },
-        {
-          "id": 193,
-          "logprob": -1.0996094,
-          "text": "\n"
-        },
-        {
-          "id": 50,
-          "logprob": -0.057159424,
-          "text": "G"
-        },
-        {
-          "id": 330,
-          "logprob": -0.13146973,
-          "text": "ir"
-        },
-        {
-          "id": 1622,
-          "logprob": -0.007095337,
-          "text": "af"
-        },
-        {
-          "id": 249,
-          "logprob": -0.008460999,
-          "text": "at"
-        },
-        {
-          "id": 1480,
-          "logprob": -0.0006928444,
-          "text": "ron"
-        },
-        {
-          "id": 37,
-          "logprob": -0.0074768066,
-          "text": ":"
+          "id": 42,
+          "logprob": -0.32958984,
+          "text": "?"
         }
       ],
       "seed": null,
       "tokens": [
-        {
-          "id": 23090,
-          "logprob": -1.8251953,
-          "special": false,
-          "text": " Hello"
-        },
-        {
-          "id": 23,
-          "logprob": -0.31713867,
-          "special": false,
-          "text": ","
-        },
-        {
-          "id": 8156,
-          "logprob": -0.23913574,
-          "special": false,
-          "text": " Daniel"
-        },
-        {
-          "id": 12,
-          "logprob": -0.5698242,
-          "special": false,
-          "text": "!"
-        },
         {
           "id": 193,
-          "logprob": -0.6152344,
+          "logprob": -0.27563477,
           "special": false,
           "text": "\n"
         },
         {
-          "id": 23626,
-          "logprob": -0.41992188,
+          "id": 27319,
+          "logprob": -1.1015625,
           "special": false,
-          "text": "Daniel"
+          "text": "Deep"
         },
         {
-          "id": 37,
-          "logprob": -0.00233078,
+          "id": 2860,
+          "logprob": -0.5292969,
           "special": false,
-          "text": ":"
+          "text": " learning"
         },
         {
-          "id": 1634,
-          "logprob": -2.0625,
+          "id": 304,
+          "logprob": -0.26538086,
           "special": false,
-          "text": " What"
+          "text": " is"
         },
         {
-          "id": 18,
-          "logprob": -1.5283203,
+          "id": 241,
+          "logprob": -0.41430664,
           "special": false,
-          "text": "'"
+          "text": " a"
         },
         {
-          "id": 94,
-          "logprob": -0.007926941,
+          "id": 30095,
+          "logprob": -1.2724609,
           "special": false,
-          "text": "s"
+          "text": " subset"
+        },
+        {
+          "id": 275,
+          "logprob": -0.0115737915,
+          "special": false,
+          "text": " of"
+        },
+        {
+          "id": 3414,
+          "logprob": -0.2409668,
+          "special": false,
+          "text": " machine"
+        },
+        {
+          "id": 2860,
+          "logprob": -0.008323669,
+          "special": false,
+          "text": " learning"
+        },
+        {
+          "id": 325,
+          "logprob": -1.1533203,
+          "special": false,
+          "text": " that"
         }
       ]
     },
-    "generated_text": " Hello, Daniel!\nDaniel: What's"
+    "generated_text": "\nDeep learning is a subset of machine learning that"
   },
   {
     "details": {
@@ -1140,375 +300,95 @@
       "generated_tokens": 10,
       "prefill": [
         {
-          "id": 50,
+          "id": 1562,
           "logprob": null,
-          "text": "G"
-        },
-        {
-          "id": 330,
-          "logprob": -5.96875,
-          "text": "ir"
-        },
-        {
-          "id": 1622,
-          "logprob": -5.6132812,
-          "text": "af"
-        },
-        {
-          "id": 249,
-          "logprob": -6.5,
-          "text": "at"
-        },
-        {
-          "id": 1480,
-          "logprob": -8.0703125,
-          "text": "ron"
+          "text": "What"
         },
         {
           "id": 304,
-          "logprob": -2.328125,
+          "logprob": -1.5371094,
           "text": " is"
         },
         {
-          "id": 23866,
-          "logprob": -9.59375,
-          "text": " obsessed"
+          "id": 12604,
+          "logprob": -8.7109375,
+          "text": " Deep"
         },
         {
-          "id": 335,
-          "logprob": -0.04840088,
-          "text": " with"
+          "id": 9078,
+          "logprob": -1.9384766,
+          "text": " Learning"
         },
         {
-          "id": 26680,
-          "logprob": -4.0039062,
-          "text": " gir"
-        },
-        {
-          "id": 1903,
-          "logprob": -0.07525635,
-          "text": "aff"
-        },
-        {
-          "id": 255,
-          "logprob": -0.006778717,
-          "text": "es"
-        },
-        {
-          "id": 23,
-          "logprob": -1.546875,
-          "text": ","
-        },
-        {
-          "id": 248,
-          "logprob": -4.328125,
-          "text": " the"
-        },
-        {
-          "id": 758,
-          "logprob": -3.734375,
-          "text": " most"
-        },
-        {
-          "id": 21735,
-          "logprob": -5.109375,
-          "text": " glorious"
-        },
-        {
-          "id": 5985,
-          "logprob": -2.09375,
-          "text": " animal"
-        },
-        {
-          "id": 313,
-          "logprob": -1.1845703,
-          "text": " on"
-        },
-        {
-          "id": 248,
-          "logprob": -0.77783203,
-          "text": " the"
-        },
-        {
-          "id": 1936,
-          "logprob": -2.3789062,
-          "text": " face"
-        },
-        {
-          "id": 275,
-          "logprob": -0.004436493,
-          "text": " of"
-        },
-        {
-          "id": 414,
-          "logprob": -1.9667969,
-          "text": " this"
-        },
-        {
-          "id": 6490,
-          "logprob": -2.0449219,
-          "text": " Earth"
-        },
-        {
-          "id": 25,
-          "logprob": -0.28198242,
-          "text": "."
-        },
-        {
-          "id": 401,
-          "logprob": -7.9140625,
-          "text": " G"
-        },
-        {
-          "id": 6013,
-          "logprob": -2.2753906,
-          "text": "ira"
-        },
-        {
-          "id": 694,
-          "logprob": -0.6230469,
-          "text": "ft"
-        },
-        {
-          "id": 1480,
-          "logprob": -0.20922852,
-          "text": "ron"
-        },
-        {
-          "id": 9369,
-          "logprob": -4.5507812,
-          "text": " believes"
-        },
-        {
-          "id": 455,
-          "logprob": -4.5664062,
-          "text": " all"
-        },
-        {
-          "id": 599,
-          "logprob": -2.7382812,
-          "text": " other"
-        },
-        {
-          "id": 5632,
-          "logprob": -0.21887207,
-          "text": " animals"
-        },
-        {
-          "id": 362,
-          "logprob": -0.76660156,
-          "text": " are"
-        },
-        {
-          "id": 23981,
-          "logprob": -4.9960938,
-          "text": " irrelevant"
-        },
-        {
-          "id": 635,
-          "logprob": -4.2304688,
-          "text": " when"
-        },
-        {
-          "id": 4354,
-          "logprob": -0.51416016,
-          "text": " compared"
-        },
-        {
-          "id": 271,
-          "logprob": -0.103637695,
-          "text": " to"
-        },
-        {
-          "id": 248,
-          "logprob": -0.58496094,
-          "text": " the"
-        },
-        {
-          "id": 21735,
-          "logprob": -3.6816406,
-          "text": " glorious"
-        },
-        {
-          "id": 64398,
-          "logprob": -1.8222656,
-          "text": " majesty"
-        },
-        {
-          "id": 275,
-          "logprob": -0.23596191,
-          "text": " of"
-        },
-        {
-          "id": 248,
-          "logprob": -0.35424805,
-          "text": " the"
-        },
-        {
-          "id": 26680,
-          "logprob": -0.24621582,
-          "text": " gir"
-        },
-        {
-          "id": 23226,
-          "logprob": -0.029632568,
-          "text": "affe"
-        },
-        {
-          "id": 25,
-          "logprob": -0.17382812,
-          "text": "."
-        },
-        {
-          "id": 193,
-          "logprob": -1.3955078,
-          "text": "\n"
-        },
-        {
-          "id": 23626,
-          "logprob": -10.0625,
-          "text": "Daniel"
-        },
-        {
-          "id": 37,
-          "logprob": -4.59375,
-          "text": ":"
-        },
-        {
-          "id": 23090,
-          "logprob": -6.9414062,
-          "text": " Hello"
-        },
-        {
-          "id": 23,
-          "logprob": -0.9926758,
-          "text": ","
-        },
-        {
-          "id": 29033,
-          "logprob": -2.2324219,
-          "text": " Gir"
-        },
-        {
-          "id": 1622,
-          "logprob": -0.10760498,
-          "text": "af"
-        },
-        {
-          "id": 249,
-          "logprob": -0.04260254,
-          "text": "at"
-        },
-        {
-          "id": 1480,
-          "logprob": -0.0024738312,
-          "text": "ron"
-        },
-        {
-          "id": 12,
-          "logprob": -1.4267578,
-          "text": "!"
-        },
-        {
-          "id": 193,
-          "logprob": -1.0996094,
-          "text": "\n"
-        },
-        {
-          "id": 50,
-          "logprob": -0.057159424,
-          "text": "G"
-        },
-        {
-          "id": 330,
-          "logprob": -0.13146973,
-          "text": "ir"
-        },
-        {
-          "id": 1622,
-          "logprob": -0.007095337,
-          "text": "af"
-        },
-        {
-          "id": 249,
-          "logprob": -0.008460999,
-          "text": "at"
-        },
-        {
-          "id": 1480,
-          "logprob": -0.0006928444,
-          "text": "ron"
-        },
-        {
-          "id": 37,
-          "logprob": -0.0074768066,
-          "text": ":"
+          "id": 42,
+          "logprob": -0.32958984,
+          "text": "?"
         }
       ],
       "seed": null,
       "tokens": [
-        {
-          "id": 23090,
-          "logprob": -1.8251953,
-          "special": false,
-          "text": " Hello"
-        },
-        {
-          "id": 23,
-          "logprob": -0.31713867,
-          "special": false,
-          "text": ","
-        },
-        {
-          "id": 8156,
-          "logprob": -0.23913574,
-          "special": false,
-          "text": " Daniel"
-        },
-        {
-          "id": 12,
-          "logprob": -0.5698242,
-          "special": false,
-          "text": "!"
-        },
         {
           "id": 193,
-          "logprob": -0.6152344,
+          "logprob": -0.27563477,
           "special": false,
           "text": "\n"
         },
         {
-          "id": 23626,
-          "logprob": -0.41992188,
+          "id": 27319,
+          "logprob": -1.1015625,
           "special": false,
-          "text": "Daniel"
+          "text": "Deep"
         },
         {
-          "id": 37,
-          "logprob": -0.00233078,
+          "id": 2860,
+          "logprob": -0.5292969,
           "special": false,
-          "text": ":"
+          "text": " learning"
         },
         {
-          "id": 1634,
-          "logprob": -2.0625,
+          "id": 304,
+          "logprob": -0.26538086,
           "special": false,
-          "text": " What"
+          "text": " is"
         },
         {
-          "id": 18,
-          "logprob": -1.5283203,
+          "id": 241,
+          "logprob": -0.41430664,
           "special": false,
-          "text": "'"
+          "text": " a"
         },
         {
-          "id": 94,
-          "logprob": -0.007926941,
+          "id": 30095,
+          "logprob": -1.2724609,
           "special": false,
-          "text": "s"
+          "text": " subset"
+        },
+        {
+          "id": 275,
+          "logprob": -0.0115737915,
+          "special": false,
+          "text": " of"
+        },
+        {
+          "id": 3414,
+          "logprob": -0.2409668,
+          "special": false,
+          "text": " machine"
+        },
+        {
+          "id": 2860,
+          "logprob": -0.008323669,
+          "special": false,
+          "text": " learning"
+        },
+        {
+          "id": 325,
+          "logprob": -1.1533203,
+          "special": false,
+          "text": " that"
         }
       ]
     },
-    "generated_text": " Hello, Daniel!\nDaniel: What's"
+    "generated_text": "\nDeep learning is a subset of machine learning that"
   }
 ]
diff --git a/integration-tests/models/test_flash_falcon.py b/integration-tests/models/test_flash_falcon.py
index eac91984..94e9eaca 100644
--- a/integration-tests/models/test_flash_falcon.py
+++ b/integration-tests/models/test_flash_falcon.py
@@ -54,7 +54,7 @@ async def test_flash_falcon_all_params(flash_falcon, response_snapshot):
 async def test_flash_falcon_load(flash_falcon, generate_load, response_snapshot):
     responses = await generate_load(
         flash_falcon,
-        "Girafatron is obsessed with giraffes, the most glorious animal on the face of this Earth. Giraftron believes all other animals are irrelevant when compared to the glorious majesty of the giraffe.\nDaniel: Hello, Girafatron!\nGirafatron:",
+        "What is Deep Learning?",
         max_new_tokens=10,
         n=4,
     )