initial commit

gen.py

@@ -2,8 +2,8 @@ import torch
 import sys
 from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
 import json
-import jsonschema  # Import the jsonschema library
-from jsonschema import validate  # Import the validate function
+import jsonschema
+from jsonschema import validate, ValidationError
 
 tokenizer = AutoTokenizer.from_pretrained('google/gemma-2-2b-it')
 
@@ -182,20 +182,31 @@ 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")  # Calculate user prompt length
+    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>")
     json_end_index = output_text.find("</json>")
 
     if json_start_index != -1 and json_end_index != -1:
-        json_string = output_text[max(json_start_index + 6, user_prompt_length):json_end_index].strip() # Trim whitespace and remove prompt
+        json_string = output_text[max(json_start_index + 6, user_prompt_length):json_end_index].strip()
 
-        # Validate JSON (you'll need to define a schema for your JSON structure)
+        # Remove any leading/trailing non-JSON characters (if present)
+        if not json_string.startswith("{"):
+            first_brace_index = json_string.find("{")
+            if first_brace_index != -1:
+                json_string = json_string[first_brace_index:]
+
+        if not json_string.endswith("}"):
+            last_brace_index = json_string.rfind("}")
+            if last_brace_index != -1:
+                json_string = json_string[:last_brace_index + 1]
+
+
+        # Validate JSON
         try:
             validate(instance=json.loads(json_string), schema=your_json_schema)
             return json_string
-        except ValidationError as e:
+        except jsonschema.exceptions.ValidationError as e:
             return f"Error: Invalid JSON - {e}"
 
     else: