o Zhod@sddlZddlmZddlmZmZmZmZmZm Z ddl Z ddl m Z ddl mZmZmZmZddlmZddlmZmZddlmZdd lmZerRdd lmZeeZeGd d d eZ Gd ddeZ!dS)N) dataclass) TYPE_CHECKINGDictListOptionalTupleUnion)GenerateDecoderOnlyOutputGenerationConfigGenerationMixinGenerationMode)LogitsProcessorList)MaxLengthCriteriaStoppingCriteriaList)GenerateNonBeamOutput)logging) BaseStreamerc@s(eZdZUdZdZeeeje d<dS)CsmGenerateOutputap Outputs of CsmForConditionalGeneration.generate. Args: sequences (`torch.LongTensor` of shape `(batch_size, sequence_length)`): The generated sequences. The second dimension (sequence_length) is either equal to `max_length` or shorter if all batches finished early due to the `eos_token_id`. scores (`tuple(torch.FloatTensor)` *optional*, returned when `output_scores=True`): Processed prediction scores of the language modeling head (scores for each vocabulary token before SoftMax) at each generation step. Tuple of `torch.FloatTensor` with up to `max_new_tokens` elements (one element for each generated token), with each tensor of shape `(batch_size, config.vocab_size)`. logits (`tuple(torch.FloatTensor)` *optional*, returned when `output_logits=True`): Unprocessed prediction scores of the language modeling head (scores for each vocabulary token before SoftMax) at each generation step. Tuple of `torch.FloatTensor` with up to `max_new_tokens` elements (one element for each generated token), with each tensor of shape `(batch_size, config.vocab_size)`. attentions (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `output_attentions=True`): Tuple (one element for each generated token) of tuples (one element for each layer of the decoder) of `torch.FloatTensor` of shape `(batch_size, num_heads, generated_length, sequence_length)`. hidden_states (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `output_hidden_states=True`): Tuple (one element for each generated token) of tuples (one element for each layer of the decoder) of `torch.FloatTensor` of shape `(batch_size, generated_length, hidden_size)`. past_key_values (`tuple(tuple(torch.FloatTensor)))`, *optional*, returned when `use_cache=True`): Returns the model cache, used to speed up decoding. Different models have a different cache format, check audio (`list(torch.FloatTensor)` of length `batch_size`): The generated audio. Naudio) __name__ __module__ __qualname____doc__rrrtorchTensor__annotations__rrU/var/www/auris/lib/python3.10/site-packages/transformers/models/csm/generation_csm.pyr*s rcseZdZdeffdd Z ddeedeedede eefffdd Z d e j d e d eded ededdeee j ffddZ         dd ee jdee jdee jdeed ee d eed eededdeedeee j fffdd ZZS)CsmGenerationMixinreturncsRtj|i|}t}|D]}t|ts!td|jjdq| |q|S)NzCsm does not support z' stopping criteria, it will be ignored.) super_get_stopping_criteriar isinstancerloggerwarning __class__rappend)selfargskwargscriteriaZ kept_criteria criterionr&rrr"Ks  z)CsmGenerationMixin._get_stopping_criteriaNgeneration_configuse_model_defaultsr*c  sdd|D}dd|D}tj||fi|\}}|jjjdi|t|jjdp5|jjd}t|jjdpB|jjd}||h|jjdhkr`t d|d|d |jjdd |jjj rot d d |jj_ ||jj_ ||jj_|jdfdd }||_||fS)z This method overrides [~generation.utils.GenerationMixin._prepare_generation_config]. It ensures that the depth decoder generation config is initialized and that passed args as depth_decoder_* are properly handled. cSs,i|]\}}|dr|tdd|qS)depth_decoder_N) startswithlen.0kvrrr ds zACsmGenerationMixin._prepare_generation_config..cSs i|] \}}|ds||qS)r0)r1r3rrrr7is min_new_tokensmax_new_tokensz2depth_decoder_generation_config's min_new_tokens (z) and max_new_tokens (z2) must be equal to self.config.num_codebooks - 1 ()zdepth_decoder_generation_config.return_dict_in_generate is set to True, but this will be ignored as the depth decoder model does not return a dictionary in generateFNcs,|}|tjtjfvrtd|d|S)NzGeneration mode zg is not supported for CSM model. Please set generation parameters to use greedy or sampling generation.)r Z GREEDY_SEARCHZSAMPLE ValueError)Zassistant_modelZgeneration_modeZoriginal_get_generation_moderrpatched_get_generation_modes  zRCsmGenerationMixin._prepare_generation_config..patched_get_generation_moderN)itemsr!_prepare_generation_config depth_decoderr.updategetattrconfigZ num_codebooksr<return_dict_in_generater$r%r8r:Zget_generation_mode) r(r.r/r*Zdepth_decoder_kwargs model_kwargsZdepth_decoder_min_new_tokensZdepth_decoder_max_new_tokensr>r-r=rrA\s>        z-CsmGenerationMixin._prepare_generation_config input_idslogits_processorstopping_criteria synced_gpusstreamerrc' Ks |jj}|jdu} |j} |j} |j} |j} |j}|j}|r!| r!dnd}|r)| r)dnd}|r1| r1dnd}|r9| r9dnd}|j dd\}}d}t j |t j |j d}|||j |}|jdkru|dduru|D]}t|trt|j|8_qf|j}|||}|rdtjd<||j}d }|j|||j d r|j|fi|}|| rd | ini|d d i|r|di|d d i}d}n |di|d d i}|||}|r|rq|jdddddf }|!|j }|||}|r| r||f7}| r||f7}| r||j"f7}| r||j#f7}|r-t$j%j&|dd} t j'| dd(d}!nt j)|dd}!|!dddf}"t$j%j*|"ddd}#|j#ddddddf}$|j+j,|#|$d}%t|%t j-rf|%n|%j.}&|&ddddf}&|&}!| r|!|/d|d|/d}!|jdkr|!dddddf}nt j0||!dddddfgdd}|dur|1|!2||dddddf|jj3k4d@}||||@}|5dk}|d7}~~%|j|||j d s|dur|6|rt7||||||ddS|S)a! This method overrides [~generation.utils.GenerationMixin._sample]. To ease maintenance, modifications are marked with the comment "Csm specific". Indeed, Csm model requires a custom generation sampling step: 1. Infer the backbone model to sample the first codebook token 2. Call generate on the depth decoder with the first codebook token as input_ids to sample the next codebook tokens 3. Use these generated codebook tokens as input_ids to sample the next first codebook token using the backbone model 4. Repeat until stopping criteria is met Csm supports two stopping criterias: - stop when the generated sequence is at max_length - stop when all the generated codebook tokens are the codebook_eos_token_id NrF)ZdtypedeviceZ inputs_embeds0ZTOKENIZERS_PARALLELISMT)rNoutput_attentionsoutput_hidden_statesZ return_dictdimr9)Z num_samples)r9rr)value)rHbackbone_last_hidden_statepast_key_values) sequencesscoreslogits attentions hidden_statesrW)8rEZcodebook_pad_token_idZ_eos_token_tensorrPrQ output_scores output_logitsrF do_sampleshaperZoneslongrNZ_get_initial_cache_positionndimgetr#r max_length__call__Z_valid_auto_compile_criteriaosenvironZget_compiled_callZcompile_configZ_has_unfinished_sequencesZprepare_inputs_for_generationrCZ#_update_model_kwargs_for_generationrZclonefloattor[r\nnZ functionalZsoftmaxZ multinomialZsqueezeZargmaxpadrBgeneraterrX unsqueezecatputcpucodebook_eos_token_idallmaxendr )'r(rHrIrJr.rKrLrGZ pad_token_idZhas_eos_stopping_criteriarPrQr]r^rFr_rYZ raw_logitsZdecoder_attentionsZdecoder_hidden_statesZ batch_sizecur_lenZthis_peer_finishedZunfinished_sequencesr,Z model_forwardZcompile_forwardZ is_prefillZ model_inputsZoutputsZnext_token_logitsZnext_token_scoresZprobsZ next_tokensZfirst_codebook_idsZdepth_decoder_input_idsrVZdepth_decoder_outputsZ codebook_idsrrr_samples               $   m zCsmGenerationMixin._sampleF input_valuesinput_values_cutoffs output_audioc  stjd ||||||||d| } t| tj } d} | rv| r"| jn| }g} tD|D]9}||jjkj dd }| dkrF| }n|j d}|d|}|j|ddd}| |jdq-Wdn1sqwY| rtd d| i| S| r| S| S) aI This method overrides [`~generation.utils.GenerationMixin.generate`] to match the specifics of the Csm model. Indeed, Csm model requires a custom generation sampling step: 1. Infer the backbone model to sample the first codebook token 2. Call generate on the depth decoder with the first codebook token as `input_ids` to sample the next codebook tokens 3. Use these generated codebook tokens as `input_ids` to sample the next first codebook token using the backbone model 4. Repeat until stopping criteria is met Most generation-controlling parameters are set in `generation_config` which, if not passed, will be set to the model's default generation configuration. You can override any `generation_config` by passing the corresponding parameters to generate(), e.g. `.generate(inputs, do_sample=True)`. Parameters: inputs_ids (`torch.Tensor` of shape (batch_size, seq_length), *optional*): The sequence used as a prompt for the backbone model. input_values (`torch.Tensor` of shape (batch_size, channels, max_concatenated_audio_length), *optional*): The batched audio input values, where each batch entry contains the concatenation of all audio segments for that entry. These values will be encoded into codebook tokens using the codec model and merged with the text input ids provided in `input_ids`. input_values_cutoffs (`torch.Tensor` of shape (batch_size, max_num_audio), *optional*): Specify the end positions of audio segments within each batch entry, relative to the concatenated audio input. If a batch entry has fewer segments than the maximum, it is padded with -1. For example, in a batch of 2 sequences where the first contains 2 audio segments of length l1, and the second contains 1 audio segment of length l2, the input_values_cutoffs would be: [[l1, 2 * l1], [l2, -1]]. generation_config ([`~generation.GenerationConfig`], *optional*): The generation configuration to be used as base parametrization for the generation call. `**kwargs` passed to generate matching the attributes of `generation_config` will override them. If `generation_config` is not provided, the default will be used, which has the following loading priority: 1) from the `generation_config.json` model file, if it exists; 2) from the model configuration. Please note that unspecified parameters will inherit [`~generation.GenerationConfig`]'s default values, whose documentation should be checked to parameterize generation. logits_processor (`LogitsProcessorList`, *optional*): Custom logits processors that complement the default logits processors built from arguments and generation config. If a logit processor is passed that is already created with the arguments or a generation config an error is thrown. This feature is intended for advanced users. stopping_criteria (`StoppingCriteriaList`, *optional*): Custom stopping criteria that complements the default stopping criteria built from arguments and a generation config. If a stopping criteria is passed that is already created with the arguments or a generation config an error is thrown. If your stopping criteria depends on the `scores` input, make sure you pass `return_dict_in_generate=True, output_scores=True` to `generate`. This feature is intended for advanced users. synced_gpus (`bool`, *optional*): Whether to continue running the while loop until max_length. Unless overridden, this flag will be set to `True` if using `FullyShardedDataParallel` or DeepSpeed ZeRO Stage 3 with multiple GPUs to avoid deadlocking if one GPU finishes generating before other GPUs. Otherwise, defaults to `False`. streamer (`BaseStreamer`, *optional*): Streamer object that will be used to stream the generated sequences. Generated tokens are passed through `streamer.put(token_ids)` and the streamer is responsible for any further processing. output_audio (`bool`, *optional*): Whether to return the generated audio. kwargs (`Dict[str, Any]`, *optional*): Ad hoc parametrization of `generation_config` and/or additional model-specific kwargs that will be forwarded to the `forward` function of the model. Depth decoder specific kwargs should be prefixed with *depth_decoder_*. Return: [`CsmGenerateOutput`] or `torch.LongTensor` or `List[torch.FloatTensor]`: A [`CsmGenerateOutput`] (if `return_dict_in_generate=True` or when `config.return_dict_in_generate=True`) or a `torch.LongTensor` when `output_audio=False` or a `List[torch.FloatTensor]` otherwise. Example: ```python >>> from transformers import CsmProcessor, CsmForConditionalGeneration >>> from datasets import load_dataset, Audio >>> model_id = "eustlb/csm-1b" >>> torch_device = "cuda" if torch.cuda.is_available() else "cpu" >>> processor = AutoProcessor.from_pretrained(model_id) >>> ds = load_dataset("hf-internal-testing/dailytalk-dummy", split="train") >>> # ensure the audio is 24kHz >>> ds = ds.cast_column("audio", Audio(sampling_rate=24000)) >>> conversation = [] >>> # prepare a conversation with text and corresponding audio >>> for text, audio, speaker_id in zip(ds[:4]["text"], ds[:4]["audio"], ds[:4]["speaker_id"]): ... conversation.append( ... { ... "role": f"{speaker_id}", ... "content": [{"type": "text", "text": text}, {"type": "audio", "path": audio["array"]}], ... } ... ) >>> # text prompt >>> conversation.append({"role": f"{ds[4]['speaker_id']}", "content": [{"type": "text", "text": ds[4]["text"]}]}) >>> inputs = processor.apply_chat_template( ... conversation, ... tokenize=True, ... return_dict=True, ... ).to(torch_device) >>> model = CsmForConditionalGeneration.from_pretrained(model_id, device_map=torch_device) >>> audio = model.generate(**inputs, output_audio=True) >>> processor.save_audio(audio, "output.wav") ``` )rHrxryr.rIrJrKrLNrRrSrr9)rrrr)r!rmr#rrrXZno_gradrErrrsZnonzeroZnumelminr`Z codec_modeldecodeZ transposernr'Z audio_valuesr)r(rHrxryr.rIrJrKrLrzr*Zgenerate_outputZgenerate_returned_dictrZgenerated_audio_codesZaudio_codes_batchZeos_idxsZ cutoff_idxZcodec_decode_outputr-rrrmRsBq       zCsmGenerationMixin.generater?) NNNNNNNNF)rrrrr"rr boolrrrArZ LongTensorrrrrwrrm __classcell__rrr-rrJsx : ?   r)"rf dataclassesrtypingrrrrrrrZtorch.nnrkZ generationr r r r Zgeneration.logits_processrZgeneration.stopping_criteriarrZgeneration.utilsrutilsrZgeneration.streamersrZ get_loggerrr$rrrrrrs