xeus / modeling_xeus.py

Convert weights to bfloat16 SafeTensors (~1.07 GB) for high efficiency and lower memory footprint

6c93afe verified about 1 month ago

10 kB

"""Pure PyTorch implementation of the XEUS E-Branchformer acoustic foundation model.
Converted to SafeTensors format (bfloat16) with zero external dependencies (no ESPnet installation required).
Original paper: XEUS - A Cross-lingual Encoder for Universal Speech (CMU WAVLab).
Weights repository: https://huggingface.co/giangndm/xeus
"""

import warnings
from typing import Optional

import torch
import torch.nn as nn
import torch.nn.functional as F
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file


class TransposeLayerNorm(nn.LayerNorm):
    """LayerNorm operating on channel dimension of (B, C, T) tensors."""

    def forward(self, input_tensor: torch.Tensor) -> torch.Tensor:
        transposed = input_tensor.transpose(-2, -1)
        normalized = F.layer_norm(
            transposed, self.normalized_shape, self.weight, self.bias, self.eps
        )
        return normalized.transpose(-2, -1)


class ConvLayerBlock(nn.Module):
    """Convolutional downsampling block used in CNNFrontend."""

    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_size: int,
        stride: int,
        bias: bool = True,
    ) -> None:
        super().__init__()
        self.kernel_size = kernel_size
        self.stride = stride
        self.conv = nn.Conv1d(
            in_channels, out_channels, kernel_size=kernel_size, stride=stride, bias=bias
        )
        self.layer_norm = TransposeLayerNorm(out_channels)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return F.gelu(self.layer_norm(self.conv(x)))


class CNNFrontend(nn.Module):
    """Wav2Vec-style 7-layer convolutional feature extractor for 16kHz audio."""

    def __init__(self) -> None:
        super().__init__()
        shapes = [
            (512, 10, 5),
            (512, 3, 2),
            (512, 3, 2),
            (512, 3, 2),
            (512, 3, 2),
            (512, 2, 2),
            (512, 2, 2),
        ]
        in_channels = 1
        blocks = []
        for out_channels, kernel_size, stride in shapes:
            blocks.append(
                ConvLayerBlock(
                    in_channels, out_channels, kernel_size=kernel_size, stride=stride
                )
            )
            in_channels = out_channels
        self.layers = nn.Sequential(*blocks)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if x.ndim == 2:
            x = F.layer_norm(x, (x.shape[-1],))
            x = x.unsqueeze(1)
        first_weight = self.layers[0].conv.weight
        if x.dtype != first_weight.dtype:
            x = x.to(dtype=first_weight.dtype)
        x = self.layers(x)
        return x.transpose(1, 2)


class Preencoder(nn.Module):
    """Linear projection from frontend 512 dim to encoder 1024 dim."""

    def __init__(self, in_dim: int = 512, out_dim: int = 1024) -> None:
        super().__init__()
        self.linear_out = nn.Linear(in_dim, out_dim)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.linear_out(x)


class ConvPositionalEmbedding(nn.Module):
    """Depthwise 1D convolutional relative positional embedding."""

    def __init__(self, hidden_dim: int = 1024, kernel_size: int = 31) -> None:
        super().__init__()
        conv = nn.Conv1d(
            hidden_dim,
            hidden_dim,
            kernel_size=kernel_size,
            padding=(kernel_size - 1) // 2,
            groups=16,
        )
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            conv = nn.utils.weight_norm(conv, name="weight", dim=2)
        self.convs = nn.ModuleList([conv])

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        residual = x
        x = x.transpose(1, 2)
        for conv in self.convs:
            x = conv(x)
            x = F.gelu(x)
        x = x.transpose(1, 2)
        return residual + x


