jasonmoo commited on
Commit
e257644
·
verified ·
1 Parent(s): e86382f

fix: preserve spatial_toks when reasoning=true

Browse files

Fixes bug where spatial_refs parameter fails when reasoning mode is enabled.

Root cause: Line 670 was replacing prompt_tokens with only the suffix,
discarding the spatial_toks that were added earlier. This caused a shape
mismatch error in _prefill_prompt() when trying to encode spatial refs.

Error was:
RuntimeError: shape mismatch: value tensor of shape [2, 2048]
cannot be broadcast to indexing result of shape [0, 2048]

Fix: Convert tensor back to list, append suffix, maintain 2D structure:
Before: prompt_tokens = [self.config.tokenizer.templates["query"]["suffix"]]
After: prompt_tokens = [prompt_tokens[0].tolist() + self.config.tokenizer.templates["query"]["suffix"]]

After _generate_reasoning() returns, prompt_tokens is a torch.Tensor, not a list.
Cannot use += operator on Tensor with list. Must convert to list first, append
suffix, then wrap in list to maintain 2D batch structure.

Tested: spatial_refs now work with reasoning=true without errors.

Files changed (1) hide show
  1. moondream.py +1 -1
moondream.py CHANGED
@@ -667,7 +667,7 @@ class MoondreamModel(nn.Module):
667
  pos, reasoning_text, reasoning_grounding = self._generate_reasoning(
668
  prompt_tokens, pos, settings, spatial_refs, attn_mask=attn_mask
669
  )
670
- prompt_tokens = [self.config.tokenizer.templates["query"]["suffix"]]
671
  reasoning_dict = {
672
  "reasoning": {"text": reasoning_text, "grounding": reasoning_grounding}
673
  }
 
667
  pos, reasoning_text, reasoning_grounding = self._generate_reasoning(
668
  prompt_tokens, pos, settings, spatial_refs, attn_mask=attn_mask
669
  )
670
+ prompt_tokens = [prompt_tokens[0].tolist() + self.config.tokenizer.templates["query"]["suffix"]]
671
  reasoning_dict = {
672
  "reasoning": {"text": reasoning_text, "grounding": reasoning_grounding}
673
  }