From 34da725872252840970b529974ea35b33078ff9f Mon Sep 17 00:00:00 2001 From: thucpham Date: Wed, 7 May 2025 10:10:25 +0200 Subject: [PATCH 1/3] add whisper nmt --- CMakeLists.txt | 1 + include/ctranslate2/layers/transformer.h | 6 + include/ctranslate2/layers/whisper.h | 20 + include/ctranslate2/models/whisper_nmt.h | 251 ++++++ python/cpp/module.cc | 1 + python/cpp/module.h | 1 + python/cpp/whisper_nmt.cc | 329 ++++++++ python/ctranslate2/converters/__init__.py | 1 + python/ctranslate2/converters/transformers.py | 331 +++++--- python/ctranslate2/converters/whisper_nmt.py | 722 +++++++++++++++++ python/ctranslate2/models/__init__.py | 2 + python/ctranslate2/specs/__init__.py | 1 + python/ctranslate2/specs/transformer_spec.py | 126 +++ python/ctranslate2/specs/whispernmt_spec.py | 533 +++++++++++++ src/layers/transformer.cc | 64 ++ src/layers/whisper.cc | 16 + src/models/model_factory.cc | 2 + src/models/whisper_nmt.cc | 735 ++++++++++++++++++ 18 files changed, 3028 insertions(+), 114 deletions(-) create mode 100644 include/ctranslate2/models/whisper_nmt.h create mode 100644 python/cpp/whisper_nmt.cc create mode 100644 python/ctranslate2/converters/whisper_nmt.py create mode 100644 python/ctranslate2/specs/whispernmt_spec.py create mode 100644 src/models/whisper_nmt.cc diff --git a/CMakeLists.txt b/CMakeLists.txt index 8eed5f18f..99bd6fcd7 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -142,6 +142,7 @@ set(SOURCES src/models/wav2vec2.cc src/models/wav2vec2bert.cc src/models/whisper.cc + src/models/whisper_nmt.cc src/ops/activation.cc src/ops/add.cc src/ops/alibi_add.cc diff --git a/include/ctranslate2/layers/transformer.h b/include/ctranslate2/layers/transformer.h index 017f9b675..9464bd0b8 100644 --- a/include/ctranslate2/layers/transformer.h +++ b/include/ctranslate2/layers/transformer.h @@ -134,6 +134,12 @@ namespace ctranslate2 { const StorageView* lengths, StorageView& output) override; + void operator()(const StorageView& input, + const std::vector& language_ids, + const std::vector& target_ids, + const StorageView* lengths, + StorageView& output); + size_t num_input_features() const override { return _embeddings.num_inputs(); } diff --git a/include/ctranslate2/layers/whisper.h b/include/ctranslate2/layers/whisper.h index 9590df8a2..92032ef5b 100644 --- a/include/ctranslate2/layers/whisper.h +++ b/include/ctranslate2/layers/whisper.h @@ -70,5 +70,25 @@ namespace ctranslate2 { StorageView& logits); }; + class WhisperConnector :public Layer { + public: + WhisperConnector(const models::Model& model, + const std::string& scope); + + void operator()(const StorageView& features, StorageView& output) const; + + DataType output_type() const override { + return _lin_2.output_type(); + } + + dim_t output_size() const override { + return _lin_2.output_size(); + } + + private: + const ops::ActivationType _activation_type; + const Dense _lin_1; + const Dense _lin_2; + }; } } diff --git a/include/ctranslate2/models/whisper_nmt.h b/include/ctranslate2/models/whisper_nmt.h new file mode 100644 index 000000000..e6b6f9271 --- /dev/null +++ b/include/ctranslate2/models/whisper_nmt.h @@ -0,0 +1,251 @@ +#pragma once + +#include "ctranslate2/generation.h" +#include "ctranslate2/layers/whisper.h" +#include "ctranslate2/models/model.h" +#include "ctranslate2/replica_pool.h" + +namespace ctranslate2 { + namespace models { + + struct WhisperNmtOptions { + // Beam size to use for beam search (set 1 to run greedy search). + size_t beam_size = 5; + + // Beam search patience factor, as described in https://arxiv.org/abs/2204.05424. + // The decoding will continue until beam_size*patience hypotheses are finished. + float patience = 1; + + // Exponential penalty applied to the length during beam search. + float length_penalty = 1; + + // Coverage penalty weight applied during beam search. + float coverage_penalty = 0; + + // Penalty applied to the score of previously generated tokens, as described in + // https://arxiv.org/abs/1909.05858 (set > 1 to penalize). + float repetition_penalty = 1; + + // Prevent repetitions of ngrams with this size (set 0 to disable). + size_t no_repeat_ngram_size = 0; + + // Maximum generation length. + size_t max_length = 448; + + // Randomly sample from the top K candidates (set 0 to sample from the full distribution). + size_t sampling_topk = 1; + + // Keep the most probable tokens whose cumulative probability exceeds this value. + float sampling_topp = 1; + + // High temperatures increase randomness. + float sampling_temperature = 1; + + // Number of hypotheses to include in the result. + size_t num_hypotheses = 1; + + // Include scores in the result. + bool return_scores = false; + + // Store attention vectors in the TranslationResult class. + bool return_attention = false; + + // Include log probs of each token in the result + bool return_logits_vocab = false; + + // Include the probability of the no speech token in the result. + bool return_no_speech_prob = false; + + // Maximum index of the first predicted timestamp. + size_t max_initial_timestamp_index = 50; + + // Suppress blank outputs at the beginning of the sampling. + bool suppress_blank = true; + + // Return alternatives at the first unconstrained decoding position. This is typically + // used with a target prefix to provide alternatives at a specifc location in the + // translation. + bool return_alternatives = false; + // Minimum probability to expand an alternative. + float min_alternative_expansion_prob = 0; + + // List of token IDs to suppress. + // -1 will suppress a default set of symbols as defined in the model config.json file. + std::vector suppress_tokens = {-1}; + + // Replace unknown target tokens by the original source token with the highest attention. + bool replace_unknowns = false; + + // Biases decoding towards a given prefix, see https://arxiv.org/abs/1912.03393 --section 4.2 + // Only activates biased-decoding when beta is in range (0, 1) and SearchStrategy is set to BeamSearch. + // The closer beta is to 1, the stronger the bias is towards the given prefix. + // + // If beta <= 0 and a non-empty prefix is given, then the prefix will be used as a + // hard-prefix rather than a soft, biased-prefix. + float prefix_bias_beta = 0; + + // Decoding length constraints. + size_t max_decoding_length = 256; + size_t min_decoding_length = 1; + + // Disable the generation of some sequences of tokens. + std::vector> suppress_sequences; + + // Disable the generation of the unknown token. + bool disable_unk = false; + + // Stop the decoding on one of these tokens (defaults to the model EOS token). + std::variant, std::vector> end_token; + + // Include the end token in the result. + bool return_end_token = false; + + // Truncate the inputs after this many tokens (set 0 to disable truncation). + size_t max_input_length = 1024; + }; + + struct WhisperNmtGenerationResult { + std::vector> sequences; + //std::vector> sequences_ids; + std::vector scores; + std::vector>> attention; + std::vector> logits; + float no_speech_prob = 0; + + size_t num_sequences() const { + return sequences.size(); + } + + bool has_scores() const { + return !scores.empty(); + } + }; + + struct WhisperNmtAlignmentResult { + std::vector> alignments; + std::vector text_token_probs; + }; + + class WhisperNmtModel : public Model { + public: + const Vocabulary& get_vocabulary() const; + size_t num_source_vocabularies() const; + const Vocabulary& get_source_vocabulary(size_t index = 0) const; + const Vocabulary& get_target_vocabulary() const; + + size_t current_spec_revision() const override; + bool is_quantizable(const std::string& variable_name) const override; + bool is_linear_weight(const std::string& variable_name) const override; + std::unique_ptr clone() const override; + + bool use_global_int16_scale() const override { + return false; + } + + bool with_source_bos() const { + return config["add_source_bos"]; + } + + bool with_source_eos() const { + return config["add_source_eos"]; + } + + const std::string* decoder_start_token() const { + auto& start_token = config["decoder_start_token"]; + return start_token.is_null() ? nullptr : start_token.get_ptr(); + } + + protected: + void initialize(ModelReader& model_reader) override; + + private: + std::shared_ptr _vocabulary; + std::vector> _source_vocabularies; + std::shared_ptr _target_vocabulary; + + void load_vocabularies(ModelReader& model_reader); + }; + + class WhisperNmtReplica : public ModelReplica { + public: + static std::unique_ptr create_from_model(const Model& model); + + WhisperNmtReplica(const std::shared_ptr& model); + + bool is_multilingual() const { + return _is_multilingual; + } + + size_t n_mels() const { + return _n_mels; + } + + size_t num_languages() const { + return _num_languages; + } + + StorageView encode(StorageView features, const bool to_cpu); + + std::vector + generate(StorageView features, + const std::vector& language, + const std::vector>& eos, + const WhisperNmtOptions& options); + + std::vector>> + make_source_ids(const std::vector>>& source_features, + size_t max_length) const; + std::vector> + make_target_ids(const std::vector>& target, + size_t max_length, + bool is_prefix) const; + + private: + std::vector + _run_translation(StorageView& features, + const std::vector& language, + const std::vector>& eos, + const WhisperNmtOptions& options); + void + _nmt_encode(StorageView& features_ids, + const std::vector>>& language_ids, + const std::vector>>& eos_ids, + StorageView& memory, + StorageView& memory_lengths); + const std::shared_ptr _model; + const std::unique_ptr _encoder; + //const std::unique_ptr _decoder; + const std::unique_ptr _transformer_encoder; + const std::unique_ptr _transformer_decoder; + const std::unique_ptr _connector; + + size_t _sot_id; + size_t _eot_id; + size_t _no_timestamps_id; + size_t _no_speech_id; + size_t _n_mels; + size_t _num_languages; + bool _is_multilingual; + + StorageView maybe_encode(StorageView features); + }; + + class WhisperNmt : public ReplicaPool { + public: + using ReplicaPool::ReplicaPool; + + bool is_multilingual() const; + size_t n_mels() const; + size_t num_languages() const; + + std::future encode(const StorageView& features, const bool to_cpu); + + std::vector> + generate(const StorageView& features, + std::vector& language, + std::vector>& eos, + WhisperNmtOptions options = {}); + }; + + } +} diff --git a/python/cpp/module.cc b/python/cpp/module.cc index 550aea5b2..85f33e343 100644 --- a/python/cpp/module.cc +++ b/python/cpp/module.cc @@ -86,6 +86,7 @@ PYBIND11_MODULE(_ext, m) ctranslate2::python::register_generator(m); ctranslate2::python::register_encoder(m); ctranslate2::python::register_whisper(m); + ctranslate2::python::register_whisper_nmt(m); ctranslate2::python::register_wav2vec2(m); ctranslate2::python::register_wav2vec2bert(m); ctranslate2::python::register_mpi(m); diff --git a/python/cpp/module.h b/python/cpp/module.h index 71d4b3b29..a3be48ff0 100644 --- a/python/cpp/module.h +++ b/python/cpp/module.h @@ -17,6 +17,7 @@ namespace ctranslate2 { void register_translation_stats(py::module& m); void register_translator(py::module& m); void register_whisper(py::module& m); + void register_whisper_nmt(py::module& m); void register_wav2vec2(py::module& m); void register_wav2vec2bert(py::module& m); void register_mpi(py::module& m); diff --git a/python/cpp/whisper_nmt.cc b/python/cpp/whisper_nmt.cc new file mode 100644 index 000000000..6f8ac858a --- /dev/null +++ b/python/cpp/whisper_nmt.cc @@ -0,0 +1,329 @@ +#include "module.h" + +#include + +#include "replica_pool.h" + +namespace ctranslate2 { + namespace python { + + class WhisperNmtWrapper : public ReplicaPoolHelper { + public: + using ReplicaPoolHelper::ReplicaPoolHelper; + + bool is_multilingual() const { + return _pool->is_multilingual(); + } + + size_t n_mels() const { + return _pool->n_mels(); + } + + size_t num_languages() const { + return _pool->num_languages(); + } + + StorageView encode(const StorageView &features, const bool to_cpu) { + return _pool->encode(features, to_cpu).get(); + } + + std::variant, + std::vector>> + generate(const StorageView &features, + Tokens language, + BatchTokens eos, + bool asynchronous, + size_t beam_size, + float patience, + size_t num_hypotheses, + float length_penalty, + float coverage_penalty, + float repetition_penalty, + size_t no_repeat_ngram_size, + size_t max_length, + bool return_scores, + bool return_attention, + bool return_logits_vocab, + bool return_no_speech_prob, + bool return_alternatives, + float min_alternative_expansion_prob, + size_t max_initial_timestamp_index, + bool suppress_blank, + const std::optional> &suppress_tokens, + size_t sampling_topk, + float sampling_topp, + float sampling_temperature, + bool replace_unknowns, + float prefix_bias_beta, + size_t max_decoding_length, + size_t min_decoding_length, + bool disable_unk, + const std::optional &end_token, + bool return_end_token, + size_t max_input_length + ) { + std::vector> futures; + + models::WhisperNmtOptions options; + options.beam_size = beam_size; + options.patience = patience; + options.length_penalty = length_penalty; + options.coverage_penalty = coverage_penalty; + options.repetition_penalty = repetition_penalty; + options.no_repeat_ngram_size = no_repeat_ngram_size; + options.sampling_topk = sampling_topk; + options.sampling_topp = sampling_topp; + options.sampling_temperature = sampling_temperature; + options.max_length = max_length; + options.num_hypotheses = num_hypotheses; + options.return_scores = return_scores; + options.return_attention = return_attention; + options.return_logits_vocab = return_logits_vocab; + options.return_no_speech_prob = return_no_speech_prob; + options.return_alternatives = return_alternatives; + options.min_alternative_expansion_prob = min_alternative_expansion_prob; + options.max_initial_timestamp_index = max_initial_timestamp_index; + options.suppress_blank = suppress_blank; + options.replace_unknowns = replace_unknowns; + options.prefix_bias_beta = prefix_bias_beta; + options.max_decoding_length = max_decoding_length; + options.min_decoding_length = min_decoding_length; + options.disable_unk = disable_unk; + options.return_end_token = return_end_token; + options.max_input_length = max_input_length; + if (end_token) + options.end_token = end_token.value(); + + if (suppress_tokens) + options.suppress_tokens = suppress_tokens.value(); + else + options.suppress_tokens.clear(); + std::shared_lock lock(_mutex); + assert_model_is_ready(); + + futures = _pool->generate(features, language, eos, options); + + return maybe_wait_on_futures(std::move(futures), asynchronous); + } + }; + + void register_whisper_nmt(py::module& m) { + py::class_(m, "WhisperNmtGenerationResult", + "A generation result from the Whisper model.") + + .def_readonly("sequences", &models::WhisperNmtGenerationResult::sequences, + "Generated sequences of tokens.") + .def_readonly("scores", &models::WhisperNmtGenerationResult::scores, + "Score of each sequence (empty if :obj:`return_scores` was disabled).") + .def_readonly("logits", &models::WhisperNmtGenerationResult::attention, + "logits in each sequence (empty if :obj:`return_logits_vocab` was disabled).") + .def_readonly("no_speech_prob", &models::WhisperNmtGenerationResult::no_speech_prob, + "Probability of the no speech token (0 if :obj:`return_no_speech_prob` was disabled).") + + .def("__repr__", [](const models::WhisperNmtGenerationResult& result) { + return "WhisperNmtGenerationResult(sequences=" + std::string(py::repr(py::cast(result.sequences))) + //+ ", sequences_ids=" + std::string(py::repr(py::cast(result.sequences_ids))) + + ", scores=" + std::string(py::repr(py::cast(result.scores))) + + ", logits=" + std::string(py::repr(py::cast(result.logits))) + + ", no_speech_prob=" + std::string(py::repr(py::cast(result.no_speech_prob))) + + ")"; + }) + ; + + declare_async_wrapper(m, "WhisperNmtGenerationResultAsync"); + + py::class_( + m, "WhisperNmt", + R"pbdoc( + Implements the Whisper speech recognition model published by OpenAI. + + See Also: + https://github.com/openai/whisper + )pbdoc") + + .def_property_readonly("is_multilingual", &WhisperNmtWrapper::is_multilingual, + "Returns ``True`` if this model is multilingual.") + + .def_property_readonly("n_mels", &WhisperNmtWrapper::n_mels, + "Returns dimension of mel input features.") + + .def_property_readonly("num_languages", &WhisperNmtWrapper::num_languages, + "Returns the number of languages supported.") + + .def(py::init>&, const StringOrMap&, size_t, size_t, long, bool, bool, py::object>(), + py::arg("model_path"), + py::arg("device")="cpu", + py::kw_only(), + py::arg("device_index")=0, + py::arg("compute_type")="default", + py::arg("inter_threads")=1, + py::arg("intra_threads")=0, + py::arg("max_queued_batches")=0, + py::arg("flash_attention")=false, + py::arg("tensor_parallel")=false, + py::arg("files")=py::none(), + R"pbdoc( + Initializes a Whisper model from a converted model. + + Arguments: + model_path: Path to the CTranslate2 model directory. + device: Device to use (possible values are: cpu, cuda, auto). + device_index: Device IDs where to place this model on. + compute_type: Model computation type or a dictionary mapping a device name + to the computation type (possible values are: default, auto, int8, int8_float32, + int8_float16, int8_bfloat16, int16, float16, bfloat16, float32). + inter_threads: Number of workers to allow executing multiple batches in parallel. + intra_threads: Number of OpenMP threads per worker (0 to use a default value). + max_queued_batches: Maximum numbers of batches in the worker queue (-1 for unlimited, + 0 for an automatic value). When the queue is full, future requests will block + until a free slot is available. + flash_attention: run model with flash attention 2 for self-attention layer + tensor_parallel: run model with tensor parallel mode + files: Load model files from the memory. This argument is a dictionary mapping + file names to file contents as file-like or bytes objects. If this is set, + :obj:`model_path` acts as an identifier for this model. + )pbdoc") + + .def_property_readonly("device", &WhisperNmtWrapper::device, + "Device this model is running on.") + .def_property_readonly("device_index", &WhisperNmtWrapper::device_index, + "List of device IDs where this model is running on.") + .def_property_readonly("compute_type", &WhisperNmtWrapper::compute_type, + "Computation type used by the model.") + .def_property_readonly("num_workers", &WhisperNmtWrapper::num_replicas, + "Number of model workers backing this instance.") + .def_property_readonly("num_queued_batches", &WhisperNmtWrapper::num_queued_batches, + "Number of batches waiting to be processed.") + .def_property_readonly("tensor_parallel", &WhisperNmtWrapper::tensor_parallel, + "Run model with tensor parallel mode.") + .def_property_readonly("num_active_batches", &WhisperNmtWrapper::num_active_batches, + "Number of batches waiting to be processed or currently processed.") + + /*.def("encode", &WhisperNmtWrapper::encode, + py::arg("features"), + py::arg("to_cpu")=false, + py::call_guard(), + R"pbdoc( + Encodes the input features. + + Arguments: + features: Mel spectogram of the audio, as a float array with shape + ``[batch_size, n_mels, chunk_length]``. + to_cpu: Copy the encoder output to the CPU before returning the value. + + Returns: + The encoder output. + )pbdoc")*/ + + .def("generate", &WhisperNmtWrapper::generate, + py::arg("features"), + py::arg("language"), + py::arg("eos"), + py::kw_only(), + py::arg("asynchronous")=false, + py::arg("beam_size")=5, + py::arg("patience")=1, + py::arg("num_hypotheses")=1, + py::arg("length_penalty")=1, + py::arg("coverage_penalty")=0, + py::arg("repetition_penalty")=1, + py::arg("no_repeat_ngram_size")=0, + py::arg("max_length")=448, + py::arg("return_scores")=false, + py::arg("return_attention")=false, + py::arg("return_logits_vocab")=false, + py::arg("return_no_speech_prob")=false, + py::arg("return_alternatives")=false, + py::arg("min_alternative_expansion_prob")=0, + py::arg("max_initial_timestamp_index")=50, + py::arg("suppress_blank")=true, + py::arg("suppress_tokens")=std::vector{-1}, + py::arg("sampling_topk")=1, + py::arg("sampling_topp")=1, + py::arg("sampling_temperature")=1, + py::arg("replace_unknowns")=false, + py::arg("prefix_bias_beta")=0, + py::arg("max_decoding_length")=256, + py::arg("min_decoding_length")=1, + py::arg("disable_unk")=false, + py::arg("end_token")=py::none(), + py::arg("return_end_token")=false, + py::arg("max_input_length")=1024, + py::call_guard(), + R"pbdoc( + Encodes the input features and generates from the given prompt. + + Arguments: + features: Mel spectogram of the audio, as a float array with shape + ``[batch_size, n_mels, chunk_length]``. This method also accepts the encoded + features returned by the method :meth:`ctranslate2.models.Whisper.encode`, + which have shape ``[batch_size, chunk_length // 2, d_model]``. + prompts: Batch of initial string tokens or token IDs. + asynchronous: Run the model asynchronously. + beam_size: Beam size (1 for greedy search). + patience: Beam search patience factor, as described in + https://arxiv.org/abs/2204.05424. The decoding will continue until + beam_size*patience hypotheses are finished. + num_hypotheses: Number of hypotheses to return. + length_penalty: Exponential penalty applied to the length during beam search. + coverage_penalty: Coverage penalty weight applied during beam search. + repetition_penalty: Penalty applied to the score of previously generated tokens + (set > 1 to penalize). + no_repeat_ngram_size: Prevent repetitions of ngrams with this size + (set 0 to disable). + max_length: Maximum generation length. + return_scores: Include the scores in the output. + return_attention: Include the attention weights in the output. + return_logits_vocab: Include the log probs in the output + return_no_speech_prob: Include the probability of the no speech token in the + result. + return_alternatives: Include alternatives in the output. + min_alternative_expansion_prob: Minimum probability to expand an alternative. + max_initial_timestamp_index: Maximum index of the first predicted timestamp. + suppress_blank: Suppress blank outputs at the beginning of the sampling. + suppress_tokens: List of token IDs to suppress. -1 will suppress a default set + of symbols as defined in the model ``config.json`` file. + sampling_topk: Randomly sample predictions from the top K candidates. + sampling_topp: Keep the most probable tokens whose cumulative probability exceeds this value. + sampling_temperature: Sampling temperature to generate more random samples. + replace_unknowns: Replace unknown target tokens by the source token with the highest attention. + prefix_bias_beta: Parameter for biasing translations towards given prefix. + max_decoding_length: Maximum prediction length. + min_decoding_length: Minimum prediction length. + disable_unk: Disable the generation of the unknown token. + end_token: Stop the decoding on one of these tokens (defaults to the model EOS token). + return_end_token: Include the end token in the results. + max_input_length: Truncate inputs after this many tokens (set 0 to disable). + + Returns: + A list of generation results. + )pbdoc") + + .def("unload_model", &WhisperNmtWrapper::unload_model, + py::arg("to_cpu")=false, + py::call_guard(), + R"pbdoc( + Unloads the model attached to this whisper but keep enough runtime context + to quickly resume whisper on the initial device. + + Arguments: + to_cpu: If ``True``, the model is moved to the CPU memory and not fully unloaded. + )pbdoc") + + .def("load_model", &WhisperNmtWrapper::load_model, + py::arg("keep_cache")=false, + py::call_guard(), + R"pbdoc( + Loads the model back to the initial device. + + Arguments: + keep_cache: If ``True``, the model cache in the CPU memory is not deleted if it exists. + )pbdoc") + + .def_property_readonly("model_is_loaded", &WhisperNmtWrapper::model_is_loaded, + "Whether the model is loaded on the initial device and ready to be used.") + ; + } + + } +} diff --git a/python/ctranslate2/converters/__init__.py b/python/ctranslate2/converters/__init__.py index eaccffee4..08cf855fa 100644 --- a/python/ctranslate2/converters/__init__.py +++ b/python/ctranslate2/converters/__init__.py @@ -6,3 +6,4 @@ from ctranslate2.converters.opennmt_tf import OpenNMTTFConverter from ctranslate2.converters.opus_mt import OpusMTConverter from ctranslate2.converters.transformers import TransformersConverter +from ctranslate2.converters.whisper_nmt import WhisperNMTConverter diff --git a/python/ctranslate2/converters/transformers.py b/python/ctranslate2/converters/transformers.py index d5f935f95..81d9aa596 100644 --- a/python/ctranslate2/converters/transformers.py +++ b/python/ctranslate2/converters/transformers.py @@ -8,12 +8,13 @@ import numpy as np -try: - import huggingface_hub - import torch - import transformers -except ImportError: - pass +import huggingface_hub +import torch +import transformers + +# try: +# except ImportError: +# pass from ctranslate2.converters import utils from ctranslate2.converters.converter import Converter @@ -207,6 +208,28 @@ def __call__(self, model, tokenizer): return spec + def set_quantization(self, model): + self.quant_type = common_spec.Quantization.CT2 + self.quant_group_size = None + self.quant_bits = None + + quantization_config = getattr(model.config, "quantization_config", None) + if quantization_config: + self.quant_type = None + if quantization_config.quant_method == "awq": + self.quant_type = _SUPPORTED_QUANTIZATION.get(quantization_config.version) + if self.quant_type is None: + raise NotImplementedError( + "Quantization type '%s' is not yet implemented. " + "The following Quantization types are currently supported: %s" + % ( + quantization_config.quant_method, + ", ".join(_SUPPORTED_QUANTIZATION.keys()), + ) + ) + self.quant_group_size = quantization_config.group_size + self.quant_bits = quantization_config.bits + def get_vocabulary(self, model, tokenizer): return [ token @@ -225,8 +248,8 @@ def set_layer_norm(self, spec, module): spec.gamma = module.weight spec.beta = module.bias - def set_linear(self, spec, module, quant_type=common_spec.Quantization.CT2): - if quant_type == common_spec.Quantization.CT2: + def set_linear(self, spec, module): + if self.quant_type == common_spec.Quantization.CT2: spec.weight = module.weight else: spec.weight = module.qweight @@ -997,6 +1020,104 @@ def set_conv1d(self, spec, module): spec.bias = module.bias +@register_loader("WhisperNmtConfig") +class WhisperLoader(BartLoader): + @property + def architecture_name(self): + return "WhisperForConditionalGeneration" + + def get_model_spec(self, model): + spec = whisper_spec.WhisperSpec( + model.config.encoder_layers, + model.config.encoder_attention_heads, + model.config.decoder_layers, + model.config.decoder_attention_heads, + ) + + self.set_encoder(spec.encoder, model.model.encoder) + + return spec + + def _get_lang_ids_from_tokenizer(self, tokenizer): + non_lang_special_tokens = [ + "<|endoftext|>", + "<|startoftranscript|>", + "<|translate|>", + "<|transcribe|>", + "<|startoflm|>", + "<|startofprev|>", + "<|nocaptions|>", + "<|notimestamps|>", + ] + return [ + token_id + for token_id, token in zip( + tokenizer.additional_special_tokens_ids, + tokenizer.additional_special_tokens, + ) + if token not in non_lang_special_tokens + ] + + def set_config(self, config, model, tokenizer): + gen_config = getattr(model, "generation_config", None) + + if gen_config is not None: + config.suppress_ids = gen_config.suppress_tokens + config.suppress_ids_begin = gen_config.begin_suppress_tokens + if hasattr(gen_config, "alignment_heads"): + config.alignment_heads = gen_config.alignment_heads + if hasattr(gen_config, "lang_to_id"): + config.lang_ids = sorted(gen_config.lang_to_id.values()) + else: + config.suppress_ids = model.config.suppress_tokens + config.suppress_ids_begin = model.config.begin_suppress_tokens + config.alignment_heads = _WHISPER_ALIGNMENT_HEADS.get(model.name_or_path) + + if getattr(config, "lang_ids", None) is None: + config.lang_ids = self._get_lang_ids_from_tokenizer(tokenizer) + + if config.alignment_heads is None: + # Use the last half layers for alignment by default. + num_layers = model.config.decoder_layers + num_heads = model.config.decoder_attention_heads + config.alignment_heads = list( + itertools.product( + range(num_layers // 2, num_layers), + range(num_heads), + ) + ) + + def get_vocabulary(self, model, tokenizer): + tokens = super().get_vocabulary(model, tokenizer) + + # Add timestamp tokens. + tokens.extend( + "<|%.2f|>" % (i * 0.02) + for i in range(model.config.vocab_size - len(tokens)) + ) + + return tokens + + def set_vocabulary(self, spec, tokens): + spec.register_vocabulary(tokens) + + def set_encoder(self, spec, encoder): + self.set_conv1d(spec.conv1, encoder.conv1) + self.set_conv1d(spec.conv2, encoder.conv2) + super().set_encoder(spec, encoder) + + def set_decoder(self, spec, decoder): + self.set_embeddings(spec.embeddings, decoder.embed_tokens) + super().set_decoder(spec, decoder) + + def set_common_layers(self, spec, module): + self.set_position_encodings(spec.position_encodings, module.embed_positions) + self.set_layer_norm(spec.layer_norm, module.layer_norm) + + def set_conv1d(self, spec, module): + spec.weight = module.weight + spec.bias = module.bias + @register_loader("Wav2Vec2Config") class Wav2Vec2Loader(BartLoader): @property @@ -1467,6 +1588,8 @@ def get_model_spec(self, model): model.config, "hidden_activation", "gelu_pytorch_tanh" ) + self.set_quantization(model) + spec = transformer_spec.TransformerDecoderModelSpec.from_config( num_layers, num_heads, @@ -1483,6 +1606,9 @@ def get_model_spec(self, model): rotary_base=getattr(model.config, "rope_theta", 10000), num_heads_kv=num_heads_kv, head_dim=model.config.head_dim, + quant_type=self.quant_type, + quant_group_size=self.quant_group_size, + quant_bits=self.quant_bits, ) self.set_decoder(spec.decoder, model.model) @@ -1528,13 +1654,20 @@ def set_decoder(self, spec, module): layer_spec.ffn.layer_norm, layer.post_attention_layernorm ) - wq = layer.self_attn.q_proj.weight - wk = layer.self_attn.k_proj.weight - wv = layer.self_attn.v_proj.weight - wo = layer.self_attn.o_proj.weight + split_layers = [common_spec.LinearSpec() for _ in range(3)] + self.set_linear(split_layers[0], layer.self_attn.q_proj) + self.set_linear(split_layers[1], layer.self_attn.k_proj) + self.set_linear(split_layers[2], layer.self_attn.v_proj) + + if self.quant_type == common_spec.Quantization.CT2: + utils.fuse_linear(layer_spec.self_attention.linear[0], split_layers) + else: + cc_dim = 1 if self.quant_type == common_spec.Quantization.AWQ_GEMM else 0 + utils.fuse_linear_prequant( + layer_spec.self_attention.linear[0], split_layers, cc_dim + ) - layer_spec.self_attention.linear[0].weight = torch.cat([wq, wk, wv]) - layer_spec.self_attention.linear[1].weight = wo + self.set_linear(layer_spec.self_attention.linear[1], layer.self_attn.o_proj) self.set_linear(layer_spec.ffn.linear_0, layer.mlp.gate_proj) self.set_linear(layer_spec.ffn.linear_0_noact, layer.mlp.up_proj) @@ -1563,6 +1696,8 @@ def get_model_spec(self, model): model.config, "hidden_activation", "gelu_pytorch_tanh" ) + self.set_quantization(model) + spec = transformer_spec.TransformerDecoderModelSpec.from_config( num_layers, num_heads, @@ -1580,6 +1715,9 @@ def get_model_spec(self, model): num_heads_kv=num_heads_kv, head_dim=model.config.head_dim, pre_post_layer_norm=True, + quant_type=self.quant_type, + quant_group_size=self.quant_group_size, + quant_bits=self.quant_bits, ) self.set_decoder(spec.decoder, model.model) @@ -1632,13 +1770,19 @@ def set_decoder(self, spec, module): layer_spec.post_feedforward_layer_norm, layer.post_feedforward_layernorm ) - wq = layer.self_attn.q_proj.weight - wk = layer.self_attn.k_proj.weight - wv = layer.self_attn.v_proj.weight - wo = layer.self_attn.o_proj.weight + split_layers = [common_spec.LinearSpec() for _ in range(3)] + self.set_linear(split_layers[0], layer.self_attn.q_proj) + self.set_linear(split_layers[1], layer.self_attn.k_proj) + self.set_linear(split_layers[2], layer.self_attn.v_proj) - layer_spec.self_attention.linear[0].weight = torch.cat([wq, wk, wv]) - layer_spec.self_attention.linear[1].weight = wo + if self.quant_type == common_spec.Quantization.CT2: + utils.fuse_linear(layer_spec.self_attention.linear[0], split_layers) + else: + cc_dim = 1 if self.quant_type == common_spec.Quantization.AWQ_GEMM else 0 + utils.fuse_linear_prequant( + layer_spec.self_attention.linear[0], split_layers, cc_dim + ) + self.set_linear(layer_spec.self_attention.linear[1], layer.self_attn.o_proj) self.set_linear(layer_spec.ffn.linear_0, layer.mlp.gate_proj) self.set_linear(layer_spec.ffn.linear_0_noact, layer.mlp.up_proj) @@ -1679,26 +1823,7 @@ def get_model_spec(self, model): rotary_scaling_type = None rotary_scaling_factor = 1 - quantization_config = getattr(model.config, "quantization_config", None) - if quantization_config: - quant_type = None - if quantization_config.quant_method == "awq": - quant_type = _SUPPORTED_QUANTIZATION.get(quantization_config.version) - if quant_type is None: - raise NotImplementedError( - "Quantization type '%s' is not yet implemented. " - "The following Quantization types are currently supported: %s" - % ( - quantization_config.quant_method, - ", ".join(_SUPPORTED_QUANTIZATION.keys()), - ) - ) - quant_group_size = quantization_config.group_size - quant_bits = quantization_config.bits - else: - quant_type = common_spec.Quantization.CT2 - quant_group_size = None - quant_bits = None + self.set_quantization(model) spec = transformer_spec.TransformerDecoderModelSpec.from_config( num_layers, @@ -1713,12 +1838,12 @@ def get_model_spec(self, model): rotary_scaling_factor=rotary_scaling_factor, rotary_base=getattr(model.config, "rope_theta", 10000), num_heads_kv=num_heads_kv, - quant_type=quant_type, - quant_group_size=quant_group_size, - quant_bits=quant_bits, + quant_type=self.quant_type, + quant_group_size=self.quant_group_size, + quant_bits=self.quant_bits, ) - self.set_decoder(spec.decoder, model.model, quant_type) + self.set_decoder(spec.decoder, model.model) self.set_linear(spec.decoder.projection, model.lm_head) # set extra RoPE parameters for Llama-3.1 @@ -1757,7 +1882,7 @@ def set_config(self, config, model, tokenizer): def set_layer_norm(self, spec, layer_norm): spec.gamma = layer_norm.weight - def set_decoder(self, spec, module, quant_type=common_spec.Quantization.CT2): + def set_decoder(self, spec, module): spec.scale_embeddings = False self.set_embeddings(spec.embeddings, module.embed_tokens) self.set_layer_norm(spec.layer_norm, module.norm) @@ -1771,38 +1896,22 @@ def set_decoder(self, spec, module, quant_type=common_spec.Quantization.CT2): ) split_layers = [common_spec.LinearSpec() for _ in range(3)] - self.set_linear( - split_layers[0], layer.self_attn.q_proj, quant_type=quant_type - ) - self.set_linear( - split_layers[1], layer.self_attn.k_proj, quant_type=quant_type - ) - self.set_linear( - split_layers[2], layer.self_attn.v_proj, quant_type=quant_type - ) + self.set_linear(split_layers[0], layer.self_attn.q_proj) + self.set_linear(split_layers[1], layer.self_attn.k_proj) + self.set_linear(split_layers[2], layer.self_attn.v_proj) - if quant_type == common_spec.Quantization.CT2: + if self.quant_type == common_spec.Quantization.CT2: utils.fuse_linear(layer_spec.self_attention.linear[0], split_layers) else: - cc_dim = 1 if quant_type == common_spec.Quantization.AWQ_GEMM else 0 + cc_dim = 1 if self.quant_type == common_spec.Quantization.AWQ_GEMM else 0 utils.fuse_linear_prequant( layer_spec.self_attention.linear[0], split_layers, cc_dim ) - self.set_linear( - layer_spec.self_attention.linear[1], - layer.self_attn.o_proj, - quant_type=quant_type, - ) + self.set_linear(layer_spec.self_attention.linear[1], layer.self_attn.o_proj) - self.set_linear( - layer_spec.ffn.linear_0, layer.mlp.gate_proj, quant_type=quant_type - ) - self.set_linear( - layer_spec.ffn.linear_0_noact, layer.mlp.up_proj, quant_type=quant_type - ) - self.set_linear( - layer_spec.ffn.linear_1, layer.mlp.down_proj, quant_type=quant_type - ) + self.set_linear(layer_spec.ffn.linear_0, layer.mlp.gate_proj) + self.set_linear(layer_spec.ffn.linear_0_noact, layer.mlp.up_proj) + self.set_linear(layer_spec.ffn.linear_1, layer.mlp.down_proj) delattr(layer, "self_attn") delattr(layer, "mlp") @@ -1840,25 +1949,7 @@ def get_model_spec(self, model): rotary_scaling_type = None rotary_scaling_factor = 1 - quantization_config = getattr(model.config, "quantization_config", None) - if quantization_config: - if quantization_config.quant_method == "awq": - quant_type = _SUPPORTED_QUANTIZATION.get(quantization_config.version) - if quant_type is None: - raise NotImplementedError( - "Quantization type '%s' is not yet implemented. " - "The following Quantization types are currently supported: %s" - % ( - quantization_config.quant_method, - ", ".join(_SUPPORTED_QUANTIZATION.keys()), - ) - ) - quant_group_size = quantization_config.group_size - quant_bits = quantization_config.bits - else: - quant_type = common_spec.Quantization.CT2 - quant_group_size = None - quant_bits = None + self.set_quantization(model) spec = transformer_spec.TransformerDecoderModelSpec.from_config( num_layers, @@ -1874,13 +1965,13 @@ def get_model_spec(self, model): rotary_base=getattr(model.config, "rope_theta", 10000), num_heads_kv=num_heads_kv, sliding_window=sliding_window, - quant_type=quant_type, - quant_group_size=quant_group_size, - quant_bits=quant_bits, + quant_type=self.quant_type, + quant_group_size=self.quant_group_size, + quant_bits=self.quant_bits, head_dim=model.config.head_dim, ) - self.set_decoder(spec.decoder, model.model, quant_type=quant_type) + self.set_decoder(spec.decoder, model.model) self.set_linear(spec.decoder.projection, model.lm_head) return spec @@ -1905,7 +1996,7 @@ def set_config(self, config, model, tokenizer): def set_layer_norm(self, spec, layer_norm): spec.gamma = layer_norm.weight - def set_decoder(self, spec, module, quant_type=common_spec.Quantization.CT2): + def set_decoder(self, spec, module): spec.scale_embeddings = False self.set_embeddings(spec.embeddings, module.embed_tokens) self.set_layer_norm(spec.layer_norm, module.norm) @@ -1919,37 +2010,27 @@ def set_decoder(self, spec, module, quant_type=common_spec.Quantization.CT2): ) split_layers = [common_spec.LinearSpec() for _ in range(3)] self.set_linear( - split_layers[0], layer.self_attn.q_proj, quant_type=quant_type + split_layers[0], layer.self_attn.q_proj ) self.set_linear( - split_layers[1], layer.self_attn.k_proj, quant_type=quant_type + split_layers[1], layer.self_attn.k_proj ) self.set_linear( - split_layers[2], layer.self_attn.v_proj, quant_type=quant_type + split_layers[2], layer.self_attn.v_proj ) - if quant_type == common_spec.Quantization.CT2: + if self.quant_type == common_spec.Quantization.CT2: utils.fuse_linear(layer_spec.self_attention.linear[0], split_layers) else: - cc_dim = 1 if quant_type == common_spec.Quantization.AWQ_GEMM else 0 + cc_dim = 1 if self.quant_type == common_spec.Quantization.AWQ_GEMM else 0 utils.fuse_linear_prequant( layer_spec.self_attention.linear[0], split_layers, cc_dim ) - self.set_linear( - layer_spec.self_attention.linear[1], - layer.self_attn.o_proj, - quant_type=quant_type, - ) + self.set_linear(layer_spec.self_attention.linear[1], layer.self_attn.o_proj) - self.set_linear( - layer_spec.ffn.linear_0, layer.mlp.gate_proj, quant_type=quant_type - ) - self.set_linear( - layer_spec.ffn.linear_0_noact, layer.mlp.up_proj, quant_type=quant_type - ) - self.set_linear( - layer_spec.ffn.linear_1, layer.mlp.down_proj, quant_type=quant_type - ) + self.set_linear(layer_spec.ffn.linear_0, layer.mlp.gate_proj) + self.set_linear(layer_spec.ffn.linear_0_noact, layer.mlp.up_proj) + self.set_linear(layer_spec.ffn.linear_1, layer.mlp.down_proj) delattr(layer, "self_attn") delattr(layer, "mlp") @@ -1986,6 +2067,8 @@ def get_model_spec(self, model): rotary_scaling_type = None rotary_scaling_factor = 1 + self.set_quantization(model) + spec = transformer_spec.TransformerDecoderModelSpec.from_config( num_layers, num_heads, @@ -1999,6 +2082,9 @@ def get_model_spec(self, model): rotary_scaling_factor=rotary_scaling_factor, rotary_base=getattr(model.config, "rope_theta", 10000), num_heads_kv=num_heads_kv, + quant_type=self.quant_type, + quant_group_size=self.quant_group_size, + quant_bits=self.quant_bits, ) self.set_decoder(spec.decoder, model.model) @@ -2049,7 +2135,14 @@ def set_decoder(self, spec, module): self.set_linear(split_layers[1], layer.self_attn.k_proj) self.set_linear(split_layers[2], layer.self_attn.v_proj) - utils.fuse_linear(layer_spec.self_attention.linear[0], split_layers) + if self.quant_type == common_spec.Quantization.CT2: + utils.fuse_linear(layer_spec.self_attention.linear[0], split_layers) + else: + cc_dim = 1 if self.quant_type == common_spec.Quantization.AWQ_GEMM else 0 + utils.fuse_linear_prequant( + layer_spec.self_attention.linear[0], split_layers, cc_dim + ) + self.set_linear( layer_spec.self_attention.linear[1], layer.self_attn.o_proj, @@ -2123,6 +2216,8 @@ def architecture_name(self): return "AutoModelForCausalLM" def get_model_spec(self, model): + self.set_quantization(model) + spec = transformer_spec.TransformerDecoderModelSpec.from_config( num_layers=model.config.n_layer, num_heads=model.config.n_head, @@ -2132,6 +2227,9 @@ def get_model_spec(self, model): rotary_interleave=False, parallel_residual=True, shared_layer_norm=True, + quant_type=self.quant_type, + quant_group_size=self.quant_group_size, + quant_bits=self.quant_bits, ) self.set_decoder(spec.decoder, model.transformer) @@ -2201,6 +2299,8 @@ def get_model_spec(self, model): rotary_scaling_type = None rotary_scaling_factor = 1 + self.set_quantization(model) + spec = transformer_spec.TransformerDecoderModelSpec.from_config( num_layers, num_heads, @@ -2216,6 +2316,9 @@ def get_model_spec(self, model): original_max_position_embeddings=original_max_position_embeddings, max_position_embeddings=max_position_embeddings, num_heads_kv=num_heads_kv, + quant_type=self.quant_type, + quant_group_size=self.quant_group_size, + quant_bits=self.quant_bits, ) self.set_decoder(spec.decoder, model.model) diff --git a/python/ctranslate2/converters/whisper_nmt.py b/python/ctranslate2/converters/whisper_nmt.py new file mode 100644 index 000000000..8607e4b5f --- /dev/null +++ b/python/ctranslate2/converters/whisper_nmt.py @@ -0,0 +1,722 @@ +import argparse +import os + +from ctranslate2.converters import utils +from ctranslate2.converters.converter import Converter +from ctranslate2.specs import ( + common_spec, + transformer_spec, + whispernmt_spec, +) + +import huggingface_hub +import transformers +import itertools +from typing import List, Optional + +class WhisperLoader: + def __call__(self, model, tokenizer, spec): + self.get_model_spec(model, spec) + self.set_config(spec.config, model, tokenizer) + #tokens = self.get_vocabulary(model, tokenizer) + #self.set_vocabulary(spec, tokens) + return spec + @property + def architecture_name(self): + return "WhisperForConditionalGeneration" + + def get_model_spec(self, model, spec): + self.set_encoder(spec.whisper_encoder, model.model.encoder) + return spec + + def _get_lang_ids_from_tokenizer(self, tokenizer): + non_lang_special_tokens = [ + "<|endoftext|>", + "<|startoftranscript|>", + "<|translate|>", + "<|transcribe|>", + "<|startoflm|>", + "<|startofprev|>", + "<|nocaptions|>", + "<|notimestamps|>", + ] + return [ + token_id + for token_id, token in zip( + tokenizer.additional_special_tokens_ids, + tokenizer.additional_special_tokens, + ) + if token not in non_lang_special_tokens + ] + + def set_config(self, config, model, tokenizer): + gen_config = getattr(model, "generation_config", None) + + if gen_config is not None: + config.suppress_ids = gen_config.suppress_tokens + config.suppress_ids_begin = gen_config.begin_suppress_tokens + if hasattr(gen_config, "alignment_heads"): + config.alignment_heads = gen_config.alignment_heads + if hasattr(gen_config, "lang_to_id"): + config.lang_ids = sorted(gen_config.lang_to_id.values()) + else: + config.suppress_ids = model.config.suppress_tokens + config.suppress_ids_begin = model.config.begin_suppress_tokens + config.alignment_heads = _WHISPER_ALIGNMENT_HEADS.get(model.name_or_path) + + if getattr(config, "lang_ids", None) is None: + config.lang_ids = self._get_lang_ids_from_tokenizer(tokenizer) + + if config.alignment_heads is None: + # Use the last half layers for alignment by default. + num_layers = model.config.decoder_layers + num_heads = model.config.decoder_attention_heads + config.alignment_heads = list( + itertools.product( + range(num_layers // 2, num_layers), + range(num_heads), + ) + ) + + def set_linear(self, spec, module): + spec.weight = module.weight + + if isinstance(module, transformers.Conv1D): + spec.weight = spec.weight.transpose(0, 1) + if module.bias is not None: + spec.bias = module.bias + + def set_attention(self, spec, attention, self_attention=False): + split_layers = [common_spec.LinearSpec() for _ in range(3)] + self.set_linear(split_layers[0], attention.q_proj) + self.set_linear(split_layers[1], attention.k_proj) + self.set_linear(split_layers[2], attention.v_proj) + + if self_attention: + utils.fuse_linear(spec.linear[0], split_layers) + else: + utils.fuse_linear(spec.linear[0], split_layers[:1]) + utils.fuse_linear(spec.linear[1], split_layers[1:]) + + self.set_linear(spec.linear[-1], attention.out_proj) + + def set_position_encodings(self, spec, module): + spec.encodings = module.weight + offset = getattr(module, "offset", 0) + if offset > 0: + spec.encodings = spec.encodings[offset:] + + def set_embeddings(self, spec, module): + spec.weight = module.weight + + def set_layer_norm(self, spec, module): + spec.gamma = module.weight + spec.beta = module.bias + + def get_vocabulary(self, model, tokenizer): + tokens = [ + token + for token, _ in sorted( + tokenizer.get_vocab().items(), key=lambda item: item[1] + ) + ] + if model.config.vocab_size < len(tokens): + tokens = tokens[: model.config.vocab_size] + + # Add timestamp tokens. + tokens.extend( + "<|%.2f|>" % (i * 0.02) + for i in range(model.config.vocab_size - len(tokens)) + ) + + return tokens + + def set_vocabulary(self, spec, tokens): + spec.register_vocabulary(tokens) + + def set_encoder(self, spec, encoder): + self.set_conv1d(spec.conv1, encoder.conv1) + self.set_conv1d(spec.conv2, encoder.conv2) + self.set_common_layers(spec, encoder) + + for layer_spec, layer in zip(spec.layer, encoder.layers): + self.set_attention( + layer_spec.self_attention, + layer.self_attn, + self_attention=True, + ) + self.set_layer_norm( + layer_spec.self_attention.layer_norm, + layer.self_attn_layer_norm, + ) + + self.set_linear(layer_spec.ffn.linear_0, layer.fc1) + self.set_linear(layer_spec.ffn.linear_1, layer.fc2) + self.set_layer_norm(layer_spec.ffn.layer_norm, layer.final_layer_norm) + + def set_decoder(self, spec, decoder): + self.set_embeddings(spec.embeddings, decoder.embed_tokens) + self.set_common_layers(spec, decoder) + + for layer_spec, layer in zip(spec.layer, decoder.layers): + self.set_attention( + layer_spec.self_attention, + layer.self_attn, + self_attention=True, + ) + self.set_layer_norm( + layer_spec.self_attention.layer_norm, + layer.self_attn_layer_norm, + ) + + if hasattr(layer, "encoder_attn"): + self.set_attention( + layer_spec.attention, + layer.encoder_attn, + self_attention=False, + ) + self.set_layer_norm( + layer_spec.attention.layer_norm, + layer.encoder_attn_layer_norm, + ) + + self.set_linear(layer_spec.ffn.linear_0, layer.fc1) + self.set_linear(layer_spec.ffn.linear_1, layer.fc2) + self.set_layer_norm(layer_spec.ffn.layer_norm, layer.final_layer_norm) + + def set_common_layers(self, spec, module): + self.set_position_encodings(spec.position_encodings, module.embed_positions) + self.set_layer_norm(spec.layer_norm, module.layer_norm) + + def set_conv1d(self, spec, module): + spec.weight = module.weight + spec.bias = module.bias + +_SUPPORTED_ACTIVATIONS = { + "gelu": common_spec.Activation.GELU, + "fast_gelu": common_spec.Activation.GELUTanh, + "relu": common_spec.Activation.RELU, + "silu": common_spec.Activation.SWISH, +} + +_SUPPORTED_FEATURES_MERGE = { + "concat": common_spec.EmbeddingsMerge.CONCAT, + "sum": common_spec.EmbeddingsMerge.ADD, +} + + +def check_opt(opt, num_source_embeddings): + with_relative_position = getattr(opt, "max_relative_positions", 0) > 0 + with_rotary = getattr(opt, "max_relative_positions", 0) == -1 + with_alibi = getattr(opt, "max_relative_positions", 0) == -2 + activation_fn = getattr(opt, "pos_ffn_activation_fn", "relu") + feat_merge = getattr(opt, "feat_merge", "concat") + self_attn_type = getattr(opt, "self_attn_type", "scaled-dot") + + check = utils.ConfigurationChecker() + check( + opt.encoder_type == opt.decoder_type + and opt.decoder_type in {"transformer", "transformer_lm"}, + "Options --encoder_type and --decoder_type must be" + " 'transformer' or 'transformer_lm", + ) + check( + self_attn_type == "scaled-dot", + "Option --self_attn_type %s is not supported (supported values are: scaled-dot)" + % self_attn_type, + ) + check( + activation_fn in _SUPPORTED_ACTIVATIONS, + "Option --pos_ffn_activation_fn %s is not supported (supported activations are: %s)" + % (activation_fn, ", ".join(_SUPPORTED_ACTIVATIONS.keys())), + ) + check( + opt.position_encoding != (with_relative_position or with_rotary or with_alibi), + "Options --position_encoding and --max_relative_positions cannot be both enabled " + "or both disabled", + ) + check( + num_source_embeddings == 1 or feat_merge in _SUPPORTED_FEATURES_MERGE, + "Option --feat_merge %s is not supported (supported merge modes are: %s)" + % (feat_merge, " ".join(_SUPPORTED_FEATURES_MERGE.keys())), + ) + check.validate() + + +def _get_model_spec_seq2seq( + opt, variables, src_vocabs, tgt_vocabs, num_source_embeddings, whisper_encoder_num_layers, whisper_encoder_num_heads +): + """Creates a model specification from the model options.""" + activation_fn = "relu" + + # Return the first head of the last layer unless the model was trained with alignments. + if getattr(opt, "lambda_align", 0) == 0: + alignment_layer = -1 + alignment_heads = 1 + else: + alignment_layer = opt.alignment_layer + alignment_heads = opt.alignment_heads + + num_heads = getattr(opt, "num_heads", 16) + num_layers = getattr(opt, "num_layers", 6) + + model_spec = whispernmt_spec.WhisperNmtSpec.from_config( + num_layers, + num_heads, + whisper_encoder_num_layers, + whisper_encoder_num_heads, + activation=_SUPPORTED_ACTIVATIONS[activation_fn], + alignment_layer=alignment_layer, + alignment_heads=alignment_heads, + num_source_embeddings=num_source_embeddings, + multi_query_attention=getattr(opt, "multiquery", False), + ) + + model_spec.config.decoder_start_token = getattr(opt, "decoder_start_token", "") + + set_transformer_spec(model_spec, variables) + for src_vocab in src_vocabs: + model_spec.register_source_vocabulary(src_vocab) + for tgt_vocab in tgt_vocabs: + model_spec.register_target_vocabulary(tgt_vocab) + + return model_spec + + +def _get_model_spec_lm(opt, variables, src_vocabs, tgt_vocabs, num_source_embeddings): + """Creates a model specification from the model options.""" + with_relative_position = getattr(opt, "max_relative_positions", 0) > 0 + with_rotary = getattr(opt, "max_relative_positions", 0) == -1 + with_alibi = getattr(opt, "max_relative_positions", 0) == -2 + activation_fn = getattr(opt, "pos_ffn_activation_fn", "relu") + num_heads = getattr(opt, "heads", 8) + num_kv = getattr(opt, "num_kv", 0) + if num_kv == num_heads or num_kv == 0: + num_kv = None + rotary_dim = 0 if with_rotary else None + rotary_interleave = getattr(opt, "rotary_interleave", True) + ffn_glu = activation_fn == "silu" + sliding_window = getattr(opt, "sliding_window", 0) + + model_spec = transformer_spec.TransformerDecoderModelSpec.from_config( + opt.dec_layers, + num_heads, + activation=_SUPPORTED_ACTIVATIONS[activation_fn], + ffn_glu=ffn_glu, + with_relative_position=with_relative_position, + alibi=with_alibi, + rms_norm=opt.layer_norm == "rms", + rotary_dim=rotary_dim, + rotary_interleave=rotary_interleave, + multi_query_attention=getattr(opt, "multiquery", False), + num_heads_kv=num_kv, + sliding_window=sliding_window, + ) + + model_spec.config.layer_norm_epsilon = getattr(opt, "norm_eps", 1e-6) + + set_transformer_decoder( + model_spec.decoder, + variables, + with_encoder_attention=False, + ) + + for tgt_vocab in tgt_vocabs: + model_spec.register_vocabulary(tgt_vocab) + + return model_spec + + +def get_vocabs(vocab_path): + with open(vocab_path, "r", encoding="utf-8") as f: + vocab = [line.strip() for line in f if line.strip()] + src_vocabs = [vocab] + tgt_vocabs = [vocab] + + return src_vocabs, tgt_vocabs + + +class WhisperNMTConverter(Converter): + """Converts models generated by OpenNMT-py.""" + + def __init__(self, model_path: str, + whisper_model: str, + vocab_path: str, + copy_files: Optional[List[str]] = None, + revision: Optional[str] = None, + low_cpu_mem_usage: bool = False, + trust_remote_code: bool = False): + """Initializes the OpenNMT-py converter. + + Arguments: + model_path: Path to the OpenNMT-py PyTorch model (.pt file). + """ + self._model_path = model_path + self._vocab_path = vocab_path + self._whisper_model = whisper_model + self._copy_files = copy_files + self._revision = revision + self._low_cpu_mem_usage = low_cpu_mem_usage + self._trust_remote_code = trust_remote_code + + def _load(self): + import torch + with torch.no_grad(): + config = transformers.AutoConfig.from_pretrained( + self._whisper_model, trust_remote_code=self._trust_remote_code + ) + loader = WhisperLoader() + model_class = getattr(transformers, loader.architecture_name) + tokenizer_class = transformers.AutoTokenizer + + kwargs = { + "torch_dtype": ( + getattr(config, "torch_dtype", None) + ) + } + + if self._revision: + kwargs["revision"] = self._revision + if self._low_cpu_mem_usage: + kwargs["low_cpu_mem_usage"] = self._low_cpu_mem_usage + if self._trust_remote_code: + kwargs["trust_remote_code"] = self._trust_remote_code + + model = self.load_model(model_class, self._whisper_model, **kwargs) + whisper_encoder_num_layers = model.config.encoder_layers + whisper_encoder_num_heads = model.config.encoder_attention_heads + tokenizer_kwargs = {} + if self._trust_remote_code: + tokenizer_kwargs["trust_remote_code"] = self._trust_remote_code + + tokenizer = self.load_tokenizer( + tokenizer_class, self._whisper_model, **tokenizer_kwargs + ) + + checkpoint = torch.load(self._model_path, map_location="cpu") + src_vocabs, tgt_vocabs = get_vocabs(self._vocab_path) + + variables = checkpoint["model"] + spec = _get_model_spec_seq2seq( + checkpoint["model_config"], + variables, + src_vocabs, + tgt_vocabs, + num_source_embeddings=len(src_vocabs), + whisper_encoder_num_layers=whisper_encoder_num_layers, + whisper_encoder_num_heads=whisper_encoder_num_heads + ) + + spec = loader(model, tokenizer, spec) + if self._copy_files: + for filename in self._copy_files: + spec.register_file(self.get_model_file(filename)) + return spec + + def load_model(self, model_class, model_name_or_path, **kwargs): + return model_class.from_pretrained(model_name_or_path, **kwargs) + + def load_tokenizer(self, tokenizer_class, model_name_or_path, **kwargs): + return tokenizer_class.from_pretrained(model_name_or_path, **kwargs) + + def get_model_file(self, filename): + if os.path.isdir(self._whisper_model): + path = os.path.join(self._whisper_model, filename) + else: + try: + path = huggingface_hub.hf_hub_download( + repo_id=self._whisper_model, filename=filename + ) + except huggingface_hub.utils.EntryNotFoundError: + path = None + + if path is None or not os.path.isfile(path): + raise ValueError( + "File %s does not exist in model %s" + % (filename, self._whisper_model) + ) + + return path + + +def set_transformer_spec(spec, variables): + #set_whisper_encoder(spec.whisper, variables) + set_whisper_reshape(spec.connector, variables) + set_transformer_encoder(spec.transformer_encoder, variables) + set_transformer_decoder(spec.transformer_decoder, variables) + +def set_whisper_reshape(spec, variables): + set_linear(spec.linear1, variables, "lin_1") + set_linear(spec.linear2, variables, "lin_2") + +def set_transformer_encoder(spec, variables): + set_input_layers(spec, variables, "src_embeddings") + set_layer_norm(spec.layer_norm, variables, "encoder.norm") + for i, layer in enumerate(spec.layer): + set_transformer_encoder_layer(layer, variables, "encoder.layers.%d" % i) + + +def set_transformer_decoder(spec, variables, with_encoder_attention=True): + set_input_layers(spec, variables, "tgt_embeddings") + set_layer_norm(spec.layer_norm, variables, "decoder.norm") + for i, layer in enumerate(spec.layer): + set_transformer_decoder_layer( + layer, + variables, + "decoder.layers.%d" % i, + with_encoder_attention=with_encoder_attention, + ) + + set_linear(spec.projection, variables, "output_layer") + #try: + # set_linear(spec.projection, variables, "generator") + #except KeyError: + # # Compatibility when the generator was a nn.Sequential module. + # set_linear(spec.projection, variables, "generator.0") + + +def set_input_layers(spec, variables, scope): + if hasattr(spec, "position_encodings"): + set_position_encodings( + spec.position_encodings, + variables, + "position_encodings", + ) + else: + # See https://github.com/OpenNMT/OpenNMT-py/issues/1722 + spec.scale_embeddings = False + + set_embeddings( + ( + spec.embeddings[0] + if isinstance(spec.embeddings, list) + else spec.embeddings + ), + variables, scope + ) + + +def set_transformer_encoder_layer(spec, variables, scope): + set_layer_norm(spec.ffn.layer_norm, variables, "%s.norm2" % scope) + set_ffn(spec.ffn, variables, "%s.ffn" % scope) + set_multi_head_attention( + spec.self_attention, + variables, + "%s.self_attention" % scope, + self_attention=True, + ) + set_layer_norm(spec.self_attention.layer_norm, variables, "%s.norm1" % scope) + + +def set_transformer_decoder_layer(spec, variables, scope, with_encoder_attention=True): + set_layer_norm(spec.ffn.layer_norm, variables, "%s.norm3" % scope) + set_ffn(spec.ffn, variables, "%s.ffn" % scope) + set_multi_head_attention( + spec.self_attention, + variables, + "%s.self_attention" % scope, + self_attention=True, + ) + set_layer_norm(spec.self_attention.layer_norm, variables, "%s.norm1" % scope) + if with_encoder_attention: + set_multi_head_attention(spec.attention, variables, "%s.attention" % scope) + set_layer_norm(spec.attention.layer_norm, variables, "%s.norm2" % scope) + + +def set_ffn(spec, variables, scope): + set_linear(spec.linear_0, variables, "%s.inner" % scope) + set_linear(spec.linear_1, variables, "%s.outer" % scope) + if hasattr(spec, "linear_0_noact"): + set_linear(spec.linear_0_noact, variables, "%s.w_3" % scope) + + +def set_multi_head_attention(spec, variables, scope, self_attention=False): + if self_attention: + set_linear(spec.linear[0], variables, "%s.in_proj" % scope) + else: + set_linear(spec.linear[0], variables, "%s.query_proj" % scope) + set_linear(spec.linear[1], variables, "%s.value_proj" % scope) + set_linear(spec.linear[-1], variables, "%s.out_proj" % scope) + if hasattr(spec, "relative_position_keys"): + spec.relative_position_keys = _get_variable( + variables, "%s.relative_positions_embeddings.weight" % scope + ) + spec.relative_position_values = spec.relative_position_keys + + +def set_layer_norm(spec, variables, scope): + try: + spec.gamma = _get_variable(variables, "%s.weight" % scope) + except KeyError: + # Compatibility with older models using a custom LayerNorm module. + spec.gamma = _get_variable(variables, "%s.a_2" % scope) + spec.beta = _get_variable(variables, "%s.b_2" % scope) + try: + spec.beta = _get_variable(variables, "%s.bias" % scope) + except KeyError: + pass + + +def set_linear(spec, variables, scope): + spec.weight = _get_variable(variables, "%s.weight" % scope) + bias = variables.get("%s.bias" % scope) + if bias is not None: + spec.bias = bias + + +def set_embeddings(spec, variables, scope): + spec.weight = _get_variable(variables, "%s.weight" % scope) + + +def set_position_encodings(spec, variables, scope): + spec.encodings = _get_variable(variables, "%s" % scope).squeeze() + + +def _get_variable(variables, name): + return variables[name] + + +def main(): + parser = argparse.ArgumentParser( + formatter_class=argparse.ArgumentDefaultsHelpFormatter + ) + parser.add_argument("--model_path", required=True, help="Model path.") + parser.add_argument("--vocab_path", required=True, help="Vocab path.") + parser.add_argument("--whisper_model", required=True, help="Whisper model name.") + Converter.declare_arguments(parser) + args = parser.parse_args() + WhisperNMTConverter(args.model_path, args.whisper_model, args.vocab_path).convert_from_args(args) + + +if __name__ == "__main__": + main() + +# Cross-attention heads that are highly correlated to the word-level timing, +# i.e. the alignment between audio and text tokens. +# Obtained from https://github.com/openai/whisper/blob/v20231106/whisper/__init__.py#L32-L47 +_WHISPER_ALIGNMENT_HEADS = { + "openai/whisper-tiny.en": [ + (1, 0), + (2, 0), + (2, 5), + (3, 0), + (3, 1), + (3, 2), + (3, 3), + (3, 4), + ], + "openai/whisper-tiny": [(2, 2), (3, 0), (3, 2), (3, 3), (3, 4), (3, 5)], + "openai/whisper-base.en": [(3, 3), (4, 7), (5, 1), (5, 5), (5, 7)], + "openai/whisper-base": [ + (3, 1), + (4, 2), + (4, 3), + (4, 7), + (5, 1), + (5, 2), + (5, 4), + (5, 6), + ], + "openai/whisper-small.en": [ + (6, 6), + (7, 0), + (7, 3), + (7, 8), + (8, 2), + (8, 5), + (8, 7), + (9, 0), + (9, 4), + (9, 8), + (9, 10), + (10, 0), + (10, 1), + (10, 2), + (10, 3), + (10, 6), + (10, 11), + (11, 2), + (11, 4), + ], + "openai/whisper-small": [ + (5, 3), + (5, 9), + (8, 0), + (8, 4), + (8, 7), + (8, 8), + (9, 0), + (9, 7), + (9, 9), + (10, 5), + ], + "openai/whisper-medium.en": [ + (11, 4), + (14, 1), + (14, 12), + (14, 14), + (15, 4), + (16, 0), + (16, 4), + (16, 9), + (17, 12), + (17, 14), + (18, 7), + (18, 10), + (18, 15), + (20, 0), + (20, 3), + (20, 9), + (20, 14), + (21, 12), + ], + "openai/whisper-medium": [(13, 15), (15, 4), (15, 15), (16, 1), (20, 0), (23, 4)], + "openai/whisper-large": [ + (9, 19), + (11, 2), + (11, 4), + (11, 17), + (22, 7), + (22, 11), + (22, 17), + (23, 2), + (23, 15), + ], + "openai/whisper-large-v2": [ + (10, 12), + (13, 17), + (16, 11), + (16, 12), + (16, 13), + (17, 15), + (17, 16), + (18, 4), + (18, 11), + (18, 19), + (19, 11), + (21, 2), + (21, 3), + (22, 3), + (22, 9), + (22, 12), + (23, 5), + (23, 7), + (23, 13), + (25, 5), + (26, 1), + (26, 12), + (27, 15), + ], + "openai/whisper-large-v3": [ + (7, 0), + (10, 17), + (12, 18), + (13, 12), + (16, 1), + (17, 14), + (19, 11), + (21, 4), + (24, 1), + (25, 6), + ], +} diff --git a/python/ctranslate2/models/__init__.py b/python/ctranslate2/models/__init__.py index 35a3dca37..80fcca6ad 100644 --- a/python/ctranslate2/models/__init__.py +++ b/python/ctranslate2/models/__init__.py @@ -7,8 +7,10 @@ Wav2Vec2, Wav2Vec2Bert, Whisper, + WhisperNmt, WhisperGenerationResult, WhisperGenerationResultAsync, + WhisperNmtGenerationResult, ) except ImportError as e: # Allow using the Python package without the compiled extension. diff --git a/python/ctranslate2/specs/__init__.py b/python/ctranslate2/specs/__init__.py index b4e53fad2..342931516 100644 --- a/python/ctranslate2/specs/__init__.py +++ b/python/ctranslate2/specs/__init__.py @@ -16,3 +16,4 @@ from ctranslate2.specs.wav2vec2_spec import Wav2Vec2Spec from ctranslate2.specs.wav2vec2bert_spec import Wav2Vec2BertSpec from ctranslate2.specs.whisper_spec import WhisperSpec +from ctranslate2.specs.whispernmt_spec import WhisperNmtSpec diff --git a/python/ctranslate2/specs/transformer_spec.py b/python/ctranslate2/specs/transformer_spec.py index 230e62cfd..fa291b426 100644 --- a/python/ctranslate2/specs/transformer_spec.py +++ b/python/ctranslate2/specs/transformer_spec.py @@ -494,6 +494,132 @@ def get_target_vocabulary_size(self): return self.decoder.embeddings.weight.shape[0] +class WhisperNmtSpec(model_spec.SequenceToSequenceModelSpec): + """Describes a WhisperNmt model. + + The specification is invariant to hidden dimensions but requires to + explicitly set the number of layers and attention heads. + """ + + def __init__( + self, encoder: TransformerEncoderSpec, decoder: TransformerDecoderSpec + ): + """Initializes a Transformer model specification. + + Args: + encoder: The encoder specification. + decoder: The decoder specification. + """ + if not isinstance(encoder, TransformerEncoderSpec): + raise TypeError("encoder argument must be a TransformerEncoderSpec") + if not isinstance(decoder, TransformerDecoderSpec): + raise TypeError("decoder argument must be a TransformerDecoderSpec") + + super().__init__() + self.encoder = encoder + self.decoder = decoder + self._config.add_attribute( + "multi_query_attention", self.encoder.multi_query_attention + ) + + @classmethod + def from_config( + cls, + num_layers: Union[int, Tuple[int, int]], + num_heads: int, + with_relative_position: bool = False, + pre_norm: bool = True, + no_final_norm: bool = False, + activation: common_spec.Activation = common_spec.Activation.RELU, + alignment_layer: int = -1, + alignment_heads: int = 1, + num_source_embeddings: int = 1, + embeddings_merge: common_spec.EmbeddingsMerge = common_spec.EmbeddingsMerge.CONCAT, + layernorm_embedding: bool = False, + relative_attention_bias: bool = False, + ffn_glu: bool = False, + rms_norm: bool = False, + multi_query_attention: bool = False, + ): + """Creates a Transformer model specification. + + Args: + num_layers: Number of encoder and decoder layers, or a 2-tuple if the + number is different. + num_heads: Number of attention heads. + with_relative_position: Use relative position representations in the self-attention + layers as described in https://arxiv.org/abs/1803.02155. + pre_norm: Enable the pre-norm Transformer architecture. + no_final_norm: Disable the final layer norm in the pre-norm architecture. + activation: Activation to apply in the feed-forward network. + alignment_layer: Layer index selected for alignment. + alignment_heads: Number of attention heads selected for alignment. + num_source_embeddings: Number of source embeddings. + embeddings_merge: When :obj:`num_source_embeddings` > 1, specify how the + embeddings are merged. + layernorm_embedding: Apply layer normalization after the embedding layer. + relative_attention_bias: Use relative attention bias in the self-attention + layers as described in the T5 paper https://arxiv.org/abs/1910.10683. + ffn_glu: Use gated linear units in the FFN layer as described in + https://arxiv.org/abs/2002.05202. + rms_norm: Use the root mean square layer normalization. + multi_query_attention: Use multi-query attention. + """ + if isinstance(num_layers, (list, tuple)): + num_encoder_layers, num_decoder_layers = num_layers + else: + num_encoder_layers, num_decoder_layers = num_layers, num_layers + + encoder = TransformerEncoderSpec( + num_encoder_layers, + num_heads, + pre_norm=pre_norm, + no_final_norm=no_final_norm, + activation=activation, + num_source_embeddings=num_source_embeddings, + embeddings_merge=embeddings_merge, + layernorm_embedding=layernorm_embedding, + relative_position=with_relative_position, + relative_attention_bias=relative_attention_bias, + ffn_glu=ffn_glu, + rms_norm=rms_norm, + multi_query_attention=multi_query_attention, + ) + + decoder = TransformerDecoderSpec( + num_decoder_layers, + num_heads, + pre_norm=pre_norm, + no_final_norm=no_final_norm, + activation=activation, + layernorm_embedding=layernorm_embedding, + relative_position=with_relative_position, + relative_attention_bias=relative_attention_bias, + alignment_layer=alignment_layer, + alignment_heads=alignment_heads, + ffn_glu=ffn_glu, + rms_norm=rms_norm, + multi_query_attention=multi_query_attention, + ) + + return cls(encoder, decoder) + + @property + def name(self): + return "TransformerSpec" + + @property + def revision(self): + return 7 + + def get_default_config(self): + return TransformerConfig() + + def get_source_vocabulary_size(self): + return [spec.weight.shape[0] for spec in self.encoder.embeddings] + + def get_target_vocabulary_size(self): + return self.decoder.embeddings.weight.shape[0] class TransformerDecoderModelConfig(model_spec.LanguageModelConfig): """Configuration for Transformer decoder models.""" diff --git a/python/ctranslate2/specs/whispernmt_spec.py b/python/ctranslate2/specs/whispernmt_spec.py new file mode 100644 index 000000000..7b2cd2225 --- /dev/null +++ b/python/ctranslate2/specs/whispernmt_spec.py @@ -0,0 +1,533 @@ +"""Declares specification of the Transformer model.""" + +from typing import Optional, Tuple, Union, List + +import numpy as np + +from ctranslate2.specs import attention_spec, common_spec, model_spec, transformer_spec + +class TransformerEncoderSpec(model_spec.LayerSpec): + def __init__( + self, + num_layers: int, + num_heads: int, + pre_norm: bool = True, + no_final_norm: bool = False, + activation: common_spec.Activation = common_spec.Activation.RELU, + num_source_embeddings: int = 1, + embeddings_merge: common_spec.EmbeddingsMerge = common_spec.EmbeddingsMerge.CONCAT, + layernorm_embedding: bool = False, + relative_position: bool = False, + relative_attention_bias: bool = False, + ffn_glu: bool = False, + rms_norm: bool = False, + multi_query_attention: bool = False, + ): + """Initializes a Transformer encoder specification. + + Args: + num_layers: Number of layers. + num_heads: Number of attention heads. + pre_norm: Enable the pre-norm Transformer architecture. + no_final_norm: Disable the final layer norm in the pre-norm architecture. + activation: Activation to apply in the feed-forward network. + num_source_embeddings: Number of source embeddings. + embeddings_merge: When :obj:`num_source_embeddings` > 1, specify how the + embeddings are merged. + layernorm_embedding: Apply layer normalization after the embedding layer. + relative_position: Use relative position representations in the self-attention + layers as described in https://arxiv.org/abs/1803.02155. + relative_attention_bias: Use relative attention bias in the self-attention + layers as described in the T5 paper https://arxiv.org/abs/1910.10683. + ffn_glu: Use gated linear units in the FFN layers as described in + https://arxiv.org/abs/2002.05202. + rms_norm: Use the root mean square layer normalization. + multi_query_attention: Use multi-query attention. + """ + self.multi_query_attention = multi_query_attention + self.num_heads = np.dtype("int16").type(num_heads) + self.pre_norm = pre_norm + self.activation = np.dtype("int8").type(activation) + self.embeddings_merge = np.dtype("int8").type(embeddings_merge) + self.embeddings = [ + common_spec.EmbeddingsSpec() for _ in range(num_source_embeddings) + ] + self.scale_embeddings = True + if not relative_position and not relative_attention_bias: + self.position_encodings = PositionEncoderSpec() + if pre_norm and not no_final_norm: + self.layer_norm = common_spec.LayerNormSpec(rms_norm=rms_norm) + if layernorm_embedding: + self.layernorm_embedding = common_spec.LayerNormSpec(rms_norm=rms_norm) + self.layer = [ + TransformerEncoderLayerSpec( + relative_position=relative_position, + relative_attention_bias=relative_attention_bias, + ffn_glu=ffn_glu, + rms_norm=rms_norm, + num_heads_kv=1 if multi_query_attention else None, + ) + for _ in range(num_layers) + ] + + +class TransformerDecoderSpec(model_spec.LayerSpec): + def __init__( + self, + num_layers: int, + num_heads: int, + pre_norm: bool = True, + activation: common_spec.Activation = common_spec.Activation.RELU, + layernorm_embedding: bool = False, + with_encoder_attention: bool = True, + no_final_norm: bool = False, + project_in_out: bool = False, + relative_position: bool = False, + relative_attention_bias: bool = False, + alignment_layer: int = -1, + alignment_heads: int = 1, + ffn_glu: bool = False, + rms_norm: bool = False, + alibi: bool = False, + alibi_use_positive_positions: bool = False, + scale_alibi: bool = False, + rotary_dim: Optional[int] = None, + rotary_interleave: bool = True, + rotary_scaling_type: Optional[attention_spec.RotaryScalingType] = None, + rotary_scaling_factor: float = 1, + rotary_base: float = 10000, + original_max_position_embeddings: int = 0, + max_position_embeddings: int = 0, + parallel_residual: bool = False, + shared_layer_norm: bool = False, + pre_post_layer_norm: bool = False, + multi_query_attention: bool = False, + num_heads_kv: Optional[int] = None, + head_dim: Optional[int] = None, + sliding_window: Optional[int] = None, + quant_type: Optional[common_spec.Quantization] = None, + quant_group_size: Optional[int] = None, + quant_bits: Optional[int] = None, + ): + """Initializes a Transformer decoder specification. + + Args: + num_layers: Number of layers. + num_heads: Number of attention heads. + pre_norm: Enable the pre-norm Transformer architecture. + activation: Activation to apply in the feed-forward network. + layernorm_embedding: Apply layer normalization after the embedding layer. + with_encoder_attention: Enable the encoder attention sublayers. + no_final_norm: Disable the final layer norm in the pre-norm architecture. + project_in_out: Add linear transformations after the embedding layer and before + the final layer. + relative_position: Use relative position representations in the self-attention + layers as described in https://arxiv.org/abs/1803.02155. + relative_attention_bias: Use relative attention bias in the self-attention + layers as described in the T5 paper https://arxiv.org/abs/1910.10683. + alignment_layer: Layer index selected for alignment. + alignment_heads: Number of attention heads selected for alignment. + ffn_glu: Use gated linear units in the FFN layers as described in + https://arxiv.org/abs/2002.05202. + rms_norm: Use the root mean square layer normalization. + alibi: Use attention with linear biases. + alibi_use_positive_positions: Use positive positions in the ALiBi definition. + scale_alibi: Apply the dot product scale factor to ALiBi. + rotary_dim: Apply rotary embeddings to these first N dimensions. If 0, rotary + embeddings are applied to all dimensions. + rotary_interleave: Interleave the head dimensions when rotary embeddings are applied. + Otherwise the head dimensions are sliced in half. + rotary_scaling_type: Type of RoPE scaling. + rotary_scaling_factor: Factor used in the RoPE scaling. + rotary_base: The base period of the rotary embeddings. + original_max_position_embeddings: The original max position embeddings + for Su rope embeddings + max_position_embeddings: The max position embeddings for Su rope embeddings + parallel_residual: Use parallel residual connections in each layer block, as used + by the GPT-J and GPT-NeoX models. + shared_layer_norm: When using parallel residual, share the input and post + attention layer norms. + pre_post_layer_norm: Add post layer norm for each pre norm layer + multi_query_attention: Use multi-query attention (alias for num_heads_kv=1). + num_heads_kv: Number of attention heads for the key and value. + sliding_window: Max sequence length to retain in KV Cache. + quant_type: quantization type used (like awq... for lower bit quantization) + quant_group_size: group size of the lower bit quantization + quant_bits: number of bit of the quantization (ex: 4bit) + """ + + self._config = dict() + if parallel_residual: + if not pre_norm: + raise ValueError("The GPT-J block expects a pre-norm architecture") + if with_encoder_attention: + raise ValueError("The GPT-J block does not have cross attention") + + if multi_query_attention: + if num_heads_kv is not None and num_heads_kv != 1: + raise ValueError( + "Enabling multi_query_attention implies num_heads_kv=1" + ) + num_heads_kv = 1 + + if with_encoder_attention and num_heads_kv not in (None, 1, num_heads): + raise ValueError( + "num_heads_kv=%d is not supported in the cross-attention layers" + % num_heads_kv + ) + + self.num_heads = np.dtype("int16").type(num_heads) + self.pre_norm = pre_norm + self.activation = np.dtype("int8").type(activation) + self.alignment_layer = np.dtype("int16").type(alignment_layer) + self.alignment_heads = np.dtype("int16").type(alignment_heads) + self.embeddings = common_spec.EmbeddingsSpec() + self.scale_embeddings = True + self.scale_outputs = model_spec.OPTIONAL + self.alibi = alibi + self.alibi_use_positive_positions = alibi_use_positive_positions + self.scale_alibi = scale_alibi + if sliding_window is not None: + self.sliding_window = np.dtype("int32").type(sliding_window) + if ( + not relative_position + and not relative_attention_bias + and not alibi + and rotary_dim is None + ): + self.position_encodings = PositionEncoderSpec() + if pre_norm and not no_final_norm: + self.layer_norm = common_spec.LayerNormSpec(rms_norm=rms_norm) + if layernorm_embedding: + self.layernorm_embedding = common_spec.LayerNormSpec(rms_norm=rms_norm) + self.projection = common_spec.LinearSpec() + self.layer = [ + TransformerDecoderLayerSpec( + with_encoder_attention=with_encoder_attention, + relative_position=relative_position, + relative_attention_bias=relative_attention_bias, + ffn_glu=ffn_glu, + rms_norm=rms_norm, + rotary_dim=rotary_dim, + rotary_interleave=rotary_interleave, + rotary_scaling_type=rotary_scaling_type, + rotary_scaling_factor=rotary_scaling_factor, + rotary_base=rotary_base, + original_max_position_embeddings=original_max_position_embeddings, + max_position_embeddings=max_position_embeddings, + parallel_residual=parallel_residual, + shared_layer_norm=shared_layer_norm, + pre_post_layer_norm=pre_post_layer_norm, + num_heads_kv=num_heads_kv, + head_dim=head_dim, + sliding_window=sliding_window, + ) + for _ in range(num_layers) + ] + self.start_from_zero_embedding = False + self._config["multi_query_attention"] = multi_query_attention or ( + num_heads_kv != num_heads + ) + + if project_in_out: + self.project_in = common_spec.LinearSpec() + self.project_out = common_spec.LinearSpec() + + if quant_type is not None: + self._config["quantization_type"] = quant_type + self._config["quantization_bits"] = quant_bits + self._config["quantization_group_size"] = quant_group_size + + @property + def config(self): + return self._config + + +class TransformerEncoderLayerSpec(model_spec.LayerSpec): + def __init__( + self, + relative_position=False, + relative_attention_bias=False, + ffn_glu=False, + rms_norm=False, + num_heads_kv=None, + sliding_window=None, + ): + self.self_attention = attention_spec.MultiHeadAttentionSpec( + self_attention=True, + relative_position=relative_position, + relative_attention_bias=relative_attention_bias, + rms_norm=rms_norm, + num_heads_kv=num_heads_kv, + sliding_window=sliding_window, + ) + self.ffn = FeedForwardSpec(glu=ffn_glu, rms_norm=rms_norm) + + +class TransformerDecoderLayerSpec(model_spec.LayerSpec): + def __init__( + self, + with_encoder_attention=True, + relative_position=False, + relative_attention_bias=False, + ffn_glu=False, + rms_norm=False, + rotary_dim=None, + rotary_interleave=True, + rotary_scaling_type=None, + rotary_scaling_factor=1, + rotary_base=10000, + original_max_position_embeddings=0, + max_position_embeddings=0, + parallel_residual=False, + shared_layer_norm=False, + pre_post_layer_norm=False, + num_heads_kv=None, + head_dim=None, + sliding_window=None, + ): + self.self_attention = attention_spec.MultiHeadAttentionSpec( + self_attention=True, + relative_position=relative_position, + relative_attention_bias=relative_attention_bias, + rms_norm=rms_norm, + rotary_dim=rotary_dim, + rotary_interleave=rotary_interleave, + rotary_scaling_type=rotary_scaling_type, + rotary_scaling_factor=rotary_scaling_factor, + rotary_base=rotary_base, + original_max_position_embeddings=original_max_position_embeddings, + max_position_embeddings=max_position_embeddings, + num_heads_kv=num_heads_kv, + head_dim=head_dim, + sliding_window=sliding_window, + ) + + if with_encoder_attention: + self.attention = attention_spec.MultiHeadAttentionSpec( + rms_norm=rms_norm, + num_heads_kv=num_heads_kv, + sliding_window=sliding_window, + ) + + self.ffn = FeedForwardSpec(glu=ffn_glu, rms_norm=rms_norm) + + if parallel_residual: + if shared_layer_norm: + self.shared_layer_norm = common_spec.LayerNormSpec() + else: + self.input_layer_norm = common_spec.LayerNormSpec() + self.post_attention_layer_norm = common_spec.LayerNormSpec() + + delattr(self.self_attention, "layer_norm") + delattr(self.ffn, "layer_norm") + + if pre_post_layer_norm: + self.input_layer_norm = common_spec.LayerNormSpec(rms_norm=rms_norm) + self.post_attention_layer_norm = common_spec.LayerNormSpec( + rms_norm=rms_norm + ) + self.pre_feedforward_layer_norm = common_spec.LayerNormSpec( + rms_norm=rms_norm + ) + self.post_feedforward_layer_norm = common_spec.LayerNormSpec( + rms_norm=rms_norm + ) + + delattr(self.self_attention, "layer_norm") + delattr(self.ffn, "layer_norm") + + +class FeedForwardSpec(model_spec.LayerSpec): + def __init__(self, glu=False, rms_norm=False): + self.layer_norm = common_spec.LayerNormSpec(rms_norm=rms_norm) + self.linear_0 = common_spec.LinearSpec() + self.linear_1 = common_spec.LinearSpec() + if glu: + self.linear_0_noact = common_spec.LinearSpec() + + +class PositionEncoderSpec(model_spec.LayerSpec): + def __init__(self): + self.encodings = model_spec.OPTIONAL + + +class TransformerConfig(model_spec.SequenceToSequenceModelConfig): + """Configuration for Transformer models.""" + + def __init__(self, layer_norm_epsilon: Optional[float] = None, **kwargs): + """Initializes the configuration for Transformer models. + + Args: + layer_norm_epsilon: The layer norm epsilon value. + **kwargs: Additional configuration. + """ + super().__init__(layer_norm_epsilon=layer_norm_epsilon, **kwargs) + +class WhisperNmtConfig(model_spec.SequenceToSequenceModelConfig): + """Configuration for the Whisper model.""" + + def __init__( + self, + suppress_ids: Optional[List[int]] = None, + suppress_ids_begin: Optional[List[int]] = None, + lang_ids: Optional[List[int]] = None, + alignment_heads: Optional[List[Tuple[int, int]]] = None, + ): + super().__init__( + suppress_ids=suppress_ids, + suppress_ids_begin=suppress_ids_begin, + lang_ids=lang_ids, + alignment_heads=alignment_heads, + ) + +class WhisperEncoderSpec(model_spec.LayerSpec): + def __init__(self, num_layers, num_heads): + self.num_heads = np.dtype("int16").type(num_heads) + self.conv1 = common_spec.Conv1DSpec() + self.conv2 = common_spec.Conv1DSpec() + self.position_encodings = transformer_spec.PositionEncoderSpec() + self.layer_norm = common_spec.LayerNormSpec() + self.layer = [ + transformer_spec.TransformerEncoderLayerSpec() for _ in range(num_layers) + ] + +class WhisperNmtConnectorSpec(model_spec.LayerSpec): + def __init__(self, activation: common_spec.Activation = common_spec.Activation.RELU): + self.linear1 = common_spec.LinearSpec() + self.linear2 = common_spec.LinearSpec() + self.activation = np.dtype("int8").type(activation) + +class WhisperNmtSpec(model_spec.SequenceToSequenceModelSpec): + """Describes a Whisper model.""" + + def __init__( + self, + transformer_encoder: TransformerEncoderSpec, + transformer_decoder: TransformerDecoderSpec, + connector: WhisperNmtConnectorSpec, + whisper_encoder: WhisperEncoderSpec, + ): + """Initializes a Transformer model specification. + + Args: + transformer_encoder: The encoder specification of nmt module. + transformer_decoder: The decoder specification of nmt module. + whisper_encoder: The encoder specification of whisper module. + """ + super().__init__() + self.transformer_encoder = transformer_encoder + self.transformer_decoder = transformer_decoder + self.whisper_encoder = whisper_encoder + self.connector = connector + self._config.add_attribute( + "multi_query_attention", self.transformer_encoder.multi_query_attention + ) + + @classmethod + def from_config( + cls, + num_layers: Union[int, Tuple[int, int]], + num_heads: int, + whisper_encoder_num_layers: int, + whisper_encoder_num_heads: int, + with_relative_position: bool = False, + pre_norm: bool = True, + no_final_norm: bool = False, + activation: common_spec.Activation = common_spec.Activation.RELU, + alignment_layer: int = -1, + alignment_heads: int = 1, + num_source_embeddings: int = 1, + embeddings_merge: common_spec.EmbeddingsMerge = common_spec.EmbeddingsMerge.CONCAT, + layernorm_embedding: bool = False, + relative_attention_bias: bool = False, + ffn_glu: bool = False, + rms_norm: bool = False, + multi_query_attention: bool = False, + ): + """Creates a Transformer model specification. + + Args: + num_layers: Number of encoder and decoder layers, or a 2-tuple if the + number is different. + num_heads: Number of attention heads. + whisper_encoder_num_layers: Number of encoder layers for whisper module. + whisper_encoder_num_heads: Number of attention heads for whisper module. + with_relative_position: Use relative position representations in the self-attention + layers as described in https://arxiv.org/abs/1803.02155. + pre_norm: Enable the pre-norm Transformer architecture. + no_final_norm: Disable the final layer norm in the pre-norm architecture. + activation: Activation to apply in the feed-forward network. + alignment_layer: Layer index selected for alignment. + alignment_heads: Number of attention heads selected for alignment. + num_source_embeddings: Number of source embeddings. + embeddings_merge: When :obj:`num_source_embeddings` > 1, specify how the + embeddings are merged. + layernorm_embedding: Apply layer normalization after the embedding layer. + relative_attention_bias: Use relative attention bias in the self-attention + layers as described in the T5 paper https://arxiv.org/abs/1910.10683. + ffn_glu: Use gated linear units in the FFN layer as described in + https://arxiv.org/abs/2002.05202. + rms_norm: Use the root mean square layer normalization. + multi_query_attention: Use multi-query attention. + """ + if isinstance(num_layers, (list, tuple)): + num_encoder_layers, num_decoder_layers = num_layers + else: + num_encoder_layers, num_decoder_layers = num_layers, num_layers + + encoder = TransformerEncoderSpec( + num_encoder_layers, + num_heads, + pre_norm=pre_norm, + no_final_norm=no_final_norm, + activation=activation, + num_source_embeddings=num_source_embeddings, + embeddings_merge=embeddings_merge, + layernorm_embedding=layernorm_embedding, + relative_position=with_relative_position, + relative_attention_bias=relative_attention_bias, + ffn_glu=ffn_glu, + rms_norm=rms_norm, + multi_query_attention=multi_query_attention, + ) + + decoder = TransformerDecoderSpec( + num_decoder_layers, + num_heads, + pre_norm=pre_norm, + no_final_norm=no_final_norm, + activation=activation, + layernorm_embedding=layernorm_embedding, + relative_position=with_relative_position, + relative_attention_bias=relative_attention_bias, + alignment_layer=alignment_layer, + alignment_heads=alignment_heads, + ffn_glu=ffn_glu, + rms_norm=rms_norm, + multi_query_attention=multi_query_attention, + ) + + connector = WhisperNmtConnectorSpec(activation=activation) + whisper_encoder = WhisperEncoderSpec(num_layers=whisper_encoder_num_layers, num_heads=whisper_encoder_num_heads) + + return cls(encoder, decoder, connector, whisper_encoder) + + @property + def name(self): + return "WhisperNmtSpec" + + @property + def revision(self): + return 7 + + def get_default_config(self): + return WhisperNmtConfig() + + def get_source_vocabulary_size(self): + return [spec.weight.shape[0] for spec in self.transformer_encoder.embeddings] + + def get_target_vocabulary_size(self): + return self.transformer_decoder.embeddings.weight.shape[0] + + diff --git a/src/layers/transformer.cc b/src/layers/transformer.cc index 5ac5bfa36..43c54e1d6 100644 --- a/src/layers/transformer.cc +++ b/src/layers/transformer.cc @@ -363,6 +363,70 @@ namespace ctranslate2 { padder->add_padding(output); } + void TransformerEncoder::operator()(const StorageView& input, + const std::vector& language_ids, + const std::vector& eos_ids, + const StorageView* lengths, + StorageView& output) { + PROFILE("TransformerEncoder"); + StorageView lang_input(output.dtype(), output.device()); + _embeddings(language_ids, lang_input); + + if (_embeddings_scale) + ops::Mul()(lang_input, *_embeddings_scale, lang_input); + if (_position_encoder) + (*_position_encoder)(lang_input); + if (_layernorm_embedding) + (*_layernorm_embedding)(lang_input, lang_input); + + StorageView eos_input(output.dtype(), output.device()); + _embeddings(eos_ids, eos_input); + + if (_embeddings_scale) + ops::Mul()(eos_input, *_embeddings_scale, eos_input); + if (_position_encoder) + (*_position_encoder)(eos_input); + if (_layernorm_embedding) + (*_layernorm_embedding)(eos_input, eos_input); + + // Concatenate the language embeddings and input and eos embeddings. + StorageView input_concat(output.dtype(), output.device()); + ops::Concat(1)({&lang_input, &input, &eos_input}, input_concat); + + StorageView hidden = input; + const dim_t max_time = input.dim(1); + + // Remove padding to reduce the amount of computation. + std::unique_ptr padder; + std::unique_ptr lengths_mask; + + if (lengths) { + if (Padder::allow_padding_removal(output.device(), _compute_type)) { + padder = std::make_unique(*lengths, max_time); + padder->remove_padding(hidden); + } + + int num_heads = _num_heads; + if (_tensor_parallel) { + num_heads = SAFE_DIVIDE(num_heads, ScopedMPISetter::getNRanks()); + } + lengths_mask = std::make_unique( + layers::MultiHeadAttention::prepare_length_mask(*lengths, num_heads, max_time)); + } + + StorageView position_bias(output.dtype(), output.device()); + + for (size_t l = 0; l < _layers.size(); ++l) { + (*_layers[l])(hidden, lengths_mask.get(), output, padder.get(), &position_bias); + if (l + 1 < _layers.size()) + hidden = std::move(output); + } + if (_output_norm) + (*_output_norm)(output, output); + if (padder) + padder->add_padding(output); + } + static std::unique_ptr make_alibi(const models::Model& model, const std::string& scope) { const bool use_alibi = model.get_flag_with_default(scope + "/alibi", false); diff --git a/src/layers/whisper.cc b/src/layers/whisper.cc index aa5dece77..6f8662551 100644 --- a/src/layers/whisper.cc +++ b/src/layers/whisper.cc @@ -81,5 +81,21 @@ namespace ctranslate2 { _proj(step_outputs, logits); } + WhisperConnector::WhisperConnector(const models::Model &model, const std::string &scope) + : _activation_type(ops::ActivationType::ReLU) + , _lin_1(model, scope + "/linear1", &_activation_type) + , _lin_2(model, scope + "/linear2") { + } + + void WhisperConnector::operator()(const StorageView &features, StorageView &output) const { + PROFILE("WhisperConnector"); + StorageView hidden(features.dtype(), features.device()); + StorageView intermediate(features.dtype(), features.device()); + ops::Transpose({0, 2, 1})(features, hidden); + _lin_1(hidden, intermediate); + _lin_2(intermediate, output); + ops::Transpose({0, 2, 1})(output, hidden); + output = std::move(hidden); + } } } diff --git a/src/models/model_factory.cc b/src/models/model_factory.cc index 059051f5d..7725c08fc 100644 --- a/src/models/model_factory.cc +++ b/src/models/model_factory.cc @@ -6,6 +6,7 @@ #include "ctranslate2/models/wav2vec2.h" #include "ctranslate2/models/wav2vec2bert.h" #include "ctranslate2/models/transformer.h" +#include "ctranslate2/models/whisper_nmt.h" namespace ctranslate2 { namespace models { @@ -22,6 +23,7 @@ namespace ctranslate2 { register_model("TransformerEncoderSpec"); register_model("WhisperSpec"); + register_model("WhisperNmtSpec"); register_model("Wav2Vec2Spec"); diff --git a/src/models/whisper_nmt.cc b/src/models/whisper_nmt.cc new file mode 100644 index 000000000..93dc87d33 --- /dev/null +++ b/src/models/whisper_nmt.cc @@ -0,0 +1,735 @@ +#include "ctranslate2/models/whisper_nmt.h" + +#include + +#include "ctranslate2/decoding.h" + +#include "dispatch.h" +#include "dtw.h" + +#ifdef CT2_WITH_CUDA +# include "cuda/utils.h" +#endif + +namespace ctranslate2 { + namespace models { + + const Vocabulary& WhisperNmtModel::get_vocabulary() const { + return *_vocabulary; + } + + size_t WhisperNmtModel::num_source_vocabularies() const { + return _source_vocabularies.size(); + } + + const Vocabulary& WhisperNmtModel::get_source_vocabulary(size_t index) const { + return *_source_vocabularies.at(index); + } + + const Vocabulary& WhisperNmtModel::get_target_vocabulary() const { + return *_target_vocabulary; + } + + size_t WhisperNmtModel::current_spec_revision() const { + return 7; + } + + void WhisperNmtModel::load_vocabularies(ModelReader& model_reader) { + { + VocabularyInfo vocab_info; + vocab_info.unk_token = config["unk_token"]; + vocab_info.bos_token = config["bos_token"]; + vocab_info.eos_token = config["eos_token"]; + + auto shared_vocabulary = load_vocabulary(model_reader, "shared_vocabulary", vocab_info); + + if (shared_vocabulary) { + _target_vocabulary = shared_vocabulary; + _source_vocabularies = {shared_vocabulary}; + + } else { + _target_vocabulary = load_vocabulary(model_reader, "target_vocabulary", vocab_info); + if (!_target_vocabulary) + throw std::runtime_error("Cannot load the target vocabulary from the model directory"); + + auto source_vocabulary = load_vocabulary(model_reader, "source_vocabulary", vocab_info); + + if (source_vocabulary) { + _source_vocabularies = {source_vocabulary}; + } else { + for (size_t i = 1;; i++) { + const std::string filename = "source_" + std::to_string(i) + "_vocabulary"; + auto vocabulary = load_vocabulary(model_reader, filename, vocab_info); + + if (!vocabulary) + break; + + _source_vocabularies.emplace_back(vocabulary); + } + } + + if (_source_vocabularies.empty()) + throw std::runtime_error("Cannot load the source vocabulary from the model directory"); + } + } + } + + void WhisperNmtModel::initialize(ModelReader& model_reader) { + //VocabularyInfo vocab_info; + //vocab_info.unk_token = "<|endoftext|>"; + //vocab_info.bos_token = "<|startoftranscript|>"; + //vocab_info.eos_token = "<|endoftext|>"; + + //_vocabulary = load_vocabulary(model_reader, "shared_vocabulary", std::move(vocab_info)); + //if (!_vocabulary) + // throw std::runtime_error("Cannot load the vocabulary from the model directory"); + // load vocabularies for nmt encoder and decoder + load_vocabularies(model_reader); + } + + bool WhisperNmtModel::is_quantizable(const std::string& variable_name) const { + return Model::is_quantizable(variable_name); + } + + bool WhisperNmtModel::is_linear_weight(const std::string& variable_name) const { + return is_quantizable(variable_name) && variable_name.find("embeddings") == std::string::npos; + } + + std::unique_ptr WhisperNmtModel::clone() const { + return std::make_unique(*this); + } + + + std::unique_ptr WhisperNmtReplica::create_from_model(const Model& model) { + if (!dynamic_cast(&model)) + throw std::invalid_argument("The model is not a Whisper model"); + + const auto scoped_device_setter = model.get_scoped_device_setter(); + const auto model_ptr = model.shared_from_this(); + const auto concrete_model = std::static_pointer_cast(model_ptr); + return std::make_unique(concrete_model); + } + + WhisperNmtReplica::WhisperNmtReplica(const std::shared_ptr& model) + : ModelReplica(model) + , _model(model) + , _encoder(std::make_unique(*model, "whisper_encoder")) + , _transformer_encoder(std::make_unique(*model, "transformer_encoder")) + , _transformer_decoder(std::make_unique(*model, "transformer_decoder")) + , _connector(std::make_unique(*model, "connector")) + { + const auto& vocabulary = model->get_source_vocabulary(0); + _sot_id = vocabulary.bos_id(); + _eot_id = vocabulary.eos_id(); + _no_timestamps_id = vocabulary.to_id("<|notimestamps|>"); + _no_speech_id = vocabulary.to_id("<|nospeech|>"); + if (_no_speech_id == vocabulary.unk_id()) + _no_speech_id = vocabulary.to_id("<|nocaptions|>"); + _is_multilingual = vocabulary.size() >= 51865; + _n_mels = _encoder->input_size(); + _num_languages = vocabulary.size() - 51765 - (_is_multilingual ? 1 : 0); + } + + StorageView WhisperNmtReplica::encode(StorageView features, const bool to_cpu) { + PROFILE("WhisperReplica::encode"); + +#ifdef CT2_WITH_CUDA + const cuda::UseTrueFp16GemmInScope use_true_fp16_gemm(false); +#endif + + const auto scoped_device_setter = _model->get_scoped_device_setter(); + const Device device = _model->device(); + const DataType dtype = _encoder->output_type(); + features.move_to(device, dtype); + + StorageView encoder_output(dtype, device); + (*_encoder)(features, encoder_output); + + if (to_cpu) { + if (device != Device::CPU) + encoder_output = encoder_output.to(Device::CPU); + return encoder_output; + } + + // Ensure all operations are finished before returning the output. + synchronize_stream(device); + + return encoder_output; + } + + StorageView WhisperNmtReplica::maybe_encode(StorageView features) { + const Device device = _model->device(); + const DataType dtype = _encoder->output_type(); + + features.move_to(device, dtype); + + if (_encoder->is_encoded(features)) + return features; + + StorageView encoder_output(dtype, device); + (*_encoder)(features, encoder_output); + return encoder_output; + } + + + static std::vector get_no_speech_probs_from_logits(const StorageView& logits, + const size_t no_speech_id) { + const Device device = logits.device(); + const DataType dtype = logits.dtype(); + + StorageView probs(dtype, device); + ops::SoftMax()(logits, probs); + + StorageView gather_ids({probs.dim(0)}, int32_t(no_speech_id), device); + StorageView no_speech_probs(dtype, device); + ops::Gather(/*axis=*/1, /*batch_dims=*/1)(probs, gather_ids, no_speech_probs); + + if (no_speech_probs.dtype() != DataType::FLOAT32) + no_speech_probs = no_speech_probs.to_float32(); + return no_speech_probs.to_vector(); + } + + static size_t get_sot_index(const std::vector& prompt, const size_t sot_id) { + const auto sot_it = std::find(prompt.begin(), prompt.end(), sot_id); + if (sot_it == prompt.end()) + throw std::invalid_argument("<|startoftranscript|> token was not found in the prompt"); + + return std::distance(prompt.begin(), sot_it); + } + + static size_t get_prompt_length(const std::vector& prompt, + const size_t sot_id, + const size_t no_timestamps_id) { + size_t index = get_sot_index(prompt, sot_id); + while (index < prompt.size() && prompt[index] >= sot_id && prompt[index] <= no_timestamps_id) + index++; + return index; + } + + static void check_prompts(const std::vector>& prompts, + const size_t sot_id, + const size_t no_timestamps_id, + size_t& sot_index, + size_t& prompt_length) { + bool first = true; + + for (const auto& prompt : prompts) { + const auto batch_sot_index = get_sot_index(prompt, sot_id); + const auto batch_prompt_length = get_prompt_length(prompt, sot_id, no_timestamps_id); + + if (first) { + sot_index = batch_sot_index; + prompt_length = batch_prompt_length; + } else if (batch_sot_index != sot_index) { + throw std::invalid_argument("The generate method currently requires the " + "<|startoftranscript|> token to be at the same position " + "in all batches. To work around this limitation, " + "simply adapt the number of previous text tokens in each " + "batch."); + } else if (batch_prompt_length != prompt_length) { + throw std::invalid_argument("The generate method currently requires each batch to have " + "the same number of task tokens after <|startoftranscript|>."); + } + + first = false; + } + } + + class ApplyTimestampRules; + + class GetNoSpeechProbs : public LogitsProcessor { + private: + const size_t _no_speech_id; + std::vector _no_speech_probs; + + public: + GetNoSpeechProbs(const size_t no_speech_id) + : _no_speech_id(no_speech_id) + { + } + + const std::vector& get_no_speech_probs() const { + return _no_speech_probs; + } + + bool apply_first() const override { + return true; + } + + void apply(dim_t step, + StorageView& logits, + DisableTokens&, + const StorageView&, + const std::vector& batch_offset, + const std::vector>*) override { + if (step == 0) { + const auto no_speech_probs = get_no_speech_probs_from_logits(logits, _no_speech_id); + + const size_t batch_size = batch_offset.size(); + const size_t beam_size = logits.dim(0) / batch_size; + + _no_speech_probs.reserve(batch_size); + for (size_t i = 0; i < batch_size; ++i) + _no_speech_probs.emplace_back(no_speech_probs[i * beam_size]); + } + } + }; + + std::vector + WhisperNmtReplica::generate(StorageView features, + const std::vector& language, + const std::vector>& eos, + const WhisperNmtOptions& options) { + PROFILE("WhisperReplica::generate"); + +#ifdef CT2_WITH_CUDA + const cuda::UseTrueFp16GemmInScope use_true_fp16_gemm(false); +#endif + + size_t sot_index = 0; + size_t prompt_length = 0; // Length of the prompt before the text tokens. + //check_prompts(prompts, _sot_id, _no_timestamps_id, sot_index, prompt_length); + + const auto& vocabulary = _model->get_vocabulary(); + const auto scoped_device_setter = _model->get_scoped_device_setter(); + + //layers::DecoderState state = _decoder->initial_state(); + StorageView encode_features = maybe_encode(std::move(features)); + + // MLP + StorageView mlp_output(_encoder->output_type(), _model->device()); + (*_connector)(encode_features, mlp_output); + // NMT + std::vector final_results = _run_translation(mlp_output, + language, + eos, + options); + + return final_results; + } + + void + WhisperNmtReplica::_nmt_encode(StorageView& features_ids, + const std::vector>>& language_ids, + const std::vector>>& eos_ids, + StorageView& memory, + StorageView& memory_lengths) { + const size_t num_langid_features = language_ids.size(); + std::vector lang_feature_ids; + lang_feature_ids.reserve(num_langid_features); + + for (size_t i = 0; i < num_langid_features; ++i) { + const auto& tokens_ids = language_ids[i]; + lang_feature_ids.emplace_back(layers::make_sequence_inputs(tokens_ids, + _model->device(), + _model->preferred_size_multiple(), + i == 0 ? &memory_lengths : nullptr)); + } + const size_t num_eosid_features = eos_ids.size(); + std::vector eos_feature_ids; + eos_feature_ids.reserve(num_eosid_features); + + for (size_t i = 0; i < num_langid_features; ++i) { + const auto& tokens_ids = language_ids[i]; + eos_feature_ids.emplace_back(layers::make_sequence_inputs(tokens_ids, + _model->device(), + _model->preferred_size_multiple(), + i == 0 ? &memory_lengths : nullptr)); + } + (*_transformer_encoder)(features_ids, lang_feature_ids, eos_feature_ids, &memory_lengths, memory); + } + + std::vector>> + WhisperNmtReplica::make_source_ids(const std::vector>>& source_features, + size_t max_length) const { + const size_t num_input_features = source_features.size(); + if (_model->num_source_vocabularies() != num_input_features) + throw std::runtime_error("The encoder expects " + + std::to_string(num_input_features) + + " input features, but " + + std::to_string(_model->num_source_vocabularies()) + + " source vocabularies are loaded"); + + std::vector>> ids; + ids.reserve(num_input_features); + + for (size_t i = 0; i < num_input_features; ++i) { + const auto& vocabulary = _model->get_source_vocabulary(i); + ids.emplace_back(vocabulary.to_ids(source_features[i], + max_length, + _model->with_source_bos(), + _model->with_source_eos())); + } + + return ids; + } + + std::vector> + WhisperNmtReplica::make_target_ids(const std::vector>& target, + size_t max_length, + bool is_prefix) const { + const auto& target_vocabulary = _model->get_target_vocabulary(); + const std::string* suffix = &target_vocabulary.eos_token(); + const std::string* prefix = _model->decoder_start_token(); + + if (is_prefix) { + suffix = nullptr; + max_length = 0; + } else if (max_length > 0) { + // The method returns the full target " a b c " but the decoder input is " a b c". + // So 1 additional token is allowed in the full target sequence. + max_length += 1; + } + + return target_vocabulary.to_ids(target, max_length, prefix, suffix); + } + + std::vector + WhisperNmtReplica::_run_translation(StorageView& features, + const std::vector& language, + const std::vector>& eos, + const WhisperNmtOptions& options) { + const auto scoped_device_setter = _model->get_scoped_device_setter(); + const auto device = _model->device(); + PROFILE("EncoderDecoderReplica::run_translation"); + + const size_t batch_size = features; + + const auto lang_features = extract_features({language}, _transformer_encoder->num_input_features()); + const auto lang_ids = make_source_ids(lang_features, options.max_input_length); + const auto eos_features = extract_features(eos, _transformer_encoder->num_input_features()); + const auto eos_ids = make_source_ids(eos_features, options.max_input_length); + + // Encode the sequence. + StorageView memory(_encoder->output_type(), device); + StorageView memory_lengths(DataType::INT32, device); + _nmt_encode(features, lang_ids, eos_ids, memory, memory_lengths); + //encode(source_ids, memory, memory_lengths); + + std::vector> start_tokens; + start_tokens = eos_ids[0]; + + layers::DecoderState state = _transformer_decoder->initial_state(); + state.emplace("memory", std::move(memory)); + state.emplace("memory_lengths", std::move(memory_lengths)); + + const auto& target_vocabulary = _model->get_target_vocabulary(); + std::vector restrict_ids; + _transformer_decoder->update_output_layer(_model->preferred_size_multiple(), restrict_ids); + + // Decode. + DecodingOptions decoding_options; + decoding_options.beam_size = options.beam_size; + decoding_options.patience = options.patience; + decoding_options.length_penalty = options.length_penalty; + decoding_options.coverage_penalty = options.coverage_penalty; + decoding_options.repetition_penalty = options.repetition_penalty; + decoding_options.no_repeat_ngram_size = options.no_repeat_ngram_size; + decoding_options.prefix_bias_beta = options.prefix_bias_beta; + decoding_options.max_length = options.max_decoding_length; + decoding_options.min_length = options.min_decoding_length; + decoding_options.sampling_topk = options.sampling_topk; + decoding_options.sampling_topp = options.sampling_topp; + decoding_options.sampling_temperature = options.sampling_temperature; + decoding_options.num_hypotheses = options.num_hypotheses; + decoding_options.return_scores = options.return_scores; + decoding_options.return_logits_vocab = options.return_logits_vocab; + decoding_options.return_attention = options.return_attention || options.replace_unknowns; + decoding_options.return_alternatives = options.return_alternatives; + decoding_options.min_alternative_expansion_prob = options.min_alternative_expansion_prob; + decoding_options.disable_sequences = target_vocabulary.to_ids(options.suppress_sequences, + /*max_length=*/0, + /*prefix=*/nullptr, + /*suffix=*/nullptr, + /*allow_unk=*/false); + + if (options.disable_unk) + decoding_options.disable_ids.push_back(target_vocabulary.unk_id()); + + const auto end_ids(std::visit(ResolveEndToken(target_vocabulary), options.end_token)); + std::vector results = decode(*_transformer_decoder, + state, + start_tokens, + end_ids, + decoding_options); + + // Convert generated ids to tokens. + std::vector final_results; + final_results.reserve(batch_size); + + for (size_t i = 0; i < batch_size; ++i) { + DecodingResult& result = results[i]; + + // Remove EOS token. + if (!options.return_end_token) { + for (size_t h = 0; h < result.hypotheses.size(); ++h) { + while (!result.hypotheses[h].empty() && is_eos(result.hypotheses[h].back(), end_ids)) { + result.hypotheses[h].pop_back(); + if (!result.attention.empty()) + result.attention[h].pop_back(); + } + } + } + + auto hypotheses = target_vocabulary.to_tokens(result.hypotheses); + + /* + if (!result.attention.empty()) { + const auto& source_original = source_features[0][i]; + const auto& source_input = source_ids[0][i]; + + for (size_t h = 0; h < result.attention.size(); ++h) { + auto& attention = result.attention[h]; + + for (auto& vector : attention) { + // Remove attenton positions for padding and implicit special tokens. + vector.resize(source_input.size()); + if (_model->with_source_bos()) + vector.erase(vector.begin()); + if (_model->with_source_eos()) + vector.pop_back(); + + // Resize to the original input size. + vector.resize(source_original.size(), 0); + } + + if (options.replace_unknowns) + replace_unknown_tokens(source_original, + hypotheses[h], + attention, + target_vocabulary.unk_token()); + } + + if (!options.return_attention) + result.attention.clear(); + }*/ + + + WhisperNmtGenerationResult final_result; + final_result.sequences = std::move(hypotheses); + //final_result.sequences_ids = std::move(result.hypotheses); + final_result.scores = std::move(result.scores); + final_result.attention = std::move(result.attention); + final_result.logits = std::move(result.logits_vocab); + //if (options.return_no_speech_prob) + // final_result.no_speech_prob = no_speech_probs[i]; + + final_results.emplace_back(std::move(final_result)); + } + + return final_results; + } + + static void remove_padding(StorageView& x, dim_t axis, dim_t size) { + const dim_t max_size = x.dim(axis); + + if (size < max_size) { + StorageView content(x.dtype(), x.device()); + StorageView padding(x.dtype(), x.device()); + + const ops::Split split_op(axis, {size, max_size - size}); + split_op(x, content, padding); + + x = std::move(content); + } + } + + static std::vector>> + compute_alignments(StorageView& attention_probs, + const std::vector& start_sequence, + const std::vector>& text_tokens, + const dim_t median_filter_width) { + const ops::MedianFilter median_filter_op(median_filter_width); + const dim_t batch_size = attention_probs.dim(0); + + // The remaining operations are not implemented on GPU, so move back to CPU. + attention_probs.move_to(Device::CPU, DataType::FLOAT32); + + ops::LayerNorm(-2, 0)(attention_probs); + + StorageView median_filter; + median_filter_op(attention_probs, median_filter); + + StorageView weights; + ops::Mean(1)(median_filter, weights); + + std::vector>> alignments; + alignments.reserve(batch_size); + + for (dim_t b = 0; b < batch_size; ++b) { + const dim_t text_length = text_tokens[b].size(); + const dim_t sot_length = start_sequence.size(); + + StorageView matrix(Shape{text_length + 1, weights.dim(2)}); + if (weights) + matrix.view(weights.index({b, sot_length, 0}), matrix.shape()); + + alignments.emplace_back(negative_dtw(matrix)); + } + + return alignments; + } + + bool WhisperNmt::is_multilingual() const { + const auto& replica = get_first_replica(); + return replica.is_multilingual(); + } + + size_t WhisperNmt::n_mels() const { + const auto& replica = get_first_replica(); + return replica.n_mels(); + } + + size_t WhisperNmt::num_languages() const { + const auto& replica = get_first_replica(); + return replica.num_languages(); + } + + std::future WhisperNmt::encode(const StorageView& features, const bool to_cpu) { + return post( + [features = features.sync_copy(), to_cpu](WhisperNmtReplica& replica) mutable { + return replica.encode(std::move(features), to_cpu); + }); + } + + std::vector> + WhisperNmt::generate(const StorageView& features, + std::vector& language, + std::vector>& eos, + WhisperNmtOptions options) { + const size_t batch_size = features.dim(0); + return post_batch( + [features = features.sync_copy(), + language = std::move(language), + eos = std::move(eos), + options = std::move(options)] + (WhisperNmtReplica& replica) mutable { + return replica.generate(std::move(features), language, eos, options); + }, + batch_size); + } + + class ApplyTimestampRules : public LogitsProcessor { + private: + const size_t _eot_id; + const size_t _no_timestamps_id; + const size_t _timestamp_begin_id; + const size_t _timestamp_end_id; + const size_t _max_initial_timestamp_id; + + public: + ApplyTimestampRules(const size_t eot_id, + const size_t no_timestamps_id, + const size_t timestamp_begin_id, + const size_t timestamp_end_id, + const size_t max_initial_timestamp_id) + : _eot_id(eot_id) + , _no_timestamps_id(no_timestamps_id) + , _timestamp_begin_id(timestamp_begin_id) + , _timestamp_end_id(timestamp_end_id) + , _max_initial_timestamp_id(max_initial_timestamp_id) + { + } + + void apply(dim_t step, + StorageView& logits, + DisableTokens& disable_tokens, + const StorageView& sequences, + const std::vector& batch_offset, + const std::vector>* prefix) override { + std::vector check_timestamps_prob_for_batch; + const dim_t batch_size = logits.dim(0); + + for (dim_t batch_id = 0; batch_id < batch_size; ++batch_id) { + const dim_t sample_begin = get_sample_begin(batch_size, batch_id, batch_offset, prefix); + + // Suppress <|notimestamps|>. + disable_tokens.add(batch_id, _no_timestamps_id); + + if (step == sample_begin && step == 0) { + // Suppress non timestamps at the beginning. + for (size_t i = 0; i < _timestamp_begin_id; ++i) + disable_tokens.add(batch_id, i); + + // Apply max_initial_timestamp option. + for (size_t i = _max_initial_timestamp_id + 1; i <= _timestamp_end_id; ++i) + disable_tokens.add(batch_id, i); + + } else if (step > sample_begin) { + // Timestamps have to appear in pairs, except directly before EOT. + const size_t last_token = sequences.at({batch_id, step - 1}); + + if (last_token >= _timestamp_begin_id) { + const size_t penultimate_token = (step - 1 > sample_begin + ? sequences.at({batch_id, step - 2}) + : last_token); + + if (penultimate_token >= _timestamp_begin_id) { // has to be non-timestamp + for (size_t i = _timestamp_begin_id; i <= _timestamp_end_id; ++i) + disable_tokens.add(batch_id, i); + } else { // cannot be normal text tokens + for (size_t i = 0; i < _eot_id; ++i) + disable_tokens.add(batch_id, i); + for (size_t i = _timestamp_begin_id; i < last_token; ++i) + disable_tokens.add(batch_id, i); + check_timestamps_prob_for_batch.push_back(batch_id); + } + } else { + check_timestamps_prob_for_batch.push_back(batch_id); + + // Timestamps shouldn't decrease: forbid timestamp tokens smaller than the last. + for (dim_t t = step - 1; t >= sample_begin; --t) { + const size_t token = sequences.at({batch_id, t}); + + if (token >= _timestamp_begin_id) { + for (size_t i = _timestamp_begin_id; i <= token; ++i) + disable_tokens.add(batch_id, i); + break; + } + } + } + } + } + + if (!check_timestamps_prob_for_batch.empty()) { + // Apply all changes to the logits before computing the log softmax. + disable_tokens.apply(); + + StorageView log_probs(logits.dtype(), logits.device()); + ops::LogSoftMax()(logits, log_probs); + + for (const dim_t batch_id : check_timestamps_prob_for_batch) { + bool sample_timestamp = false; + + DEVICE_AND_FLOAT_DISPATCH( + "ApplyTimestampRules", log_probs.device(), log_probs.dtype(), + (sample_timestamp = should_sample_timestamp(log_probs, batch_id))); + + if (sample_timestamp) { + for (size_t i = 0; i < _timestamp_begin_id; ++i) + disable_tokens.add(batch_id, i); + } + } + } + } + + template + bool should_sample_timestamp(const StorageView& log_probs, const dim_t batch_id) { + const dim_t num_text_tokens = _timestamp_begin_id; + const dim_t num_timestamp_tokens = _timestamp_end_id - _timestamp_begin_id + 1; + + const T* text_log_probs = log_probs.index({batch_id, 0}); + const T* timestamp_log_probs = text_log_probs + num_text_tokens; + + // If sum of probability over timestamps is above any other token, sample timestamp. + const float max_text_token_log_prob = primitives::max(text_log_probs, num_text_tokens); + const float timestamp_log_prob = primitives::logsumexp(timestamp_log_probs, + num_timestamp_tokens); + + return timestamp_log_prob > max_text_token_log_prob; + } + + }; + + } +} From fe034ee0457639b6bab011215f5650d40b277e0b Mon Sep 17 00:00:00 2001 From: thucpham Date: Wed, 7 May 2025 13:40:45 +0200 Subject: [PATCH 2/3] fix bugs --- include/ctranslate2/models/whisper_nmt.h | 1 - src/layers/transformer.cc | 4 ++-- src/models/whisper_nmt.cc | 4 ++-- 3 files changed, 4 insertions(+), 5 deletions(-) diff --git a/include/ctranslate2/models/whisper_nmt.h b/include/ctranslate2/models/whisper_nmt.h index e6b6f9271..d824d146b 100644 --- a/include/ctranslate2/models/whisper_nmt.h +++ b/include/ctranslate2/models/whisper_nmt.h @@ -214,7 +214,6 @@ namespace ctranslate2 { StorageView& memory_lengths); const std::shared_ptr _model; const std::unique_ptr _encoder; - //const std::unique_ptr _decoder; const std::unique_ptr _transformer_encoder; const std::unique_ptr _transformer_decoder; const std::unique_ptr _connector; diff --git a/src/layers/transformer.cc b/src/layers/transformer.cc index 43c54e1d6..a8236cdea 100644 --- a/src/layers/transformer.cc +++ b/src/layers/transformer.cc @@ -393,8 +393,8 @@ namespace ctranslate2 { StorageView input_concat(output.dtype(), output.device()); ops::Concat(1)({&lang_input, &input, &eos_input}, input_concat); - StorageView hidden = input; - const dim_t max_time = input.dim(1); + StorageView hidden = input_concat; + const dim_t max_time = input_concat.dim(1); // Remove padding to reduce the amount of computation. std::unique_ptr padder; diff --git a/src/models/whisper_nmt.cc b/src/models/whisper_nmt.cc index 93dc87d33..4f6865383 100644 --- a/src/models/whisper_nmt.cc +++ b/src/models/whisper_nmt.cc @@ -329,8 +329,8 @@ namespace ctranslate2 { std::vector eos_feature_ids; eos_feature_ids.reserve(num_eosid_features); - for (size_t i = 0; i < num_langid_features; ++i) { - const auto& tokens_ids = language_ids[i]; + for (size_t i = 0; i < num_eosid_features; ++i) { + const auto& tokens_ids = eos_ids[i]; eos_feature_ids.emplace_back(layers::make_sequence_inputs(tokens_ids, _model->device(), _model->preferred_size_multiple(), From ffb4e474e4c838c7e9a8fd7a03d2a9ffd0f58a12 Mon Sep 17 00:00:00 2001 From: thucpham Date: Wed, 7 May 2025 15:20:14 +0200 Subject: [PATCH 3/3] add examples --- examples/whisper_nmt/whisper_nmt.py | 17 +++++++++++++++ src/models/whisper_nmt.cc | 32 ----------------------------- 2 files changed, 17 insertions(+), 32 deletions(-) create mode 100644 examples/whisper_nmt/whisper_nmt.py diff --git a/examples/whisper_nmt/whisper_nmt.py b/examples/whisper_nmt/whisper_nmt.py new file mode 100644 index 000000000..54e412911 --- /dev/null +++ b/examples/whisper_nmt/whisper_nmt.py @@ -0,0 +1,17 @@ +import ctranslate2 +import librosa +import transformers + +# Load and resample the audio file. +audio, _ = librosa.load("testaudio_16000_test01_20s.wav", sr=16000, mono=True) + +# Compute the features of the first 30 seconds of audio. +processor = transformers.WhisperProcessor.from_pretrained("openai/whisper-medium") +inputs = processor(audio, return_tensors="np", sampling_rate=16000) +features = ctranslate2.StorageView.from_array(inputs.input_features) + +# Load the model on CPU. +model = ctranslate2.models.WhisperNmt("models/whisper-nmt-ct2", device="cuda") + +results = model.generate(features, ["en"], [[""]]) +print(results) diff --git a/src/models/whisper_nmt.cc b/src/models/whisper_nmt.cc index 4f6865383..d4dfa2786 100644 --- a/src/models/whisper_nmt.cc +++ b/src/models/whisper_nmt.cc @@ -473,38 +473,6 @@ namespace ctranslate2 { auto hypotheses = target_vocabulary.to_tokens(result.hypotheses); - /* - if (!result.attention.empty()) { - const auto& source_original = source_features[0][i]; - const auto& source_input = source_ids[0][i]; - - for (size_t h = 0; h < result.attention.size(); ++h) { - auto& attention = result.attention[h]; - - for (auto& vector : attention) { - // Remove attenton positions for padding and implicit special tokens. - vector.resize(source_input.size()); - if (_model->with_source_bos()) - vector.erase(vector.begin()); - if (_model->with_source_eos()) - vector.pop_back(); - - // Resize to the original input size. - vector.resize(source_original.size(), 0); - } - - if (options.replace_unknowns) - replace_unknown_tokens(source_original, - hypotheses[h], - attention, - target_vocabulary.unk_token()); - } - - if (!options.return_attention) - result.attention.clear(); - }*/ - - WhisperNmtGenerationResult final_result; final_result.sequences = std::move(hypotheses); //final_result.sequences_ids = std::move(result.hypotheses);