Transformer encoder consisting of *config.encoder_layers* self attention layers. Each layer is a [`FunAudioChatDecoderLayer`]. Args: config: FunAudioChatAudioEncoderConfig
| 561 | |
| 562 | |
| 563 | class FunAudioChatDecoder(FunAudioChatPreTrainedModel): |
| 564 | """ |
| 565 | Transformer encoder consisting of *config.encoder_layers* self attention layers. Each layer is a |
| 566 | [`FunAudioChatDecoderLayer`]. |
| 567 | |
| 568 | Args: |
| 569 | config: FunAudioChatAudioEncoderConfig |
| 570 | """ |
| 571 | |
| 572 | config_class = FunAudioChatAudioEncoderConfig |
| 573 | main_input_name = "audio_ids" |
| 574 | _tied_weights_keys = ["lm_head.weight"] |
| 575 | |
| 576 | def __init__(self, config: FunAudioChatAudioEncoderConfig): |
| 577 | super().__init__(config) |
| 578 | self.group_size = config.group_size |
| 579 | self.hidden_size = config.output_dim |
| 580 | self.pre_matching = nn.Linear(self.hidden_size, self.hidden_size * self.group_size, bias=True) |
| 581 | |
| 582 | crq_transformer_config = AutoConfig.for_model(**config.crq_transformer_config) |
| 583 | self.crq_transformer = AutoModel.from_config(crq_transformer_config) |
| 584 | del self.crq_transformer.embed_tokens |
| 585 | self.input_matching = nn.Linear(self.hidden_size, crq_transformer_config.hidden_size, bias=False) |
| 586 | self.output_matching = nn.Linear(crq_transformer_config.hidden_size, self.hidden_size, bias=False) |
| 587 | |
| 588 | self.lm_head = nn.Linear(config.output_dim, config.codebook_size, bias=False) |
| 589 | self.config = config |
| 590 | |
| 591 | # Initialize weights and apply final processing |
| 592 | self.post_init() |
| 593 | |
| 594 | def get_embeddings(self, audio_tokens): |
| 595 | return self.lm_head.weight.data[audio_tokens] |
| 596 | |
| 597 | def sampling_step(self, logits): |
| 598 | # Copy is needed to avoid keeping a hanging ref to outputs.logits which may be very large for first iteration |
| 599 | # (the clone itself is always small) |
| 600 | next_token_logits = logits[:, -1, :].to(copy=True, dtype=torch.float32, device=logits.device) |
| 601 | |
| 602 | # pre-process distribution |
| 603 | next_token_scores = self.crq_logits_processor(torch.cat([self.crq_speech_ids, *self.crq_generate_tokens], dim=-1), next_token_logits) |
| 604 | |
| 605 | # token selection |
| 606 | if self.crq_do_sample: |
| 607 | probs = nn.functional.softmax(next_token_scores, dim=-1) |
| 608 | # TODO (joao): this OP throws "skipping cudagraphs due to ['incompatible ops']", find solution |
| 609 | next_tokens = torch.multinomial(probs, num_samples=1).squeeze(1) |
| 610 | else: |
| 611 | next_tokens = torch.argmax(next_token_scores, dim=-1) |
| 612 | |
| 613 | return next_tokens, logits |
| 614 | |
| 615 | def crq_generate_forward( |
| 616 | self, |
| 617 | inputs_embeds=None, |
| 618 | audio_embeds=None, |
| 619 | labels=None, |
| 620 | attention_mask=None, |