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 #10 array generation

Browse Source
This commit is contained in:
rahulgs12
2023-05-03 11:24:40 -04:00
parent d38be16dd4
commit 4c710431de
1 changed files with 12 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
jsonformer/main.py
+12 -5
View File
@@ -122,13 +122,22 @@ class Jsonformer:
) -> Any:
schema_type = schema["type"]
if schema_type == "number":
obj[key if key else -1] = self.generation_marker
if key:
obj[key] = self.generation_marker
else:
obj.append(self.generation_marker)
return self.generate_number()
elif schema_type == "boolean":
obj[key if key else -1] = self.generation_marker
if key:
obj[key] = self.generation_marker
else:
obj.append(self.generation_marker)
return self.generate_boolean()
elif schema_type == "string":
obj[key if key else -1] = self.generation_marker
if key:
obj[key] = self.generation_marker
else:
obj.append(self.generation_marker)
return self.generate_string()
elif schema_type == "array":
new_array = []
@@ -157,9 +166,7 @@ class Jsonformer:
comma_token_id = self.tokenizer.convert_tokens_to_ids(", ")
close_bracket_logits = logits[close_bracket_token_id]
comma_logits = logits[comma_token_id]
if close_bracket_logits > comma_logits:
break