Fixed anyOf handling in Jsonformer

Author: armand0e

jsonformer/main.py

@@ -23,10 +23,10 @@ def __init__(
         prompt: str,
         *,
         debug: bool = False,
-        max_array_length: int = 128,
-        max_number_tokens: int = 2048,
+        max_array_length: int = 10,
+        max_number_tokens: int = 6,
         temperature: float = 1.0,
-        max_string_token_length: int = 1024,
+        max_string_token_length: int = 10,
     ):
         self.model = model
         self.tokenizer = tokenizer
@@ -71,15 +71,14 @@ def generate_number(self, temperature: Union[float, None] = None, iterations=0):
         )
         response = self.tokenizer.decode(response[0], skip_special_tokens=True)
 
-        response = response[len(prompt) :]
+        response = response[len(prompt):]
         response = response.strip().rstrip(".")
         self.debug("[generate_number]", response)
         try:
             return float(response)
         except ValueError:
             if iterations > 3:
                 raise ValueError("Failed to generate a valid number")
-
             return self.generate_number(temperature=self.temperature * 1.3, iterations=iterations+1)
 
     def generate_boolean(self) -> bool:
@@ -90,16 +89,11 @@ def generate_boolean(self) -> bool:
         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")
 
         result = logits[true_token_id] > logits[false_token_id]
-
         self.debug("[generate_boolean]", result)
-
         return result.item()
 
     def generate_string(self) -> str:
@@ -120,8 +114,6 @@ def generate_string(self) -> str:
             pad_token_id=self.tokenizer.eos_token_id,
         )
 
-        # Some models output the prompt as part of the response
-        # This removes the prompt from the response if it is present
         if (
             len(response[0]) >= len(input_tokens[0])
             and (response[0][: len(input_tokens[0])] == input_tokens).all()
@@ -131,7 +123,6 @@ def generate_string(self) -> str:
             response = response[0]
 
         response = self.tokenizer.decode(response, skip_special_tokens=True)
-
         self.debug("[generate_string]", "|" + response + "|")
 
         if response.count('"') < 1:
@@ -153,6 +144,12 @@ def generate_value(
         obj: Union[Dict[str, Any], List[Any]],
         key: Union[str, None] = None,
     ) -> Any:
+        if "anyOf" in schema:
+            options = [option for option in schema["anyOf"] if "type" in option]
+            if options:
+                schema = options[0]
+            else:
+                raise ValueError("No valid type in anyOf for key: " + str(key))
         schema_type = schema["type"]
         if schema_type == "number":
             if key:
@@ -174,7 +171,10 @@ def generate_value(
             return self.generate_string()
         elif schema_type == "array":
             new_array = []
-            obj[key] = new_array
+            if key:
+                obj[key] = new_array
+            else:
+                obj.append(new_array)
             return self.generate_array(schema["items"], new_array)
         elif schema_type == "object":
             new_obj = {}
@@ -186,60 +186,7 @@ def generate_value(
         else:
             raise ValueError(f"Unsupported schema type: {schema_type}")
 
-    def generate_array(self, item_schema: Dict[str, Any], obj: Dict[str, Any]) -> list:
-        for _ in range(self.max_array_length):
-            # forces array to have at least one element
-            element = self.generate_value(item_schema, obj)
-            obj[-1] = element
-
-            obj.append(self.generation_marker)
-            input_prompt = self.get_prompt()
-            obj.pop()
-            input_tensor = self.tokenizer.encode(input_prompt, return_tensors="pt")
-            output = self.model.forward(input_tensor.to(self.model.device))
-            logits = output.logits[0, -1]
-
-
-            top_indices = logits.topk(30).indices
-            sorted_token_ids = top_indices[logits[top_indices].argsort(descending=True)]
-
-            found_comma = False
-            found_close_bracket = False
-
-            for token_id in sorted_token_ids:
-                decoded_token = self.tokenizer.decode(token_id)
-                if ',' in decoded_token:
-                    found_comma = True
-                    break
-                if ']' in decoded_token:
-                    found_close_bracket = True
-                    break
-
-            if found_close_bracket or not found_comma:
-                break
-
-        return obj
-
-    def get_prompt(self):
-        template = """{prompt}\nOutput result in the following JSON schema format:\n{schema}\nResult: {progress}"""
-        progress = json.dumps(self.value)
-        gen_marker_index = progress.find(f'"{self.generation_marker}"')
-        if gen_marker_index != -1:
-            progress = progress[:gen_marker_index]
-        else:
-            raise ValueError("Failed to find generation marker")
-
-        prompt = template.format(
-            prompt=self.prompt,
-            schema=json.dumps(self.json_schema),
-            progress=progress,
-        )
-
-        return prompt
-
     def __call__(self) -> Dict[str, Any]:
         self.value = {}
-        generated_data = self.generate_object(
-            self.json_schema["properties"], self.value
-        )
-        return generated_data
+        generated_data = self.generate_object(self.json_schema["properties"], self.value)
+        return generated_data