adt-str / model.py

Upload folder using huggingface_hub

162fa19 verified about 2 months ago

18.9 kB

import math
from typing import Optional

import torch
import torch.nn as nn
import torch.nn.functional as F
from config import ADTModelConfig
from transformers import PreTrainedModel
import torchaudio.transforms as T
from utils.utils import create_mask_plain


class TokenEmbedding(nn.Module):
    def __init__(self, vocab_size, emb_size):
        """
        Args:
            vocab_size: Vocabulary size
            emb_size: Dimension of embedding vectors
        """
        super(TokenEmbedding, self).__init__()
        self.embedding = nn.Embedding(vocab_size, emb_size)
        self.emb_size = emb_size

    def forward(self, token_vector: torch.Tensor):
        """
        Args:
            token_vector: One-hot or multi-hot tensor (batch, seq_len, vocab_size)
        Returns:
            Embedded vectors (batch, seq_len, emb_size)
        """
        # Get embedding weights (vocab_size, emb_size)
        embedding_weights = self.embedding.weight  # Get weights

        # Multiply one-hot or multi-hot vector with embedding weights
        embeddings = torch.matmul(token_vector, embedding_weights) * math.sqrt(
            self.emb_size
        )

        return embeddings


class TokenEmbedding_plain(nn.Module):
    def __init__(self, vocab_size, emb_size):
        super(TokenEmbedding_plain, self).__init__()
        self.embedding = nn.Embedding(vocab_size, emb_size)
        self.emb_size = emb_size

    def forward(self, tokens: torch.Tensor):
        return self.embedding(tokens.long()) * math.sqrt(self.emb_size)


class PositionalEncoding(nn.Module):
    def __init__(self, emb_size: int, maxlen: int = 2048):
        super(PositionalEncoding, self).__init__()
        den = torch.exp(-torch.arange(0, emb_size, 2) * math.log(10000) / emb_size)
        pos = torch.arange(0, maxlen).reshape(maxlen, 1)
        pos_embedding = torch.zeros((maxlen, emb_size))
        pos_embedding[:, 0::2] = torch.sin(pos * den)
        pos_embedding[:, 1::2] = torch.cos(pos * den)
        pos_embedding = pos_embedding.unsqueeze(0)

        self.register_buffer("pos_embedding", pos_embedding)

    def forward(self, token_embedding: torch.Tensor):
        return token_embedding + self.pos_embedding[:, : token_embedding.size(1), :]


