ctranslate2 4.6.1__cp314-cp314-win_amd64.whl → 4.6.3__cp314-cp314-win_amd64.whl

ctranslate2/extensions.py

@@ -556,12 +556,28 @@ def _process_iterable(process_func, iterables, max_batch_size, batch_type, **kwa
 
 def _batch_iterator(iterable, batch_size, batch_type):
     streams = None
-    cur_batch_size = 0
+    max_length = 0
 
     for example in iterable:
         if not isinstance(example, tuple):
             example = (example,)
 
+        if batch_type == "examples":
+            if streams and len(streams[0]) == batch_size:
+                yield streams
+                streams = None
+
+        elif batch_type == "tokens":
+            max_length = max(max_length, len(example[0]))
+
+            if streams and (len(streams[0]) + 1) * max_length > batch_size:
+                yield streams
+                streams = None
+                max_length = len(example[0])
+
+        else:
+            raise ValueError("Invalid batch type %s" % batch_type)
+
         if streams is None:
             streams = tuple([] for _ in example)
         for batch, element in zip(streams, example):
@@ -569,17 +585,5 @@ def _batch_iterator(iterable, batch_size, batch_type):
                 raise ValueError("Input iterables do not have the same length")
             batch.append(element)
 
-        if batch_type == "examples":
-            cur_batch_size += 1
-        elif batch_type == "tokens":
-            cur_batch_size += len(example[0])
-        else:
-            raise ValueError("Invalid batch type %s" % batch_type)
-
-        if cur_batch_size >= batch_size:
-            yield streams
-            streams = None
-            cur_batch_size = 0
-
     if streams is not None:
         yield streams

ctranslate2/specs/attention_spec.py

@@ -32,14 +32,22 @@ class MultiHeadAttentionSpec(model_spec.LayerSpec):
         num_heads_kv=None,
         head_dim=None,
         sliding_window=None,
+        qk_norm=False,
+        qk_norm_rms=True,
+        has_norm=True,
     ):
         self.queries_scale = model_spec.OPTIONAL
 
-        self.layer_norm = common_spec.LayerNormSpec(rms_norm=rms_norm)
+        if has_norm:
+            self.layer_norm = common_spec.LayerNormSpec(rms_norm=rms_norm)
         self.linear = [
             common_spec.LinearSpec() for _ in range(2 if self_attention else 3)
         ]
 
+        if qk_norm:
+            self.q_norm = common_spec.LayerNormSpec(rms_norm=qk_norm_rms)
+            self.k_norm = common_spec.LayerNormSpec(rms_norm=qk_norm_rms)
+
         if relative_position:
             self.relative_position_keys = None
             self.relative_position_values = None

ctranslate2/specs/transformer_spec.py

@@ -23,6 +23,16 @@ class TransformerEncoderSpec(model_spec.LayerSpec):
         ffn_glu: bool = False,
         rms_norm: bool = False,
         multi_query_attention: bool = False,
+        num_heads_kv: Optional[int] = None,
+        head_dim: Optional[int] = None,
+        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,
+        sliding_window: Optional[int] = None,
+        qk_norm: Optional[bool] = False,
+        pre_post_layer_norm: bool = False,
     ):
         """Initializes a Transformer encoder specification.
 
@@ -43,8 +53,28 @@ class TransformerEncoderSpec(model_spec.LayerSpec):
           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.
+          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.
+          head_dim: Number of dimensions per attention head.
+          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.
+          sliding_window: Max sequence length to retain in KV Cache.
+          qk_norm: Apply layer normalization to the query and key projections.
+          pre_post_layer_norm: Add post layer norm for each pre norm layer.
         """
+
+        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
+
         self.multi_query_attention = multi_query_attention
         self.num_heads = np.dtype("int16").type(num_heads)
         self.pre_norm = pre_norm
@@ -60,13 +90,24 @@ class TransformerEncoderSpec(model_spec.LayerSpec):
             self.layer_norm = common_spec.LayerNormSpec(rms_norm=rms_norm)
         if layernorm_embedding:
             self.layernorm_embedding = common_spec.LayerNormSpec(rms_norm=rms_norm)
+        if sliding_window is not None:
+            self.sliding_window = np.dtype("int32").type(sliding_window)
+
         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,
