import torch
from transformers import AutoModel, AutoTokenizer
model_path = "zenlm/zen-guard-stream"
tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
model = AutoModel.from_pretrained(
model_path,
device_map="auto",
torch_dtype=torch.bfloat16,
trust_remote_code=True,
).eval()
user_message = "Hello, how are you?"
assistant_message = "I'm doing well, thank you for asking!"
messages = [
{"role": "user", "content": user_message},
{"role": "assistant", "content": assistant_message}
]
text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=False)
inputs = tokenizer(text, return_tensors="pt")
token_ids = inputs.input_ids[0]
stream_state = None
for i, token_id in enumerate(token_ids):
result, stream_state = model.stream_moderate_from_ids(
token_id,
role="assistant" if i > len(user_message) else "user",
stream_state=stream_state
)
token_str = tokenizer.decode([token_id])
risk = result['risk_level'][-1]
print(f"Token: {repr(token_str)} → Risk: {risk}")
model.close_stream(stream_state)