update

jsonformer/example.py

@@ -0,0 +1,47 @@
+from jsonformer.format import highlight_values
+from jsonformer.main import Jsonformer
+from transformers import (
+    AutoModelForCausalLM,
+    AutoTokenizer,
+)
+import torch
+
+
+weather_schema = {
+    "type": "object",
+    "properties": {
+        "temperature": {"type": "number"},
+        "humidity": {
+            "type": "number",
+        },
+        "wind_speed": {
+            "type": "object",
+            "properties": {
+                "value": {"type": "number"},
+                "unit": {"type": "string"},
+            },
+        },
+    },
+}
+
+print("Loading model and tokenizer...")
+model_name = "databricks/dolly-v2-12b"
+model = AutoModelForCausalLM.from_pretrained(model_name, use_cache=True)
+tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True, use_cache=True)
+print("Loaded model and tokenizer")
+
+
+builder = Jsonformer(
+    model=model,
+    tokenizer=tokenizer,
+    json_schema=weather_schema,
+    prompt="generate the weather",
+    debug=True,
+)
+
+print("Generating...")
+output = builder()
+
+highlight_values(
+    output,
+)