initial commit

gen.py

@@ -2,8 +2,6 @@ import torch
 import sys
 from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
 import json
-import jsonschema
-from jsonschema import validate, ValidationError
 
 tokenizer = AutoTokenizer.from_pretrained('google/gemma-2-2b-it')
 
@@ -182,6 +180,7 @@ def generate(event):
 
 
     output_text = tokenizer.decode(tokens[0], skip_special_tokens=False)
+    print(output_text)
     user_prompt_length = len(f"<bos><start_of_turn>user\n{prompt}\n{event}<end_of_turn>\n<start_of_turn>model\n")
 
     json_start_index = output_text.find("<json>")
@@ -201,12 +200,11 @@ def generate(event):
             if last_brace_index != -1:
                 json_string = json_string[:last_brace_index + 1]
 
-
-        # Validate JSON
+        # Load JSON without validation
         try:
-            validate(instance=json.loads(json_string), schema=your_json_schema)
-            return json_string
-        except jsonschema.exceptions.ValidationError as e:
+            data = json.loads(json_string)
+            return data  # Return the parsed JSON data
+        except json.JSONDecodeError as e:
             return f"Error: Invalid JSON - {e}"
 
     else: