From 5dbfc07c6e4c53ac49ce6980327d24b416e087d8 Mon Sep 17 00:00:00 2001 From: OlivierDehaene <23298448+OlivierDehaene@users.noreply.github.com> Date: Wed, 8 Mar 2023 17:55:21 +0100 Subject: [PATCH] fix(server): fix index out of range for watermarking --- server/text_generation_server/models/model.py | 11 +++-------- 1 file changed, 3 insertions(+), 8 deletions(-) diff --git a/server/text_generation_server/models/model.py b/server/text_generation_server/models/model.py index e0ce6686..d0645959 100644 --- a/server/text_generation_server/models/model.py +++ b/server/text_generation_server/models/model.py @@ -16,13 +16,8 @@ class Model(ABC): self.device = device # see `decode_token` method - self.tokenizer.add_special_tokens( - {"additional_special_tokens": [""]} - ) - self.special_decode_token_id = self.tokenizer.convert_tokens_to_ids( - "" - ) - self.special_decode_token_length = len("") + self.special_decode_token_id = self.tokenizer.pad_token_id + self.special_decode_token_length = len(self.tokenizer.pad_token) @property @abstractmethod @@ -40,4 +35,4 @@ class Model(ABC): [self.special_decode_token_id, token_id], skip_special_tokens=False ) # slice to remove special decode token - return result[self.special_decode_token_length :] + return result[self.special_decode_token_length:]