vikhyatk commited on
Commit
2b705ee
1 Parent(s): 010e389

Upload Moondream

Browse files
Files changed (1) hide show
  1. moondream.py +4 -2
moondream.py CHANGED
@@ -49,9 +49,11 @@ class Moondream(PreTrainedModel):
49
  else:
50
  assert prompt.count("<image>") == 1
51
  before, after = prompt.split("<image>")
52
- embeds.append(text_emb(_tokenize(f"{before}<image>")))
 
53
  embeds.append(image_embeds.to(self.device))
54
- embeds.append(text_emb(_tokenize(f"</image>{after}")))
 
55
 
56
  return torch.cat(embeds, dim=1)
57
 
 
49
  else:
50
  assert prompt.count("<image>") == 1
51
  before, after = prompt.split("<image>")
52
+ if len(before) > 0:
53
+ embeds.append(text_emb(_tokenize(before)))
54
  embeds.append(image_embeds.to(self.device))
55
+ if len(after) > 0:
56
+ embeds.append(text_emb(_tokenize(after)))
57
 
58
  return torch.cat(embeds, dim=1)
59