From 188ceb2fef4696dd122326b0ef866899da314304 Mon Sep 17 00:00:00 2001
From: Nicolas Patry <patry.nicolas@protonmail.com>
Date: Fri, 16 Aug 2024 13:16:51 +0200
Subject: [PATCH] Update the conftest to allow NaN (first logprob).

---
 integration-tests/conftest.py                             | 2 ++
 .../test_flash_llama_fp8/test_flash_llama_fp8_load.json   | 8 ++++----
 2 files changed, 6 insertions(+), 4 deletions(-)

diff --git a/integration-tests/conftest.py b/integration-tests/conftest.py
index 46a8769f..e283ee50 100644
--- a/integration-tests/conftest.py
+++ b/integration-tests/conftest.py
@@ -118,6 +118,8 @@ class ResponseComparator(JSONSnapshotExtension):
                 and token.text == other.text
                 and (
                     self.ignore_logprob
+                    or token.logprob == other.logprob
+                    and token.logprob is None
                     or math.isclose(token.logprob, other.logprob, rel_tol=self.rtol)
                 )
                 and token.special == other.special
diff --git a/integration-tests/models/__snapshots__/test_flash_llama_fp8/test_flash_llama_fp8_load.json b/integration-tests/models/__snapshots__/test_flash_llama_fp8/test_flash_llama_fp8_load.json
index bf10cc44..1e9ff6c0 100644
--- a/integration-tests/models/__snapshots__/test_flash_llama_fp8/test_flash_llama_fp8_load.json
+++ b/integration-tests/models/__snapshots__/test_flash_llama_fp8/test_flash_llama_fp8_load.json
@@ -7,7 +7,7 @@
       "prefill": [
         {
           "id": 128000,
-          "logprob": 0.0,
+          "logprob": null,
           "text": "<|begin_of_text|>"
         },
         {
@@ -96,7 +96,7 @@
       "prefill": [
         {
           "id": 128000,
-          "logprob": 0.0,
+          "logprob": null,
           "text": "<|begin_of_text|>"
         },
         {
@@ -185,7 +185,7 @@
       "prefill": [
         {
           "id": 128000,
-          "logprob": 0.0,
+          "logprob": null,
           "text": "<|begin_of_text|>"
         },
         {
@@ -274,7 +274,7 @@
       "prefill": [
         {
           "id": 128000,
-          "logprob": 0.0,
+          "logprob": null,
           "text": "<|begin_of_text|>"
         },
         {