mirror of
https://github.com/wassname/prob_jsonformer.git
synced 2026-09-07 17:00:24 +08:00
48 lines
1.0 KiB
Python
48 lines
1.0 KiB
Python
from jsonformer.format import highlight_values
|
|
from jsonformer.main import Jsonformer
|
|
from transformers import (
|
|
AutoModelForCausalLM,
|
|
AutoTokenizer,
|
|
)
|
|
|
|
|
|
|
|
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,
|
|
)
|