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-GSM8k"
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 = "Natalia sold clips to 48 of her friends in April, and then she sold half as many clips in May. How many clips did Natalia sell altogether in April and May?"
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-GSM8k}
6}