Upload processing_kosmos2.py

processing_kosmos2.py

@@ -331,10 +331,12 @@ class Kosmos2Processor(ProcessorMixin):
         """
         return self.tokenizer.decode(*args, **kwargs)
 
-    def post_processor_generation(self, text):
+    def post_processor_generation(self, text, cleanup_and_extract=True):
 
         caption = text.split("</image>")[-1]
-        return clean_text_and_extract_entities_with_bboxes(caption)
+        if cleanup_and_extract:
+            return clean_text_and_extract_entities_with_bboxes(caption)
+        return caption
 
     @property
     # Copied from transformers.models.blip.processing_blip.BlipProcessor.model_input_names