Spaces:
Runtime error
Runtime error
File size: 3,569 Bytes
1547a56 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 |
# --------------------------------------------------------
# ArTST: Arabic Text and Speech Transformer (https://arxiv.org/abs/2310.16621)
# Github source: https://github.com/mbzuai-nlp/ArTST
# Based on speecht5, fairseq and espnet code bases
# https://github.com/microsoft/SpeechT5/tree/main/SpeechT5; https://github.com/pytorch/fairseq; https://github.com/espnet/espnet
# --------------------------------------------------------
import torch.nn as nn
import torch
import contextlib
from fairseq import utils
from fairseq.modules import (
AdaptiveSoftmax,
)
class TextDecoderPostnet(nn.Module):
"""
Args:
in_channels (int): the number of input channels
mid_channels (int): the number of intermediate channels
out_channels (int): the number of output channels
kernel_sizes (List[int]): the kernel size for each convolutional layer
"""
def __init__(self, embed_tokens, dictionary, args, output_projection=None,):
super(TextDecoderPostnet, self).__init__()
self.output_embed_dim = args.decoder_output_dim
self.output_projection = output_projection
self.adaptive_softmax = None
self.share_input_output_embed = args.share_input_output_embed
if self.output_projection is None:
self.build_output_projection(args, dictionary, embed_tokens)
self.freeze_decoder_updates = args.freeze_decoder_updates
self.num_updates = 0
def output_layer(self, features):
"""Project features to the vocabulary size."""
if self.adaptive_softmax is None:
# project back to size of vocabulary
return self.output_projection(features)
else:
return features
def build_output_projection(self, args, dictionary, embed_tokens):
if args.adaptive_softmax_cutoff is not None:
self.adaptive_softmax = AdaptiveSoftmax(
len(dictionary),
self.output_embed_dim,
utils.eval_str_list(args.adaptive_softmax_cutoff, type=int),
dropout=args.adaptive_softmax_dropout,
adaptive_inputs=embed_tokens if args.tie_adaptive_weights else None,
factor=args.adaptive_softmax_factor,
tie_proj=args.tie_adaptive_proj,
)
elif self.share_input_output_embed:
self.output_projection = nn.Linear(
embed_tokens.weight.shape[1],
embed_tokens.weight.shape[0],
bias=False,
)
self.output_projection.weight = embed_tokens.weight
else:
self.output_projection = nn.Linear(
self.output_embed_dim, len(dictionary), bias=False
)
nn.init.normal_(
self.output_projection.weight, mean=0, std=self.output_embed_dim ** -0.5
)
# num_base_layers = getattr(args, "base_layers", 0)
# for i in range(num_base_layers):
# self.layers.insert(
# ((i + 1) * args.decoder_layers) // (num_base_layers + 1),
# BaseLayer(args),
# )
def forward(self, x):
ft = self.freeze_decoder_updates <= self.num_updates
with torch.no_grad() if not ft else contextlib.ExitStack():
return self._forward(x)
def _forward(self, x):
# embed positions
x = self.output_layer(x)
return x
def set_num_updates(self, num_updates):
"""Set the number of parameters updates."""
self.num_updates = num_updates
|