o }oi @s ddlZddlZddlZddlZddlZddlZddlmZddlm Z m Z m Z m Z m Z mZddlZddlZddlmmZddlZddlmZddlmZddlmZddlmZmZddlm Z m!Z!ddl"m#m$m%m&m'Z(dd l)m*Z*dd l"m+Z+m,Z,m-Z-m.Z.m/Z/m0Z0m1Z1m2Z2m3Z3m4Z4m5Z5dd l6m7Z7d Z8d Z9zddl:Z:Wne;yeZ>dZ?Wn e@eAfydZ?YnwGdddZBGdddZCdeeDe eDfde deEdeEdejFf ddZGddZHdEdeId eId!eEd"eJde f d#d$ZKdFd&d'ZLdeIde fd(d)ZM dEdeId*eJde fd+d,ZNdeIde fd-d.ZOdeIde fd/d0ZPdeIde fd1d2ZQdeIde fd3d4ZRdeIde fd5d6ZSde fd7d8ZTd9d:ZUGd;d<dd>eVZWeGd?d@d@eXZYdGde fdAdBZZGdCdDdDeZ[dS)HN) dataclass)AnyDictListSequenceTupleUnion) rearrange) DictConfig)Image)Datasetdefault_collate)CLIPImageProcessorSiglipImageProcessor)image_transform) DEFAULT_BOS_TOKENDEFAULT_EOS_TOKENDEFAULT_IM_END_TOKENDEFAULT_IM_START_TOKENDEFAULT_IMAGE_PATCH_TOKENDEFAULT_IMAGE_TOKENDEFAULT_LABELS_TOKENDEFAULT_PAD_TOKENDEFAULT_VID_END_TOKENDEFAULT_VID_START_TOKENDEFAULT_VIDEO_TOKEN)get_ltor_masks_and_position_idsz;The package `decord` was not installed in this environment.)IndexedDatasetTFc@s(eZdZdZddZddZddZdS) TarOrFolderImageLoadera A class for loading images from a tar archive or a regular folder. This class provides functionality to open and read images from either a tar archive (.tar file) or a standard directory with image files. It builds an index of images if the source is a tar archive for efficient access. Attributes: image_folder (str): The path to the tar archive or image folder. tar_index (dict): A dictionary that maps file names to their tarfile member objects if the image source is a tar archive. Methods: __init__(self, image_folder): Initializes the loader with the specified image folder. build_index(self): Builds an index of image file names and their corresponding tarfile member objects for a tar archive. open_image(self, file_name): Opens and returns an image by its file name. The image is returned as an RGB PIL Image object. cCs(||_i|_|jdr|dSdSN.tar) image_folder tar_indexendswith build_index)selfr#r(f/home/ubuntu/.local/lib/python3.10/site-packages/nemo/collections/multimodal/data/neva/neva_dataset.py__init__Ys   zTarOrFolderImageLoader.__init__cCNt|jd}|D]}||j|j<q WddS1s wYdSNr)tarfileopenr# getmembersr$namer'tarmemberr(r(r)r&_  "z"TarOrFolderImageLoader.build_indexcCs|jdrt|jd'}|j|}|r,||}t|}| |WdSWddS1s7wYdStt j |j|}| |S)Nr"r-) rEr%r.r/r$r7r8decord VideoReaderflatten_framesr:r;r<)r'r=r3r4r>capr(r(r) open_videos        z!TarOrFolderVideoLoader.open_videocs>|jddkrd}t|dS|jddkr.td}t|dS|jddkrCd}t|dS|jd dkrdg}D]}|}t|d}||qN|Stt|jd }tj dtd |t d }fd d |D}t||jd kr||dt||jd ks|S)Nsplice_single_framefirstrr6middlelastr num_framesrdtypecs$g|]}t|dqS)r6)r fromarrayasnumpyr9).0irJr(r) s$z9TarOrFolderVideoLoader.flatten_frames..) rFrUr rTr9lenappendminnplinspaceint)r'rJframeframes rgb_frameimgrQindicesr(rXr)rIs.   z%TarOrFolderVideoLoader.flatten_framesN)r@rArBrCr*r&rKrIr(r(r(r)rDps  rDtexts tokenizercontext_lengthadd_extra_tokenreturnc s|dks |dks Jdd}t|tr|g}d}fdd|D}tdd|D}t|||}tjt|||tjd }t|D]!\}} t| ||krV| d ||} t | ||d t| f<qB|rj|d}|S) a Returns the tokenized representation of given input string(s). If the list of tokens exceeds the context length plus the number of extra tokens, it gets truncated. If it's smaller, it gets padded with zeros. Parameters ---------- texts : Union[str, List[str]] An input string or a list of input strings to tokenize. tokenizer : Any A tokenizer to be used for tokenization. context_length : int The context length to be used for the output tensor. add_extra_token : int Number of extra tokens to add, should be either 0 or 1. Returns ------- torch.LongTensor A two-dimensional tensor containing the resulting tokens, shape = [number of input strings, context_length + add_extra_token]. rrz*`add_extra_token` should be either 0 or 1.FTcg|]}|qSr() text_to_ids)rVtrfr(r)rYztokenize..cSsg|]}t|qSr()rZrVtokenr(r(r)rYrRN) isinstancestrmaxr\torchzerosrZlong enumeratetensor) rerfrgrh texts_is_strtokensmax_lenresultrWrpr(rmr)tokenizes  r~csfdd|DS)a Returns the token id for a given token. Parameters ---------- tokenizer : nemo tokenizer A tokenizer to be used for tokenization. tokens : list A list of tokens to get the token id for. Returns ------- List The token ids. crjr() token_to_idrormr(r)rYrnz"get_tokens_ids..r()rfr{r(rmr)get_tokens_idssrsourcesmultimodal_cfg cur_token_len use_plaincCs6|d}|d}|d}|}|dkrt}n |dkrt}n|S|s"|S|d} |ddkr0|d }|} |dkr:| | 9} |d rEt|| } nt|| d } t|| t|} |dkr |d d r |ddsmtd|dd} | d} }| dkrt|dd| }||}n|}| |}|ddd}t|}t|}t|}|dkr t|t|}}tt }}|d r||||g|}|||| |g7}d |} n,|||d|g}|||||g|d7}|||| d|g7}d |} || |} n|d r||} n||d } || |} |D]v}|d}|dr`||ddvs8J|dd |d |dd<|t jjt jjdd|dd|dd<|rt||ddvsnJ||dd<|ddkr| || }||d<q"|D] }|d || |d<qq"|S) a Preprocesses multimodal sources based on the provided configuration. This function modifies the sources for multimodal data processing. It checks if the data is multimodal and adjusts the token lengths accordingly. It also handles the start and end tokens for images and replaces image tokens in conversations. Parameters: - sources (dict): A dictionary containing the multimodal sources to be processed. - multimodal_cfg (dict): A configuration dictionary specifying various options for multimodal processing. It includes keys like 'is_multimodal', 'use_im_start_end', and 'sep_image_conv_front'. - cur_token_len (int): The current length of tokens to be considered for image processing. - use_plain (bool, optional): A boolean flag to use plain image token replacement without additional processing. Defaults to False. Returns: - dict: The processed sources dictionary after applying multimodal preprocessing steps. is_multimodal model_type media_typeimagevideorQmm_mlp_adapter_typemlp_downsampleuse_im_start_endrOuse_litaFlitaNzLITA config is missinglita_video_archrtemporal_all_resolution sample_framesvisual_token_formatv1im_vid_start_endr conversationssep_image_conv_frontvalue: conv_template interleaved)rrrrrr7 ValueErrorr\rrr<replacestripconversation_libdefault_conversationseproles)rrrrrrrimage_token_len default_tokenrQ num_patches replace_tokenrnum_temporal_tokensnum_spatial_tokensr num_tokensr media_start media_end image_patch image_start image_end vid_startvid_endreplace_token_listsource conversationupdated_conversationturnr(r(r)preprocess_multimodals                     rsquarec Cst|ts t|trj|dkr;t|jt|j}}||}d\}}tt|||}|j|ddd|iddd}|S|d kr]d d } | |td d |j D}|j|dddd}|S|j|dddd}|S|dksrJd||}|S)NkeepiptF shortest_edgereturn_tensorsdo_center_cropsize pixel_valuesrpadcS~|j\}}||kr |S||kr't|j||f|}||d||df|St|j||f|}||||ddf|SNrrOrr newmodepastepil_imgbackground_colorwidthheightr}r(r(r) expand2square} z$process_image..expand2squarecs|] }t|dVqdSNr_rVxr(r(r) z process_image..rrCNeMo image transform with setting `image_aspect_ratio` to `square`.) rrrrrtrr\r_ preprocesstuple image_mean) processorrimage_aspect_ratiomax_hwmin_hw aspect_ratior|min_lenrrr(r(r) process_imageps0  rcCsZtj}|jd|jdd}g}t|D]J\}}|d}||dd|jdkr1|dd}g|_t|D] \}} || d} | |j|dksPJ||| | dq8||q| d } t ||| d | d } | } d }d }t || D]z\}}||}||ddg|dd}d}t|D]T\}}|dkrnK||}t|dkrn>|d|7<|dkrt||}t||d}nt|||}t|||d}t||||<||7}qt||d<q| r| ddddf} | ddddf} ntj| ddd} t| dddf<t| | dS)a  Preprocesses sources for the LLaMA 3 model configuration. The function applies prompt templates and tokenizes the conversations according to the LLaMA 2 model specifications. It involves special handling of tokens, masking of labels, and adjustments based on configuration settings. Parameters: - sources (dict): A dictionary of sources containing conversations to be processed. - tokenizer: The tokenizer to be used for processing the text. - cfg: Configuration settings for preprocessing, including context length and additional tokens. Returns: - Dict: A dictionary containing tokenized and labeled data suitable for the LLaMA 2 model. This includes tokens, labels, and any special processing as defined in the configuration. rrhumangptrfromNrOrrhrgrerfrgrhz/<|start_header_id|>assistant<|end_header_id|> z*<|start_header_id|>user<|end_header_id|> rrshiftsdimsr{labels)rconv_llava_llama_3copyrrxmessagesappend_messager[ get_promptr7r~clonedetachzipsplitr<rZrk IGNORE_INDEX contiguousrurolldict)rrfcfgconvrrrWrjsentencerolerhr{rr round_seprtargetroundscur_lenrouparts round_leninstruction_lenr(r(r)preprocess_llama_3sf          r  is_mistralcCs|rtj}ntj}|jd|jdd}g}t|D]J\}}|d}||dd|jdkr9|dd}g|_t|D] \} } || d} | |j| dksXJ||| | dq@|| q| d } t ||| d | d } t d }|t tg}t||\}}}d| | |k<d| | |k<d| | |k<| }|rd }nd}t||D]s\}}||j}d}t|D]\\}}|dkrnS||}t|dkrnF|d|7<t|||j}|rt||dd}n t||dd}|dkr |d8}n|d7}t||||<||7}qt||d<q| r>| ddddf} |ddddf}ntj|ddd}t|dddf<t| |dS)a  Preprocesses sources for the LLaMA 2 model configuration. The function applies prompt templates and tokenizes the conversations according to the LLaMA 2 model specifications. It involves special handling of tokens, masking of labels, and adjustments based on configuration settings. Parameters: - sources (dict): A dictionary of sources containing conversations to be processed. - tokenizer: The tokenizer to be used for processing the text. - cfg: Configuration settings for preprocessing, including context length and additional tokens. Returns: - Dict: A dictionary containing tokenized and labeled data suitable for the LLaMA 2 model. This includes tokens, labels, and any special processing as defined in the configuration. rrrrrNrOrrhrgrllama_2z[/INST]z[/INST] rrrr)r conv_mistralrconv_llava_llama_2rrxrrr[rr7r~rrrrrrrrsep2rZrkrrrurr)rrfrr rrrrWrrrrrhr{image_patch_tokenDEFAULT_TOKENS img_patch_idbos_ideos_idrrrrrrrrrrr(r(r)preprocess_llama_2sz                rcsZ tj}g}t|D]_\}}|d}d}t|D]5\}}|ddkr>|jd|d<|d} ||d| |ds=d}q|jd|d<||d|dq|} |rf| d rf| d td  d } | | q | d } t ||| d | d} | } dd}t|| D]o\}}|}|d dgfdd|dd D}t|ttt|ksJd}t|D]5\}}|dkrn,||}t|dkrnt||d|}t||}t||||<||7}qt||d <q| r| d d d df} | d d dd f} ntj| ddd} t| d d df<t| | dS)a Preprocess sources for Yi-1.5 34b model configuration. The function applies prompt templates and tokenizes the conversations according to the Yi-1.5 34b model specifications. It involves special handling of tokens, masking of labels, and adjustments based on configuration settings. This template works with the following tokenizer configs: - model.tokenizer.library='huggingface' - model.tokenizer.type='01-ai/Yi-1.5-34B' - model.tokenizer.additional_special_tokens='{additional_special_tokens: ["", "", "", "", "", ""]}' At inference time, add end string to stop sampling: - inference.end_strings='["<|im_end|>"]' Parameters: - sources (dict): A dictionary of sources containing conversations to be processed. - tokenizer: The tokenizer to be used for processing the text. - cfg: Configuration settings for preprocessing, including context length and additional tokens. Returns: - Dict: A dictionary containing tokenized and labeled data suitable for the LLaMA 2 model. This includes tokens, labels, and any special processing as defined in the configuration. rFrOrrrTrz <|im_end|>N rhrgrz<|im_start|>user z<|im_start|>assistant csg|]}|qSr(r(rrr(r)rYrqz%preprocess_yi_34b..rrrr)r conv_yi_34brrxrrrr%rZr[r7r~rrrrr<summaprkrrrurr)rrfrrrrWrstrip_end_for_inferencerrcontextrhr{rrrrrrrrrrr(rr)preprocess_yi_34b`sn      .   rcCs~tj}|jd|jdd}g}t|D]J\}}|d}||dd|jdkr1|dd}g|_t|D] \}} || d} | |j|dksPJ||| | dq8||q| d } t ||| d | d } t d } | t t g}t||\}}}d| | |k<d| | |k<d| | |k<| }|j|jdd }t||D]c\}}||j}d}t|D]L\}}|dkrnC||}t|dkrn6|d|7<t|||j}t||dd}|dkr|d8}|d8}t||||<||7}qt||d<q| r)| ddddf} |ddddf}ntj|ddd}t|dddf<t| |dS)a Preprocesses sources for the Vicuna V1 model configuration. Similar to `preprocess_llama_2`, this function applies prompt templates and performs tokenization, but it is tailored for the Vicuna V1 model. It includes specific handling for token translations, label masking, and tokenizer configuration. Parameters: - sources (dict): A dictionary of sources containing conversations to be processed. - tokenizer: The tokenizer to be used for processing the text. - cfg: Configuration settings for preprocessing, which may include context length and additional tokens. Returns: - Dict: A dictionary containing the processed data, including tokens and labels, formatted for the Vicuna V1 model. rrrrrNrOrrhrgrr rrrrr)rconv_vicuna_v1rrrxrrr[rr7r~rrrrrrrrrrrZrkrrrurr)rrfrrrrrWrrrrrhr{rrrrrrrrrrrrrrrr(r(r) preprocess_v1sn             rcCs$g}|D] }||dq|d}t|||d|d}|d}t|}t|} t|} || | tg} t|| \} } }}d||| k<| }t ||| k<t |||k<t ||dk<t |||k<|r||ddddf }|ddd df }nt j |ddd }t |dddf<t||d S) zDtokenize the interleaved prompt and mask the text part of the promptrrhrgrrrNrrrr)r[r7r~rrrrrrrrrrurr)rrfrrrrhr{rrimage_start_tokenimage_end_tokenrr img_start_id img_end_idpad_idrr(r(r)preprocess_interleaved_prompt!s>        r$c Cs tj}g}|D]m}g|_|d|j|_d}t|dD]E\}}|ddkrS|jd|d<d|vr8d|d<t|dd |d } | |d| |d sRd }q|jd |d<| |d|d q| } |rr| d d } | | q |d} t |||d| d} | } |j|jdd }ttd}t|| D]\}}||j}|j|ddg}tdt|dD]}| |j|||dqd }t|D]I\}}|dkrn@||}t|dkrn3||d}|r|nd}t||d ||}t|||j}t||||<||7}qt||d<q| rG| ddddf} | ddddf} ntj| ddd} t| dddf<t| | dS)! Preprocess a given set of conversational sources using nvgpt conversation template This function processes conversations by first ensuring the conversation starts with a 'human' role, then tokenizes the conversations, applies specific token replacements, and finally masks labels for training purposes. Parameters: - sources: A dictionary containing conversational data. Expected format is a dict of conversations, where each conversation is a list of messages, and each message is a dict with 'from' (role) and 'value' (message text). - tokenizer: A tokenizer from the Hugging Face Transformers library used for tokenizing the conversations. - cfg: Configuration settings which include 'add_extra_token' (bool) to determine if an extra token should be added to the tokenized output, and 'context_length' for specifying the tokenization context length. Returns: - Dict: A dictionary containing two keys: - 'tokens': A tensor of tokenized conversation data. - 'labels': A tensor of labels for the conversation data, used for training models. Labels are masked based on the conversation structure. Note: - The function includes specific token replacements (e.g., DEFAULT_IMAGE_PATCH_TOKEN, , ) and masking techniques for labels. - It is designed to work with conversational data where messages alternate between a 'human' and a 'gpt' role. - The function asserts that each message in a conversation alternates between the defined roles and skips messages not starting with the 'human' role. systemFrrOrrlabelzjquality:4,toxicity:0,humor:0,creativity:0,helpfulness:4,correctness:4,coherence:4,complexity:4,verbosity:4rrTr rhrgrz quality:.* Nrrrr) r conv_nvgptrrr7r&rxrrrrrstripr[r~rrrrecompilerrr<rangerZsearchgrouprkrrrurr)rrfrrrrrrWrrrrhr{rrlabels_str_regexprrr re_roundsconv_idxrrrmatch labels_strrrr(r(r)preprocess_nvgptOs~          r6c Cs tj}g}|D]w}g|_|d|j|_d}t|dD]F\}}|ddkrT|jd|d<d|vrAt|dd|d } n|d } | |d| |d sSd }q|jd |d<| |d|d q| } |r|| d r|| d t d  d} | | q |d} t|||d| d} | } |j|jdd}t|| D]\}}||j}|j|d dg}tdt |dD]}| |j|||dqd }t|D]U\}}|dkrnL||}t |dkrn?tttd|d}t ||d ||r|dd |nd}t |||j}t||||<||7}qt||d <q| rU| d d d df} | d d dd f} ntj| ddd} t| d d df<t| | dS)r%r&FrrOrrr'rrTrr(Nrhrgrr)rz.*?\nrrr) r conv_nv_dporrr7r&rxrrrrr%rZr[r~rrrrrr<r.r,r/escaperkendrrrurr)rrfrrrrrrWrrrrhr{rrrrrr2r3rrr labels_matchrrr(r(r)preprocess_nv_dpos~          . r;c Csg}|D]!}|d}t|dksJ|dd|ddd}||q|d}t|||d|d }|}t||D]\} }|d}t||dd} t| d | <qA|rv|d d d d f }|d d dd f }nt j |d d d }t|d d d f<t ||d S)aV Preprocesses plain text sources (no template) for tokenization and label generation. This function concatenates conversations with an end signal, tokenizes them, and prepares labels for training. It handles sources with a specific structure (expecting two elements in 'conversations') and includes the option to add an extra token as specified in the configuration. The function also applies masking to the labels. Parameters: - sources: A list of source dictionaries. Each source dictionary should have a key 'conversations' containing a list of conversation parts. - tokenizer: The tokenizer to be used for converting text to tokens. - cfg: Configuration dictionary which may include 'context_length' and 'add_extra_token' settings. Returns: - Dict: A dictionary containing tokenized data and corresponding labels. This includes 'tokens' which are the tokenized representations of the conversations, and 'labels' which are used for training the model. The labels have specific indices masked with IGNORE_INDEX as per the preprocessing logic. rrOrrrrrhrgrNrrr) rZr[r7r~rrrrkrrrurr) rrfrrrrrhr{rr tokenized_lenr(r(r)preprocess_plain1s6   r=cCs|jdvr t||j|jS|jdkrt||j|jS|jdkr't||j|jS|jdkr4t||j|jS|jdkrAt||j|jS|jdkrPt||j|jddS|jd kr]t||j|jS|jd krjt ||j|jS|jd krwt ||j|jSt d |jd )N)nvgpt nv_steerlmnv_dporr llama_3mistralT)r yi_34bplainrzConversation template `z` is not supported in Neva now.) rr6rfrr;rrr rr=r$r)r'rr(r(r)preprocess_conversationsks&         rEcsNeZdZdZdededeffdd ZddZd eee j ffd d Z Z S) LazySupervisedDataset#Dataset for supervised fine-tuning. data_pathrrFcstt||dur&t|d }t|}Wdn1s wYng}td||_||_ ||_ |d|_ |d|_ |d|_ |d|_|j rRt|j nd|_|j rat|j ||_dSd|_dS)Nr-z%Formatting inputs...Skip in lazy moderr#rEimage_processor)superrFr*r/jsonloadloggingwarningrflist_data_dictrrr#rErr image_loaderrD video_loader)r'rHrfrrFfilerO __class__r(r)r*s"       "zLazySupervisedDataset.__init__cCs t|jSN)rZrOr'r(r(r)__len__s zLazySupervisedDataset.__len__ricst|tjr t|}j|}t|tr|g}t|dks!Jdd|dvrtj|dts>j|dgj|d<g}j|dD]$}j|}|dur\t d|dt j |j d}||qGtg}|rt|}j d}|djd|}|djd |} j d d kr|d dkr|d7}| d dkr| d7} || } tt|j | jd kd }nd|dvrtj|dts݈j|dgj|d<g} j|dD]} j| } | durt d| dtj ts tj trhj ddkr.expand2squarecs(g|]}|tddjjDqS)csrrrrr(r(r)rrz?LazySupervisedDataset.__getitem__...)rrr)rVr`rr'r(r)rYsz5LazySupervisedDataset.__getitem__..rrrrr{rrrrr crop_sizer)rRdimr)&rrr]integerr_rOrZlistrPr?rMrNrrrr[rurystackshaperrdeepcopyrrQrKrrrtrr\rrErr[MAX_NUM_IMAGESrvfloatcat)r'rWrimages image_filer media_tensorsrXheight_num_patcheswidth_num_patchesrvideos video_filerarrrr|rr data_dictr[ padding_size zero_paddingr(rYr) __getitem__s                         z!LazySupervisedDataset.__getitem__) r@rArBrCrsrr*rWrruTensorrp __classcell__r(r(rSr)rFs  rFcs.eZdZdZdededeffdd ZZS) NevaDatasetrGrHrrFc s4|drtt|||||dS|drtt|d|||td|ddkr|d}t|dD]Y}t|}g|d<|dD]A}t d |d } | D])} | d  d d } t j|| } t j| sstd| qQ|d| qQt dt|d |d <qE|j|q6dSdStd|d)Nz.jsonz.jsonlz)Loading image inputs from SteerLM Datasetrrr#r-rzzFormatting of z is not supported in Neva.)r%rJrsr*rMrNr/rKloadsr,finditerr0rr:r;r<isfiler[subrrOr) r'rHrfrrFr#linerecordrmatchesr4 image_name image_pathrSr(r)r*,s0       zNevaDataset.__init__)r@rArBrCrsrr*rrr(r(rSr)rs)s"rsc@sDeZdZUdZeed<ejed<dee de e e j ffddZ dS) DataCollatorForSupervisedDatasetz,Collate examples for supervised fine-tuning. model_cfgrf instancesricCs d|dv}tdd|D}|dddd}|D]A}||djd}t|dd|fdd|d<t|d d|fdd |d <|r\|dd |kr\t|dt|gfd|d<q|rtd d|D}td d|D}|D]C}||djd}t|dd|fd||d<|d } || jd} tj| g| jddR| j| j d} tj| | fdd|d <qst |} |j } |j }| d}| d }|j dd }|d kr| d }n|dkr| d}ntd||r>| d}g}|D])}|gtdt|dD]}||d||}|d tt|qqt|}tj|tj|j d}tj|tj|j d}nt|| j|j ddddd\}}}d||d k<d||d k<d||d k<|durit|d krtt|d}n |dkr~t|d}||||||d} |r|| d<| S)N cu_seqlensrcs|] }|djdVqdS)r{rNrarVinstancer(r(r)rWz.r@r{constantrrcsr)rrNrrr(r(r)rarcsr)rrNrrr(r(r)rbrr)rSdevicer\rrzUnsupported media type eod_mask_lossF)data eod_tokenrreset_attention_maskreset_position_idsgzb T c h w -> b T 1 c h wzb T F c h w -> b T F c h w)r{rattention_mask loss_mask position_idsmedia)rtraFrrure IntTensorrvrSrr rfrrr7rr[r.rZextendr_ LongTensoronesrrdrwrrNotImplementedErrorr )r'rpacked_sequencer|rpad_len max_len_cu max_len_image pad_len_curnum_pad pad_tensorbatchrfrr{rrrrr cu_seqlenindseqlenrrr(r(r)__call__Us  (             z)DataCollatorForSupervisedDataset.__call__N)r@rArBrCr __annotations__ transformersPreTrainedTokenizerrrrsrurqrr(r(r(r)r~Ns  $r~c CsX|j}|j}d}t|ddrd}|jdd}|dur|n|j}t||td"id|jd |j d |j d d d |d d d |jjj d|d|ddd|ddd|j dt|jddd|d|d|j d|ddd|dddt|jdddt|jdid|jddt|dd|dd|dddd } t| | d!S)#z5Make dataset and collator for supervised fine-tuning.rno_seqlen_plus_one_input_tokensFrr[rrNrrrr>rrXr#rErrrIrhrgrrrQrrrrlinearrLsep_token_between_frames)rLrQr)rfrHrrF) train_dataset eval_datasetr()rmm_cfggetattrvision_encoderr7rHrsrrrllmrXrencoder_seq_length) rfrIreach_file_from_pathrFrrhr[rHrr(r(r)make_supervised_data_modulesl                rc@sHeZdZd dedeeeffddZddZdeee j ffd d Z d S) NevaPackedSeqDatatsetrrHr[cCst||_||_dSrU)rdsr[)r'rHr[r(r(r)r*s  zNevaPackedSeqDatatset.__init__cCst|jjdS)Nr)rZrdocument_indicesrVr(r(r)rWszNevaPackedSeqDatatset.__len__ricCsl|jj|}t|j|t|j|dt|j|dt|j|djddg|jRd}|S)NrrOr)r)rr{rr)rrrurr FloatTensorreshaper[)r'rW doc_startrr(r(r)rps $z!NevaPackedSeqDatatset.__getitem__N)r) r@rArBrsrr_r*rWrrurqrpr(r(r(r)rsr)F)rrU)\rrKrMr:r,r. dataclassesrtypingrrrrrrnumpyr]rutorch.nn.functionalnn functionalrreinopsr omegaconfr PILr torch.utils.datar r rr2nemo.collections.multimodal.data.neva.conversation collections multimodalrnevarrAnemo.collections.multimodal.data.clip.augmentations.augmentationsrrrrrrrrrrrr2nemo.collections.nlp.modules.common.megatron.utilsrrcrrG ExceptionrN&megatron.core.datasets.indexed_datasetrHAVE_MEGATRON_CORE ImportErrorModuleNotFoundErrorr rDrsr_rr~rrboolrrr rrrr$r6r;r=rErFrsobjectr~rrr(r(r(r)s      $ 4    ,M /  q$ ` p b _ . p r :(%X*