From f33fccfb13f720f0a84503b32b95bacc4874c8b8 Mon Sep 17 00:00:00 2001
From: Nicolas Patry <patry.nicolas@protonmail.com>
Date: Thu, 25 Apr 2024 15:41:18 +0200
Subject: [PATCH] Handle None.

---
 server/text_generation_server/utils/tokens.py | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/server/text_generation_server/utils/tokens.py b/server/text_generation_server/utils/tokens.py
index e1d00aa1..a661eed9 100644
--- a/server/text_generation_server/utils/tokens.py
+++ b/server/text_generation_server/utils/tokens.py
@@ -185,7 +185,13 @@ class StoppingCriteria:
             StopSequenceCriteria(sequence) for sequence in pb.stop_sequences
         ]
         eos_token_id = tokenizer.eos_token_id
-        eos_token_ids: Set[int] = eos_token_id if isinstance(eos_token_id, set) else {eos_token_id}
+
+        if eos_token_id is None:
+            eos_token_ids = {}
+        elif isinstance(eos_token_id, set):
+            eos_token_ids = eos_token_id
+        else:
+            eos_token_ids = eos_token_id
         return StoppingCriteria(
             eos_token_ids,
             stop_sequence_criterias,