+                num_heads_kv=num_heads_kv,
+                head_dim=head_dim,
+                rotary_dim=rotary_dim,
+                rotary_interleave=rotary_interleave,
+                rotary_scaling_type=rotary_scaling_type,
+                rotary_scaling_factor=rotary_scaling_factor,
+                rotary_base=rotary_base,
+                qk_norm=qk_norm,
+                pre_post_layer_norm=pre_post_layer_norm,
             )
             for _ in range(num_layers)
         ]
@@ -109,6 +150,8 @@ class TransformerDecoderSpec(model_spec.LayerSpec):
         quant_type: Optional[common_spec.Quantization] = None,
         quant_group_size: Optional[int] = None,
         quant_bits: Optional[int] = None,
+        qk_norm: bool = False,
+        external_pre_post_encoder_layers: Optional[bool] = False,
     ):
         """Initializes a Transformer decoder specification.
 
@@ -155,6 +198,8 @@ class TransformerDecoderSpec(model_spec.LayerSpec):
           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)
+          external_pre_post_encoder_layers: if the encoder attention pre and processing
+            is done outside the attention.
         """
 
         self._config = dict()
@@ -171,12 +216,6 @@ class TransformerDecoderSpec(model_spec.LayerSpec):
                 )
             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)
@@ -222,6 +261,8 @@ class TransformerDecoderSpec(model_spec.LayerSpec):
                 num_heads_kv=num_heads_kv,
                 head_dim=head_dim,
                 sliding_window=sliding_window,
+                qk_norm=qk_norm,
+                external_pre_post_encoder_layers=external_pre_post_encoder_layers,
             )
             for _ in range(num_layers)
         ]
@@ -252,7 +293,15 @@ class TransformerEncoderLayerSpec(model_spec.LayerSpec):
         ffn_glu=False,
         rms_norm=False,
         num_heads_kv=None,
+        head_dim=None,
         sliding_window=None,
