From 4c710431de60ec6b3979c49c78df103cf7cd4bdd Mon Sep 17 00:00:00 2001 From: rahulgs12 Date: Wed, 3 May 2023 11:24:40 -0400 Subject: [PATCH] fix #10 array generation --- jsonformer/main.py | 17 ++++++++++++----- 1 file changed, 12 insertions(+), 5 deletions(-) diff --git a/jsonformer/main.py b/jsonformer/main.py index 0a08152..fed8189 100644 --- a/jsonformer/main.py +++ b/jsonformer/main.py @@ -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