From ab4ad40a4c8ed74bd93e4e48f5a09972108160f9 Mon Sep 17 00:00:00 2001 From: Ryan <18477649+Ryul0rd@users.noreply.github.com> Date: Mon, 15 May 2023 00:30:47 -0700 Subject: [PATCH] Fix tokenizer edgecase todo in generate_boolean --- jsonformer/main.py | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/jsonformer/main.py b/jsonformer/main.py index 2dfab48..61c8ef8 100644 --- a/jsonformer/main.py +++ b/jsonformer/main.py @@ -124,11 +124,8 @@ class Jsonformer: output = self.model.forward(input_tensor.to(self.model.device)) logits = output.logits[0, -1] - # todo: this assumes that "true" and "false" are both tokenized to a single token - # this is probably not true for all tokenizers - # this can be fixed by looking at only the first token of both "true" and "false" - true_token_id = self.tokenizer.convert_tokens_to_ids("true") - false_token_id = self.tokenizer.convert_tokens_to_ids("false") + true_token_id = self.tokenizer.encode("true", return_tensors="pt")[0, 0] + false_token_id = self.tokenizer.encode("false", return_tensors="pt")[0, 0] result = logits[true_token_id] > logits[false_token_id]