Repository: huggingface/transformers
File: src/transformers/models/higgs_audio_v2/generation_higgs_audio_v2.py (Python)

Code before:

        # keep track of which sequences are already finished
        batch_size, cur_len = input_ids.shape[:2]
        this_peer_finished = False
        unfinished_sequences = torch.ones(batch_size, dtype=torch.long, device=input_ids.device)

        model_forward = (
            self.get_compiled_call(generation_config.compile_config)
            if self._valid_auto_compile_criteria(model_kwargs, generation_config)
            else self.__call__
        )

        # Assisted generation completes the prefill stage in candidate generator so that
        # we don't have several `prefill` calls in one generation loop. Skip `_prefill` for assistants
        if not generation_config.is_assistant:
            outputs = self._prefill(input_ids, generation_config, model_kwargs)
            prefill_consumed = False
        else:
            model_kwargs = self._get_initial_cache_position(input_ids.shape[1], input_ids.device, model_kwargs)
            prefill_consumed = True

        while self._has_unfinished_sequences(this_peer_finished, synced_gpus, device=input_ids.device):
            if prefill_consumed:
                use_cache = model_kwargs.get("use_cache", True)
                new_input_ids = input_ids[:, -1:] if use_cache else input_ids
                model_inputs = self.prepare_inputs_for_generation(
                    new_input_ids, full_input_ids=input_ids, **model_kwargs
                )
                with self._optimize_model_for_decode():
                    outputs = model_forward(**model_inputs, return_dict=True)
            prefill_consumed = True
            model_kwargs = self._update_model_kwargs_for_generation(
                outputs,
                model_kwargs,
                is_encoder_decoder=self.config.is_encoder_decoder,
            )
            if synced_gpus and this_peer_finished:
                continue

            # Copy is needed to avoid keeping a hanging ref to outputs.logits which may be very large for first iteration
            # (the clone itself is always small)
            next_token_logits = outputs.logits[:, -1, :].to(copy=True, dtype=torch.float32, device=input_ids.device)

            # pre-process distribution
            next_token_scores = logits_processor(input_ids, next_token_logits)

            # ===========================
            # BELOW DIFFERENCES WITH GenerationMixin._sample()
            # Store scores, attentions and hidden_states when required
            if return_dict_in_generate:
                if output_scores:

Review comment (suggestion):
I guess we need the full input ids for the custom mask logic? It's a bit weird but otherwise aligns with the rest of the new changes on generate

Code after:

        # keep track of which sequences are already finished
        batch_size, cur_len = input_ids.shape[:2]
        this_peer_finished = False
        unfinished_sequences = torch.ones(batch_size, dtype=torch.long, device=input_ids.device)

        model_forward = (
            self.get_compiled_call(generation_config.compile_config)
            if self._valid_auto_compile_criteria(model_kwargs, generation_config)
            else self.__call__
        )

        prefill_consumed = False
        outputs = self._prefill(
            input_ids,
            generation_config,
            model_kwargs,
            is_first_iteration=not generation_config.is_assistant,
        )

        while self._has_unfinished_sequences(this_peer_finished, synced_gpus, device=input_ids.device):
            if prefill_consumed:
                use_cache = model_kwargs.get("use_cache", True)
                new_input_ids = input_ids[:, -1:] if use_cache else input_ids
                model_inputs = self.prepare_inputs_for_generation(
                    new_input_ids, full_input_ids=input_ids, **model_kwargs
                )
                with self._optimize_model_for_decode():
                    outputs = model_forward(**model_inputs, return_dict=True)
            prefill_consumed = True
            model_kwargs = self._update_model_kwargs_for_generation(
                outputs,
                model_kwargs,
                is_encoder_decoder=self.config.is_encoder_decoder,
            )
            if synced_gpus and this_peer_finished:
                continue

            # Copy is needed to avoid keeping a hanging ref to outputs.logits which may be very large for first iteration
            # (the clone itself is always small)
            next_token_logits = outputs.logits[:, -1, :].to(copy=True, dtype=torch.float32, device=input_ids.device)

            # pre-process distribution
            next_token_scores = logits_processor(input_ids, next_token_logits)

            # ===========================
            # BELOW DIFFERENCES WITH GenerationMixin._sample()
            # Store scores, attentions and hidden_states when required
            if return_dict_in_generate:
                if output_scores:
                    scores += (