From 7b8cec8b3d5c0b28499b4240b59bacd0cfd34cf1 Mon Sep 17 00:00:00 2001 From: wassname Date: Sat, 11 May 2024 14:49:05 +0800 Subject: [PATCH] fix len --- prob_jsonformer/logits_processors.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/prob_jsonformer/logits_processors.py b/prob_jsonformer/logits_processors.py index dc2608f..a6e334c 100644 --- a/prob_jsonformer/logits_processors.py +++ b/prob_jsonformer/logits_processors.py @@ -26,8 +26,14 @@ class StringStoppingCriteria(StoppingCriteria): if self.max_length is not None: # because of tokens this wont work pefectly - str_l = len(self.tokenizer.decode(input_ids[0], skip_special_tokens=True)) - if str_l > self.max_length: + gen_ids = input_ids[0][self.prompt_length :] + o = self.tokenizer.decode(gen_ids, skip_special_tokens=True) + # print("o", o) + str_l = len(o) + # print( + # f"stopping? self.max_length={self.max_length} str_l={str_l},self.prompt_length={self.prompt_length}" + # ) + if (str_l - -self.prompt_length) > self.max_length: return True return result