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

Refactor token decoding to skip special tokens

Browse Source
This commit is contained in:
wassname
2024-05-11 13:18:43 +08:00
parent c0e929b814
commit aff64e4a8c
2 changed files with 7 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
prob_jsonformer/logits_processors.py
+3 -2
View File
@@ -25,6 +25,7 @@ class StringStoppingCriteria(StoppingCriteria):
result = '"' in last_token
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:
return True
@@ -86,7 +87,7 @@ class OutputNumbersTokens(LogitsWarper):
self.allowed_mask = torch.zeros(vocab_size, dtype=torch.bool)
for _, token_id in tokenizer.get_vocab().items():
token_str = tokenizer.decode(token_id).strip()
token_str = tokenizer.decode(token_id, skip_special_tokens=True).strip()
if (
token_str == ""
@@ -157,7 +158,7 @@ class OutputIntegersTokens(LogitsWarper):
self.allowed_mask = torch.zeros(vocab_size, dtype=torch.bool)
for _, token_id in tokenizer.get_vocab().items():
token_str = tokenizer.decode(token_id).strip()
token_str = tokenizer.decode(token_id, skip_special_tokens=True).strip()
if (
token_str == ""
prob_jsonformer/main.py
+4 -2
View File
@@ -386,7 +386,9 @@ class Jsonformer:
found_close_bracket = False
for token_id in sorted_token_ids:
decoded_token = self.tokenizer.decode(token_id)
decoded_token = self.tokenizer.decode(
token_id, skip_special_tokens=True
)
if "," in decoded_token:
found_comma = True
break
@@ -400,7 +402,7 @@ class Jsonformer:
return obj
def get_prompt(self):
template = """{prompt}\nOutput result in the following JSON schema format:\n{schema}\nResult: {progress}"""
template = """{prompt}\nOutput result in the following JSON schema format:\n```json{schema}```\nResult: ```json\n{progress}"""
progress = json.dumps(self.value)
gen_marker_index = progress.find(f'"{self.generation_marker}"')
if gen_marker_index != -1: