Upload processing_kosmos2.py
Browse files- processing_kosmos2.py +4 -2
processing_kosmos2.py
CHANGED
@@ -331,10 +331,12 @@ class Kosmos2Processor(ProcessorMixin):
|
|
331 |
"""
|
332 |
return self.tokenizer.decode(*args, **kwargs)
|
333 |
|
334 |
-
def post_processor_generation(self, text):
|
335 |
|
336 |
caption = text.split("</image>")[-1]
|
337 |
-
|
|
|
|
|
338 |
|
339 |
@property
|
340 |
# Copied from transformers.models.blip.processing_blip.BlipProcessor.model_input_names
|
|
|
331 |
"""
|
332 |
return self.tokenizer.decode(*args, **kwargs)
|
333 |
|
334 |
+
def post_processor_generation(self, text, cleanup_and_extract=True):
|
335 |
|
336 |
caption = text.split("</image>")[-1]
|
337 |
+
if cleanup_and_extract:
|
338 |
+
return clean_text_and_extract_entities_with_bboxes(caption)
|
339 |
+
return caption
|
340 |
|
341 |
@property
|
342 |
# Copied from transformers.models.blip.processing_blip.BlipProcessor.model_input_names
|