class CSGU(nn.Module):
    """Convolutional Spatial Gating Unit."""

    def __init__(self, size: int = 3072, kernel_size: int = 31) -> None:
        super().__init__()
        self.norm = nn.LayerNorm(size // 2)
        self.conv = nn.Conv1d(
            size // 2,
            size // 2,
            kernel_size=kernel_size,
            padding=(kernel_size - 1) // 2,
            groups=size // 2,
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        u, v = x.chunk(2, dim=-1)
        v = self.norm(v)
        v = v.transpose(1, 2)
        v = self.conv(v)
        v = v.transpose(1, 2)
        return u * v


class ConvolutionalGatingMLP(nn.Module):
    """Branch containing CSGU gated linear unit and depthwise conv."""

    def __init__(
        self, size: int = 1024, linear_units: int = 3072, kernel_size: int = 31
    ) -> None:
        super().__init__()
        self.channel_proj1 = nn.Linear(size, linear_units)
        self.csgu = CSGU(linear_units, kernel_size)
        self.channel_proj2 = nn.Linear(linear_units // 2, size)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.channel_proj1(x)
        x = F.gelu(x)
        x = self.csgu(x)
        return self.channel_proj2(x)


class MultiHeadedAttention(nn.Module):
    """Multi-Head Self-Attention with projection matrices."""

    def __init__(self, n_head: int = 8, n_feat: int = 1024) -> None:
        super().__init__()
        self.d_k = n_feat // n_head
        self.h = n_head
        self.linear_q = nn.Linear(n_feat, n_feat)
        self.linear_k = nn.Linear(n_feat, n_feat)
        self.linear_v = nn.Linear(n_feat, n_feat)
        self.linear_out = nn.Linear(n_feat, n_feat)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        b, t, d = x.shape
        q = self.linear_q(x).view(b, t, self.h, self.d_k).transpose(1, 2)
        k = self.linear_k(x).view(b, t, self.h, self.d_k).transpose(1, 2)
        v = self.linear_v(x).view(b, t, self.h, self.d_k).transpose(1, 2)
        out = F.scaled_dot_product_attention(q, k, v)
        out = out.transpose(1, 2).contiguous().view(b, t, d)
        return self.linear_out(out)


class EBranchformerEncoderLayer(nn.Module):
    """Single E-Branchformer layer combining Attn, CGMLP and FFN branches."""

    def __init__(
        self,
        size: int = 1024,
        linear_units: int = 4096,
        cgmlp_units: int = 3072,
        heads: int = 8,
        kernel: int = 31,
    ) -> None:
        super().__init__()
        self.norm_mha = nn.LayerNorm(size)
        self.norm_mlp = nn.LayerNorm(size)
        self.mha = MultiHeadedAttention(heads, size)
        self.cgmlp = ConvolutionalGatingMLP(size, cgmlp_units, kernel)
        self.norm_merge = nn.LayerNorm(size)
        self.merge_proj = nn.Linear(size + size, size)
        self.norm_ffn = nn.LayerNorm(size)
        self.ffn = nn.Sequential(
            nn.Linear(size, linear_units),
            nn.GELU(),
            nn.Linear(linear_units, size),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # Branch 1: MHA
        branch1 = self.mha(self.norm_mha(x))
        # Branch 2: CGMLP
        branch2 = self.cgmlp(self.norm_mlp(x))
        # Merge branches
        merged = torch.cat([branch1, branch2], dim=-1)
        merged = self.merge_proj(merged)
        x = x + self.norm_merge(merged)
        # FFN sublayer
        x = x + self.ffn(self.norm_ffn(x))
        return x


class EBranchformerEncoder(nn.Module):
    """Stack of E-Branchformer blocks with relative positional embeddings."""

    def __init__(
        self,
        num_blocks: int = 19,
        size: int = 1024,
        linear_units: int = 4096,
        heads: int = 8,
        kernel: int = 31,
    ) -> None:
        super().__init__()
        self.embed = nn.Sequential(ConvPositionalEmbedding(size, kernel))
        self.encoders = nn.ModuleList([
            EBranchformerEncoderLayer(size, linear_units, 3072, heads, kernel)
            for _ in range(num_blocks)
        ])
        self.after_norm = nn.LayerNorm(size)
        self.gradient_checkpointing = False

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.embed(x)
        for layer in self.encoders:
            if self.gradient_checkpointing and self.training:
                x = torch.utils.checkpoint.checkpoint(layer, x, use_reentrant=False)
            else:
                x = layer(x)
        return self.after_norm(x)


class XeusTrunk(nn.Module):
    """Complete XEUS acoustic foundation trunk (Frontend + Preencoder + EBranchformer)."""

    DEFAULT_REPO_ID = "giangndm/xeus"
    DEFAULT_FILENAME = "model.safetensors"

    def __init__(
        self,
        num_blocks: int = 19,
        hidden_dim: int = 1024,
        linear_units: int = 4096,
        heads: int = 8,
        kernel: int = 31,
    ) -> None:
        super().__init__()
        self.frontend = CNNFrontend()
        self.preencoder = Preencoder(512, hidden_dim)
        self.encoder = EBranchformerEncoder(
            num_blocks, hidden_dim, linear_units, heads, kernel
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.frontend(x)
        x = self.preencoder(x)
        return self.encoder(x)

    @classmethod
    def from_pretrained(
        cls,
        repo_id: str = DEFAULT_REPO_ID,
        revision: Optional[str] = None,
        filename: str = DEFAULT_FILENAME,
        dtype: Optional[torch.dtype] = torch.bfloat16,
    ) -> "XeusTrunk":
        model = cls()
        if dtype is not None:
            model = model.to(dtype=dtype)
        ckpt_path = hf_hub_download(repo_id=repo_id, revision=revision, filename=filename)
        if filename.endswith(".safetensors"):
            state_dict = load_file(ckpt_path)
        else:
            state_dict = torch.load(ckpt_path, map_location="cpu", weights_only=True)
        filtered_state = {
            k: v
            for k, v in state_dict.items()
            if k.startswith(("frontend.", "preencoder.", "encoder."))
        }
        model.load_state_dict(filtered_state, strict=True)
        return model

	"""Pure PyTorch implementation of the XEUS E-Branchformer acoustic foundation model.

	Converted to SafeTensors format (bfloat16) with zero external dependencies (no ESPnet installation required).
	Original paper: XEUS - A Cross-lingual Encoder for Universal Speech (CMU WAVLab).
	Weights repository: https://huggingface.co/giangndm/xeus
	"""

	import warnings
	from typing import Optional

	import torch
	import torch.nn as nn
	import torch.nn.functional as F
	from huggingface_hub import hf_hub_download
	from safetensors.torch import load_file


	class TransposeLayerNorm(nn.LayerNorm):
	"""LayerNorm operating on channel dimension of (B, C, T) tensors."""

	def forward(self, input_tensor: torch.Tensor) -> torch.Tensor:
	transposed = input_tensor.transpose(-2, -1)
	normalized = F.layer_norm(
	transposed, self.normalized_shape, self.weight, self.bias, self.eps
	)
	return normalized.transpose(-2, -1)


	class ConvLayerBlock(nn.Module):
	"""Convolutional downsampling block used in CNNFrontend."""

	def __init__(
	self,
	in_channels: int,
	out_channels: int,
	kernel_size: int,
	stride: int,
	bias: bool = True,
	) -> None:
	super().__init__()
	self.kernel_size = kernel_size
	self.stride = stride
	self.conv = nn.Conv1d(
	in_channels, out_channels, kernel_size=kernel_size, stride=stride, bias=bias
	)
	self.layer_norm = TransposeLayerNorm(out_channels)

	def forward(self, x: torch.Tensor) -> torch.Tensor:
	return F.gelu(self.layer_norm(self.conv(x)))


	class CNNFrontend(nn.Module):
	"""Wav2Vec-style 7-layer convolutional feature extractor for 16kHz audio."""

	def __init__(self) -> None:
	super().__init__()
	shapes = [
	(512, 10, 5),
	(512, 3, 2),
	(512, 3, 2),
	(512, 3, 2),
	(512, 3, 2),
	(512, 2, 2),
	(512, 2, 2),
	]
	in_channels = 1
	blocks = []
	for out_channels, kernel_size, stride in shapes:
	blocks.append(
	ConvLayerBlock(
	in_channels, out_channels, kernel_size=kernel_size, stride=stride
	)
	)
	in_channels = out_channels
	self.layers = nn.Sequential(*blocks)

	def forward(self, x: torch.Tensor) -> torch.Tensor:
	if x.ndim == 2:
	x = F.layer_norm(x, (x.shape[-1],))
	x = x.unsqueeze(1)
	first_weight = self.layers[0].conv.weight
	if x.dtype != first_weight.dtype:
	x = x.to(dtype=first_weight.dtype)
	x = self.layers(x)
	return x.transpose(1, 2)


	class Preencoder(nn.Module):
	"""Linear projection from frontend 512 dim to encoder 1024 dim."""

	def __init__(self, in_dim: int = 512, out_dim: int = 1024) -> None:
	super().__init__()
	self.linear_out = nn.Linear(in_dim, out_dim)

	def forward(self, x: torch.Tensor) -> torch.Tensor:
	return self.linear_out(x)


	class ConvPositionalEmbedding(nn.Module):
	"""Depthwise 1D convolutional relative positional embedding."""

	def __init__(self, hidden_dim: int = 1024, kernel_size: int = 31) -> None:
	super().__init__()
	conv = nn.Conv1d(
	hidden_dim,
	hidden_dim,
	kernel_size=kernel_size,
	padding=(kernel_size - 1) // 2,
	groups=16,
	)
	with warnings.catch_warnings():
	warnings.simplefilter("ignore")
	conv = nn.utils.weight_norm(conv, name="weight", dim=2)
	self.convs = nn.ModuleList([conv])

	def forward(self, x: torch.Tensor) -> torch.Tensor:
	residual = x
	x = x.transpose(1, 2)
	for conv in self.convs:
	x = conv(x)
	x = F.gelu(x)
	x = x.transpose(1, 2)
	return residual + x


	class CSGU(nn.Module):
	"""Convolutional Spatial Gating Unit."""

	def __init__(self, size: int = 3072, kernel_size: int = 31) -> None:
	super().__init__()
	self.norm = nn.LayerNorm(size // 2)
	self.conv = nn.Conv1d(
	size // 2,
	size // 2,
	kernel_size=kernel_size,
	padding=(kernel_size - 1) // 2,
	groups=size // 2,
	)

	def forward(self, x: torch.Tensor) -> torch.Tensor:
	u, v = x.chunk(2, dim=-1)
	v = self.norm(v)
	v = v.transpose(1, 2)
	v = self.conv(v)
	v = v.transpose(1, 2)
	return u * v


	class ConvolutionalGatingMLP(nn.Module):
	"""Branch containing CSGU gated linear unit and depthwise conv."""

	def __init__(
	self, size: int = 1024, linear_units: int = 3072, kernel_size: int = 31
	) -> None:
	super().__init__()
	self.channel_proj1 = nn.Linear(size, linear_units)
	self.csgu = CSGU(linear_units, kernel_size)
	self.channel_proj2 = nn.Linear(linear_units // 2, size)

	def forward(self, x: torch.Tensor) -> torch.Tensor:
	x = self.channel_proj1(x)
	x = F.gelu(x)
	x = self.csgu(x)
	return self.channel_proj2(x)


	class MultiHeadedAttention(nn.Module):
	"""Multi-Head Self-Attention with projection matrices."""

	def __init__(self, n_head: int = 8, n_feat: int = 1024) -> None:
	super().__init__()
	self.d_k = n_feat // n_head
	self.h = n_head
	self.linear_q = nn.Linear(n_feat, n_feat)
	self.linear_k = nn.Linear(n_feat, n_feat)
	self.linear_v = nn.Linear(n_feat, n_feat)
	self.linear_out = nn.Linear(n_feat, n_feat)

	def forward(self, x: torch.Tensor) -> torch.Tensor:
	b, t, d = x.shape
	q = self.linear_q(x).view(b, t, self.h, self.d_k).transpose(1, 2)
	k = self.linear_k(x).view(b, t, self.h, self.d_k).transpose(1, 2)
	v = self.linear_v(x).view(b, t, self.h, self.d_k).transpose(1, 2)
	out = F.scaled_dot_product_attention(q, k, v)
	out = out.transpose(1, 2).contiguous().view(b, t, d)
	return self.linear_out(out)


	class EBranchformerEncoderLayer(nn.Module):
	"""Single E-Branchformer layer combining Attn, CGMLP and FFN branches."""

	def __init__(
	self,
	size: int = 1024,
	linear_units: int = 4096,
	cgmlp_units: int = 3072,
	heads: int = 8,
	kernel: int = 31,
	) -> None:
	super().__init__()
	self.norm_mha = nn.LayerNorm(size)
	self.norm_mlp = nn.LayerNorm(size)
	self.mha = MultiHeadedAttention(heads, size)
	self.cgmlp = ConvolutionalGatingMLP(size, cgmlp_units, kernel)
	self.norm_merge = nn.LayerNorm(size)
	self.merge_proj = nn.Linear(size + size, size)
	self.norm_ffn = nn.LayerNorm(size)
	self.ffn = nn.Sequential(
	nn.Linear(size, linear_units),
	nn.GELU(),
	nn.Linear(linear_units, size),
	)

	def forward(self, x: torch.Tensor) -> torch.Tensor:
	# Branch 1: MHA
	branch1 = self.mha(self.norm_mha(x))
	# Branch 2: CGMLP
	branch2 = self.cgmlp(self.norm_mlp(x))
	# Merge branches
	merged = torch.cat([branch1, branch2], dim=-1)
	merged = self.merge_proj(merged)
	x = x + self.norm_merge(merged)
	# FFN sublayer
	x = x + self.ffn(self.norm_ffn(x))
	return x


	class EBranchformerEncoder(nn.Module):
	"""Stack of E-Branchformer blocks with relative positional embeddings."""

	def __init__(
	self,
	num_blocks: int = 19,
	size: int = 1024,
	linear_units: int = 4096,
	heads: int = 8,
	kernel: int = 31,
	) -> None:
	super().__init__()
	self.embed = nn.Sequential(ConvPositionalEmbedding(size, kernel))
	self.encoders = nn.ModuleList([
	EBranchformerEncoderLayer(size, linear_units, 3072, heads, kernel)
	for _ in range(num_blocks)
	])
	self.after_norm = nn.LayerNorm(size)
	self.gradient_checkpointing = False

	def forward(self, x: torch.Tensor) -> torch.Tensor:
	x = self.embed(x)
	for layer in self.encoders:
	if self.gradient_checkpointing and self.training:
	x = torch.utils.checkpoint.checkpoint(layer, x, use_reentrant=False)
	else:
	x = layer(x)
	return self.after_norm(x)


	class XeusTrunk(nn.Module):
	"""Complete XEUS acoustic foundation trunk (Frontend + Preencoder + EBranchformer)."""

	DEFAULT_REPO_ID = "giangndm/xeus"
	DEFAULT_FILENAME = "model.safetensors"

	def __init__(
	self,
	num_blocks: int = 19,
	hidden_dim: int = 1024,
	linear_units: int = 4096,
	heads: int = 8,
	kernel: int = 31,
	) -> None:
	super().__init__()
	self.frontend = CNNFrontend()
	self.preencoder = Preencoder(512, hidden_dim)
	self.encoder = EBranchformerEncoder(
	num_blocks, hidden_dim, linear_units, heads, kernel
	)

	def forward(self, x: torch.Tensor) -> torch.Tensor:
	x = self.frontend(x)
	x = self.preencoder(x)
	return self.encoder(x)

	@classmethod
	def from_pretrained(
	cls,
	repo_id: str = DEFAULT_REPO_ID,
	revision: Optional[str] = None,
	filename: str = DEFAULT_FILENAME,
	dtype: Optional[torch.dtype] = torch.bfloat16,
	) -> "XeusTrunk":
	model = cls()
	if dtype is not None:
	model = model.to(dtype=dtype)
	ckpt_path = hf_hub_download(repo_id=repo_id, revision=revision, filename=filename)
	if filename.endswith(".safetensors"):
	state_dict = load_file(ckpt_path)
	else:
	state_dict = torch.load(ckpt_path, map_location="cpu", weights_only=True)
	filtered_state = {
	k: v
	for k, v in state_dict.items()
	if k.startswith(("frontend.", "preencoder.", "encoder."))
	}
	model.load_state_dict(filtered_state, strict=True)
	return model