a hM@sddlmZmZddlmZddlmZmZddlm Z m Z m Z m Z ddl mZmZddlmZernddlZGd d d e d d ZGd dde ZdgZdS))OptionalUnion) BatchFeature) ImageInputis_valid_image)MultiModalDataProcessingKwargsProcessorMixinUnpack)PreTokenizedInput TextInput)is_torch_availableNc@s&eZdZddidddddidZd S) ColQwen2ProcessorKwargspaddingZlongestZchannels_firstT)Z data_formatZdo_convert_rgbZreturn_tensorspt) text_kwargs images_kwargsZ common_kwargsN)__name__ __module__ __qualname__ _defaultsrrl/var/www/html/assistant/venv/lib/python3.9/site-packages/transformers/models/colqwen2/processing_colqwen2.pyr#srF)totalcseZdZdZddgZdZdZd"eeeedfdd Z d#e e e e ee ee feeed d d Zd$d dZeedddZd%e eeedddZe e ee feeedddZd&e dedfe dedfeede defddddZed d!ZZS)'ColQwen2Processora Constructs a ColQwen2 processor which wraps a Qwen2VLProcessor and special methods to process images and queries, as well as to compute the late-interaction retrieval score. [`ColQwen2Processor`] offers all the functionalities of [`Qwen2VLProcessor`]. See the [`~Qwen2VLProcessor.__call__`] for more information. Args: image_processor ([`Qwen2VLImageProcessor`], *optional*): The image processor is a required input. tokenizer ([`Qwen2TokenizerFast`], *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*): A string that gets tokenized and prepended to the image tokens. query_prefix (`str`, *optional*): A prefix to be used for the query. image_processor tokenizerZAutoImageProcessor)ZQwen2TokenizerZQwen2TokenizerFastN)visual_prompt_prefix query_prefixc sftj|||dt|ds dn|j|_t|ds6dn|j|_|durJd}||_|dur\d}||_dS)N) chat_template image_tokenz <|image_pad|> video_tokenz <|video_pad|>zf<|im_start|>user <|vision_start|><|image_pad|><|vision_end|>Describe the image.<|im_end|><|endoftext|>zQuery: )super__init__hasattrr!r"rr)selfrrr rrkwargs __class__rrr$Hs zColQwen2Processor.__init__)imagestextr'returncKs|jtfd|jji|}|ddd}|du}|durJ|durJtd|durb|durbtd|dur0t|r||g}nHt|trt|drn0t|trt|dtrt|ddstd|j gt |} |j fd |i|d } | d } | dur||j j d } d} t t | D]`}|j| |vrb| ||jd | | | d| |<| d7} q| |d |j| |<q|j| fddi|d}ti|| d}|d dddf|d ddd f}tt|d|}tjjjj|dd|d<|r,|d|ddkd}|d|i|S|durt|trN|g}n$t|trjt|dtsrtd|dur|jd}g}|D]}|j||}||q|j|fddi|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 Qwen2VLProcessor's [`~Qwen2VLProcessor.__call__`] method adapted for the ColQwen2 model. It cannot process both text and images at the same time. When preparing the the text(s), this method forwards the `text` and `kwargs` arguments to Qwen2TokenizerFast's [`~Qwen2TokenizerFast.__call__`]. When preparing the the image(s), this method forwards the `images` and `kwargs` arguments to Qwen2VLImageProcessor's [`~Qwen2VLImageProcessor.__call__`]. Please refer to the doctsring 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 timerzAimages must be an image, list of images or list of list of imagesr*rimage_grid_thwz<|placeholder|>return_token_type_idsF)data pixel_valuesT) batch_firstZ input_idsZtoken_type_idsilabelsz*Text must be a string or a list of strings )Z _merge_kwargsrrZ init_kwargspop ValueErrorr isinstancelistrlenr merge_sizeranger!replaceprodrtorchsplittolistnnutilsrnn pad_sequenceZ masked_fillupdatestrquery_augmentation_tokenrappend)r&r*r+ZaudioZvideosr'Z output_kwargsr-r1Z texts_docZ image_inputsr.Z merge_lengthindexiZ text_inputsZ return_dataoffsetsr3r5Z texts_queryqueryZaugmented_queryZ batch_queryrrr__call__]s- (   (      zColQwen2Processor.__call__c s|i}|durntjdi|ddp6jjfdd|D}fdd|D}|||dtfi|S)a Computes the number of placeholder tokens needed for multimodal inputs with the given sizes. Args: image_sizes (`list[list[int]]`, *optional*): The input sizes formatted as (height, width) per each image. Returns: `MultiModalData`: A `MultiModalData` object holding number of tokens per each of the provided input modalities, along with other useful data. Nrr<cs"g|]}jjg|RqSr)rZget_number_of_image_patches).0Z image_size)rr&rr sz@ColQwen2Processor._get_num_multimodal_tokens..csg|]}|dqS)r/r)rPZ num_patches)r<rrrQ)num_image_tokensnum_image_patches)rrgetrGrr<r)r&Z image_sizesr'Z vision_datarTrSr)rr<r&r_get_num_multimodal_tokenss   z,ColQwen2Processor._get_num_multimodal_tokens)r,cCs|jjS)z Return the query augmentation token. Query augmentation buffers are used as reasoning buffers during inference. )rZ pad_token)r&rrrrIsz*ColQwen2Processor.query_augmentation_token)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 ColQwen2Processor's [`ColQwen2Processor.__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!z ColQwen2Processor.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 ColQwen2Processor's [`ColQwen2Processor.__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+rW)r&r+r'rrrprocess_queries(s z!ColQwen2Processor.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)a[ Compute the late-interaction/MaxSim score (ColBERT-like) for the given multi-vector query embeddings (`qs`) and passage embeddings (`ps`). For ColQwen2, 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)r4Z padding_valuez bnd,csd->bcnsr)dimr/r0)r;r8ZdeviceZdtyper=r@rCrDrErFrJZeinsummaxsumcatto) r&r\r]r^r_r`ZscoresrLZ batch_scoresZ batch_queriesjZbatch_passagesrrrscore_retrievalJs2      "z!ColQwen2Processor.score_retrievalcCs&|jj}|jj}dd|D}||S)NcSsg|]}|dvr|qS))Zpixel_values_videosZvideo_grid_thwr)rPnamerrrrQsz7ColQwen2Processor.model_input_names..)rmodel_input_namesr)r&Ztokenizer_input_namesZimage_processor_input_namesrrrris z#ColQwen2Processor.model_input_names)NNNNN)NNNN)N)N)rZNr[)rrr__doc__ attributesZimage_processor_classZtokenizer_classrrHr$rrr r r:r rrrOrVpropertyrIrXrYintrgri __classcell__rrr(rr0sd    % &  @r)typingrrZfeature_extraction_utilsrZ image_utilsrrZprocessing_utilsrr r r Ztokenization_utils_baser r rDrr@rr__all__rrrrs   i