from outformer import Jsonformer, highlight_values
from transformers import AutoModelForCausalLM, AutoTokenizer

# Initialize model and tokenizer
model_name = "Qwen/Qwen3-1.7B"
model = AutoModelForCausalLM.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)

# Create Jsonformer instance
jsonformer = Jsonformer(model, tokenizer, max_tokens_string=30)

# Define your JSON schema
json_schema = {
    "type": "object",
    "properties": {
        "brand": {
            "type": "string",
            "description": "Brand of the product",
        },
        "model": {
            "type": "string",
            "description": "Model of the product",
        },
        "product_type": {
            "type": "string",
            "description": "Type of the product",
        },
        "gender": {
            "type": "string",
            "enum": ["Female", "Male", "Unisex"],
        },
        "color": {
            "type": "string",
            "description": "Color of the product if specified, otherwise return 'Unknown'",
        },
        "material": {
            "type": "string",
            "description": "Material of the product if specified, otherwise return 'Unknown'",
        },
        "features": {
            "type": "array",
            "minItems": 3,
            "items": {
                "type": "string",
                "description": "Features of the product that may be relevant for the customer. Extract as much as possible.",
            },
        },
    },
}

# Your input prompt
prompt = """
Extract key information from the product description:

adidas Men's Powerlift.3 Cross-Trainer Shoes
A powerful shoe with lockdown fit. Made with an extra-wide design that allows the foot to spread, these men's lifting/weight-training shoes pair a snug-fitting upper with a wide midfoot strap for extra support. A high-density die-cut wedge midsole keeps you close to the ground.
100% Synthetic leather
Imported
Rubber sole
Removable Insole
"""

# Generate structured output
generated_data = jsonformer.generate(schema=json_schema, prompt=prompt)

# Highlight generated values
highlight_values(generated_data)

The code above will generate a structured JSON output and display it with highlighted values. Here's what you'll get:

{
    "brand": "Adidas",
    "model": "Powerlift.3 Cross-Trainer Shoes",
    "product_type": "Cross-Trainer Shoes",
    "gender": "Male",
    "color": "Unknown",
    "material": "Synthetic leather",
    "features": [
        "Lockdown fit",
        "Extra-wide design",
        "High-density die-cut wedge midsole",
    ],
}

When using highlight_values(), the output will be displayed in your terminal with the generated values highlighted in color, making it easy to distinguish between the structure and the generated content.