+        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,
+        qk_norm=False,
+        pre_post_layer_norm: bool = False,
     ):
         self.self_attention = attention_spec.MultiHeadAttentionSpec(
             self_attention=True,
@@ -260,10 +309,32 @@ class TransformerEncoderLayerSpec(model_spec.LayerSpec):
             relative_attention_bias=relative_attention_bias,
             rms_norm=rms_norm,
             num_heads_kv=num_heads_kv,
+            head_dim=head_dim,
             sliding_window=sliding_window,
+            rotary_dim=rotary_dim,
+            rotary_interleave=rotary_interleave,
+            rotary_scaling_type=rotary_scaling_type,
+            rotary_scaling_factor=rotary_scaling_factor,
+            rotary_base=rotary_base,
+            qk_norm=qk_norm,
         )
         self.ffn = FeedForwardSpec(glu=ffn_glu, rms_norm=rms_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 TransformerDecoderLayerSpec(model_spec.LayerSpec):
     def __init__(
@@ -286,6 +357,8 @@ class TransformerDecoderLayerSpec(model_spec.LayerSpec):
         num_heads_kv=None,
         head_dim=None,
         sliding_window=None,
+        qk_norm=False,
+        external_pre_post_encoder_layers=False,
     ):
         self.self_attention = attention_spec.MultiHeadAttentionSpec(
             self_attention=True,
@@ -302,13 +375,17 @@ class TransformerDecoderLayerSpec(model_spec.LayerSpec):
             num_heads_kv=num_heads_kv,
             head_dim=head_dim,
             sliding_window=sliding_window,
+            qk_norm=qk_norm,
         )
 
         if with_encoder_attention:
             self.attention = attention_spec.MultiHeadAttentionSpec(
                 rms_norm=rms_norm,
                 num_heads_kv=num_heads_kv,
+                head_dim=head_dim,
                 sliding_window=sliding_window,
+                qk_norm=qk_norm,
+                has_norm=external_pre_post_encoder_layers is False,
             )
 
         self.ffn = FeedForwardSpec(glu=ffn_glu, rms_norm=rms_norm)
@@ -324,10 +401,21 @@ class TransformerDecoderLayerSpec(model_spec.LayerSpec):
             delattr(self.ffn, "layer_norm")
 
         if pre_post_layer_norm:
+            # Self-attention layer norms
             self.input_layer_norm = common_spec.LayerNormSpec(rms_norm=rms_norm)
             self.post_attention_layer_norm = common_spec.LayerNormSpec(
                 rms_norm=rms_norm
             )
+
+            if with_encoder_attention and external_pre_post_encoder_layers:
+                self.external_post_encoder_attention_layer_norm = (
+                    common_spec.LayerNormSpec(rms_norm=rms_norm)
+                )
+                self.external_pre_encoder_attention_layer_norm = (
+                    common_spec.LayerNormSpec(rms_norm=rms_norm)
+                )
+
+            # Feed-forward layer norms
             self.pre_feedforward_layer_norm = common_spec.LayerNormSpec(
                 rms_norm=rms_norm
             )
@@ -557,6 +645,7 @@ class TransformerDecoderModelSpec(model_spec.LanguageModelSpec):
         quant_type: Optional[common_spec.Quantization] = None,
         quant_group_size: Optional[int] = None,
         quant_bits: Optional[int] = None,
+        qk_norm: bool = False,
     ):
         """Creates a Transformer decoder model specification.
 
@@ -631,6 +720,7 @@ class TransformerDecoderModelSpec(model_spec.LanguageModelSpec):
             quant_type=quant_type,
             quant_group_size=quant_group_size,
             quant_bits=quant_bits,
+            qk_norm=qk_norm,
         )
 
         return cls(decoder)

ctranslate2/version.py

@@ -1,3 +1,3 @@
 """Version information."""
 
-__version__ = "4.6.1"
+__version__ = "4.6.3"

{ctranslate2-4.6.1.dist-info → ctranslate2-4.6.3.dist-info}/METADATA

@@ -1,9 +1,10 @@
 Metadata-Version: 2.4
 Name: ctranslate2
-Version: 4.6.1
+Version: 4.6.3
 Summary: Fast inference engine for Transformer models
 Home-page: https://opennmt.net
 Author: OpenNMT
+License: MIT
 Project-URL: Documentation, https://opennmt.net/CTranslate2
 Project-URL: Forum, https://forum.opennmt.net
 Project-URL: Gitter, https://gitter.im/OpenNMT/CTranslate2
@@ -13,7 +14,6 @@ Classifier: Development Status :: 5 - Production/Stable
 Classifier: Environment :: GPU :: NVIDIA CUDA :: 12 :: 12.4
 Classifier: Intended Audience :: Developers
 Classifier: Intended Audience :: Science/Research
-Classifier: License :: OSI Approved :: MIT License
 Classifier: Programming Language :: Python :: 3
 Classifier: Programming Language :: Python :: 3 :: Only
 Classifier: Programming Language :: Python :: 3.9
@@ -34,6 +34,7 @@ Dynamic: description
 Dynamic: description-content-type
 Dynamic: home-page
 Dynamic: keywords
+Dynamic: license
 Dynamic: project-url
 Dynamic: requires-dist
 Dynamic: requires-python
@@ -49,7 +50,7 @@ The project implements a custom runtime that applies many performance optimizati
 
 The following model types are currently supported:
 
-* Encoder-decoder models: Transformer base/big, M2M-100, NLLB, BART, mBART, Pegasus, T5, Whisper
+* Encoder-decoder models: Transformer base/big, M2M-100, NLLB, BART, mBART, Pegasus, T5, Whisper T5Gemma
 * Decoder-only models: GPT-2, GPT-J, GPT-NeoX, OPT, BLOOM, MPT, Llama, Mistral, Gemma, CodeGen, GPTBigCode, Falcon, Qwen2
 * Encoder-only models: BERT, DistilBERT, XLM-RoBERTa
 
@@ -160,6 +161,16 @@ Executed with 4 threads on a [*c5.2xlarge*](https://aws.amazon.com/ec2/instance-
 
 Executed with CUDA 11 on a [*g5.xlarge*](https://aws.amazon.com/ec2/instance-types/g5/) Amazon EC2 instance equipped with a NVIDIA A10G GPU (driver version: 510.47.03).
 
+## Contributing
+
+CTranslate2 is a community-driven project. We welcome contributions of all kinds:
+* **New Model Support:** Help us implement more Transformer architectures.
+* **Performance:** Propose optimizations for CPU or GPU kernels.
+* **Bug Reports:** Open an issue if you find something not working as expected.
+* **Documentation:** Improve our guides or add new examples.
+
+Check out our [Contributing Guide](CONTRIBUTING.md) to learn how to set up your development environment.
+
 ## Additional resources
 
 * [Documentation](https://opennmt.net/CTranslate2)

{ctranslate2-4.6.1.dist-info → ctranslate2-4.6.3.dist-info}/RECORD

@@ -1,33 +1,33 @@
-ctranslate2/__init__.py,sha256=HmklwU3BCnUgQAAGQeCBOr44QS5X0npr_WDkBh2e6Lo,1508
-ctranslate2/_ext.cp314-win_amd64.pyd,sha256=gqC7yPHNxHKG_7GTH0K2Vajp8yjGvYWdjkZe5xjQqHg,719872
-ctranslate2/ctranslate2.dll,sha256=GJzPqItIFHs3sD9PwQm1yRHJyASjW3xszKIxUW3A0oc,58389504
-ctranslate2/cudnn64_9.dll,sha256=wHzEfy-kpWZZPHr0qn5X7fCamFoP3dFMuNb0VuJSrwU,438840
-ctranslate2/extensions.py,sha256=axO2FI8ddiFmlko2AzQ6VcdtF-3hDA7VmPGnTIkrPkI,21782
+ctranslate2/__init__.py,sha256=CGqShDaFxQ-u-aCtVq99T4HKuBdMB8b49l2KSxnQb8M,1735
+ctranslate2/_ext.cp314-win_amd64.pyd,sha256=7ls8m8Bo_fA3itDnwgbjTUvsVJa0IG8Z4D6eMDUh_9Q,720384
+ctranslate2/ctranslate2.dll,sha256=9TRGHvoyNSzXe9eEc3jKJa5-21-AeSENgp3DXvhCZ4M,58590720
+ctranslate2/cudnn64_9.dll,sha256=ntvN_3OwrwcOsWCyzmbln-ygSqAXNR2O7cxejhSZZ9I,266288
+ctranslate2/extensions.py,sha256=kDNt0H9KvfNCc3PrRGzfkj9Fkvna84i2O5Y-rav6UkU,21940
 ctranslate2/libiomp5md.dll,sha256=mCIzNmsK_NoeD1WgsTQJfjW3eWE_VN22nmhebNBrdV8,1614192
 ctranslate2/logging.py,sha256=P9evHdxuMx_iHvwJjEASEq-j5062H64Pl5-fJjxEuHk,1221
-ctranslate2/version.py,sha256=4YXv9jQt8K_nFjSqbVGYqjJBzqTZM2P6siU7xOzZDwY,53
+ctranslate2/version.py,sha256=TboXlbA67GNmSOm1v2u_U8AKgYh5iminMMLTvi3Xho4,53
 ctranslate2/converters/__init__.py,sha256=ufYjcXf2sK4fiXAUU6tIJyWmNuLjKFf_KH3GWLXe4ls,507
 ctranslate2/converters/converter.py,sha256=Qkb8NGLLmgqMT6HZkFq61zwbxyq3NlWcaxLZ6Ap-YOQ,3601
 ctranslate2/converters/eole_ct2.py,sha256=RUcDJH_2AUt0jDs5oAqccE6tQPbO9LQ6JmVriC1DTy8,12564
-ctranslate2/converters/fairseq.py,sha256=uQpd-ftYSO4c6WdEwCUyuZWhzWX1UTG7dGOC6EtcDVE,12765
+ctranslate2/converters/fairseq.py,sha256=2vlBk4AVCHwXxKkwPHVmcjyfo1dAV0_DJS1i6q-44NE,12822
 ctranslate2/converters/marian.py,sha256=1_7P3EbIDPOdyJbtb_Lp-LCBPBb9A8E9OhzoyFwTb64,11274
 ctranslate2/converters/openai_gpt2.py,sha256=1rXKM2ZURZHWRv4XZ135fPkVWpM4rTG-q7VR7OD6d-A,3304
-ctranslate2/converters/opennmt_py.py,sha256=Vva60az6tGqlQXs0UgC09r_fCD3u2u6wUJB-8V4OUFQ,13183
+ctranslate2/converters/opennmt_py.py,sha256=zex4TbHiiJMy0tkqQg39oNjxmSZKf8dnRLH3iQ1H4z0,13227
 ctranslate2/converters/opennmt_tf.py,sha256=uBRp2wz5xriSQcA_c0S0ekY7ws6RpRX_0EKeMRdM7-s,16222
 ctranslate2/converters/opus_mt.py,sha256=5KbPaTiBhhorPzMpTugIfIJ8SgcqHfJUbJrWKBN-Djs,1254
-ctranslate2/converters/transformers.py,sha256=Wzih7qmqNpen_EdxKaAoEWLMV4RR16GVsrwrbB9lj8A,114233
+ctranslate2/converters/transformers.py,sha256=VRal3vKSQrAOvcNPwewjVMtgvWskz0KD5bdIrpNrZNA,142380
 ctranslate2/converters/utils.py,sha256=w7NG39lx-9dOdL57OqKVTdC__opkuP8RACg1TLlUJwM,3817
 ctranslate2/models/__init__.py,sha256=53p98uemtuvVPz8xK7_LbOhBiUJJu-c-NdmOHJgdXus,497
 ctranslate2/specs/__init__.py,sha256=9GabtSyczznYqiqUS6XvULi8pQ3_3RNRogXobGP0G80,653
-ctranslate2/specs/attention_spec.py,sha256=ios3aZRWbZ8PmcYi9pXIad52lMweqOUgV5ZJbkFOKmE,3218
+ctranslate2/specs/attention_spec.py,sha256=FnaSiQREWQw_cURgsCb9_aIpGOCxyVGTCpIOdd-08v8,3492
 ctranslate2/specs/common_spec.py,sha256=freTDhQMy5PYofBrij4_FDgrKokMYApWSPIpASZIlJc,1608
 ctranslate2/specs/model_spec.py,sha256=atCAYzDEIzyJ1TCayFGZVutHqSWa1ww-vbZ0OiIJqh8,25736
-ctranslate2/specs/transformer_spec.py,sha256=vBTnBaZ8nslREF8FSJKP6VveyYAvS2_L0h8yqPAhpY0,30124
+ctranslate2/specs/transformer_spec.py,sha256=-GJ0oSjI3ns-Ei_-xXIM_P2GaZxt5Z-g03zJ0m_4ciU,34317
 ctranslate2/specs/wav2vec2_spec.py,sha256=NITsuOuf2F5bU1-aXit8-WEtWV9fH2Eq7A7857UyYho,2106
 ctranslate2/specs/wav2vec2bert_spec.py,sha256=UgtsJWC9mMgJ7bn4T_xg1uXK0rqA4-9tT2KMGVgPKnw,3529
 ctranslate2/specs/whisper_spec.py,sha256=_vm1sc5yOowOJ4iyvcxMXrgt-UcLJrZT8OtPscUXcQQ,2447
-ctranslate2-4.6.1.dist-info/METADATA,sha256=80HyCFtLjBdWxYcQtHLdyoLJwYOvGoRPk3r_arkWMsg,10354
-ctranslate2-4.6.1.dist-info/WHEEL,sha256=7k6Wcy588iJYe5lf5K095NLg-uoBTnE-T8eHJ92G4_4,101
-ctranslate2-4.6.1.dist-info/entry_points.txt,sha256=ZHkojut_TmVRHl0bJIGm2b9wqr98GAJqxN9rlJtQshs,466
-ctranslate2-4.6.1.dist-info/top_level.txt,sha256=1hUaWzcFIuSo2BAIUHFA3Osgsu6S1giq0y6Rosv8HOQ,12
-ctranslate2-4.6.1.dist-info/RECORD,,
+ctranslate2-4.6.3.dist-info/METADATA,sha256=awoc6t4JSxpv51lmfAG28ZG91FhGQ8DHspyLzLqLo_Q,10839
+ctranslate2-4.6.3.dist-info/WHEEL,sha256=7k6Wcy588iJYe5lf5K095NLg-uoBTnE-T8eHJ92G4_4,101
+ctranslate2-4.6.3.dist-info/entry_points.txt,sha256=ZHkojut_TmVRHl0bJIGm2b9wqr98GAJqxN9rlJtQshs,466
+ctranslate2-4.6.3.dist-info/top_level.txt,sha256=1hUaWzcFIuSo2BAIUHFA3Osgsu6S1giq0y6Rosv8HOQ,12
+ctranslate2-4.6.3.dist-info/RECORD,,