import json, math, torchfrom transformers import AutoProcessor, AutoModelForMultimodalLM model_id = "mohamedlotfy50/sifr-0.8b-v3"processor = AutoProcessor.from_pretrained(model_id)tok = processor.tokenizermodel = AutoModelForMultimodalLM.from_pretrained( model_id, dtype=torch.bfloat16, device_map="cuda:0").eval() SYSTEM = ("You are a decision engine. Read the state, then answer the question " "by replying with exactly one option key from the list. Do not explain.") def render(state, q): s = state if isinstance(state, str) else json.dumps( state, ensure_ascii=False, separators=(",", ":")) body = ["State:", s if s else "(empty)", "", "Question:", q["instructions"], "", "Options:"] for key, desc in q["criteria"].items(): body.append(f"{key}: {key if desc is None else desc}") body += ["", "Reply with exactly one option key."] prompt = tok.apply_chat_template( [{"role": "system", "content": SYSTEM}, {"role": "user", "content": "\n".join(body)}], tokenize=False, add_generation_prompt=True) return tok(prompt, add_special_tokens=False)["input_ids"] @torch.inference_mode()def score(prompt_ids, criteria): """Length-normalized P(key | prompt) for every option.""" keys, sums, lens = list(criteria), [], [] for k in keys: kids = tok(k, add_special_tokens=False)["input_ids"] out = model(input_ids=torch.tensor([prompt_ids + kids], device=model.device), logits_to_keep=len(kids) + 1) lsm = torch.log_softmax(out.logits[0].float(), dim=-1) sums.append(sum(lsm[t, kt].item() for t, kt in enumerate(kids))) lens.append(len(kids)) norm = [s / l for s, l in zip(sums, lens)] top = max(norm); w = [math.exp(x - top) for x in norm]; z = sum(w) return {k: wi / z for k, wi in zip(keys, w)} state = {"task": "sentiment classification"}q = {"instructions": "Classify the sentiment of this review:\n" "One star. The screen arrived cracked and support never replied.", "criteria": {"A": "negative", "B": "positive"}}probs = score(render(state, q), q["criteria"])# -> {'A': 0.755, 'B': 0.245}print("choice:", max(probs, key=probs.get)) # -> A