import pickle
from huggingface_hub import hf_hub_download

# load probe
path = hf_hub_download(
repo_id="jammjelly/role-probes",
filename="probes/gpt-oss-20b/layer_12/probe.pkl"
)

probe = pickle.load(open(path, "rb"))

## get relevant (token-level) hidden states ##

# obtain (token-level) role predictions
probs = probe.predict_proba(layer_hs)