1import torch
2import torch.nn as nn
3from transformers import AutoModelForCausalLM, AutoTokenizer
4from huggingface_hub import hf_hub_download
5
6class ProjectionHead(nn.Module):
7 def __init__(self, hidden_size=2048):
8 super().__init__()
9 self.mlp = nn.Sequential(
10 nn.Linear(hidden_size, hidden_size),
11 nn.GELU(),
12 nn.Linear(hidden_size, hidden_size),
13 nn.LayerNorm(hidden_size),
14 )
15 def forward(self, x):
16 return self.mlp(x)
17
18# Load model and projection head
19repo_id = "LakshyAAAgrawal/QThink-Qwen3-1.7B-Tooluse"
20model = AutoModelForCausalLM.from_pretrained(repo_id, torch_dtype=torch.bfloat16, device_map="auto")
21tokenizer = AutoTokenizer.from_pretrained(repo_id)
22proj = ProjectionHead(2048).to(model.device).to(torch.bfloat16)
23proj.load_state_dict(torch.load(
24 hf_hub_download(repo_id, "projection_head.pt"), map_location=model.device
25))
26proj.eval()
27model.eval()
28
29# Prepare input
30question = "What's the weather like in San Francisco?"
31messages = [{"role": "user", "content": question}]
32text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True, enable_thinking=False)
33inputs = tokenizer(text, return_tensors="pt").to(model.device)
34
35with torch.no_grad():
36 # Step 1: Process prompt
37 out = model(**inputs, output_hidden_states=True, use_cache=True)
38 past_kv = out.past_key_values
39 latent = out.hidden_states[-1][:, -1, :]
40
41 # Step 2: K=6 latent reasoning steps
42 mask = inputs["attention_mask"].clone()
43 for k in range(6):
44 latent = proj(latent)
45 mask = torch.cat([mask, mask.new_ones(1, 1)], dim=1)
46 out = model(inputs_embeds=latent.unsqueeze(1), attention_mask=mask,
47 past_key_values=past_kv, output_hidden_states=True, use_cache=True)
48 past_kv = out.past_key_values
49 latent = out.hidden_states[-1][:, -1, :]
50
51 # Step 3: Greedy decode answer
52 next_token = out.logits[:, -1, :].argmax(dim=-1)
53 tokens = [next_token]
54 eos_id = tokenizer.eos_token_id
55 for _ in range(2047):
56 if next_token.item() == eos_id:
57 break
58 mask = torch.cat([mask, mask.new_ones(1, 1)], dim=1)
59 out = model(input_ids=next_token.unsqueeze(0), attention_mask=mask,
60 past_key_values=past_kv, use_cache=True)
61 past_kv = out.past_key_values
62 next_token = out.logits[:, -1, :].argmax(dim=-1)
63 tokens.append(next_token)
64
65print(tokenizer.decode(torch.cat(tokens), skip_special_tokens=True))
1@misc{qthink2025,
2 title={QThink: Parallel Latent Reasoning via Per-Step Distillation of Multiple Rollouts},
3 author={Lakshya Agrawal},
4 year={2025},
5 url={https://huggingface.co/LakshyAAAgrawal/QThink-Qwen3-1.7B-Tooluse}
6}