{
    "name": "string",
    "company": "string",
    "job_title": "string"
}

{
    "name": "John Smith",
    "company": "Acme Corporation",
    "job_title": "Senior Software Engineer"
}

import torch
import json
from transformers import AutoModelForCausalLM, AutoTokenizer

MODEL_PATH = "FrontiersMind/Lumma-0.6B-Extract"

tokenizer = AutoTokenizer.from_pretrained(
    MODEL_PATH,
    trust_remote_code=True
)

model = AutoModelForCausalLM.from_pretrained(
    MODEL_PATH,
    trust_remote_code=True,
    torch_dtype=torch.bfloat16,
    device_map="auto"
)

model.eval()


input_text = "John Smith joined Acme Corporation as a Senior Software Engineer."

template = {
    "name": "string",
    "company": "string",
    "job_title": "string",
}

prompt = tokenizer.apply_chat_template(
input_text=input_text,
template=template,
tokenize=False,
add_generation_prompt=True
)


inputs = tokenizer(prompt, return_tensors="pt", add_special_tokens=False).to(model.device)

with torch.inference_mode():
    out = model.generate(
        **inputs,
        max_new_tokens=512,
        do_sample=False,
        eos_token_id=tokenizer.eos_token_id,
        pad_token_id=tokenizer.pad_token_id,
    )

raw = tokenizer.decode(out[0, inputs["input_ids"].shape[1]:], skip_special_tokens=False)
pred = raw.split("<|endoftext|>")[0].strip()
print(json.dumps(json.loads(pred), indent=4))