ctranslate2 4.7.0__cp314-cp314-macosx_11_0_arm64.whl

ctranslate2/converters/utils.py

@@ -0,0 +1,127 @@
+import numpy as np
+
+
+def fuse_linear(spec, layers):
+    if not layers:
+        raise ValueError("Cannot fuse linear layers: at least one layer is required")
+
+    if isinstance(layers[0].weight, np.ndarray):
+        concatenate = np.concatenate
+        zeros = np.zeros
+    else:
+        import torch
+
+        concatenate = torch.cat
+        zeros = torch.zeros
+
+    spec.weight = concatenate([layer.weight for layer in layers])
+
+    bias_dtype = None
+    for layer in layers:
+        if layer.has_bias():
+            bias_dtype = layer.bias.dtype
+            break
+
+    if bias_dtype is not None:
+        spec.bias = concatenate(
+            [
+                (
+                    layer.bias
+                    if layer.has_bias()
+                    else zeros([layer.weight.shape[0]], dtype=bias_dtype)
+                )
+                for layer in layers
+            ]
+        )
+
+
+def fuse_linear_prequant(spec, layers, axis):
+    if not layers:
+        raise ValueError("Cannot fuse linear layers: at least one layer is required")
+    params = ["weight", "weight_scale", "weight_zero"]
+    if isinstance(layers[0].weight, np.ndarray):
+        concatenate = np.concatenate
+    else:
+        import torch
+
+        concatenate = torch.cat
+
+    for param in params:
+        setattr(
+            spec,
+            param,
+            concatenate([getattr(layer, param) for layer in layers], axis=axis),
+        )
+
+
+def permute_for_sliced_rotary(weight, num_heads, rotary_dim=None):
+    """Permutes the weight to use the sliced rotary implementation."""
+    if rotary_dim is not None:
+        weight = weight.reshape(num_heads, weight.shape[0] // num_heads, -1)
+
+        rotary_weight = weight[:, :rotary_dim]
+        rotary_weight = permute_for_sliced_rotary(
+            rotary_weight.reshape(num_heads * rotary_dim, -1), num_heads
+        ).reshape(num_heads, rotary_dim, -1)
+
+        weight[:, :rotary_dim] = rotary_weight
+
+        return weight.reshape(-1, weight.shape[-1])
+
+    return (
+        weight.reshape(num_heads, weight.shape[0] // num_heads // 2, 2, weight.shape[1])
+        .swapaxes(1, 2)
+        .reshape(weight.shape[0], weight.shape[1])
+    )
+
+
+def smooth_activation(layer_norm, linear, activation_scales):
+    """Applies the activation smoothing technique described in
+    https://github.com/mit-han-lab/smoothquant.
+    """
+    if not isinstance(linear.weight, np.ndarray):
+        linear_weight = linear.weight.numpy()
+        activation_scales = activation_scales.numpy()
+    else:
+        linear_weight = linear.weight
+
+    weight_scales = np.amax(np.absolute(linear_weight), axis=0)
+    weight_scales = np.maximum(weight_scales, 1e-5)
+
+    activation_scales = activation_scales.astype(weight_scales.dtype)
+
+    scales = np.sqrt(activation_scales / weight_scales)
+    scales = np.maximum(scales, 1e-5)
+
+    if not isinstance(linear.weight, np.ndarray):
+        import torch
+
+        scales = torch.from_numpy(scales)
+
+    layer_norm.gamma /= scales
+    layer_norm.beta /= scales
+
+    linear.weight *= scales.reshape(1, -1)
+
+
+def raise_unsupported(reasons):
+    message = (
+        "The model you are trying to convert is not supported by CTranslate2. "
+        "We identified the following reasons:\n"
+    )
+    for reason in reasons:
+        message += "\n- " + reason
+    raise ValueError(message)
+
+
+class ConfigurationChecker:
+    def __init__(self):
+        self._unsupported_reasons = []
+
+    def __call__(self, assert_condition, error_message):
+        if not assert_condition:
+            self._unsupported_reasons.append(error_message)
+
+    def validate(self):
+        if self._unsupported_reasons:
+            raise_unsupported(self._unsupported_reasons)