wassname/prob_jsonformer
1
0
Fork 0
mirror of https://github.com/wassname/prob_jsonformer.git synced 2026-09-09 11:29:57 +08:00
Code Issues Packages Projects Releases Wiki Activity

Fix tokenizer edgecase todo in generate_boolean

Browse Source
This commit is contained in:
Ryan
2023-05-15 00:30:47 -07:00
parent 76bec08f7f
commit ab4ad40a4c
1 changed files with 2 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
jsonformer/main.py
+2 -5
View File
@@ -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]