bbarazandeh2026

Add working readout example script

d91ffd3 verified 21 days ago

4.72 kB

"""Minimal working example: cribl-decision structured readout.
Loads the LoRA adapter, builds a prompt in the training format,
extracts hidden-state logits for candidate tokens, and returns
a probability distribution.  No text is generated.
"""

import json
import torch
from transformers import AutoTokenizer, Qwen3_5ForConditionalGeneration
from peft import PeftModel

# ── 0.  Triton workaround (needed on some systems) ──────────────
# Uncomment the two lines below if you hit a Triton / Python.h
# compilation error on first inference:
#
# from torch._native import triton_utils
# triton_utils.deregister_op_overrides()

# ── 1.  Load model ──────────────────────────────────────────────
BASE_ID = "Qwen/Qwen3.5-4B-Base"
BASE_REV = "1001bb4d826a52d1f399e183466143f4da7b741b"
ADAPTER_ID = "cribl-ai/cribl-decision-1.0"

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
dtype = torch.bfloat16 if device.type == "cuda" else torch.float32

tokenizer = AutoTokenizer.from_pretrained(BASE_ID, revision=BASE_REV)

# IMPORTANT: use Qwen3_5ForConditionalGeneration, not AutoModelForCausalLM.
# AutoModelForCausalLM resolves to Qwen3_5ForCausalLM which silently
# ignores the LoRA adapter in transformers >= 5.19.
base = Qwen3_5ForConditionalGeneration.from_pretrained(
    BASE_ID, revision=BASE_REV, torch_dtype=dtype,
)
model = PeftModel.from_pretrained(base, ADAPTER_ID, is_trainable=False).to(device)
model.eval()

# Split into language_model (transformer) and lm_head (projection).
base_model = model.get_base_model()
language_model = base_model.model.language_model
lm_head = base_model.lm_head

# ── 2.  Label table ────────────────────────────────────────────
# A-Z for up to 26 options.  For >26, continues with two-letter
# single-token labels (AA, AB, ...) up to 255 total.
LETTERS = "ABCDEFGHIJ"


def letter_token_ids(tok, n):
    """Token IDs for the first n candidate letters."""
    return [tok.encode(LETTERS[i], add_special_tokens=False)[0] for i in range(n)]


# ── 3.  Build prompt (training format) ─────────────────────────
def build_prompt(state, question, options):
    """Build a prompt in the exact format the model was trained on.
    Format:
        Context:
        <JSON or plain-text state>
        Question: <text>
        Options:
        (A) <option 1>
        (B) <option 2>
        ...
        Answer: (
    """
    if isinstance(state, dict):
        state_str = json.dumps(state, ensure_ascii=False, separators=(", ", ": "))
    else:
        state_str = str(state)

    lines = [f"Context:\n{state_str}", "",
             f"Question: {question}", "Options:"]
    for i, opt in enumerate(options):
        lines.append(f"({LETTERS[i]}) {opt}")
    lines.append("Answer: (")
    return "\n".join(lines)


# ── 4.  Score candidates ───────────────────────────────────────
@torch.inference_mode()
def score(prompt, n_options):
    """Run one forward pass, return {option_index: probability}."""
    ids = tokenizer.encode(prompt, add_special_tokens=False)
    input_ids = torch.tensor([ids], dtype=torch.long, device=device)
    mask = torch.ones_like(input_ids)

    hidden = language_model(input_ids=input_ids, attention_mask=mask).last_hidden_state
    h = hidden[0, -1]  # hidden state at the "(" token

    cand_ids = torch.tensor(letter_token_ids(tokenizer, n_options),
                            dtype=torch.long, device=device)
    weights = lm_head.weight[cand_ids]
    logits = (weights @ h).float()
    probs = torch.softmax(logits, dim=-1).cpu().tolist()
    return probs


