mirror of https://github.com/hpcaitech/ColossalAI
98 lines
3.6 KiB
Python
98 lines
3.6 KiB
Python
import torch
|
|
import transformers
|
|
|
|
from ..registry import ModelAttribute, model_zoo
|
|
|
|
# ===============================
|
|
# Register single-sentence Whisper
|
|
# ===============================
|
|
|
|
|
|
# define data gen function
|
|
def data_gen():
|
|
# Generated from following code snippet
|
|
#
|
|
# from transformers import AutoFeatureExtractor, WhisperModel
|
|
# from datasets import load_dataset
|
|
|
|
# model = WhisperModel.from_pretrained("openai/whisper-base")
|
|
# feature_extractor = AutoFeatureExtractor.from_pretrained("openai/whisper-base")
|
|
# ds = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation")
|
|
# inputs = feature_extractor(ds[0]["audio"]["array"], return_tensors="pt")
|
|
# input_features = inputs.input_features
|
|
# decoder_input_ids = torch.tensor([[1, 1]]) * model.config.decoder_start_token_id
|
|
|
|
input_features = torch.rand(1, 80, 3000)
|
|
decoder_input_ids = torch.tensor([[1, 1]]) * 50258
|
|
return dict(input_features=input_features, decoder_input_ids=decoder_input_ids)
|
|
|
|
|
|
def data_gen_for_conditional_generation():
|
|
# labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
|
|
# Labels for computing the language modeling loss. Indices should either be in `[0, ..., config.vocab_size]`
|
|
# or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored (masked), the loss is
|
|
# only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
|
|
data = data_gen()
|
|
data["labels"] = torch.tensor([[0, 1]], dtype=torch.int64)
|
|
return data
|
|
|
|
|
|
def data_gen_for_audio_classification():
|
|
# labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
|
|
# Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
|
|
# config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
|
|
# `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
|
|
# `WhisperForAudioClassification` does not need `decoder_input_ids`
|
|
data = data_gen()
|
|
data.pop("decoder_input_ids")
|
|
data["labels"] = torch.tensor([1], dtype=torch.int64)
|
|
return data
|
|
|
|
|
|
# define output transform function
|
|
output_transform_fn = lambda x: x
|
|
|
|
# define loss funciton
|
|
loss_fn = lambda x: torch.nn.functional.mse_loss(x.last_hidden_state, torch.ones_like(x.last_hidden_state))
|
|
loss_fn_attr = lambda x: x.loss
|
|
|
|
config = transformers.WhisperConfig(
|
|
classifier_proj_size=256,
|
|
d_model=256,
|
|
decoder_attention_heads=4,
|
|
decoder_ffn_dim=1536,
|
|
decoder_layers=2,
|
|
encoder_attention_heads=4,
|
|
encoder_ffn_dim=1536,
|
|
encoder_layers=2,
|
|
vocab_size=51866,
|
|
)
|
|
|
|
# register the Whisper variants
|
|
model_zoo.register(
|
|
name="transformers_whisper",
|
|
model_fn=lambda: transformers.WhisperModel(config),
|
|
data_gen_fn=data_gen,
|
|
output_transform_fn=output_transform_fn,
|
|
loss_fn=loss_fn,
|
|
model_attribute=ModelAttribute(has_control_flow=True),
|
|
)
|
|
|
|
model_zoo.register(
|
|
name="transformers_whisper_for_conditional_generation",
|
|
model_fn=lambda: transformers.WhisperForConditionalGeneration(config),
|
|
data_gen_fn=data_gen_for_conditional_generation,
|
|
output_transform_fn=output_transform_fn,
|
|
loss_fn=loss_fn_attr,
|
|
model_attribute=ModelAttribute(has_control_flow=True),
|
|
)
|
|
|
|
model_zoo.register(
|
|
name="transformers_whisper_for_audio_classification",
|
|
model_fn=lambda: transformers.WhisperForAudioClassification(config),
|
|
data_gen_fn=data_gen_for_audio_classification,
|
|
output_transform_fn=output_transform_fn,
|
|
loss_fn=loss_fn_attr,
|
|
model_attribute=ModelAttribute(has_control_flow=True),
|
|
)
|