Views
No views yet

1from transformers import AutoModelForCausalLM, AutoTokenizer
2
3model_name = "ValiantLabs/Qwen3-4B-ShiningValiant3"
4
5# load the tokenizer and the model
6tokenizer = AutoTokenizer.from_pretrained(model_name)
7model = AutoModelForCausalLM.from_pretrained(
8 model_name,
9 torch_dtype="auto",
10 device_map="auto"
11)
12
13# prepare the model input
14prompt = "Propose a novel cognitive architecture where the primary memory component is a Graph Neural Network (GNN). How would this GNN represent working, declarative, and procedural memory? How would the \"cognitive cycle\" be implemented as operations on this graph?"
15messages = [
16 {"role": "user", "content": prompt}
17]
18text = tokenizer.apply_chat_template(
19 messages,
20 tokenize=False,
21 add_generation_prompt=True,
22 enable_thinking=True # Switches between thinking and non-thinking modes. Default is True.
23)
24model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
25
26# conduct text completion
27generated_ids = model.generate(
28 **model_inputs,
29 max_new_tokens=32768
30)
31output_ids = generated_ids[0][len(model_inputs.input_ids[0]):].tolist()
32
33# parsing thinking content
34try:
35 # rindex finding 151668 (</think>)
36 index = len(output_ids) - output_ids[::-1].index(151668)
37except ValueError:
38 index = 0
39
40thinking_content = tokenizer.decode(output_ids[:index], skip_special_tokens=True).strip("\n")
41content = tokenizer.decode(output_ids[index:], skip_special_tokens=True).strip("\n")
42
43print("thinking content:", thinking_content)
44print("content:", content)