S^(hLHddlmZmZmZmZddlmZddlmZm Z m Z ddl m Z m Z mZddlmZmZmZddlmZerddlZGd d e d Zd ZedDcgc] }d|dd c}edDcgc] }d|dd c}zZdZGdde ZdgZycc}wcc}w))ClassVarListOptionalUnion) BatchFeature) ImageInputis_valid_imagemake_flat_list_of_images)ProcessingKwargsProcessorMixinUnpack) AddedTokenPreTokenizedInput TextInput)is_torch_availableNc&eZdZddidddddidZy ) ColPaliProcessorKwargspaddinglongestchannels_firstT) data_formatdo_convert_rgbreturn_tensorspt) text_kwargs images_kwargs common_kwargsN)__name__ __module__ __qualname__ _defaultsl/var/www/skyplay_api_hub/venv/lib/python3.12/site-packages/transformers/models/colpali/processing_colpali.pyrr$s, y ," +D1 Ir$rF)totalziz4>z3c ||z|z||dS)aZ Builds a string from the input prompt and image tokens. For example, for the call: build_string_from_input( prompt="Prefix str" bos_token="", image_seq_len=3, image_token="", ) The output will be: "Initial str" Args: prompt (`List[Union[str, ImageInput]]`): The input prompt. bos_token (`str`): The beginning of sentence token. image_seq_len (`int`): The length of the image sequence. image_token (`str`): The image token. num_images (`int`): Number of images in the prompt.  r#prompt bos_token image_seq_len image_token num_imagess r%build_string_from_inputr15s$&M)J6 7 {6(" MMr$c ~eZdZUdZddgZdgZdZdZdZe e e d<d Z e e e d < d fd Z d!d ed eeeeeeefdeedefdZdZdZedZede fdZ d"d edeedefdZd eeeefdeedefdZ d#dededfdededfdededdede fddf dZ xZ!S)$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. image_processor tokenizer chat_template)SiglipImageProcessorSiglipImageProcessorFast)GemmaTokenizerGemmaTokenizerFastzDescribe the image.visual_prompt_prefixz Question: query_prefixc | td| tdt|ds td|j|_t|dsCttdd}d |gi}|j ||j t|_n|j|_|jtd|_ d|_ t|5||| y) Nz)You need to specify an `image_processor`.z"You need to specify a `tokenizer`.image_seq_lengthz;Image processor is missing an `image_seq_length` attribute.r/FT) normalizedspecialadditional_special_tokens)r6) ValueErrorhasattrr>r IMAGE_TOKENadd_special_tokensconvert_tokens_to_idsimage_token_id add_tokens EXTRA_TOKENS add_bos_token add_eos_tokensuper__init__)selfr4r5r6kwargsr/ tokens_to_add __class__s r%rMzColPaliProcessor.__init__ds  "HI I  AB B(:;Z[ [ / @ @y-0$[UDQK8;-HM  ( ( 7"+"A"A+"ND "+":":D \*"' "'  )=Qr$imagestextrOreturnc |jtfd|jji|}|dj dd}|dnd}| | t d| | t d|t |r|g}n^t|trt |d rn?t|tr$t|d trt |d d s t d |jgt|z} |D cgc]} | jd }} t| |D cgc]R\} } t| |jj|jt t| tr t| nd T} } } t#|}|j$|fi|dd}|dj'dd|ddxx|jz cc<|j| fddi|d}i|d|i}|r.|dj)|dd k(d}|j+d|it-|S|t|t.r|g}n.t|trt|d t.s t d||j0dz}g}|D]C}|jj|j2z|z}||z }|dz }|j5|E|dj'dd|dd<|j|fddi|d}|Sycc} wcc} } w)a Main method to prepare for the model either (1) one or several texts, either (2) one or several image(s). This method is 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`. tokenizer_init_kwargsrsuffixNTFz&Either text or images must be providedz5Only one of text or images can be processed at a timerzAimages must be an image, list of images or list of list of imagesRGBr+r pixel_values max_lengthreturn_token_type_ids input_idstoken_type_idsilabels)dataz*Text must be a string or a list of strings r*2) _merge_kwargsrr5 init_kwargspoprBr isinstancelistr;lenconvertzipr1r-r>rDr r4get masked_fillupdaterstrquery_augmentation_tokenr<append)rNrRrSaudiovideosrO output_kwargsrWr\ texts_docimager, image_list input_stringsrZinputs return_datar_ texts_queryquery batch_querys r%__call__zColPaliProcessor.__call__sqZ+** " "&.."<"<  }-11(DA(.(: ?uemmE*?F?+.i*@ 'FJ(!"nn66"&"7"7 +2S>SS:#T^^&+ .F CVB^\BK$ ,88@P9QUV9VX\]""Hf#56[1 1  $$v t,DGS1I !MNN~66;%'K *0043D3DDuL ""5)  * :G}9U9Y9YZfhj9kM- ( 6($..&+ .K  1C@ s 4K;ALc:|jj|i|S)z This method forwards all its arguments to GemmaTokenizerFast's [`~PreTrainedTokenizer.batch_decode`]. Please refer to the docstring of this method for more information. )r5 batch_decoderNargsrOs r%rzColPaliProcessor.batch_decodes +t~~**D;F;;r$c:|jj|i|S)z This method forwards all its arguments to GemmaTokenizerFast's [`~PreTrainedTokenizer.decode`]. Please refer to the docstring of this method for more information. )r5decoders r%rzColPaliProcessor.decodes %t~~$$d5f55r$c|jj}|jj}ttj ||zSN)r5model_input_namesr4rgdictfromkeys)rNtokenizer_input_namesimage_processor_input_namess r%rz"ColPaliProcessor.model_input_namess? $ @ @&*&:&:&L&L#DMM"7:U"UVWWr$c.|jjS)z Return the query augmentation token. Query augmentation buffers are used as reasoning buffers during inference. )r5 pad_token)rNs r%roz)ColPaliProcessor.query_augmentation_tokens~~'''r$c *|jdd|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 SiglipImageProcessor's [`~SiglipImageProcessor.__call__`]. 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`. rRr#r})rNrRrOs r%process_imageszColPaliProcessor.process_imagessBt}}5F5f55r$c *|jdd|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 LlamaTokenizerFast's [`~LlamaTokenizerFast.__call__`]. 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`). rSr#r)rNrSrOs r%process_queriesz ColPaliProcessor.process_queriesAs@t}}1$1&11r$query_embeddingsz torch.Tensorpassage_embeddings batch_size output_dtypez torch.dtype output_devicez torch.devicec t|dk(r tdt|dk(r td|dj|djk7r td|dj|djk7r td||dj}g}t dt||D]%}g}t j jjj||||zdd} t dt||D]} t j jjj|| | |zdd} |jt jd| | jd djd |jt j|d j|j|(t j|d 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 dtypeT) batch_first padding_valuez bnd,csd->bcnsr)dimrY)rhrBdevicedtyperangetorchnnutilsrnn pad_sequencerpeinsummaxsumcatto) rNrrrrrscoresi batch_scores batch_queriesjbatch_passagess r%score_retrievalz ColPaliProcessor.score_retrievalcs@  A %23 3 ! "a '34 4 A  % %);A)>)E)E ENO O A  $ $(:1(=(C(C CLM M  +A.44L%'q#./< ]A/1L!HHNN..;; Q^4$VW<M1c"45zB !&!3!3!@!@&q1z>:\]"A"##LL-PTTYZT[\]^bbghbi   MM%))La8;;LILL][ \ ]yyQ''r$)NNN)NNNNr)r(Ncpu)"rr r!__doc__ attributes valid_kwargsimage_processor_classtokenizer_classr;rrn__annotations__r<rMr rrrrrrrr}rrpropertyrrorrintrr __classcell__)rQs@r%r3r3Ks $[1J#$LP>O*?(3-?".L(3-. R@"^b }}I0$y/4HYCZZ[} /0 } }~<6XX (#(("!6!6/0!6  !6F 2ItI./ 2/0 2  2L0449 >(^0D DE>(".$~2F"FG>( >( }- >( ^S01 >( >(r$r3)typingrrrrfeature_extraction_utilsr image_utilsr r r processing_utilsr r rtokenization_utils_baserrrrrrrrDrrIr1r3__all__)rs0r%rs.324OOHHOO' -U  ).t5A$qgQ5RWX[R\8]Q4#wa8]] N,V(~V(r  e 68]s B5B