@@ -3901,6 +3901,93 @@ def __call__(self, **kwargs):
39013901 return super ().__call__ (** kwargs )
39023902
39033903
3904+ class GraniteDoclingChatHandler (Llava15ChatHandler ):
3905+ """
3906+ Handler for Granite-Docling models.
3907+
3908+ Format(512x512): <loc_xmin><loc_ymin><loc_xmax><loc_ymax>Content
3909+
3910+ Note(JamePeng): The GGUF files for Model and MMPROJ should be BF16 version !!!
3911+ Since the model does not have special tokens for the start and end of an image,
3912+ it is recommended to process only one image at a time.
3913+ You can iterate through the images individually for recognition.
3914+
3915+ """
3916+ GRANITE_BOS_TOKEN = "<|start_of_role|>"
3917+ GRANITE_EOS_TOKEN = "<|end_of_text|>"
3918+ GRANITE_PAD_TOKEN = "<|end_of_text|>"
3919+ GRANITE_IMAGE_TOKEN = "<image>"
3920+
3921+ CHAT_FORMAT = (
3922+ "{%- for message in messages -%}"
3923+ "{{- '<|start_of_role|>' + message['role'] + '<|end_of_role|>' -}}"
3924+ "{%- if message['content'] is string -%}"
3925+ "{{- message['content'] -}}"
3926+ "{%- else -%}"
3927+ "{%- for part in message['content'] -%}"
3928+ "{%- if part['type'] == 'text' -%}"
3929+ "{{- part['text'] -}}"
3930+ "{%- elif part['type'] == 'image_url' -%}"
3931+ "{%- if part.image_url is string -%}"
3932+ "{{- part.image_url -}}"
3933+ "{%- else -%}"
3934+ "{{- part.image_url.url -}}"
3935+ "{%- endif -%}"
3936+ "{%- endif -%}"
3937+ "{%- endfor -%}"
3938+ "{%- endif -%}"
3939+ "{{- '<|end_of_text|>\n ' -}}"
3940+ "{%- endfor -%}"
3941+ "{%- if add_generation_prompt -%}"
3942+ "{{- '<|start_of_role|>assistant' -}}"
3943+ # Support the 'controls' parameter if present in generation arguments
3944+ "{%- if controls -%}{{- ' ' + controls | tojson() -}}{%- endif -%}"
3945+ "{{- '<|end_of_role|>' -}}"
3946+ "{%- endif -%}"
3947+ )
3948+
3949+ def __init__ (self , controls : dict = None , ** kwargs ):
3950+ """
3951+ Granite-Docling Handler
3952+ Args:
3953+ controls (dict, optional): Operational parameters passed to the assistant role.
3954+
3955+ The 'controls' parameter is used to guide the model's behavior or output format.
3956+ Common examples for 'controls' include:
3957+ - Document Parsing: {"mode": "document_parsing", "format": "json"}
3958+ """
3959+ self .controls = controls
3960+ super ().__init__ (** kwargs )
3961+
3962+ def __call__ (self , ** kwargs ):
3963+ # Inject controls into the template environment
3964+ self .extra_template_arguments ["controls" ] = self .controls
3965+ self .DEFAULT_SYSTEM_MESSAGE = None
3966+ kwargs ['stop' ] = [self .GRANITE_EOS_TOKEN ]
3967+
3968+ llama = kwargs ['llama' ]
3969+ llama .reset ()
3970+ llama ._ctx .memory_clear (True )
3971+ llama .n_tokens = 0
3972+
3973+ if hasattr (llama , 'input_ids' ):
3974+ llama .input_ids .fill (0 )
3975+
3976+ if hasattr (self , '_last_image_embed' ):
3977+ self ._last_image_embed = None
3978+ self ._last_image_hash = None
3979+
3980+ if self .verbose :
3981+ messages = kwargs .get ('messages' , [])
3982+ try :
3983+ image_count = len (self .get_image_urls (messages ))
3984+ print (f"GraniteDoclingChatHandler - Cleared state, processing { image_count } images" , file = sys .stderr )
3985+ except Exception :
3986+ print (f"GraniteDoclingChatHandler - Cleared state" , file = sys .stderr )
3987+
3988+ return super ().__call__ (** kwargs )
3989+
3990+
39043991class LFM2VLChatHandler (Llava15ChatHandler ):
39053992 LFM2VL_BOS_TOKEN = "<|startoftext|>"
39063993 LFM2VL_EOS_TOKEN = "<|im_end|>"
0 commit comments