a h=@sddlmZmZddlmZmZmZddlmZddl m Z m Z ddl m Z mZddlmZmZddlmZmZer~dd lZeeZGd d d e d d ZGdddeZdgZd S))OptionalUnion) IMAGE_TOKENPaliGemmaProcessorbuild_string_from_input) BatchFeature) ImageInputmake_flat_list_of_images)ProcessingKwargsUnpack)PreTokenizedInput TextInput)is_torch_availableloggingNc@s&eZdZddidddddidZd S) ColPaliProcessorKwargspaddingZlongestZchannels_firstT)Z data_formatZdo_convert_rgbZreturn_tensorspt) text_kwargs images_kwargsZ common_kwargsN)__name__ __module__ __qualname__ _defaultsrrg/var/www/html/assistant/venv/lib/python3.9/site-packages/transformers/models/colpali/modular_colpali.pyr#srF)totalcseZdZdZdeedfdd Zeedd d Zdee e e e e e e fe eed d d Zdee eedddZe e e e fe eedddZde de dfe de dfeede defddddZZS) ColPaliProcessora Constructs a ColPali processor which wraps a PaliGemmaProcessor and special methods to process images and queries, as well as to compute the late-interaction retrieval score. [`ColPaliProcessor`] offers all the functionalities of [`PaliGemmaProcessor`]. See the [`~PaliGemmaProcessor.__call__`] for more information. Args: image_processor ([`SiglipImageProcessor`], *optional*): The image processor is a required input. tokenizer ([`LlamaTokenizerFast`], *optional*): The tokenizer is a required input. chat_template (`str`, *optional*): A Jinja template which will be used to convert lists of messages in a chat into a tokenizable string. visual_prompt_prefix (`str`, *optional*, defaults to `"Describe the image."`): A string that gets tokenized and prepended to the image tokens. query_prefix (`str`, *optional*, defaults to `"Question: "`): A prefix to be used for the query. NDescribe the image. Question: )visual_prompt_prefix query_prefixcs"tj|||d||_||_dS)N)image_processor tokenizer chat_template)super__init__r r!)selfr"r#r$r r! __class__rrr&EszColPaliProcessor.__init__)returncCs|jjS)z Return the query augmentation token. Query augmentation buffers are used as reasoning buffers during inference. )r#Z pad_tokenr'rrrquery_augmentation_tokenQsz)ColPaliProcessor.query_augmentation_token)imagestextkwargsr*c sjtfdjji|}|ddd}|du}|durJ|durJtd|durb|durbtd|durZj|}t|}j gt |} dd|D}fd dt | |D} j|fi|d d } |d d ddur|dd j 7<j| fd di|d} i| d | i} |rP| d| ddkd}| d|it| dS|durt|trx|g}n$t|trt|dtstd|durjd}g}|D]*}jjj||d}||q|d d d|dd <j|fd di|d}|SdS)a Main method to prepare for the model either (1) one or several texts, either (2) one or several image(s). This method is a custom wrapper around the PaliGemmaProcessor's [`~PaliGemmaProcessor.__call__`] method adapted for the ColPali model. It cannot process both text and images at the same time. When preparing the text(s), this method forwards the `text` and `kwargs` arguments to LlamaTokenizerFast's [`~LlamaTokenizerFast.__call__`]. When preparing the image(s), this method forwards the `images` and `kwargs` arguments to SiglipImageProcessor's [`~SiglipImageProcessor.__call__`]. Please refer to the docstring of the above two methods for more information. Args: images (`PIL.Image.Image`, `np.ndarray`, `torch.Tensor`, `list[PIL.Image.Image]`, `list[np.ndarray]`, `list[torch.Tensor]`): The image or batch of images to be prepared. Each image can be a PIL image, NumPy array or PyTorch tensor. In case of a NumPy array/PyTorch tensor, each image should be of shape (C, H, W), where C is a number of channels, H and W are image height and width. text (`str`, `list[str]`, `list[list[str]]`): The sequence or batch of sequences to be encoded. Each sequence can be a string or a list of strings (pretokenized string). If the sequences are provided as list of strings (pretokenized), you must set `is_split_into_words=True` (to lift the ambiguity with a batch of sequences). return_tensors (`str` or [`~utils.TensorType`], *optional*): If set, will return tensors of a particular framework. Acceptable values are: - `'tf'`: Return TensorFlow `tf.constant` objects. - `'pt'`: Return PyTorch `torch.Tensor` objects. - `'np'`: Return NumPy `np.ndarray` objects. - `'jax'`: Return JAX `jnp.ndarray` objects. Returns: [`BatchFeature`]: A [`BatchFeature`] with the following fields: - **input_ids** -- List of token ids to be fed to a model. - **attention_mask** -- List of indices specifying which tokens should be attended to by the model (when `return_attention_mask=True` or if *"attention_mask"* is in `self.model_input_names` and if `text` is not `None`). - **pixel_values** -- Pixel values to be fed to a model. Returned when `images` is not `None`. Ztokenizer_init_kwargsrsuffixNz&Either text or images must be providedz5Only one of text or images can be processed at a timecSsg|]}|dqS)RGB)convert).0imagerrr z-ColPaliProcessor.__call__..c s:g|]2\}}t|jjjtt|tr.t|nddqS))prompt bos_tokenZ image_seq_lenZ image_tokenZ num_images)rr#r9image_seq_lengthr isinstancelistlen)r3r8Z image_listr+rrr5sr pixel_values max_lengthreturn_token_type_idsFZ input_idsZtoken_type_idsrilabels)dataz*Text must be a string or a list of strings  2)Z _merge_kwargsrr#Z init_kwargspop ValueErrorr"Z fetch_imagesr r r=zipgetr:Z masked_fillupdaterr;strr<r,r9r!append)r'r-r.ZaudioZvideosr/Z output_kwargsr0r@Z texts_docZ input_stringsr>inputsZ return_datarAZ texts_queryqueryZ batch_queryrr+r__call__Zsp-         zColPaliProcessor.__call__)r-r/r*cKs|jfd|i|S)a Prepare for the model one or several image(s). This method is a wrapper around the `__call__` method of the ColPaliProcessor's [`ColPaliProcessor.__call__`]. This method forwards the `images` and `kwargs` arguments to the image processor. Args: images (`PIL.Image.Image`, `np.ndarray`, `torch.Tensor`, `list[PIL.Image.Image]`, `list[np.ndarray]`, `list[torch.Tensor]`): The image or batch of images to be prepared. Each image can be a PIL image, NumPy array or PyTorch tensor. In case of a NumPy array/PyTorch tensor, each image should be of shape (C, H, W), where C is a number of channels, H and W are image height and width. return_tensors (`str` or [`~utils.TensorType`], *optional*): If set, will return tensors of a particular framework. Acceptable values are: - `'tf'`: Return TensorFlow `tf.constant` objects. - `'pt'`: Return PyTorch `torch.Tensor` objects. - `'np'`: Return NumPy `np.ndarray` objects. - `'jax'`: Return JAX `jnp.ndarray` objects. Returns: [`BatchFeature`]: A [`BatchFeature`] with the following fields: - **input_ids** -- List of token ids to be fed to a model. - **attention_mask** -- List of indices specifying which tokens should be attended to by the model (when `return_attention_mask=True` or if *"attention_mask"* is in `self.model_input_names` and if `text` is not `None`). - **pixel_values** -- Pixel values to be fed to a model. Returned when `images` is not `None`. r-rO)r'r-r/rrrprocess_imagess!zColPaliProcessor.process_images)r.r/r*cKs|jfd|i|S)a Prepare for the model one or several texts. This method is a wrapper around the `__call__` method of the ColPaliProcessor's [`ColPaliProcessor.__call__`]. This method forwards the `text` and `kwargs` arguments to the tokenizer. Args: text (`str`, `list[str]`, `list[list[str]]`): The sequence or batch of sequences to be encoded. Each sequence can be a string or a list of strings (pretokenized string). If the sequences are provided as list of strings (pretokenized), you must set `is_split_into_words=True` (to lift the ambiguity with a batch of sequences). return_tensors (`str` or [`~utils.TensorType`], *optional*): If set, will return tensors of a particular framework. Acceptable values are: - `'tf'`: Return TensorFlow `tf.constant` objects. - `'pt'`: Return PyTorch `torch.Tensor` objects. - `'np'`: Return NumPy `np.ndarray` objects. - `'jax'`: Return JAX `jnp.ndarray` objects. Returns: [`BatchFeature`]: A [`BatchFeature`] with the following fields: - **input_ids** -- List of token ids to be fed to a model. - **attention_mask** -- List of indices specifying which tokens should be attended to by the model (when `return_attention_mask=True` or if *"attention_mask"* is in `self.model_input_names` and if `text` is not `None`). r.rP)r'r.r/rrrprocess_queriess z ColPaliProcessor.process_queriescpuz torch.Tensorz torch.dtypez torch.device)query_embeddingspassage_embeddings batch_size output_dtype output_devicer*c Cs@t|dkrtdt|dkr(td|dj|djkrDtd|dj|djkr`td|durr|dj}g}tdt||D]}g}tjjjj ||||ddd} tdt||D]N} tjjjj || | |ddd} | t d | | j d d dj d d q| tj|d d ||qtj|dd S)aZ Compute the late-interaction/MaxSim score (ColBERT-like) for the given multi-vector query embeddings (`qs`) and passage embeddings (`ps`). For ColPali, a passage is the image of a document page. Because the embedding tensors are multi-vector and can thus have different shapes, they should be fed as: (1) a list of tensors, where the i-th tensor is of shape (sequence_length_i, embedding_dim) (2) a single tensor of shape (n_passages, max_sequence_length, embedding_dim) -> usually obtained by padding the list of tensors. Args: query_embeddings (`Union[torch.Tensor, list[torch.Tensor]`): Query embeddings. passage_embeddings (`Union[torch.Tensor, list[torch.Tensor]`): Passage embeddings. batch_size (`int`, *optional*, defaults to 128): Batch size for computing scores. output_dtype (`torch.dtype`, *optional*, defaults to `torch.float32`): The dtype of the output tensor. If `None`, the dtype of the input embeddings is used. output_device (`torch.device` or `str`, *optional*, defaults to "cpu"): The device of the output tensor. Returns: `torch.Tensor`: A tensor of shape `(n_queries, n_passages)` containing the scores. The score tensor is saved on the "cpu" device. rzNo queries providedzNo passages providedz/Queries and passages must be on the same devicez-Queries and passages must have the same dtypeNT)Z batch_firstZ padding_valuez bnd,csd->bcnsr)dimr7)r=rGZdeviceZdtyperangetorchnnutilsZrnnZ pad_sequencerLZeinsummaxsumcatto) r'rUrVrWrXrYZscoresiZ batch_scoresZ batch_queriesjZbatch_passagesrrrscore_retrievals2      "z ColPaliProcessor.score_retrieval)NNNrr)NNNN)N)rSNrT)rrr__doc__rKr&propertyr,r rrr r<r rrrOrQrRintrrf __classcell__rrr(rr0sV   y % & r)typingrrZ2transformers.models.paligemma.processing_paligemmarrrZfeature_extraction_utilsrZ image_utilsr r Zprocessing_utilsr r Ztokenization_utils_baser rr_rrr]Z get_loggerrloggerrr__all__rrrrs   *