o }oiXf@sddlZddlmZmZddlZddlmZddlmZm Z m Z ddl m Z m Z ddlmZGdddeZGd d d eZGd d d eZGd ddeZGdddeZdS)N)ABCabstractmethod) rearrange)InterleavedSampleSimilarityInterleavedSample VQASample)ImageTextSampleMultiModalSampleConfig)loggingc@s eZdZddZeddZdS) SampleEncodercCsdS)an Initialize the SampleEncoder class. This class serves as an abstract base class for encoding samples. It provides a common interface for different types of sample encoders. Subclasses should implement the encode method to perform the actual encoding process. Parameters: None Returns: None Nselfr r k/home/ubuntu/.local/lib/python3.10/site-packages/nemo/collections/multimodal/data/energon/sample_encoder.py__init__szSampleEncoder.__init__cCtd)a; Abstract method to encode a sample. Must be implemented by subclasses. This method is responsible for encoding a given input sample into a format suitable for further processing. The encoded sample is then stored in the output_sample object. Parameters: input_sample (object): The input sample to be encoded. The type and structure of this object depend on the specific subclass. output_sample (object): The object where the encoded sample will be stored. The type and structure of this object depend on the specific subclass. Returns: None: The method does not return any value. Raises: NotImplementedError: If the method is called directly on the abstract class, it will raise this exception. Subclasses must implement this method. ,Subclasses must implement the encode method.NotImplementedErrorr input_sample output_sampler r rencode*szSampleEncoder.encodeN)__name__ __module__ __qualname__rrrr r r rr sr csjeZdZdZedffdd ZdejdejfddZd ejdejfd d Z d e d e ddfddZ Z S)BaseSampleEncodera# Base class for encoding multimodal samples, specifically for handling text and image data. This class provides basic functionality for preprocessing images, computing loss masks, and managing sample configuration. It serves as a base class for more specialized encoders. Attributes: tokenizer (Tokenizer): The HF tokenizer used for tokenizing input text. image_processor (ImageProcessor): The HF image processor used for preprocessing input images. multimodal_sample_config (MultiModalSampleConfig): Configuration for multimodal samples, including tokens and placeholders. ignore_place_holder (int): Token ID used to ignore certain tokens during loss computation. image_token (Token): Token dataclass representing image placeholders in the tokenized sequence. NcsJtt|dr|j|_n||_||_||_|j|_|j|_||_dS)a Initialize the BaseSampleEncoder. Parameters: tokenizer (Tokenizer): The tokenizer used for processing text. image_processor (ImageProcessor): The image processor used for preprocessing images. multimodal_sample_config (MultiModalSampleConfig, optional): Configuration object for multimodal samples. Defaults to MultiModalSampleConfig(). tokenizerN) superrhasattrrimage_processormultimodal_sample_configignore_place_holder image_tokenimage_tag_typerrr r!r$ __class__r rrNs   zBaseSampleEncoder.__init__imagereturncCs4|jj|dddd}t|tjsJt|d}|S)a Preprocess and reshape an input image for encoding. The function preprocesses an image using the specified image processor and reshapes it to the expected format for further processing. Parameters: image (torch.Tensor): A tensor representing the input image with dimensions (channels, height, width). Returns: torch.Tensor: A preprocessed and reshaped image tensor with dimensions (1, 1, channels, height, width). ptF)return_tensors do_rescale pixel_valueszF c h w -> 1 F c h w)r preprocess isinstancetorchTensorr)rr(r r r process_imagees  zBaseSampleEncoder.process_imagelabelscCs&tj|tjd}d|||jk<|S)a2 Compute a binary loss mask based on the provided labels. The function generates a mask where the loss is computed only for tokens that are not equal to the `ignore_place_holder` token. Parameters: labels (torch.Tensor): A tensor containing labels for which the loss mask needs to be generated. Returns: torch.Tensor: A binary mask tensor with the same shape as the input labels. The mask has ones for tokens where loss should be computed and zeros for `ignore_place_holder` tokens. dtypeg)r0onessizefloatr")rr3 loss_maskr r rcompute_loss_maskwsz#BaseSampleEncoder.compute_loss_maskrrcCr)a Abstract method to encode an input sample. Subclasses must implement this method to encode input samples into the desired format. Parameters: input_sample (ImageTextSample): The sample to be encoded. output_sample (ImageTextSample): The object to store the encoded sample. Returns: None Raises: NotImplementedError: If the method is called directly on the abstract class. rrrr r rrszBaseSampleEncoder.encode) rrr__doc__r rr0r1r2r:rr __classcell__r r r&rr?srcseZdZdZedffdd ZddefddZd ed e j fd d Z d e j ded e j fddZ dede fddZddZZS)VQASampleEncodera Encoder specifically designed for Visual Question Answering (VQA) samples. This class extends the BaseSampleEncoder to handle VQA tasks, applying a specific prompt template and computing labels and loss masks based on the VQA input. Attributes: conversation_template_config (ConversationTemplateConfig): Configuration for conversation templates used in VQA. Nct|||||j|_dS)a Initialize the VQASampleEncoder. Parameters: tokenizer (Tokenizer): The HF tokenizer used for processing text. image_processor (ImageProcessor): The HF image processor used for preprocessing images. multimodal_sample_config (MultiModalSampleConfig, optional): Configuration object for multimodal samples. Defaults to MultiModalSampleConfig(). N)rrconversation_template_configr%r&r rr zVQASampleEncoder.__init__F input_textcCstd|jd|jg}|jjr|d|jjdt|jtrZt|jtrZt t |jt |j}t |D] }||jj d|j|d||jj d|j|dq8n+t|jt rt|jt r||jj d|jd||jj d|jdntd|jjr|jj|j_n |jjdurtd |jj|d d d }|jd kr|d d}n |jdusJd|jtd||S)a Apply a conversation template to the input text for VQA. This method generates a templated prompt by combining system, user, and assistant messages. Parameters: input_text (VQASample): The VQA sample containing the context and answer. use_plain (bool, optional): Whether to use a plain format for the prompt. Defaults to False. Returns: str: The generated templated prompt as a string. z$apply_conversation_template context z answer system)rolecontentrzYVQA Sample context/answers should either be a List[str] or str. Other types not supportedNzBoth tokenizer and conversation template does not have chat template defined. Refer to https://huggingface.co/docs/transformers/main/en/chat_templatingF)tokenizeadd_generation_promptinternvlzzzNot supported image_tag_type z'apply prompt template templated_prompt )r debugcontextanswersr?rBappendr/listminlenrangerolesstr ValueError chat_templaterapply_chat_templater$replace)rrA use_plainmessages min_lengthitemplated_promptr r rapply_prompt_templates8     z&VQASampleEncoder.apply_prompt_templatepromptr)cCsdddd|jjfDd}t||}g}|D]!}||jjkr+||jjqt|dkr<||j |ddj qt j |t j d S) a Tokenize the input prompt, replacing special tokens with their IDs. This method splits the prompt into chunks based on the presence of special tokens (like ) and tokenizes each chunk. Special tokens are replaced with their corresponding token IDs. Parameters: prompt (str): The prompt string to tokenize. Returns: torch.Tensor: A tensor containing the tokenized prompt. (|css|]}t|VqdS)N)reescape).0tokenr r r sz,VQASampleEncoder.tokenize..)rFadd_special_tokensr4)joinr# token_strr`splitrLtoken_idrOextendr input_idsr0tensorlong)rr] regex_patternchunkstokenized_chunkschunkr r rrFs"   zVQASampleEncoder.tokenizetokenssamplec Csddlm}t||j}d}t|jdd}t|jt r |jn|jg}|D]H}| ||}|j j |dddd} |j | ddkrI| d d} ||| |\} } | dkrbtd |j|| ||S|| | || | <| }q&|S) a) Compute labels for the tokenized prompt based on the answers in the VQA sample. This method generates a label tensor where the tokens corresponding to the answers are marked with their token IDs, while other tokens are marked with the `ignore_place_holder` ID. Parameters: tokens (torch.Tensor): A tensor containing the tokenized prompt. sample (VQASample): The VQA sample containing the answers. Returns: torch.Tensor: A tensor containing the labels for the tokenized prompt. r)_find_pattern_indices stop_stringNFr*)rgr+rEzUnable to find a valid answer in the conversation. Details: - Messages: %s - Tokens: %s - Answer Tokens: %s - Search Start Index: %d)nemo.collections.vlm.data.utilsrvr0 ones_liker"getattrr?r/rKrMprocess_answer_strrrdecoder warning) rrtrurvr3search_start_indexstop_strrKanswer answer_tokens answer_start answer_endr r rcompute_labelss0    zVQASampleEncoder.compute_labelsrrcCs||}td|||}|||}|dd}|dd}td|td|||}||j}|j |_ |j dg|_ | d|_ ||_||_||_|S)a Encode a VQA sample into a format suitable for further processing. This method applies a prompt template, tokenizes the prompt, computes labels and a loss mask, and processes the image. The encoded sample is then stored in the output_sample object. Parameters: input_sample (VQASample): The VQA sample to be encoded. output_sample (ImageTextSample): The object to store the encoded sample. Returns: ImageTextSample: The encoded sample stored in output_sample. z/task encoder encode_sample conversation_prompt NrE8task encoder encode_sample after tokenize prompt tokens z"task encoder encode_sample labels r)r\r rIrFr contiguousr:r2r(__key__shapenum_image_tilessqueezeimagesrtr3r9)rrrconversation_promptrtr3r9processed_imager r rrCs"      zVQASampleEncoder.encodecCsd||dur dS|S)N rxr )rrrr r rr|fsz#VQASampleEncoder.process_answer_str)F)rrrr;r rrr\rRr0r1rFrrrr|r<r r r&rr=s 6=#r=csjeZdZdZedffdd ZdeejejffddZ dejdejfd d Z d e d e fd dZ ZS)InterleavedSampleEncodera Encoder for handling interleaved sequences of text and images (InterleavedSample from energon). This class extends the BaseSampleEncoder to handle interleaved samples, where the input consists of a sequence of text strings and image tensors. The text and images are processed and encoded into a format suitable for further processing. Attributes: tokenizer (Tokenizer): The tokenizer used for processing text. image_processor (ImageProcessor): The image processor used for preprocessing images. multimodal_sample_config (MultiModalSampleConfig): Configuration for multimodal samples, including tokens and placeholders. Ncst||||dS)a Initialize the InterleavedSampleEncoder. Parameters: tokenizer (Tokenizer): The tokenizer used for processing text. image_processor (ImageProcessor): The image processor used for preprocessing images. multimodal_sample_config (MultiModalSampleConfig, optional): Configuration object for multimodal samples. Defaults to MultiModalSampleConfig(). N)rrr%r&r rrxs z!InterleavedSampleEncoder.__init__r)cCsg}g}|jD]<}t|tjr!||jj||}||qt|dkr;t d|| |j |ddj qtdt|tj|tjd}t d|tj|dd }||fS) a Tokenize the input sequence and process images in an interleaved sample. This method processes a sequence that consists of text strings and image tensors. The text is tokenized, and the images are processed. The method returns a tensor of tokenized text and a concatenated tensor of processed images. Parameters: sample (InterleavedSample): The interleaved sample containing a sequence of text strings and image tensors. Returns: tuple[torch.Tensor, torch.Tensor]: A tuple containing: - A tensor with tokenized text and image token IDs. - A concatenated tensor of processed images. rz)a Initialize the SimilarityInterleavedEncoder. Parameters: tokenizer (Tokenizer): The tokenizer used for processing text. image_processor (ImageProcessor): The image processor used for preprocessing images. multimodal_sample_config (MultiModalSampleConfig, optional): Configuration object for multimodal samples. Defaults to MultiModalSampleConfig(). N)rrimage_following_textr%r&r rrr@z%SimilarityInterleavedEncoder.__init__rur)cs|j}|j}|j}g}d}ddtt||D}fdd|D}t|}tt|D]5} |t|krZ||| krZjsD|j j ||| jrU|j j |d7}q,||| q,g} |D] } | rt | dt rt | t r| dd| 7<qf| | qfg} | D]} | j j kr| | q| j| dd jqtj| tjd }td | tj|dd }||fS) a& Tokenize the input sequence and process images based on similarity indices. This method processes a sequence of text strings and images, interleaving them based on similarity indices (matched_text_indices). The text is tokenized, and the images are processed. The method returns a tensor of tokenized text and a concatenated tensor of processed images. Parameters: sample (SimilarityInterleavedSample): The sample containing a sequence of text strings and images, along with similarity indices that determine the interleaving order. Returns: tuple[torch.Tensor, torch.Tensor]: A tuple containing: - A tensor with tokenized text and image token IDs. - A concatenated tensor of processed images. rcSsg|]\}}|qSr r )rb_imgr r r sz9SimilarityInterleavedEncoder.tokenize..csg|]}|qSr )r2)rbrsr r rr srErrFrfr4zLMultimodal dataloader encode similarity interleaved sample tokenized chunks r)rtextsmatched_text_indicessortedziprPrOrrLr#rkr/rRrlrrmr0rnror rIr)rrurrrinterleaved_list image_idx sorted_imagessorted_indicestext_idxfinal_sequenceitemrrrsrtrr r rrFsB    z%SimilarityInterleavedEncoder.tokenize) rrrr;r rrrr0r1rFr<r r r&rrs  &r)r`abcrrr0einopsrmegatron.energonrrr/nemo.collections.multimodal.data.energon.configrr nemo.utilsr r rr=rrr r r rs   &]Om