import torchfrom 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()Ā # Simulate streaming moderationuser_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]Ā # Initialize stream statestream_state = NoneĀ # Process tokens one by one (simulating streaming)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)