class ComputeMelSpectrogram(torch.nn.Module):
    def __init__(self, sample_rate, win_length, time_res, n_mels):
        super(ComputeMelSpectrogram, self).__init__()
        with torch.device("cpu"):
            self.compute_spec = T.MelSpectrogram(
                sample_rate=sample_rate,
                n_fft=win_length,
                hop_length=int(time_res * sample_rate),
                n_mels=n_mels,
                f_min=20.0,
                power=2,
            )
        self.window_pad_idxs = int((win_length / 2) // int(time_res * sample_rate) + 1)

    def forward(self, wave):
        wave_device = wave.device
        self.compute_spec = self.compute_spec.to(wave_device)
        # Esegui in float32: torchaudio MelSpectrogram + bf16 autocast causano CUBLAS_STATUS_INVALID_VALUE
        with torch.amp.autocast(
            device_type="cuda" if wave_device.type == "cuda" else "cpu", enabled=False
        ):
            wave_f32 = wave.float()
            mel_spec = self.compute_spec(wave_f32)

        logmel_spec = torch.log(mel_spec + 1e-10)
        logmel_spec = torch.clamp(logmel_spec, -23, 12)
        logmel_spec = (logmel_spec + 23) / (12 + 23)

        return logmel_spec.permute(0, 2, 1)[
            :, self.window_pad_idxs : -(self.window_pad_idxs + 1), :
        ]  # (batch, len, emb)


class Encoder(nn.Module):
    def __init__(
        self,
        enc_layers,
        d_query,
        nhead,
        ffn_hid_dim,
        dropout,
    ):
        super(Encoder, self).__init__()
        self.num_features = d_query * nhead
        self.dense_layer = nn.Linear(self.num_features, self.num_features, bias=False)
        self.positional_encoding = PositionalEncoding(self.num_features)
        self.layer_norm = nn.LayerNorm(
            normalized_shape=self.num_features, elementwise_affine=True
        )
        self.dropout_layer = nn.Dropout(p=dropout)

        layer = nn.TransformerEncoderLayer(
            d_model=self.num_features,
            nhead=nhead,
            dim_feedforward=ffn_hid_dim,
            dropout=dropout,
            activation="gelu",
            batch_first=True,
            norm_first=False,
        )
        self.encoder = nn.TransformerEncoder(layer, num_layers=enc_layers)

    def forward(self, src_emb):
        src_emb = self.dense_layer(src_emb)
        src_emb = self.positional_encoding(src_emb)
        src_emb = self.dropout_layer(src_emb)
        encoder_output = self.encoder(src_emb)
        encoder_output = self.dropout_layer(self.layer_norm(encoder_output))
        return encoder_output


class Decoder(nn.Module):
    def __init__(
        self,
        dec_layers,
        d_query,
        nhead,
        ffn_hid_dim,
        tgt_vocab_size,
        dropout,
        plain=False,
    ):
        super(Decoder, self).__init__()
        self.num_features = d_query * nhead
        if plain:
            self.tgt_tok_emb = TokenEmbedding_plain(tgt_vocab_size, self.num_features)
        if not plain:
            self.tgt_tok_emb = TokenEmbedding(tgt_vocab_size, self.num_features)
        self.positional_encoding = PositionalEncoding(self.num_features)
        self.dropout_layer = nn.Dropout(p=dropout)
        self.generator = nn.Linear(self.num_features, tgt_vocab_size)

        layer = nn.TransformerDecoderLayer(
            d_model=self.num_features,
            nhead=nhead,
            dim_feedforward=ffn_hid_dim,
            dropout=dropout,
            activation="gelu",
            batch_first=True,
            norm_first=False,
        )
        self.decoder = nn.TransformerDecoder(layer, num_layers=dec_layers)

    def forward(self, tgt, encoder_output, tgt_mask, tgt_padding_mask):
        tgt_emb = self.positional_encoding(self.tgt_tok_emb(tgt))
        tgt_emb = self.dropout_layer(tgt_emb)
        # bf16-safe and PyTorch same-type: use additive float masks (0 / -1e4) instead of bool (-inf)
        if tgt_mask is not None:
            tgt_mask = torch.zeros_like(
                tgt_mask, dtype=tgt_emb.dtype, device=tgt_mask.device
            ).masked_fill_(tgt_mask, -1e4)
        if tgt_padding_mask is not None and tgt_padding_mask.dtype == torch.bool:
            tgt_padding_mask = torch.zeros_like(
                tgt_padding_mask, dtype=tgt_emb.dtype, device=tgt_padding_mask.device
            ).masked_fill_(tgt_padding_mask, -1e4)
        decoder_outputs = self.decoder(
            tgt_emb,
            encoder_output,
            tgt_mask=tgt_mask,
            tgt_key_padding_mask=tgt_padding_mask,
            memory_mask=None,
            memory_key_padding_mask=None,
        )
        return self.generator(decoder_outputs)


class ADTModel(PreTrainedModel):
    config_class = ADTModelConfig
    _supports_assign_param_buffer = False
    _tied_weights_keys = []
    all_tied_weights_keys = {}

    def __init__(
        self,
        config: ADTModelConfig,
    ) -> None:
        super(ADTModel, self).__init__(config)
        self.config = config
        self.encoder = Encoder(
            enc_layers=config.enc_layers,
            d_query=config.d_query,
            nhead=config.nhead,
            ffn_hid_dim=int(config.d_query * config.nhead * 4),
            dropout=config.dropout,
        )
        self.decoder = Decoder(
            dec_layers=config.dec_layers,
            d_query=config.d_query,
            nhead=config.nhead,
            ffn_hid_dim=int(config.d_query * config.nhead * 4),
            dropout=config.dropout,
            tgt_vocab_size=config.tgt_vocab_size,
            plain=config.plain,
        )
        self.compute_spectrogram = ComputeMelSpectrogram(
            sample_rate=config.sample_rate,
            win_length=config.win_length,
            time_res=config.time_res,
            n_mels=config.n_mels,
        )
        self.project_to_mel = nn.Linear(
            config.n_mels, int(config.d_query * config.nhead)
        )

    def _loss_fn(
        self, logits: torch.FloatTensor, tgt: torch.LongTensor
    ) -> torch.FloatTensor:
        # Use fp32 for stable loss and ignore PAD token (id=1)
        logits = logits.float()
        logits = torch.nan_to_num(logits, nan=0.0, posinf=1e4, neginf=-1e4)
        return F.cross_entropy(
            logits.reshape(-1, logits.shape[-1]),
            tgt.long().reshape(-1),
            ignore_index=1,
        )

    def forward(
        self,
        src: torch.FloatTensor,
        tgt: torch.LongTensor,
        tgt_mask: Optional[torch.BoolTensor],
        tgt_padding_mask: Optional[torch.BoolTensor],
        labels: torch.LongTensor,
    ) -> torch.FloatTensor:
        src_emb = self.compute_spectrogram(src)
        src_emb = self.project_to_mel(src_emb)
        encoder_out = self.encoder(src_emb)

        # Build causal mask inside model so DataParallel does not split (T,T) across devices
        if tgt_mask is None:
            tgt_seq_len = tgt.size(1)
            tgt_mask, _ = create_mask_plain(tgt_seq_len, None, tgt.device)
        logits = self.decoder(tgt, encoder_out, tgt_mask, tgt_padding_mask)
        loss = self._loss_fn(logits, labels)
        return loss

    def sample(
        self,
        src: torch.FloatTensor,
        src_mask: torch.BoolTensor,
        tgt_mask: torch.BoolTensor,
        max_length: int = 1000,
        start_token: int = 2,
        end_token: int = 3,
    ) -> torch.LongTensor:
        """
        Basic autoregressive greedy sampling for ADT model.
        Args:
            src: Input audio tensor (batch_size, audio_length)
            src_mask: Source mask (not used in current implementation)
            tgt_mask: Target mask (not used in current implementation)
            max_length: Maximum sequence length to generate
            start_token: Token ID to start generation (default: 0)
            end_token: Token ID to end generation (default: 2, assuming EOS)
        Returns:
            Generated token sequence (batch_size, seq_length)
        """
        if not self.config.plain:
            raise NotImplementedError("Non-plain mode is not implemented")
        self.eval()
        device = src.device
        batch_size = src.shape[0]

        # Process input through encoder
        src_emb = self.compute_spectrogram(src)
        src_emb = self.project_to_mel(src_emb)
        memory = self.encoder(src_emb)

        # Initialize with start token
        generated = torch.full(
            (batch_size, 1), start_token, dtype=torch.long, device=device
        )
        finished = torch.zeros(batch_size, dtype=torch.bool, device=device)
        # Greedy decoding loop (per-row finished mask so batch items can end at different steps)
        for step in range(max_length - 1):
            # Create causal mask for current sequence length (no padding in decoding)
            seq_len = generated.shape[1]
            tgt_mask, tgt_padding_mask = create_mask_plain(seq_len, None, device)

            # Get decoder output
            logits = self.decoder(
                generated,
                memory,
                tgt_mask,
                tgt_padding_mask=None,
            )

            # Get next token (greedy); freeze finished rows at EOS so they do not keep consuming argmax noise
            next_token_logits = logits[:, -1, :]  # Last position
            next_token = torch.argmax(next_token_logits, dim=-1)  # Greedy selection
            next_token = torch.where(
                finished, torch.full_like(next_token, end_token), next_token
            )
            generated = torch.cat([generated, next_token.unsqueeze(1)], dim=1)
            finished = finished | (next_token == end_token)
            if torch.all(finished):
                break

        return generated

    def beam_search(
        self,
        src: torch.FloatTensor,
        src_mask: torch.BoolTensor,
        tgt_mask: torch.BoolTensor,
        beam_size: int = 5,
        max_length: int = 1000,
        start_token: int = 2,
        end_token: int = 3,
        length_penalty: float = 1.0,
    ) -> torch.LongTensor:
        """
        Beam search sampling for ADT model.
        Args:
            src: Input audio tensor (batch_size, audio_length)
            src_mask: Source mask (not used in current implementation)
            tgt_mask: Target mask (not used in current implementation)
            beam_size: Number of beams to keep at each step
            max_length: Maximum sequence length to generate
            start_token: Token ID to start generation (default: 0)
            end_token: Token ID to end generation (default: 2, assuming EOS)
            length_penalty: Length penalty factor for scoring (default: 1.0)
        Returns:
            Best generated token sequence (batch_size, seq_length)
        """
        if not self.config.plain:
            raise NotImplementedError("Non-plain mode is not implemented")

        self.eval()
        device = src.device
        batch_size = src.shape[0]

        # Process input through encoder
        src_emb = self.compute_spectrogram(src)
        src_emb = self.project_to_mel(src_emb)
        memory = self.encoder(src_emb)

        # Initialize beams for each batch
        # Each beam contains: [batch_idx, sequence, log_prob, finished]
        all_beams = []

        for batch_idx in range(batch_size):
            # Initialize beams for this batch item
            initial_sequence = torch.full(
                (1, 1), start_token, dtype=torch.long, device=device
            )
            beams = [{"sequence": initial_sequence, "log_prob": 0.0, "finished": False}]
            all_beams.append(beams)

        # Beam search loop
        for step in range(max_length - 1):
            new_all_beams = []

            for batch_idx in range(batch_size):
                beams = all_beams[batch_idx]
                new_beams = []

                # Skip if all beams are finished
                if all(beam["finished"] for beam in beams):
                    new_all_beams.append(beams)
                    continue

                # Collect all active sequences for this batch
                active_beams = [beam for beam in beams if not beam["finished"]]
                if not active_beams:
                    new_all_beams.append(beams)
                    continue

                # Stack sequences for parallel processing
                sequences = torch.cat(
                    [beam["sequence"] for beam in active_beams], dim=0
                )
                batch_memory = memory[batch_idx : batch_idx + 1].expand(
                    len(active_beams), -1, -1
                )

                # Create causal mask
                seq_len = sequences.shape[1]
                tgt_mask = torch.triu(
                    torch.ones(seq_len, seq_len, device=device), diagonal=1
                ).bool()

                # Get decoder output
                logits = self.decoder(
                    sequences,
                    batch_memory,
                    tgt_mask,
                    None,
                )

                # Get next token probabilities
                next_token_logits = logits[:, -1, :]  # Last position
                log_probs = torch.log_softmax(next_token_logits, dim=-1)

                # Generate candidates for each active beam
                candidates = []
                for beam_idx, beam in enumerate(active_beams):
                    beam_log_probs = log_probs[beam_idx]
                    top_k_log_probs, top_k_tokens = torch.topk(
                        beam_log_probs, beam_size
                    )

                    for k in range(beam_size):
                        token = top_k_tokens[k].item()
                        token_log_prob = top_k_log_probs[k].item()
                        new_log_prob = beam["log_prob"] + token_log_prob

                        # Create new sequence
                        new_sequence = torch.cat(
                            [beam["sequence"], torch.tensor([[token]], device=device)],
                            dim=1,
                        )

                        # Check if finished
                        finished = token == end_token

                        candidates.append(
                            {
                                "sequence": new_sequence,
                                "log_prob": new_log_prob,
                                "finished": finished,
                            }
                        )

                # Add finished beams (they don't expand further)
                finished_beams = [beam for beam in beams if beam["finished"]]
                candidates.extend(finished_beams)

                # Select top beam_size candidates
                candidates.sort(
                    key=lambda x: self._score_sequence(x, length_penalty), reverse=True
                )
                new_beams = candidates[:beam_size]

                new_all_beams.append(new_beams)

            all_beams = new_all_beams

            # Check if all beams in all batches are finished
            all_finished = all(
                all(beam["finished"] for beam in beams) for beams in all_beams
            )
            if all_finished:
                break

        # Select best sequence for each batch
        results = []
        for batch_idx in range(batch_size):
            beams = all_beams[batch_idx]
            # Sort by score and take the best
            beams.sort(
                key=lambda x: self._score_sequence(x, length_penalty), reverse=True
            )
            best_sequence = beams[0]["sequence"]
            results.append(best_sequence)

        # Stack results and pad to same length
        max_len = max(seq.shape[1] for seq in results)
        padded_results = []
        for seq in results:
            if seq.shape[1] < max_len:
                padding = torch.full(
                    (1, max_len - seq.shape[1]),
                    end_token,
                    dtype=torch.long,
                    device=device,
                )
                padded_seq = torch.cat([seq, padding], dim=1)
            else:
                padded_seq = seq
            padded_results.append(padded_seq)

        return torch.cat(padded_results, dim=0)

    def _score_sequence(self, beam: dict, length_penalty: float) -> float:
        seq_len = beam["sequence"].shape[1]
        if length_penalty == 0.0:
            return beam["log_prob"]
        else:
            # Length normalization: log_prob / (length^length_penalty)
            return beam["log_prob"] / (seq_len**length_penalty)

	import math
	from typing import Optional

	import torch
	import torch.nn as nn
	import torch.nn.functional as F
	from config import ADTModelConfig
	from transformers import PreTrainedModel
	import torchaudio.transforms as T
	from utils.utils import create_mask_plain


	class TokenEmbedding(nn.Module):
	def __init__(self, vocab_size, emb_size):
	"""
	Args:
	vocab_size: Vocabulary size
	emb_size: Dimension of embedding vectors
	"""
	super(TokenEmbedding, self).__init__()
	self.embedding = nn.Embedding(vocab_size, emb_size)
	self.emb_size = emb_size

	def forward(self, token_vector: torch.Tensor):
	"""
	Args:
	token_vector: One-hot or multi-hot tensor (batch, seq_len, vocab_size)
	Returns:
	Embedded vectors (batch, seq_len, emb_size)
	"""
	# Get embedding weights (vocab_size, emb_size)
	embedding_weights = self.embedding.weight # Get weights

	# Multiply one-hot or multi-hot vector with embedding weights
	embeddings = torch.matmul(token_vector, embedding_weights) * math.sqrt(
	self.emb_size
	)

	return embeddings


	class TokenEmbedding_plain(nn.Module):
	def __init__(self, vocab_size, emb_size):
	super(TokenEmbedding_plain, self).__init__()
	self.embedding = nn.Embedding(vocab_size, emb_size)
	self.emb_size = emb_size

	def forward(self, tokens: torch.Tensor):
	return self.embedding(tokens.long()) * math.sqrt(self.emb_size)


	class PositionalEncoding(nn.Module):
	def __init__(self, emb_size: int, maxlen: int = 2048):
	super(PositionalEncoding, self).__init__()
	den = torch.exp(-torch.arange(0, emb_size, 2) * math.log(10000) / emb_size)
	pos = torch.arange(0, maxlen).reshape(maxlen, 1)
	pos_embedding = torch.zeros((maxlen, emb_size))
	pos_embedding[:, 0::2] = torch.sin(pos * den)
	pos_embedding[:, 1::2] = torch.cos(pos * den)
	pos_embedding = pos_embedding.unsqueeze(0)

	self.register_buffer("pos_embedding", pos_embedding)

	def forward(self, token_embedding: torch.Tensor):
	return token_embedding + self.pos_embedding[:, : token_embedding.size(1), :]


	class ComputeMelSpectrogram(torch.nn.Module):
	def __init__(self, sample_rate, win_length, time_res, n_mels):
	super(ComputeMelSpectrogram, self).__init__()
	with torch.device("cpu"):
	self.compute_spec = T.MelSpectrogram(
	sample_rate=sample_rate,
	n_fft=win_length,
	hop_length=int(time_res * sample_rate),
	n_mels=n_mels,
	f_min=20.0,
	power=2,
	)
	self.window_pad_idxs = int((win_length / 2) // int(time_res * sample_rate) + 1)

	def forward(self, wave):
	wave_device = wave.device
	self.compute_spec = self.compute_spec.to(wave_device)
	# Esegui in float32: torchaudio MelSpectrogram + bf16 autocast causano CUBLAS_STATUS_INVALID_VALUE
	with torch.amp.autocast(
	device_type="cuda" if wave_device.type == "cuda" else "cpu", enabled=False
	):
	wave_f32 = wave.float()
	mel_spec = self.compute_spec(wave_f32)

	logmel_spec = torch.log(mel_spec + 1e-10)
	logmel_spec = torch.clamp(logmel_spec, -23, 12)
	logmel_spec = (logmel_spec + 23) / (12 + 23)

	return logmel_spec.permute(0, 2, 1)[
	:, self.window_pad_idxs : -(self.window_pad_idxs + 1), :
	] # (batch, len, emb)


	class Encoder(nn.Module):
	def __init__(
	self,
	enc_layers,
	d_query,
	nhead,
	ffn_hid_dim,
	dropout,
	):
	super(Encoder, self).__init__()
	self.num_features = d_query * nhead
	self.dense_layer = nn.Linear(self.num_features, self.num_features, bias=False)
	self.positional_encoding = PositionalEncoding(self.num_features)
	self.layer_norm = nn.LayerNorm(
	normalized_shape=self.num_features, elementwise_affine=True
	)
	self.dropout_layer = nn.Dropout(p=dropout)

	layer = nn.TransformerEncoderLayer(
	d_model=self.num_features,
	nhead=nhead,
	dim_feedforward=ffn_hid_dim,
	dropout=dropout,
	activation="gelu",
	batch_first=True,
	norm_first=False,
	)
	self.encoder = nn.TransformerEncoder(layer, num_layers=enc_layers)

	def forward(self, src_emb):
	src_emb = self.dense_layer(src_emb)
	src_emb = self.positional_encoding(src_emb)
	src_emb = self.dropout_layer(src_emb)
	encoder_output = self.encoder(src_emb)
	encoder_output = self.dropout_layer(self.layer_norm(encoder_output))
	return encoder_output


	class Decoder(nn.Module):
	def __init__(
	self,
	dec_layers,
	d_query,
	nhead,
	ffn_hid_dim,
	tgt_vocab_size,
	dropout,
	plain=False,
	):
	super(Decoder, self).__init__()
	self.num_features = d_query * nhead
	if plain:
	self.tgt_tok_emb = TokenEmbedding_plain(tgt_vocab_size, self.num_features)
	if not plain:
	self.tgt_tok_emb = TokenEmbedding(tgt_vocab_size, self.num_features)
	self.positional_encoding = PositionalEncoding(self.num_features)
	self.dropout_layer = nn.Dropout(p=dropout)
	self.generator = nn.Linear(self.num_features, tgt_vocab_size)

	layer = nn.TransformerDecoderLayer(
	d_model=self.num_features,
	nhead=nhead,
	dim_feedforward=ffn_hid_dim,
	dropout=dropout,
	activation="gelu",
	batch_first=True,
	norm_first=False,
	)
	self.decoder = nn.TransformerDecoder(layer, num_layers=dec_layers)

	def forward(self, tgt, encoder_output, tgt_mask, tgt_padding_mask):
	tgt_emb = self.positional_encoding(self.tgt_tok_emb(tgt))
	tgt_emb = self.dropout_layer(tgt_emb)
	# bf16-safe and PyTorch same-type: use additive float masks (0 / -1e4) instead of bool (-inf)
	if tgt_mask is not None:
	tgt_mask = torch.zeros_like(
	tgt_mask, dtype=tgt_emb.dtype, device=tgt_mask.device
	).masked_fill_(tgt_mask, -1e4)
	if tgt_padding_mask is not None and tgt_padding_mask.dtype == torch.bool:
	tgt_padding_mask = torch.zeros_like(
	tgt_padding_mask, dtype=tgt_emb.dtype, device=tgt_padding_mask.device
	).masked_fill_(tgt_padding_mask, -1e4)
	decoder_outputs = self.decoder(
	tgt_emb,
	encoder_output,
	tgt_mask=tgt_mask,
	tgt_key_padding_mask=tgt_padding_mask,
	memory_mask=None,
	memory_key_padding_mask=None,
	)
	return self.generator(decoder_outputs)


	class ADTModel(PreTrainedModel):
	config_class = ADTModelConfig
	_supports_assign_param_buffer = False
	_tied_weights_keys = []
	all_tied_weights_keys = {}

	def __init__(
	self,
	config: ADTModelConfig,
	) -> None:
	super(ADTModel, self).__init__(config)
	self.config = config
	self.encoder = Encoder(
	enc_layers=config.enc_layers,
	d_query=config.d_query,
	nhead=config.nhead,
	ffn_hid_dim=int(config.d_query * config.nhead * 4),
	dropout=config.dropout,
	)
	self.decoder = Decoder(
	dec_layers=config.dec_layers,
	d_query=config.d_query,
	nhead=config.nhead,
	ffn_hid_dim=int(config.d_query * config.nhead * 4),
	dropout=config.dropout,
	tgt_vocab_size=config.tgt_vocab_size,
	plain=config.plain,
	)
	self.compute_spectrogram = ComputeMelSpectrogram(
	sample_rate=config.sample_rate,
	win_length=config.win_length,
	time_res=config.time_res,
	n_mels=config.n_mels,
	)
	self.project_to_mel = nn.Linear(
	config.n_mels, int(config.d_query * config.nhead)
	)

	def _loss_fn(
	self, logits: torch.FloatTensor, tgt: torch.LongTensor
	) -> torch.FloatTensor:
	# Use fp32 for stable loss and ignore PAD token (id=1)
	logits = logits.float()
	logits = torch.nan_to_num(logits, nan=0.0, posinf=1e4, neginf=-1e4)
	return F.cross_entropy(
	logits.reshape(-1, logits.shape[-1]),
	tgt.long().reshape(-1),
	ignore_index=1,
	)

	def forward(
	self,
	src: torch.FloatTensor,
	tgt: torch.LongTensor,
	tgt_mask: Optional[torch.BoolTensor],
	tgt_padding_mask: Optional[torch.BoolTensor],
	labels: torch.LongTensor,
	) -> torch.FloatTensor:
	src_emb = self.compute_spectrogram(src)
	src_emb = self.project_to_mel(src_emb)
	encoder_out = self.encoder(src_emb)

	# Build causal mask inside model so DataParallel does not split (T,T) across devices
	if tgt_mask is None:
	tgt_seq_len = tgt.size(1)
	tgt_mask, _ = create_mask_plain(tgt_seq_len, None, tgt.device)
	logits = self.decoder(tgt, encoder_out, tgt_mask, tgt_padding_mask)
	loss = self._loss_fn(logits, labels)
	return loss

	def sample(
	self,
	src: torch.FloatTensor,
	src_mask: torch.BoolTensor,
	tgt_mask: torch.BoolTensor,
	max_length: int = 1000,
	start_token: int = 2,
	end_token: int = 3,
	) -> torch.LongTensor:
	"""
	Basic autoregressive greedy sampling for ADT model.

	Args:
	src: Input audio tensor (batch_size, audio_length)
	src_mask: Source mask (not used in current implementation)
	tgt_mask: Target mask (not used in current implementation)
	max_length: Maximum sequence length to generate
	start_token: Token ID to start generation (default: 0)
	end_token: Token ID to end generation (default: 2, assuming EOS)

	Returns:
	Generated token sequence (batch_size, seq_length)
	"""
	if not self.config.plain:
	raise NotImplementedError("Non-plain mode is not implemented")
	self.eval()
	device = src.device
	batch_size = src.shape[0]

	# Process input through encoder
	src_emb = self.compute_spectrogram(src)
	src_emb = self.project_to_mel(src_emb)
	memory = self.encoder(src_emb)

	# Initialize with start token
	generated = torch.full(
	(batch_size, 1), start_token, dtype=torch.long, device=device
	)
	finished = torch.zeros(batch_size, dtype=torch.bool, device=device)
	# Greedy decoding loop (per-row finished mask so batch items can end at different steps)
	for step in range(max_length - 1):
	# Create causal mask for current sequence length (no padding in decoding)
	seq_len = generated.shape[1]
	tgt_mask, tgt_padding_mask = create_mask_plain(seq_len, None, device)

	# Get decoder output
	logits = self.decoder(
	generated,
	memory,
	tgt_mask,
	tgt_padding_mask=None,
	)

	# Get next token (greedy); freeze finished rows at EOS so they do not keep consuming argmax noise
	next_token_logits = logits[:, -1, :] # Last position
	next_token = torch.argmax(next_token_logits, dim=-1) # Greedy selection
	next_token = torch.where(
	finished, torch.full_like(next_token, end_token), next_token
	)
	generated = torch.cat([generated, next_token.unsqueeze(1)], dim=1)
	finished = finished \| (next_token == end_token)
	if torch.all(finished):
	break

	return generated

	def beam_search(
	self,
	src: torch.FloatTensor,
	src_mask: torch.BoolTensor,
	tgt_mask: torch.BoolTensor,
	beam_size: int = 5,
	max_length: int = 1000,
	start_token: int = 2,
	end_token: int = 3,
	length_penalty: float = 1.0,
	) -> torch.LongTensor:
	"""
	Beam search sampling for ADT model.

	Args:
	src: Input audio tensor (batch_size, audio_length)
	src_mask: Source mask (not used in current implementation)
	tgt_mask: Target mask (not used in current implementation)
	beam_size: Number of beams to keep at each step
	max_length: Maximum sequence length to generate
	start_token: Token ID to start generation (default: 0)
	end_token: Token ID to end generation (default: 2, assuming EOS)
	length_penalty: Length penalty factor for scoring (default: 1.0)

	Returns:
	Best generated token sequence (batch_size, seq_length)
	"""
	if not self.config.plain:
	raise NotImplementedError("Non-plain mode is not implemented")

	self.eval()
	device = src.device
	batch_size = src.shape[0]

	# Process input through encoder
	src_emb = self.compute_spectrogram(src)
	src_emb = self.project_to_mel(src_emb)
	memory = self.encoder(src_emb)

	# Initialize beams for each batch
	# Each beam contains: [batch_idx, sequence, log_prob, finished]
	all_beams = []

	for batch_idx in range(batch_size):
	# Initialize beams for this batch item
	initial_sequence = torch.full(
	(1, 1), start_token, dtype=torch.long, device=device
	)
	beams = [{"sequence": initial_sequence, "log_prob": 0.0, "finished": False}]
	all_beams.append(beams)

	# Beam search loop
	for step in range(max_length - 1):
	new_all_beams = []

	for batch_idx in range(batch_size):
	beams = all_beams[batch_idx]
	new_beams = []

	# Skip if all beams are finished
	if all(beam["finished"] for beam in beams):
	new_all_beams.append(beams)
	continue

	# Collect all active sequences for this batch
	active_beams = [beam for beam in beams if not beam["finished"]]
	if not active_beams:
	new_all_beams.append(beams)
	continue

	# Stack sequences for parallel processing
	sequences = torch.cat(
	[beam["sequence"] for beam in active_beams], dim=0
	)
	batch_memory = memory[batch_idx : batch_idx + 1].expand(
	len(active_beams), -1, -1
	)

	# Create causal mask
	seq_len = sequences.shape[1]
	tgt_mask = torch.triu(
	torch.ones(seq_len, seq_len, device=device), diagonal=1
	).bool()

	# Get decoder output
	logits = self.decoder(
	sequences,
	batch_memory,
	tgt_mask,
	None,
	)

	# Get next token probabilities
	next_token_logits = logits[:, -1, :] # Last position
	log_probs = torch.log_softmax(next_token_logits, dim=-1)

	# Generate candidates for each active beam
	candidates = []
	for beam_idx, beam in enumerate(active_beams):
	beam_log_probs = log_probs[beam_idx]
	top_k_log_probs, top_k_tokens = torch.topk(
	beam_log_probs, beam_size
	)

	for k in range(beam_size):
	token = top_k_tokens[k].item()
	token_log_prob = top_k_log_probs[k].item()
	new_log_prob = beam["log_prob"] + token_log_prob

	# Create new sequence
	new_sequence = torch.cat(
	[beam["sequence"], torch.tensor([[token]], device=device)],
	dim=1,
	)

	# Check if finished
	finished = token == end_token

	candidates.append(
	{
	"sequence": new_sequence,
	"log_prob": new_log_prob,
	"finished": finished,
	}
	)

	# Add finished beams (they don't expand further)
	finished_beams = [beam for beam in beams if beam["finished"]]
	candidates.extend(finished_beams)

	# Select top beam_size candidates
	candidates.sort(
	key=lambda x: self._score_sequence(x, length_penalty), reverse=True
	)
	new_beams = candidates[:beam_size]

	new_all_beams.append(new_beams)

	all_beams = new_all_beams

	# Check if all beams in all batches are finished
	all_finished = all(
	all(beam["finished"] for beam in beams) for beams in all_beams
	)
	if all_finished:
	break

	# Select best sequence for each batch
	results = []
	for batch_idx in range(batch_size):
	beams = all_beams[batch_idx]
	# Sort by score and take the best
	beams.sort(
	key=lambda x: self._score_sequence(x, length_penalty), reverse=True
	)
	best_sequence = beams[0]["sequence"]
	results.append(best_sequence)

	# Stack results and pad to same length
	max_len = max(seq.shape[1] for seq in results)
	padded_results = []
	for seq in results:
	if seq.shape[1] < max_len:
	padding = torch.full(
	(1, max_len - seq.shape[1]),
	end_token,
	dtype=torch.long,
	device=device,
	)
	padded_seq = torch.cat([seq, padding], dim=1)
	else:
	padded_seq = seq
	padded_results.append(padded_seq)

	return torch.cat(padded_results, dim=0)

	def _score_sequence(self, beam: dict, length_penalty: float) -> float:
	seq_len = beam["sequence"].shape[1]
	if length_penalty == 0.0:
	return beam["log_prob"]
	else:
	# Length normalization: log_prob / (length^length_penalty)
	return beam["log_prob"] / (seq_len**length_penalty)