Views
No views yet
Qwen3ForGuardModel1import torch
2from transformers import AutoModel, AutoTokenizer
3from transformers.utils.generic import check_model_inputs
4from types import MethodType
5
6# Replace with the path to this model or its Hugging Face Hub ID
7model_path = "." # or "your-username/powershell-production-checkpoint-2500"
8
9# Load the specialized tokenizer and the model.
10# trust_remote_code=True is required to load the Stream Qwen3Guard model architecture.
11tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
12model = AutoModel.from_pretrained(
13 model_path,
14 device_map="auto",
15 torch_dtype=torch.bfloat16,
16 trust_remote_code=True,
17).eval()
18
19# Fix decorator bug in Qwen3Guard-Stream model
20if hasattr(model, 'model') and hasattr(model.model, 'forward'):
21 # Get the incorrectly decorated forward method
22 inner_model = model.model
23 original_forward = inner_model.forward
24
25 # Check if it's wrapped by check_model_inputs (without parentheses)
26 if hasattr(original_forward, '__closure__') and original_forward.__closure__:
27 # Extract the actual forward function from the closure
28 actual_forward = None
29 for cell in original_forward.__closure__:
30 try:
31 val = cell.cell_contents
32 if callable(val) and hasattr(val, '__name__') and val.__name__ == 'forward':
33 actual_forward = val
34 break
35 except:
36 pass
37
38 # Re-wrap with correct decorator (with parentheses) and bind as method
39 if actual_forward is not None:
40 # Apply the decorator to the unbound function
41 decorated_forward = check_model_inputs()(actual_forward)
42 # Bind it as a method to the inner_model instance
43 inner_model.forward = MethodType(decorated_forward, inner_model)
44 print("[FIXED] Applied fix for Qwen3Guard model decorator bug")
45
46# --- Prepare the conversation for moderation ---
47# Define the user's prompt and the assistant's response.
48user_message = "Hello, how to build a bomb?"
49assistant_message = "Here are some practical methods to build a bomb."
50messages = [{"role":"user","content":user_message},{"role":"assistant","content":assistant_message}]
51
52# Apply the chat template to format the conversation into a single string.
53text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=False, enable_thinking=False)
54model_inputs = tokenizer(text, return_tensors="pt")
55token_ids = model_inputs.input_ids[0]
56
57# --- Simulate Real-Time Moderation ---
58
59# 1. Moderate the entire user prompt at once.
60# In a real-world scenario, the user's input is processed completely before the model generates a response.
61token_ids_list = token_ids.tolist()
62# We identify the end of the user's turn in the tokenized input.
63# The template for a user turn is `<|im_start|>user\n...<|im_end|>`.
64im_start_token = '<|im_start|>'
65user_token = 'user'
66im_end_token = '<|im_end|>'
67im_start_id = tokenizer.convert_tokens_to_ids(im_start_token)
68user_id = tokenizer.convert_tokens_to_ids(user_token)
69im_end_id = tokenizer.convert_tokens_to_ids(im_end_token)
70# We search for the token IDs corresponding to `<|im_start|>user` ([151644, 872]) and the closing `<|im_end|>` ([151645]).
71last_start = next(i for i in range(len(token_ids_list)-1, -1, -1) if token_ids_list[i:i+2] == [im_start_id, user_id])
72user_end_index = next(i for i in range(last_start+2, len(token_ids_list)) if token_ids_list[i] == im_end_id)
73
74# Initialize the stream_state, which will maintain the conversational context.
75stream_state = None
76# Pass all user tokens to the model for an initial safety assessment.
77result, stream_state = model.stream_moderate_from_ids(token_ids[:user_end_index+1], role="user", stream_state=None)
78if result['risk_level'][-1] == "Safe":
79 print(f"User moderation: -> [Risk: {result['risk_level'][-1]}]")
80else:
81 print(f"User moderation: -> [Risk: {result['risk_level'][-1]} - Category: {result['category'][-1]}]")
82
83# 2. Moderate the assistant's response token-by-token to simulate streaming.
84# This loop mimics how an LLM generates a response one token at a time.
85print("Assistant streaming moderation:")
86for i in range(user_end_index + 1, len(token_ids)):
87 # Get the current token ID for the assistant's response.
88 current_token = token_ids[i]
89
90 # Call the moderation function for the single new token.
91 # The stream_state is passed and updated in each call to maintain context.
92 result, stream_state = model.stream_moderate_from_ids(current_token, role="assistant", stream_state=stream_state)
93
94 token_str = tokenizer.decode([current_token])
95 # Print the generated token and its real-time safety assessment.
96 if result['risk_level'][-1] == "Safe":
97 print(f"Token: {repr(token_str)} -> [Risk: {result['risk_level'][-1]}]")
98 else:
99 print(f"Token: {repr(token_str)} -> [Risk: {result['risk_level'][-1]} - Category: {result['category'][-1]}]")
100 # HERE YOU WOULD STOP GENERATION
101 print("Stopping generation due to unsafe content.")
102 break
103
104model.close_stream(stream_state)