# ── 5.  Demo ────────────────────────────────────────────────────
if __name__ == "__main__":
    # Example: route a support ticket
    state = {"message": "My card was charged twice for the same purchase."}
    question = "Which department should handle this?"
    options = ["billing", "technical support", "sales"]

    prompt = build_prompt(state, question, options)
    print("=== Prompt ===")
    print(prompt)
    print()

    probs = score(prompt, len(options))
    print("=== Probabilities ===")
    for i, (opt, p) in enumerate(zip(options, probs)):
        print(f"  ({LETTERS[i]}) {opt:20s}  {p:.4f}")
    print(f"\n  Pick: ({LETTERS[probs.index(max(probs))]}) {options[probs.index(max(probs))]}")

	"""Minimal working example: cribl-decision structured readout.

	Loads the LoRA adapter, builds a prompt in the training format,
	extracts hidden-state logits for candidate tokens, and returns
	a probability distribution. No text is generated.
	"""

	import json
	import torch
	from transformers import AutoTokenizer, Qwen3_5ForConditionalGeneration
	from peft import PeftModel

	# ── 0. Triton workaround (needed on some systems) ──────────────
	# Uncomment the two lines below if you hit a Triton / Python.h
	# compilation error on first inference:
	#
	# from torch._native import triton_utils
	# triton_utils.deregister_op_overrides()

	# ── 1. Load model ──────────────────────────────────────────────
	BASE_ID = "Qwen/Qwen3.5-4B-Base"
	BASE_REV = "1001bb4d826a52d1f399e183466143f4da7b741b"
	ADAPTER_ID = "cribl-ai/cribl-decision-1.0"

	device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
	dtype = torch.bfloat16 if device.type == "cuda" else torch.float32

	tokenizer = AutoTokenizer.from_pretrained(BASE_ID, revision=BASE_REV)

	# IMPORTANT: use Qwen3_5ForConditionalGeneration, not AutoModelForCausalLM.
	# AutoModelForCausalLM resolves to Qwen3_5ForCausalLM which silently
	# ignores the LoRA adapter in transformers >= 5.19.
	base = Qwen3_5ForConditionalGeneration.from_pretrained(
	BASE_ID, revision=BASE_REV, torch_dtype=dtype,
	)
	model = PeftModel.from_pretrained(base, ADAPTER_ID, is_trainable=False).to(device)
	model.eval()

	# Split into language_model (transformer) and lm_head (projection).
	base_model = model.get_base_model()
	language_model = base_model.model.language_model
	lm_head = base_model.lm_head

	# ── 2. Label table ────────────────────────────────────────────
	# A-Z for up to 26 options. For >26, continues with two-letter
	# single-token labels (AA, AB, ...) up to 255 total.
	LETTERS = "ABCDEFGHIJ"


	def letter_token_ids(tok, n):
	"""Token IDs for the first n candidate letters."""
	return [tok.encode(LETTERS[i], add_special_tokens=False)[0] for i in range(n)]


	# ── 3. Build prompt (training format) ─────────────────────────
	def build_prompt(state, question, options):
	"""Build a prompt in the exact format the model was trained on.

	Format:
	Context:
	<JSON or plain-text state>

	Question: <text>
	Options:
	(A) <option 1>
	(B) <option 2>
	...
	Answer: (
	"""
	if isinstance(state, dict):
	state_str = json.dumps(state, ensure_ascii=False, separators=(", ", ": "))
	else:
	state_str = str(state)

	lines = [f"Context:\n{state_str}", "",
	f"Question: {question}", "Options:"]
	for i, opt in enumerate(options):
	lines.append(f"({LETTERS[i]}) {opt}")
	lines.append("Answer: (")
	return "\n".join(lines)


	# ── 4. Score candidates ───────────────────────────────────────
	@torch.inference_mode()
	def score(prompt, n_options):
	"""Run one forward pass, return {option_index: probability}."""
	ids = tokenizer.encode(prompt, add_special_tokens=False)
	input_ids = torch.tensor([ids], dtype=torch.long, device=device)
	mask = torch.ones_like(input_ids)

	hidden = language_model(input_ids=input_ids, attention_mask=mask).last_hidden_state
	h = hidden[0, -1] # hidden state at the "(" token

	cand_ids = torch.tensor(letter_token_ids(tokenizer, n_options),
	dtype=torch.long, device=device)
	weights = lm_head.weight[cand_ids]
	logits = (weights @ h).float()
	probs = torch.softmax(logits, dim=-1).cpu().tolist()
	return probs


	# ── 5. Demo ────────────────────────────────────────────────────
	if __name__ == "__main__":
	# Example: route a support ticket
	state = {"message": "My card was charged twice for the same purchase."}
	question = "Which department should handle this?"
	options = ["billing", "technical support", "sales"]

	prompt = build_prompt(state, question, options)
	print("=== Prompt ===")
	print(prompt)
	print()

	probs = score(prompt, len(options))
	print("=== Probabilities ===")
	for i, (opt, p) in enumerate(zip(options, probs)):
	print(f" ({LETTERS[i]}) {opt:20s} {p:.4f}")
	print(f"\n Pick: ({LETTERS[probs.index(max(probs))]}) {options[probs.index(max(probs))]}")