import json, sys, torchfrom huggingface_hub import snapshot_downloadfrom safetensors.torch import load_filefrom transformers import AutoModelForCausalLM, AutoTokenizer path = snapshot_download("FluidInference/clef-text-0.6b")sys.path.insert(0, path)from joint_schema_model import ClefModel, JointSchemaHead, collate_records, encode_record tokenizer = AutoTokenizer.from_pretrained(path)backbone = AutoModelForCausalLM.from_pretrained(path, dtype=torch.float32)head = JointSchemaHead(**json.loads(open(f"{path}/joint_head_config.json").read()))head.load_state_dict(load_file(f"{path}/joint_head.safetensors"))model = ClefModel(backbone, head).eval() record = {"state": "Our checkout is down and customers can't pay.", "questions": {"team": {"type": "choice", "instructions": "Which team should handle this?", "criteria": {"billing": "Payments", "engineering": "Bugs or outages"}}, "urgency": {"type": "score", "criteria": ["Low", "Normal", "High", "Critical"]}}}encoded = encode_record(tokenizer, record)with torch.inference_mode(): logits = model(collate_records([encoded], tokenizer.pad_token_id, torch.device("cpu")))[0]for question, question_logits in zip(encoded.questions, logits): print(question.question_id, dict(zip(question.option_ids, question_logits.softmax(-1).tolist())))