vikhyatk
/

moondream2

Image-Text-to-Text

text-generation

Model card Files Files and versions

vikhyatk commited on Mar 13, 2024

Commit

2b705ee

·

verified ·

1 Parent(s): 010e389

Upload Moondream

Files changed (1) hide show

moondream.py +4 -2

moondream.py CHANGED Viewed

@@ -49,9 +49,11 @@ class Moondream(PreTrainedModel):
         else:
             assert prompt.count("<image>") == 1
             before, after = prompt.split("<image>")
-            embeds.append(text_emb(_tokenize(f"{before}<image>")))
             embeds.append(image_embeds.to(self.device))
-            embeds.append(text_emb(_tokenize(f"</image>{after}")))
         return torch.cat(embeds, dim=1)

         else:
             assert prompt.count("<image>") == 1
             before, after = prompt.split("<image>")
+            if len(before) > 0:
+                embeds.append(text_emb(_tokenize(before)))
             embeds.append(image_embeds.to(self.device))
+            if len(after) > 0:
+                embeds.append(text_emb(_tokenize(after)))
         return torch.cat(embeds, dim=1)