Upload moondream.py

moondream.py

@@ -87,6 +87,7 @@ class Moondream(PreTrainedModel):
 
         return tokenizer.batch_decode(output_ids, skip_special_tokens=True)
 
+    """
     def answer_question(
         self,
         image_embeds,
@@ -105,14 +106,32 @@ class Moondream(PreTrainedModel):
             **kwargs,
         )[0]
         cleaned_answer = answer.strip()
-        print("DONE")
 
         # Use the result_queue to pass the result if it is provided
         if result_queue:
             result_queue.put(cleaned_answer)
         else:
             return cleaned_answer
+    """
+    def answer_question(
+        self,
+        image_embeds,
+        question,
+        tokenizer,
+        chat_history="",
+        result_queue=None,
+        **kwargs,
+    ):
+        prompt = f"<image>\n\n{chat_history}Question: {question}\n\nAnswer:"
+        streamer = TextStreamer(tokenizer)
+        output_ids = self.text_model.generate(
+            inputs_embeds=self.input_embeds(prompt, image_embeds, tokenizer),
+            streamer=streamer,
+            **kwargs,
+        )
 
+        for output_id in output_ids:
+            yield tokenizer.decode(output_id, skip_special_tokens=True)
     def batch_answer(
         self,
         images,