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

Fix bug resulting in overly long numbers/ints

Browse Source
This commit is contained in:
Ryan
2023-05-18 02:33:46 -07:00
parent c0e51b367b
commit 58e8270537
2 changed files with 38 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files
jsonformer/logits_processors.py
+32 -7
View File
@@ -48,14 +48,21 @@ class NumberStoppingCriteria(StoppingCriteria):
if (
decoded.count(".") == 1
and len(decoded.strip().split(".")[1]) > self.precision
and len(decoded.replace(" ", "").split(".")[1]) > self.precision
):
return True
if (
len(decoded) > 1
and "," in decoded
and any(c.isdigit() for c in decoded.split(",")[0])
):
return True
if (
len(decoded) > 1
and any(c.isdigit() for c in decoded)
and decoded[-1] in [" ", "\n"]
and ("," in decoded or decoded[-1] in (" ", "\n"))
):
return True
@@ -71,9 +78,16 @@ class OutputNumbersTokens(LogitsWarper):
for _, token_id in tokenizer.get_vocab().items():
token_str = tokenizer.decode(token_id).strip()
if token_str == "" or (
all(c.isdigit() or c == "." for c in token_str)
and token_str.count(".") <= 1
if (
token_str == ""
or (
all(c.isdigit() or c == "." for c in token_str)
and token_str.count(".") <= 1
) or (
"," in token_str
and all(c.isdigit() or c == "." for c in token_str.split(",")[0])
and token_str.count(".") <= 1
)
):
self.allowed_mask[token_id] = True
@@ -106,10 +120,17 @@ class IntegerStoppingCriteria(StoppingCriteria):
if len(decoded.strip()) > self.max_digits:
return True
if (
len(decoded) > 1
and "," in decoded
and any(c.isdigit() for c in decoded.split(",")[0])
):
return True
if (
len(decoded) > 1
and any(c.isdigit() for c in decoded)
and decoded[-1] in [" ", "\n"]
and decoded[-1] in (" ", "\n")
):
return True
@@ -125,7 +146,11 @@ class OutputIntegersTokens(LogitsWarper):
for _, token_id in tokenizer.get_vocab().items():
token_str = tokenizer.decode(token_id).strip()
if token_str == "" or all(c.isdigit() for c in token_str):
if (
token_str == ""
or all(c.isdigit() for c in token_str)
or "," in token_str and all(c.isdigit() for c in token_str.split(",")[0])
):
self.allowed_mask[token_id] = True
def __call__(self, _, scores):
jsonformer/main.py
+6 -2
View File
@@ -76,7 +76,9 @@ class Jsonformer:
response = self.tokenizer.decode(response[0], skip_special_tokens=True)
response = response[len(prompt) :]
response = response.strip().rstrip(".")
if "," in response:
response = response.split(",")[0]
response = response.replace(" ", "").rstrip(".")
self.debug("[generate_number]", response)
try:
return float(response)
@@ -106,7 +108,9 @@ class Jsonformer:
response = self.tokenizer.decode(response[0], skip_special_tokens=True)
response = response[len(prompt) :]
response = response.strip()
if "," in response:
response = response.split(",")[0]
response = response.replace(" ", "")
self.debug("[generate_integer]", response)
try:
return int(response)