Commit
•
1d2ae3c
1
Parent(s):
a55fa23
fix deps
Browse files- requirements.txt +5 -4
- whisper_jax/__init__.py +21 -0
- whisper_jax/layers.py +1310 -0
- whisper_jax/modeling_flax_whisper.py +1686 -0
- whisper_jax/partitioner.py +939 -0
- whisper_jax/pipeline.py +506 -0
- whisper_jax/train_state.py +130 -0
requirements.txt
CHANGED
@@ -1,7 +1,8 @@
|
|
1 |
--find-links https://storage.googleapis.com/jax-releases/libtpu_releases.html
|
2 |
jax[tpu]
|
3 |
-
git+https://github.com/sanchit-gandhi/whisper-jax.git
|
4 |
-
requests
|
5 |
-
yt-dlp>=2023.3.4
|
6 |
torch
|
7 |
-
transformers>=4.40.0
|
|
|
|
|
|
|
|
|
|
1 |
--find-links https://storage.googleapis.com/jax-releases/libtpu_releases.html
|
2 |
jax[tpu]
|
|
|
|
|
|
|
3 |
torch
|
4 |
+
transformers>=4.40.0
|
5 |
+
flax
|
6 |
+
cached-property
|
7 |
+
requests
|
8 |
+
yt-dlp>=2023.3.4
|
whisper_jax/__init__.py
ADDED
@@ -0,0 +1,21 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# coding=utf-8
|
2 |
+
# Copyright 2023 The HuggingFace Inc. team.
|
3 |
+
#
|
4 |
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
5 |
+
# you may not use this file except in compliance with the License.
|
6 |
+
# You may obtain a copy of the License at
|
7 |
+
#
|
8 |
+
# http://www.apache.org/licenses/LICENSE-2.0
|
9 |
+
#
|
10 |
+
# Unless required by applicable law or agreed to in writing, software
|
11 |
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
12 |
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
13 |
+
# See the License for the specific language governing permissions and
|
14 |
+
# limitations under the License.
|
15 |
+
|
16 |
+
__version__ = "0.0.1"
|
17 |
+
|
18 |
+
from .modeling_flax_whisper import FlaxWhisperForConditionalGeneration
|
19 |
+
from .partitioner import PjitPartitioner
|
20 |
+
from .pipeline import FlaxWhisperPipline
|
21 |
+
from .train_state import InferenceState
|
whisper_jax/layers.py
ADDED
@@ -0,0 +1,1310 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# coding=utf-8
|
2 |
+
# Copyright 2023 The T5X Authors and The HuggingFace Inc. team.
|
3 |
+
#
|
4 |
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
5 |
+
# you may not use this file except in compliance with the License.
|
6 |
+
# You may obtain a copy of the License at
|
7 |
+
#
|
8 |
+
# http://www.apache.org/licenses/LICENSE-2.0
|
9 |
+
#
|
10 |
+
# Unless required by applicable law or agreed to in writing, software
|
11 |
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
12 |
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
13 |
+
# See the License for the specific language governing permissions and
|
14 |
+
# limitations under the License.
|
15 |
+
|
16 |
+
"""Dense attention classes and mask/weighting functions."""
|
17 |
+
|
18 |
+
# pylint: disable=attribute-defined-outside-init,g-bare-generic
|
19 |
+
|
20 |
+
import dataclasses
|
21 |
+
import functools
|
22 |
+
import operator
|
23 |
+
from typing import Any, Callable, Iterable, List, Optional, Sequence, Tuple, Union
|
24 |
+
|
25 |
+
import jax
|
26 |
+
import jax.numpy as jnp
|
27 |
+
import numpy as np
|
28 |
+
from flax import linen as nn
|
29 |
+
from flax.linen import partitioning as nn_partitioning
|
30 |
+
from flax.linen.dtypes import promote_dtype
|
31 |
+
from jax import lax, random
|
32 |
+
|
33 |
+
|
34 |
+
# from flax.linen.partitioning import param_with_axes, with_sharding_constraint
|
35 |
+
param_with_axes = nn_partitioning.param_with_axes
|
36 |
+
with_sharding_constraint = nn_partitioning.with_sharding_constraint
|
37 |
+
|
38 |
+
|
39 |
+
# Type annotations
|
40 |
+
Array = jnp.ndarray
|
41 |
+
DType = jnp.dtype
|
42 |
+
PRNGKey = jnp.ndarray
|
43 |
+
Shape = Iterable[int]
|
44 |
+
Activation = Callable[..., Array]
|
45 |
+
PrecisionLike = Union[None, str, lax.Precision, Tuple[str, str], Tuple[lax.Precision, lax.Precision]]
|
46 |
+
DotGeneralT = Callable[..., Array]
|
47 |
+
ConvGeneralDilatedT = Callable[..., Array]
|
48 |
+
PaddingLike = Union[str, int, Sequence[Union[int, Tuple[int, int]]]]
|
49 |
+
LaxPadding = Union[str, Sequence[Tuple[int, int]]]
|
50 |
+
|
51 |
+
# Parameter initializers.
|
52 |
+
Initializer = Callable[[PRNGKey, Shape, DType], Array]
|
53 |
+
InitializerAxis = Union[int, Tuple[int, ...]]
|
54 |
+
NdInitializer = Callable[[PRNGKey, Shape, DType, InitializerAxis, InitializerAxis], Array]
|
55 |
+
|
56 |
+
default_embed_init = nn.initializers.variance_scaling(1.0, "fan_in", "normal", out_axis=0)
|
57 |
+
|
58 |
+
|
59 |
+
# ------------------------------------------------------------------------------
|
60 |
+
# Temporary inlined JAX N-d initializer code
|
61 |
+
# TODO(levskaya): remove once new JAX release is out.
|
62 |
+
# ------------------------------------------------------------------------------
|
63 |
+
def _compute_fans(shape: jax.core.NamedShape, in_axis=-2, out_axis=-1):
|
64 |
+
"""Inlined JAX `nn.initializer._compute_fans`."""
|
65 |
+
if isinstance(in_axis, int):
|
66 |
+
in_size = shape[in_axis]
|
67 |
+
else:
|
68 |
+
in_size = int(np.prod([shape[i] for i in in_axis]))
|
69 |
+
if isinstance(out_axis, int):
|
70 |
+
out_size = shape[out_axis]
|
71 |
+
else:
|
72 |
+
out_size = int(np.prod([shape[i] for i in out_axis]))
|
73 |
+
receptive_field_size = shape.total / in_size / out_size
|
74 |
+
fan_in = in_size * receptive_field_size
|
75 |
+
fan_out = out_size * receptive_field_size
|
76 |
+
return fan_in, fan_out
|
77 |
+
|
78 |
+
|
79 |
+
def variance_scaling(scale, mode, distribution, in_axis=-2, out_axis=-1, dtype=jnp.float_):
|
80 |
+
"""Inlined JAX `nn.initializer.variance_scaling`."""
|
81 |
+
|
82 |
+
def init(key, shape, dtype=dtype):
|
83 |
+
return jnp.zeros(shape, dtype=dtype)
|
84 |
+
dtype = jax.dtypes.canonicalize_dtype(dtype)
|
85 |
+
shape = jax.core.as_named_shape(shape)
|
86 |
+
fan_in, fan_out = _compute_fans(shape, in_axis, out_axis)
|
87 |
+
if mode == "fan_in":
|
88 |
+
denominator = fan_in
|
89 |
+
elif mode == "fan_out":
|
90 |
+
denominator = fan_out
|
91 |
+
elif mode == "fan_avg":
|
92 |
+
denominator = (fan_in + fan_out) / 2
|
93 |
+
else:
|
94 |
+
raise ValueError("invalid mode for variance scaling initializer: {}".format(mode))
|
95 |
+
variance = jnp.array(scale / denominator, dtype=dtype)
|
96 |
+
|
97 |
+
if distribution == "truncated_normal":
|
98 |
+
# constant is stddev of standard normal truncated to (-2, 2)
|
99 |
+
stddev = jnp.sqrt(variance) / jnp.array(0.87962566103423978, dtype)
|
100 |
+
return random.truncated_normal(key, -2, 2, shape, dtype) * stddev
|
101 |
+
elif distribution == "normal":
|
102 |
+
return random.normal(key, shape, dtype) * jnp.sqrt(variance)
|
103 |
+
elif distribution == "uniform":
|
104 |
+
return random.uniform(key, shape, dtype, -1) * jnp.sqrt(3 * variance)
|
105 |
+
else:
|
106 |
+
raise ValueError("invalid distribution for variance scaling " "initializer: {}".format(distribution))
|
107 |
+
|
108 |
+
return init
|
109 |
+
|
110 |
+
|
111 |
+
# ------------------------------------------------------------------------------
|
112 |
+
|
113 |
+
|
114 |
+
def nd_dense_init(scale, mode, distribution):
|
115 |
+
"""Initializer with in_axis, out_axis set at call time."""
|
116 |
+
|
117 |
+
def init_fn(key, shape, dtype, in_axis, out_axis):
|
118 |
+
fn = variance_scaling(scale, mode, distribution, in_axis, out_axis)
|
119 |
+
return fn(key, shape, dtype)
|
120 |
+
|
121 |
+
return init_fn
|
122 |
+
|
123 |
+
|
124 |
+
def dot_product_attention(
|
125 |
+
query: Array,
|
126 |
+
key: Array,
|
127 |
+
value: Array,
|
128 |
+
bias: Optional[Array] = None,
|
129 |
+
dropout_rng: Optional[PRNGKey] = None,
|
130 |
+
dropout_rate: float = 0.0,
|
131 |
+
deterministic: bool = False,
|
132 |
+
dtype: DType = jnp.float32,
|
133 |
+
float32_logits: bool = False,
|
134 |
+
):
|
135 |
+
"""Computes dot-product attention given query, key, and value.
|
136 |
+
|
137 |
+
This is the core function for applying attention based on
|
138 |
+
https://arxiv.org/abs/1706.03762. It calculates the attention weights given
|
139 |
+
query and key and combines the values using the attention weights.
|
140 |
+
|
141 |
+
Args:
|
142 |
+
query: queries for calculating attention with shape of `[batch, q_length,
|
143 |
+
num_heads, qk_depth_per_head]`.
|
144 |
+
key: keys for calculating attention with shape of `[batch, kv_length,
|
145 |
+
num_heads, qk_depth_per_head]`.
|
146 |
+
value: values to be used in attention with shape of `[batch, kv_length,
|
147 |
+
num_heads, v_depth_per_head]`.
|
148 |
+
bias: bias for the attention weights. This should be broadcastable to the
|
149 |
+
shape `[batch, num_heads, q_length, kv_length]` This can be used for
|
150 |
+
incorporating causal masks, padding masks, proximity bias, etc.
|
151 |
+
dropout_rng: JAX PRNGKey: to be used for dropout
|
152 |
+
dropout_rate: dropout rate
|
153 |
+
deterministic: bool, deterministic or not (to apply dropout)
|
154 |
+
dtype: the dtype of the computation (default: float32)
|
155 |
+
float32_logits: bool, if True then compute logits in float32 to avoid
|
156 |
+
numerical issues with bfloat16.
|
157 |
+
|
158 |
+
Returns:
|
159 |
+
Output of shape `[batch, length, num_heads, v_depth_per_head]`.
|
160 |
+
"""
|
161 |
+
assert key.ndim == query.ndim == value.ndim, "q, k, v must have same rank."
|
162 |
+
assert query.shape[:-3] == key.shape[:-3] == value.shape[:-3], "q, k, v batch dims must match."
|
163 |
+
assert query.shape[-2] == key.shape[-2] == value.shape[-2], "q, k, v num_heads must match."
|
164 |
+
assert key.shape[-3] == value.shape[-3], "k, v lengths must match."
|
165 |
+
assert query.shape[-1] == key.shape[-1], "q, k depths must match."
|
166 |
+
|
167 |
+
# Casting logits and softmax computation for float32 for model stability.
|
168 |
+
if float32_logits:
|
169 |
+
query = query.astype(jnp.float32)
|
170 |
+
key = key.astype(jnp.float32)
|
171 |
+
|
172 |
+
# `attn_weights`: [batch, num_heads, q_length, kv_length]
|
173 |
+
attn_weights = jnp.einsum("bqhd,bkhd->bhqk", query, key)
|
174 |
+
|
175 |
+
# Apply attention bias: masking, dropout, proximity bias, etc.
|
176 |
+
if bias is not None:
|
177 |
+
attn_weights = attn_weights + bias.astype(attn_weights.dtype)
|
178 |
+
|
179 |
+
# Normalize the attention weights across `kv_length` dimension.
|
180 |
+
attn_weights = jax.nn.softmax(attn_weights).astype(dtype)
|
181 |
+
|
182 |
+
# Apply attention dropout.
|
183 |
+
if not deterministic and dropout_rate > 0.0:
|
184 |
+
keep_prob = 1.0 - dropout_rate
|
185 |
+
# T5 broadcasts along the "length" dim, but unclear which one that
|
186 |
+
# corresponds to in positional dimensions here, assuming query dim.
|
187 |
+
dropout_shape = list(attn_weights.shape)
|
188 |
+
dropout_shape[-2] = 1
|
189 |
+
keep = random.bernoulli(dropout_rng, keep_prob, dropout_shape)
|
190 |
+
keep = jnp.broadcast_to(keep, attn_weights.shape)
|
191 |
+
multiplier = keep.astype(attn_weights.dtype) / jnp.asarray(keep_prob, dtype=dtype)
|
192 |
+
attn_weights = attn_weights * multiplier
|
193 |
+
|
194 |
+
# Take the linear combination of `value`.
|
195 |
+
return jnp.einsum("bhqk,bkhd->bqhd", attn_weights, value)
|
196 |
+
|
197 |
+
|
198 |
+
dynamic_vector_slice_in_dim = jax.vmap(lax.dynamic_slice_in_dim, in_axes=(None, 0, None, None))
|
199 |
+
|
200 |
+
|
201 |
+
class MultiHeadDotProductAttention(nn.Module):
|
202 |
+
"""Multi-head dot-product attention.
|
203 |
+
|
204 |
+
Attributes:
|
205 |
+
num_heads: number of attention heads. Features (i.e. inputs_q.shape[-1])
|
206 |
+
should be divisible by the number of heads.
|
207 |
+
head_dim: dimension of each head.
|
208 |
+
dtype: the dtype of the computation.
|
209 |
+
dropout_rate: dropout rate
|
210 |
+
kernel_init: initializer for the kernel of the Dense layers.
|
211 |
+
float32_logits: bool, if True then compute logits in float32 to avoid
|
212 |
+
numerical issues with bfloat16.
|
213 |
+
"""
|
214 |
+
|
215 |
+
num_heads: int
|
216 |
+
head_dim: int
|
217 |
+
dtype: DType = jnp.float32
|
218 |
+
dropout_rate: float = 0.0
|
219 |
+
kernel_init: NdInitializer = nd_dense_init(1.0, "fan_in", "normal")
|
220 |
+
float32_logits: bool = False # computes logits in float32 for stability.
|
221 |
+
|
222 |
+
@nn.compact
|
223 |
+
def __call__(
|
224 |
+
self,
|
225 |
+
inputs_q: Array,
|
226 |
+
inputs_kv: Array,
|
227 |
+
mask: Optional[Array] = None,
|
228 |
+
bias: Optional[Array] = None,
|
229 |
+
*,
|
230 |
+
decode: bool = False,
|
231 |
+
deterministic: bool = False,
|
232 |
+
) -> Array:
|
233 |
+
"""Applies multi-head dot product attention on the input data.
|
234 |
+
|
235 |
+
Projects the inputs into multi-headed query, key, and value vectors,
|
236 |
+
applies dot-product attention and project the results to an output vector.
|
237 |
+
|
238 |
+
There are two modes: decoding and non-decoding (e.g., training). The mode is
|
239 |
+
determined by `decode` argument. For decoding, this method is called twice,
|
240 |
+
first to initialize the cache and then for an actual decoding process. The
|
241 |
+
two calls are differentiated by the presence of 'cached_key' in the variable
|
242 |
+
dict. In the cache initialization stage, the cache variables are initialized
|
243 |
+
as zeros and will be filled in the subsequent decoding process.
|
244 |
+
|
245 |
+
In the cache initialization call, `inputs_q` has a shape [batch, length,
|
246 |
+
q_features] and `inputs_kv`: [batch, length, kv_features]. During the
|
247 |
+
incremental decoding stage, query, key and value all have the shape [batch,
|
248 |
+
1, qkv_features] corresponding to a single step.
|
249 |
+
|
250 |
+
Args:
|
251 |
+
inputs_q: input queries of shape `[batch, q_length, q_features]`.
|
252 |
+
inputs_kv: key/values of shape `[batch, kv_length, kv_features]`.
|
253 |
+
mask: attention mask of shape `[batch, num_heads, q_length, kv_length]`.
|
254 |
+
bias: attention bias of shape `[batch, num_heads, q_length, kv_length]`.
|
255 |
+
decode: Whether to prepare and use an autoregressive cache.
|
256 |
+
deterministic: Disables dropout if set to True.
|
257 |
+
|
258 |
+
Returns:
|
259 |
+
output of shape `[batch, length, q_features]`.
|
260 |
+
"""
|
261 |
+
projection = functools.partial(
|
262 |
+
DenseGeneral,
|
263 |
+
axis=-1,
|
264 |
+
features=(self.num_heads, self.head_dim),
|
265 |
+
kernel_axes=("embed", "heads", "kv"),
|
266 |
+
dtype=self.dtype,
|
267 |
+
)
|
268 |
+
|
269 |
+
# NOTE: T5 does not explicitly rescale the attention logits by
|
270 |
+
# 1/sqrt(depth_kq)! This is folded into the initializers of the
|
271 |
+
# linear transformations, which is equivalent under Adafactor.
|
272 |
+
depth_scaling = jnp.sqrt(self.head_dim).astype(self.dtype)
|
273 |
+
|
274 |
+
def query_init(*args):
|
275 |
+
return self.kernel_init(*args) / depth_scaling
|
276 |
+
|
277 |
+
# Project inputs_q to multi-headed q/k/v
|
278 |
+
# dimensions are then [batch, length, num_heads, head_dim]
|
279 |
+
query = projection(kernel_init=query_init, name="query")(inputs_q)
|
280 |
+
key = projection(kernel_init=self.kernel_init, name="key")(inputs_kv)
|
281 |
+
value = projection(kernel_init=self.kernel_init, name="value")(inputs_kv)
|
282 |
+
|
283 |
+
query = with_sharding_constraint(query, ("batch", "length", "heads", "kv"))
|
284 |
+
key = with_sharding_constraint(key, ("batch", "length", "heads", "kv"))
|
285 |
+
value = with_sharding_constraint(value, ("batch", "length", "heads", "kv"))
|
286 |
+
|
287 |
+
if decode:
|
288 |
+
# Detect if we're initializing by absence of existing cache data.
|
289 |
+
is_initialized = self.has_variable("cache", "cached_key")
|
290 |
+
|
291 |
+
# The key and value have dimension [batch, length, num_heads, head_dim],
|
292 |
+
# but we cache them as [batch, num_heads, head_dim, length] as a TPU
|
293 |
+
# fusion optimization. This also enables the "scatter via one-hot
|
294 |
+
# broadcast" trick, which means we do a one-hot broadcast instead of a
|
295 |
+
# scatter/gather operations, resulting in a 3-4x speedup in practice.
|
296 |
+
def swap_dims(x):
|
297 |
+
return x[:-3] + tuple(x[i] for i in [-2, -1, -3])
|
298 |
+
|
299 |
+
cached_key = self.variable("cache", "cached_key", jnp.zeros, swap_dims(key.shape), key.dtype)
|
300 |
+
cached_value = self.variable("cache", "cached_value", jnp.zeros, swap_dims(value.shape), value.dtype)
|
301 |
+
cache_index = self.variable("cache", "cache_index", lambda: jnp.array(0, dtype=jnp.int32))
|
302 |
+
if is_initialized:
|
303 |
+
batch, num_heads, head_dim, length = cached_key.value.shape
|
304 |
+
# During fast autoregressive decoding, we feed one position at a time,
|
305 |
+
# and cache the keys and values step by step.
|
306 |
+
# Sanity shape check of cached key against input query.
|
307 |
+
expected_shape = (batch, 1, num_heads, head_dim)
|
308 |
+
if expected_shape != query.shape:
|
309 |
+
raise ValueError(
|
310 |
+
"Autoregressive cache shape error, "
|
311 |
+
"expected query shape %s instead got %s." % (expected_shape, query.shape)
|
312 |
+
)
|
313 |
+
|
314 |
+
# Create a OHE of the current index. NOTE: the index is increased below.
|
315 |
+
cur_index = cache_index.value
|
316 |
+
one_hot_indices = jax.nn.one_hot(cur_index, length, dtype=key.dtype)
|
317 |
+
# In order to update the key, value caches with the current key and
|
318 |
+
# value, we move the length axis to the back, similar to what we did for
|
319 |
+
# the cached ones above.
|
320 |
+
# Note these are currently the key and value of a single position, since
|
321 |
+
# we feed one position at a time.
|
322 |
+
one_token_key = jnp.moveaxis(key, -3, -1)
|
323 |
+
one_token_value = jnp.moveaxis(value, -3, -1)
|
324 |
+
# Update key, value caches with our new 1d spatial slices.
|
325 |
+
# We implement an efficient scatter into the cache via one-hot
|
326 |
+
# broadcast and addition.
|
327 |
+
key = cached_key.value + one_token_key * one_hot_indices
|
328 |
+
value = cached_value.value + one_token_value * one_hot_indices
|
329 |
+
cached_key.value = key
|
330 |
+
cached_value.value = value
|
331 |
+
cache_index.value = cache_index.value + 1
|
332 |
+
# Move the keys and values back to their original shapes.
|
333 |
+
key = jnp.moveaxis(key, -1, -3)
|
334 |
+
value = jnp.moveaxis(value, -1, -3)
|
335 |
+
|
336 |
+
# Causal mask for cached decoder self-attention: our single query
|
337 |
+
# position should only attend to those key positions that have already
|
338 |
+
# been generated and cached, not the remaining zero elements.
|
339 |
+
mask = combine_masks(
|
340 |
+
mask,
|
341 |
+
jnp.broadcast_to(
|
342 |
+
jnp.arange(length) <= cur_index,
|
343 |
+
# (1, 1, length) represent (head dim, query length, key length)
|
344 |
+
# query length is 1 because during decoding we deal with one
|
345 |
+
# index.
|
346 |
+
# The same mask is applied to all batch elements and heads.
|
347 |
+
(batch, 1, 1, length),
|
348 |
+
),
|
349 |
+
)
|
350 |
+
|
351 |
+
# Grab the correct relative attention bias during decoding. This is
|
352 |
+
# only required during single step decoding.
|
353 |
+
if bias is not None:
|
354 |
+
# The bias is a full attention matrix, but during decoding we only
|
355 |
+
# have to take a slice of it.
|
356 |
+
# This is equivalent to bias[..., cur_index:cur_index+1, :].
|
357 |
+
bias = dynamic_vector_slice_in_dim(jnp.squeeze(bias, axis=0), jnp.reshape(cur_index, (-1)), 1, -2)
|
358 |
+
|
359 |
+
# Convert the boolean attention mask to an attention bias.
|
360 |
+
if mask is not None:
|
361 |
+
# attention mask in the form of attention bias
|
362 |
+
attention_bias = lax.select(
|
363 |
+
mask > 0, jnp.full(mask.shape, 0.0).astype(self.dtype), jnp.full(mask.shape, -1e10).astype(self.dtype)
|
364 |
+
)
|
365 |
+
else:
|
366 |
+
attention_bias = None
|
367 |
+
|
368 |
+
# Add provided bias term (e.g. relative position embedding).
|
369 |
+
if bias is not None:
|
370 |
+
attention_bias = combine_biases(attention_bias, bias)
|
371 |
+
|
372 |
+
dropout_rng = None
|
373 |
+
if not deterministic and self.dropout_rate > 0.0:
|
374 |
+
dropout_rng = self.make_rng("dropout")
|
375 |
+
|
376 |
+
# Apply attention.
|
377 |
+
x = dot_product_attention(
|
378 |
+
query,
|
379 |
+
key,
|
380 |
+
value,
|
381 |
+
bias=attention_bias,
|
382 |
+
dropout_rng=dropout_rng,
|
383 |
+
dropout_rate=self.dropout_rate,
|
384 |
+
deterministic=deterministic,
|
385 |
+
dtype=self.dtype,
|
386 |
+
float32_logits=self.float32_logits,
|
387 |
+
)
|
388 |
+
|
389 |
+
# Back to the original inputs dimensions.
|
390 |
+
out = DenseGeneral(
|
391 |
+
features=inputs_q.shape[-1], # output dim is set to the input dim.
|
392 |
+
axis=(-2, -1),
|
393 |
+
kernel_init=self.kernel_init,
|
394 |
+
kernel_axes=("heads", "kv", "embed"),
|
395 |
+
dtype=self.dtype,
|
396 |
+
name="out",
|
397 |
+
)(x)
|
398 |
+
return out
|
399 |
+
|
400 |
+
|
401 |
+
def _normalize_axes(axes: Iterable[int], ndim: int) -> Tuple[int]:
|
402 |
+
# A tuple by convention. len(axes_tuple) then also gives the rank efficiently.
|
403 |
+
return tuple([ax if ax >= 0 else ndim + ax for ax in axes])
|
404 |
+
|
405 |
+
|
406 |
+
def _canonicalize_tuple(x):
|
407 |
+
if isinstance(x, Iterable):
|
408 |
+
return tuple(x)
|
409 |
+
else:
|
410 |
+
return (x,)
|
411 |
+
|
412 |
+
|
413 |
+
# ------------------------------------------------------------------------------
|
414 |
+
# DenseGeneral for attention layers.
|
415 |
+
# ------------------------------------------------------------------------------
|
416 |
+
class DenseGeneral(nn.Module):
|
417 |
+
"""A linear transformation (without bias) with flexible axes.
|
418 |
+
|
419 |
+
Attributes:
|
420 |
+
features: tuple with numbers of output features.
|
421 |
+
axis: tuple with axes to apply the transformation on.
|
422 |
+
dtype: the dtype of the computation (default: float32).
|
423 |
+
kernel_init: initializer function for the weight matrix.
|
424 |
+
"""
|
425 |
+
|
426 |
+
features: Union[Iterable[int], int]
|
427 |
+
axis: Union[Iterable[int], int] = -1
|
428 |
+
dtype: DType = jnp.float32
|
429 |
+
params_dtype: DType = jnp.float32
|
430 |
+
kernel_init: NdInitializer = nd_dense_init(1.0, "fan_in", "normal")
|
431 |
+
kernel_axes: Tuple[str, ...] = ()
|
432 |
+
use_bias: bool = True
|
433 |
+
bias_init: Any = nn.initializers.zeros
|
434 |
+
|
435 |
+
@nn.compact
|
436 |
+
def __call__(self, inputs: Array) -> Array:
|
437 |
+
"""Applies a linear transformation to the inputs along multiple dimensions.
|
438 |
+
|
439 |
+
Args:
|
440 |
+
inputs: The nd-array to be transformed.
|
441 |
+
|
442 |
+
Returns:
|
443 |
+
The transformed input.
|
444 |
+
"""
|
445 |
+
features = _canonicalize_tuple(self.features)
|
446 |
+
axis = _canonicalize_tuple(self.axis)
|
447 |
+
|
448 |
+
inputs = jnp.asarray(inputs, self.dtype)
|
449 |
+
axis = _normalize_axes(axis, inputs.ndim)
|
450 |
+
|
451 |
+
kernel_shape = tuple([inputs.shape[ax] for ax in axis]) + features
|
452 |
+
kernel_in_axis = np.arange(len(axis))
|
453 |
+
kernel_out_axis = np.arange(len(axis), len(axis) + len(features))
|
454 |
+
kernel = param_with_axes(
|
455 |
+
"kernel",
|
456 |
+
self.kernel_init,
|
457 |
+
kernel_shape,
|
458 |
+
self.params_dtype,
|
459 |
+
kernel_in_axis,
|
460 |
+
kernel_out_axis,
|
461 |
+
axes=self.kernel_axes,
|
462 |
+
)
|
463 |
+
if self.use_bias:
|
464 |
+
bias = param_with_axes("bias", self.bias_init, features, self.params_dtype, axes=(self.kernel_axes[-1],))
|
465 |
+
kernel = jnp.asarray(kernel, self.dtype)
|
466 |
+
|
467 |
+
contract_ind = tuple(range(0, len(axis)))
|
468 |
+
y = lax.dot_general(inputs, kernel, ((axis, contract_ind), ((), ())))
|
469 |
+
if self.use_bias:
|
470 |
+
bias = jnp.asarray(bias, self.dtype)
|
471 |
+
# y += jnp.reshape(bias, (1,) * (y.ndim - 1) + (-1,))
|
472 |
+
y += jnp.reshape(bias, (1,) * (len(features) - y.ndim) + bias.shape[:])
|
473 |
+
return y
|
474 |
+
|
475 |
+
|
476 |
+
def _convert_to_activation_function(fn_or_string: Union[str, Callable]) -> Callable:
|
477 |
+
"""Convert a string to an activation function."""
|
478 |
+
if fn_or_string == "linear":
|
479 |
+
return lambda x: x
|
480 |
+
elif isinstance(fn_or_string, str):
|
481 |
+
return getattr(nn, fn_or_string)
|
482 |
+
elif callable(fn_or_string):
|
483 |
+
return fn_or_string
|
484 |
+
else:
|
485 |
+
raise ValueError("don't know how to convert %s to an activation function" % (fn_or_string,))
|
486 |
+
|
487 |
+
|
488 |
+
class MlpBlock(nn.Module):
|
489 |
+
"""Transformer MLP / feed-forward block.
|
490 |
+
|
491 |
+
Attributes:
|
492 |
+
intermediate_dim: Shared dimension of hidden layers.
|
493 |
+
activations: Type of activations for each layer. Each element is either
|
494 |
+
'linear', a string function name in flax.linen, or a function.
|
495 |
+
kernel_init: Kernel function, passed to the dense layers.
|
496 |
+
deterministic: Whether the dropout layers should be deterministic.
|
497 |
+
intermediate_dropout_rate: Dropout rate used after the intermediate layers.
|
498 |
+
dtype: Type for the dense layer.
|
499 |
+
"""
|
500 |
+
|
501 |
+
intermediate_dim: int = 2048
|
502 |
+
activations: Sequence[Union[str, Callable]] = ("relu",)
|
503 |
+
kernel_init: NdInitializer = nd_dense_init(1.0, "fan_in", "truncated_normal")
|
504 |
+
intermediate_dropout_rate: float = 0.1
|
505 |
+
dtype: Any = jnp.float32
|
506 |
+
|
507 |
+
@nn.compact
|
508 |
+
def __call__(self, inputs, decode: bool = False, deterministic: bool = False):
|
509 |
+
"""Applies Transformer MlpBlock module."""
|
510 |
+
# Iterate over specified MLP input activation functions.
|
511 |
+
# e.g. ('relu',) or ('gelu', 'linear') for gated-gelu.
|
512 |
+
activations = []
|
513 |
+
for idx, act_fn in enumerate(self.activations):
|
514 |
+
dense_name = "wi" if len(self.activations) == 1 else f"wi_{idx}"
|
515 |
+
x = DenseGeneral(
|
516 |
+
self.intermediate_dim,
|
517 |
+
dtype=self.dtype,
|
518 |
+
kernel_init=self.kernel_init,
|
519 |
+
kernel_axes=("embed", "mlp"),
|
520 |
+
name=dense_name,
|
521 |
+
)(inputs)
|
522 |
+
x = _convert_to_activation_function(act_fn)(x)
|
523 |
+
activations.append(x)
|
524 |
+
|
525 |
+
# Take elementwise product of above intermediate activations.
|
526 |
+
x = functools.reduce(operator.mul, activations)
|
527 |
+
# Apply dropout and final dense output projection.
|
528 |
+
x = nn.Dropout(rate=self.intermediate_dropout_rate, broadcast_dims=(-2,))(
|
529 |
+
x, deterministic=deterministic
|
530 |
+
) # Broadcast along length.
|
531 |
+
x = with_sharding_constraint(x, ("batch", "length", "mlp"))
|
532 |
+
output = DenseGeneral(
|
533 |
+
inputs.shape[-1], dtype=self.dtype, kernel_init=self.kernel_init, kernel_axes=("mlp", "embed"), name="wo"
|
534 |
+
)(x)
|
535 |
+
return output
|
536 |
+
|
537 |
+
|
538 |
+
class Embed(nn.Module):
|
539 |
+
"""A parameterized function from integers [0, n) to d-dimensional vectors.
|
540 |
+
|
541 |
+
Attributes:
|
542 |
+
num_embeddings: number of embeddings.
|
543 |
+
features: number of feature dimensions for each embedding.
|
544 |
+
dtype: the dtype of the embedding vectors (default: float32).
|
545 |
+
embedding_init: embedding initializer.
|
546 |
+
one_hot: performs the gather with a one-hot contraction rather than a true
|
547 |
+
gather. This is currently needed for SPMD partitioning.
|
548 |
+
"""
|
549 |
+
|
550 |
+
num_embeddings: int
|
551 |
+
features: int
|
552 |
+
cast_input_dtype: Optional[DType] = None
|
553 |
+
dtype: DType = jnp.float32
|
554 |
+
params_dtype: DType = jnp.float32
|
555 |
+
attend_dtype: Optional[DType] = None
|
556 |
+
embedding_init: Initializer = default_embed_init
|
557 |
+
one_hot: bool = True
|
558 |
+
embedding: Array = dataclasses.field(init=False)
|
559 |
+
|
560 |
+
def setup(self):
|
561 |
+
self.embedding = param_with_axes(
|
562 |
+
"embedding",
|
563 |
+
self.embedding_init,
|
564 |
+
(self.num_embeddings, self.features),
|
565 |
+
self.params_dtype,
|
566 |
+
axes=("vocab", "embed"),
|
567 |
+
)
|
568 |
+
|
569 |
+
def __call__(self, inputs: Array) -> Array:
|
570 |
+
"""Embeds the inputs along the last dimension.
|
571 |
+
|
572 |
+
Args:
|
573 |
+
inputs: input data, all dimensions are considered batch dimensions.
|
574 |
+
|
575 |
+
Returns:
|
576 |
+
Output which is embedded input data. The output shape follows the input,
|
577 |
+
with an additional `features` dimension appended.
|
578 |
+
"""
|
579 |
+
if self.cast_input_dtype:
|
580 |
+
inputs = inputs.astype(self.cast_input_dtype)
|
581 |
+
if not jnp.issubdtype(inputs.dtype, jnp.integer):
|
582 |
+
raise ValueError("Input type must be an integer or unsigned integer.")
|
583 |
+
if self.one_hot:
|
584 |
+
iota = lax.iota(jnp.int32, self.num_embeddings)
|
585 |
+
one_hot = jnp.array(inputs[..., jnp.newaxis] == iota, dtype=self.dtype)
|
586 |
+
output = jnp.dot(one_hot, jnp.asarray(self.embedding, self.dtype))
|
587 |
+
else:
|
588 |
+
output = jnp.asarray(self.embedding, self.dtype)[inputs]
|
589 |
+
output = with_sharding_constraint(output, ("batch", "length", "embed"))
|
590 |
+
return output
|
591 |
+
|
592 |
+
def attend(self, query: Array) -> Array:
|
593 |
+
"""Attend over the embedding using a query array.
|
594 |
+
|
595 |
+
Args:
|
596 |
+
query: array with last dimension equal the feature depth `features` of the
|
597 |
+
embedding.
|
598 |
+
|
599 |
+
Returns:
|
600 |
+
An array with final dim `num_embeddings` corresponding to the batched
|
601 |
+
inner-product of the array of query vectors against each embedding.
|
602 |
+
Commonly used for weight-sharing between embeddings and logit transform
|
603 |
+
in NLP models.
|
604 |
+
"""
|
605 |
+
dtype = self.attend_dtype if self.attend_dtype is not None else self.dtype
|
606 |
+
return jnp.dot(query, jnp.asarray(self.embedding, dtype).T)
|
607 |
+
|
608 |
+
|
609 |
+
class RelativePositionBiases(nn.Module):
|
610 |
+
"""Adds T5-style relative positional embeddings to the attention logits.
|
611 |
+
|
612 |
+
Attributes:
|
613 |
+
num_buckets: Number of buckets to bucket distances between key and query
|
614 |
+
positions into.
|
615 |
+
max_distance: Maximum distance before everything is lumped into the last
|
616 |
+
distance bucket.
|
617 |
+
num_heads: Number of heads in the attention layer. Each head will get a
|
618 |
+
different relative position weighting.
|
619 |
+
dtype: Type of arrays through this module.
|
620 |
+
embedding_init: initializer for relative embedding table.
|
621 |
+
"""
|
622 |
+
|
623 |
+
num_buckets: int
|
624 |
+
max_distance: int
|
625 |
+
num_heads: int
|
626 |
+
dtype: Any
|
627 |
+
embedding_init: Callable[..., Array] = nn.linear.default_embed_init
|
628 |
+
|
629 |
+
@staticmethod
|
630 |
+
def _relative_position_bucket(relative_position, bidirectional=True, num_buckets=32, max_distance=128):
|
631 |
+
"""Translate relative position to a bucket number for relative attention.
|
632 |
+
|
633 |
+
The relative position is defined as memory_position - query_position, i.e.
|
634 |
+
the distance in tokens from the attending position to the attended-to
|
635 |
+
position. If bidirectional=False, then positive relative positions are
|
636 |
+
invalid.
|
637 |
+
We use smaller buckets for small absolute relative_position and larger
|
638 |
+
buckets for larger absolute relative_positions. All relative
|
639 |
+
positions >=max_distance map to the same bucket. All relative
|
640 |
+
positions <=-max_distance map to the same bucket. This should allow for
|
641 |
+
more graceful generalization to longer sequences than the model has been
|
642 |
+
trained on.
|
643 |
+
|
644 |
+
Args:
|
645 |
+
relative_position: an int32 array
|
646 |
+
bidirectional: a boolean - whether the attention is bidirectional
|
647 |
+
num_buckets: an integer
|
648 |
+
max_distance: an integer
|
649 |
+
|
650 |
+
Returns:
|
651 |
+
a Tensor with the same shape as relative_position, containing int32
|
652 |
+
values in the range [0, num_buckets)
|
653 |
+
"""
|
654 |
+
ret = 0
|
655 |
+
n = -relative_position
|
656 |
+
if bidirectional:
|
657 |
+
num_buckets //= 2
|
658 |
+
ret += (n < 0).astype(np.int32) * num_buckets
|
659 |
+
n = np.abs(n)
|
660 |
+
else:
|
661 |
+
n = np.maximum(n, 0)
|
662 |
+
# now n is in the range [0, inf)
|
663 |
+
max_exact = num_buckets // 2
|
664 |
+
is_small = n < max_exact
|
665 |
+
val_if_large = max_exact + (
|
666 |
+
np.log(n.astype(np.float32) / max_exact + np.finfo(np.float32).eps)
|
667 |
+
/ np.log(max_distance / max_exact)
|
668 |
+
* (num_buckets - max_exact)
|
669 |
+
).astype(np.int32)
|
670 |
+
val_if_large = np.minimum(val_if_large, num_buckets - 1)
|
671 |
+
ret += np.where(is_small, n, val_if_large)
|
672 |
+
return ret
|
673 |
+
|
674 |
+
@nn.compact
|
675 |
+
def __call__(self, qlen, klen, bidirectional=True):
|
676 |
+
"""Produce relative position embedding attention biases.
|
677 |
+
|
678 |
+
Args:
|
679 |
+
qlen: attention query length.
|
680 |
+
klen: attention key length.
|
681 |
+
bidirectional: whether to allow positive memory-query relative position
|
682 |
+
embeddings.
|
683 |
+
|
684 |
+
Returns:
|
685 |
+
output: `(1, len, q_len, k_len)` attention bias
|
686 |
+
"""
|
687 |
+
# TODO(levskaya): should we be computing this w. numpy as a program
|
688 |
+
# constant?
|
689 |
+
context_position = np.arange(qlen, dtype=jnp.int32)[:, None]
|
690 |
+
memory_position = np.arange(klen, dtype=jnp.int32)[None, :]
|
691 |
+
relative_position = memory_position - context_position # shape (qlen, klen)
|
692 |
+
rp_bucket = self._relative_position_bucket(
|
693 |
+
relative_position,
|
694 |
+
bidirectional=bidirectional,
|
695 |
+
num_buckets=self.num_buckets,
|
696 |
+
max_distance=self.max_distance,
|
697 |
+
)
|
698 |
+
relative_attention_bias = param_with_axes(
|
699 |
+
"rel_embedding",
|
700 |
+
self.embedding_init,
|
701 |
+
(self.num_heads, self.num_buckets),
|
702 |
+
jnp.float32,
|
703 |
+
axes=("heads", "relpos_buckets"),
|
704 |
+
)
|
705 |
+
|
706 |
+
relative_attention_bias = jnp.asarray(relative_attention_bias, self.dtype)
|
707 |
+
# Instead of using a slow gather, we create a leading-dimension one-hot
|
708 |
+
# array from rp_bucket and use it to perform the gather-equivalent via a
|
709 |
+
# contraction, i.e.:
|
710 |
+
# (num_head, num_buckets) x (num_buckets one-hot, qlen, klen).
|
711 |
+
# This is equivalent to relative_attention_bias[:, rp_bucket]
|
712 |
+
bcast_iota = lax.broadcasted_iota(jnp.int32, (self.num_buckets, 1, 1), 0)
|
713 |
+
rp_bucket_one_hot = jnp.array(rp_bucket[jnp.newaxis, ...] == bcast_iota, dtype=self.dtype)
|
714 |
+
# --> shape (qlen, klen, num_heads)
|
715 |
+
values = lax.dot_general(
|
716 |
+
relative_attention_bias, rp_bucket_one_hot, (((1,), (0,)), ((), ())) # rhs, lhs contracting dims
|
717 |
+
) # no batched dims
|
718 |
+
# Add a singleton batch dimension.
|
719 |
+
# --> shape (1, num_heads, qlen, klen)
|
720 |
+
return values[jnp.newaxis, ...]
|
721 |
+
|
722 |
+
|
723 |
+
# ------------------------------------------------------------------------------
|
724 |
+
# T5 Layernorm - no subtraction of mean or bias.
|
725 |
+
# ------------------------------------------------------------------------------
|
726 |
+
# class LayerNorm(nn.Module):
|
727 |
+
# """T5 Layer normalization operating on the last axis of the input data."""
|
728 |
+
# epsilon: float = 1e-6
|
729 |
+
# dtype: Any = jnp.float32
|
730 |
+
# scale_init: Initializer = nn.initializers.ones
|
731 |
+
|
732 |
+
# @nn.compact
|
733 |
+
# def __call__(self, x: jnp.ndarray) -> jnp.ndarray:
|
734 |
+
# """Applies layer normalization on the input."""
|
735 |
+
# x = jnp.asarray(x, jnp.float32)
|
736 |
+
# features = x.shape[-1]
|
737 |
+
# mean2 = jnp.mean(lax.square(x), axis=-1, keepdims=True)
|
738 |
+
# y = jnp.asarray(x * lax.rsqrt(mean2 + self.epsilon), self.dtype)
|
739 |
+
# scale = param_with_axes(
|
740 |
+
# 'scale', self.scale_init, (features,), jnp.float32, axes=('embed',))
|
741 |
+
|
742 |
+
# scale = jnp.asarray(scale, self.dtype)
|
743 |
+
# return y * scale
|
744 |
+
|
745 |
+
|
746 |
+
class LayerNorm(nn.Module):
|
747 |
+
"""Layer normalization (https://arxiv.org/abs/1607.06450).
|
748 |
+
Operates on the last axis of the input data.
|
749 |
+
It normalizes the activations of the layer for each given example in a
|
750 |
+
batch independently, rather than across a batch like Batch Normalization.
|
751 |
+
i.e. applies a transformation that maintains the mean activation within
|
752 |
+
each example close to 0 and the activation standard deviation close to 1.
|
753 |
+
Attributes:
|
754 |
+
epsilon: A small float added to variance to avoid dividing by zero.
|
755 |
+
dtype: the dtype of the computation (default: float32).
|
756 |
+
use_bias: If True, bias (beta) is added.
|
757 |
+
use_scale: If True, multiply by scale (gamma). When the next layer is linear
|
758 |
+
(also e.g. nn.relu), this can be disabled since the scaling will be done
|
759 |
+
by the next layer.
|
760 |
+
bias_init: Initializer for bias, by default, zero.
|
761 |
+
scale_init: Initializer for scale, by default, one.
|
762 |
+
"""
|
763 |
+
|
764 |
+
epsilon: float = 1e-6
|
765 |
+
dtype: Any = jnp.float32
|
766 |
+
params_dtype: DType = jnp.float32
|
767 |
+
use_bias: bool = True
|
768 |
+
use_scale: bool = True
|
769 |
+
bias_init: Callable[[PRNGKey, Shape, Any], Array] = nn.initializers.zeros
|
770 |
+
scale_init: Callable[[PRNGKey, Shape, Any], Array] = nn.initializers.ones
|
771 |
+
|
772 |
+
@nn.compact
|
773 |
+
def __call__(self, x):
|
774 |
+
"""Applies layer normalization on the input.
|
775 |
+
Args:
|
776 |
+
x: the inputs
|
777 |
+
Returns:
|
778 |
+
Normalized inputs (the same shape as inputs).
|
779 |
+
"""
|
780 |
+
x = jnp.asarray(x, jnp.float32)
|
781 |
+
features = x.shape[-1]
|
782 |
+
mean = jnp.mean(x, axis=-1, keepdims=True)
|
783 |
+
mean2 = jnp.mean(lax.square(x), axis=-1, keepdims=True)
|
784 |
+
var = mean2 - lax.square(mean)
|
785 |
+
mul = lax.rsqrt(var + self.epsilon)
|
786 |
+
if self.use_scale:
|
787 |
+
scale = param_with_axes("scale", self.scale_init, (features,), self.params_dtype, axes=("embed",))
|
788 |
+
mul = mul * jnp.asarray(scale, self.dtype)
|
789 |
+
y = (x - mean) * mul
|
790 |
+
if self.use_bias:
|
791 |
+
bias = param_with_axes("bias", self.bias_init, (features,), self.params_dtype, axes=("embed",))
|
792 |
+
y = y + jnp.asarray(bias, self.dtype)
|
793 |
+
return jnp.asarray(y, self.dtype)
|
794 |
+
|
795 |
+
|
796 |
+
# ------------------------------------------------------------------------------
|
797 |
+
# Mask-making utility functions.
|
798 |
+
# ------------------------------------------------------------------------------
|
799 |
+
def make_attention_mask(
|
800 |
+
query_input: Array,
|
801 |
+
key_input: Array,
|
802 |
+
pairwise_fn: Callable = jnp.multiply,
|
803 |
+
extra_batch_dims: int = 0,
|
804 |
+
dtype: DType = jnp.float32,
|
805 |
+
) -> Array:
|
806 |
+
"""Mask-making helper for attention weights.
|
807 |
+
|
808 |
+
In case of 1d inputs (i.e., `[batch, len_q]`, `[batch, len_kv]`, the
|
809 |
+
attention weights will be `[batch, heads, len_q, len_kv]` and this
|
810 |
+
function will produce `[batch, 1, len_q, len_kv]`.
|
811 |
+
|
812 |
+
Args:
|
813 |
+
query_input: a batched, flat input of query_length size
|
814 |
+
key_input: a batched, flat input of key_length size
|
815 |
+
pairwise_fn: broadcasting elementwise comparison function
|
816 |
+
extra_batch_dims: number of extra batch dims to add singleton axes for, none
|
817 |
+
by default
|
818 |
+
dtype: mask return dtype
|
819 |
+
|
820 |
+
Returns:
|
821 |
+
A `[batch, 1, len_q, len_kv]` shaped mask for 1d attention.
|
822 |
+
"""
|
823 |
+
# [batch, len_q, len_kv]
|
824 |
+
mask = pairwise_fn(
|
825 |
+
# [batch, len_q] -> [batch, len_q, 1]
|
826 |
+
jnp.expand_dims(query_input, axis=-1),
|
827 |
+
# [batch, len_q] -> [batch, 1, len_kv]
|
828 |
+
jnp.expand_dims(key_input, axis=-2),
|
829 |
+
)
|
830 |
+
|
831 |
+
# [batch, 1, len_q, len_kv]. This creates the head dim.
|
832 |
+
mask = jnp.expand_dims(mask, axis=-3)
|
833 |
+
mask = jnp.expand_dims(mask, axis=tuple(range(extra_batch_dims)))
|
834 |
+
return mask.astype(dtype)
|
835 |
+
|
836 |
+
|
837 |
+
def make_causal_mask(x: Array, extra_batch_dims: int = 0, dtype: DType = jnp.float32) -> Array:
|
838 |
+
"""Make a causal mask for self-attention.
|
839 |
+
|
840 |
+
In case of 1d inputs (i.e., `[batch, len]`, the self-attention weights
|
841 |
+
will be `[batch, heads, len, len]` and this function will produce a
|
842 |
+
causal mask of shape `[batch, 1, len, len]`.
|
843 |
+
|
844 |
+
Note that a causal mask does not depend on the values of x; it only depends on
|
845 |
+
the shape. If x has padding elements, they will not be treated in a special
|
846 |
+
manner.
|
847 |
+
|
848 |
+
Args:
|
849 |
+
x: input array of shape `[batch, len]`
|
850 |
+
extra_batch_dims: number of batch dims to add singleton axes for, none by
|
851 |
+
default
|
852 |
+
dtype: mask return dtype
|
853 |
+
|
854 |
+
Returns:
|
855 |
+
A `[batch, 1, len, len]` shaped causal mask for 1d attention.
|
856 |
+
"""
|
857 |
+
idxs = jnp.broadcast_to(jnp.arange(x.shape[-1], dtype=jnp.int32), x.shape)
|
858 |
+
return make_attention_mask(idxs, idxs, jnp.greater_equal, extra_batch_dims=extra_batch_dims, dtype=dtype)
|
859 |
+
|
860 |
+
|
861 |
+
def combine_masks(*masks: Optional[Array], dtype: DType = jnp.float32):
|
862 |
+
"""Combine attention masks.
|
863 |
+
|
864 |
+
Args:
|
865 |
+
*masks: set of attention mask arguments to combine, some can be None.
|
866 |
+
dtype: final mask dtype
|
867 |
+
|
868 |
+
Returns:
|
869 |
+
Combined mask, reduced by logical and, returns None if no masks given.
|
870 |
+
"""
|
871 |
+
masks = [m for m in masks if m is not None]
|
872 |
+
if not masks:
|
873 |
+
return None
|
874 |
+
assert all(
|
875 |
+
(x.ndim == masks[0].ndim for x in masks)
|
876 |
+
), f"masks must have same rank: {tuple((x.ndim for x in masks))}"
|
877 |
+
mask, *other_masks = masks
|
878 |
+
for other_mask in other_masks:
|
879 |
+
mask = jnp.logical_and(mask, other_mask)
|
880 |
+
return mask.astype(dtype)
|
881 |
+
|
882 |
+
|
883 |
+
def combine_biases(*masks: Optional[Array]):
|
884 |
+
"""Combine attention biases.
|
885 |
+
|
886 |
+
Args:
|
887 |
+
*masks: set of attention bias arguments to combine, some can be None.
|
888 |
+
|
889 |
+
Returns:
|
890 |
+
Combined mask, reduced by summation, returns None if no masks given.
|
891 |
+
"""
|
892 |
+
masks = [m for m in masks if m is not None]
|
893 |
+
if not masks:
|
894 |
+
return None
|
895 |
+
assert all(
|
896 |
+
(x.ndim == masks[0].ndim for x in masks)
|
897 |
+
), f"masks must have same rank: {tuple((x.ndim for x in masks))}"
|
898 |
+
mask, *other_masks = masks
|
899 |
+
for other_mask in other_masks:
|
900 |
+
mask = mask + other_mask
|
901 |
+
return mask
|
902 |
+
|
903 |
+
|
904 |
+
def make_decoder_mask(
|
905 |
+
decoder_target_tokens: Array,
|
906 |
+
dtype: DType,
|
907 |
+
decoder_causal_attention: Optional[Array] = None,
|
908 |
+
decoder_segment_ids: Optional[Array] = None,
|
909 |
+
) -> Array:
|
910 |
+
"""Compute the self-attention mask for a decoder.
|
911 |
+
|
912 |
+
Decoder mask is formed by combining a causal mask, a padding mask and an
|
913 |
+
optional packing mask. If decoder_causal_attention is passed, it makes the
|
914 |
+
masking non-causal for positions that have value of 1.
|
915 |
+
|
916 |
+
A prefix LM is applied to a dataset which has a notion of "inputs" and
|
917 |
+
"targets", e.g., a machine translation task. The inputs and targets are
|
918 |
+
concatenated to form a new target. `decoder_target_tokens` is the concatenated
|
919 |
+
decoder output tokens.
|
920 |
+
|
921 |
+
The "inputs" portion of the concatenated sequence can attend to other "inputs"
|
922 |
+
tokens even for those at a later time steps. In order to control this
|
923 |
+
behavior, `decoder_causal_attention` is necessary. This is a binary mask with
|
924 |
+
a value of 1 indicating that the position belonged to "inputs" portion of the
|
925 |
+
original dataset.
|
926 |
+
|
927 |
+
Example:
|
928 |
+
|
929 |
+
Suppose we have a dataset with two examples.
|
930 |
+
|
931 |
+
ds = [{"inputs": [6, 7], "targets": [8]},
|
932 |
+
{"inputs": [3, 4], "targets": [5]}]
|
933 |
+
|
934 |
+
After the data preprocessing with packing, the two examples are packed into
|
935 |
+
one example with the following three fields (some fields are skipped for
|
936 |
+
simplicity).
|
937 |
+
|
938 |
+
decoder_target_tokens = [[6, 7, 8, 3, 4, 5, 0]]
|
939 |
+
decoder_segment_ids = [[1, 1, 1, 2, 2, 2, 0]]
|
940 |
+
decoder_causal_attention = [[1, 1, 0, 1, 1, 0, 0]]
|
941 |
+
|
942 |
+
where each array has [batch, length] shape with batch size being 1. Then,
|
943 |
+
this function computes the following mask.
|
944 |
+
|
945 |
+
mask = [[[[1, 1, 0, 0, 0, 0, 0],
|
946 |
+
[1, 1, 0, 0, 0, 0, 0],
|
947 |
+
[1, 1, 1, 0, 0, 0, 0],
|
948 |
+
[0, 0, 0, 1, 1, 0, 0],
|
949 |
+
[0, 0, 0, 1, 1, 0, 0],
|
950 |
+
[0, 0, 0, 1, 1, 1, 0],
|
951 |
+
[0, 0, 0, 0, 0, 0, 0]]]]
|
952 |
+
|
953 |
+
mask[b, 1, :, :] represents the mask for the example `b` in the batch.
|
954 |
+
Because mask is for a self-attention layer, the mask's shape is a square of
|
955 |
+
shape [query length, key length].
|
956 |
+
|
957 |
+
mask[b, 1, i, j] = 1 means that the query token at position i can attend to
|
958 |
+
the key token at position j.
|
959 |
+
|
960 |
+
Args:
|
961 |
+
decoder_target_tokens: decoder output tokens. [batch, length]
|
962 |
+
dtype: dtype of the output mask.
|
963 |
+
decoder_causal_attention: a binary mask indicating which position should
|
964 |
+
only attend to earlier positions in the sequence. Others will attend
|
965 |
+
bidirectionally. [batch, length]
|
966 |
+
decoder_segment_ids: decoder segmentation info for packed examples. [batch,
|
967 |
+
length]
|
968 |
+
|
969 |
+
Returns:
|
970 |
+
the combined decoder mask.
|
971 |
+
"""
|
972 |
+
masks = []
|
973 |
+
# The same mask is applied to all attention heads. So the head dimension is 1,
|
974 |
+
# i.e., the mask will be broadcast along the heads dim.
|
975 |
+
# [batch, 1, length, length]
|
976 |
+
causal_mask = make_causal_mask(decoder_target_tokens, dtype=dtype)
|
977 |
+
|
978 |
+
# Positions with value 1 in `decoder_causal_attneition` can attend
|
979 |
+
# bidirectionally.
|
980 |
+
if decoder_causal_attention is not None:
|
981 |
+
# [batch, 1, length, length]
|
982 |
+
inputs_mask = make_attention_mask(
|
983 |
+
decoder_causal_attention, decoder_causal_attention, jnp.logical_and, dtype=dtype
|
984 |
+
)
|
985 |
+
masks.append(jnp.logical_or(causal_mask, inputs_mask).astype(dtype))
|
986 |
+
else:
|
987 |
+
masks.append(causal_mask)
|
988 |
+
|
989 |
+
# Padding mask.
|
990 |
+
masks.append(make_attention_mask(decoder_target_tokens > 0, decoder_target_tokens > 0, dtype=dtype))
|
991 |
+
|
992 |
+
# Packing mask
|
993 |
+
if decoder_segment_ids is not None:
|
994 |
+
masks.append(make_attention_mask(decoder_segment_ids, decoder_segment_ids, jnp.equal, dtype=dtype))
|
995 |
+
|
996 |
+
return combine_masks(*masks, dtype=dtype)
|
997 |
+
|
998 |
+
|
999 |
+
def canonicalize_padding(padding: PaddingLike, rank: int) -> LaxPadding:
|
1000 |
+
""" "Canonicalizes conv padding to a jax.lax supported format."""
|
1001 |
+
if isinstance(padding, str):
|
1002 |
+
return padding
|
1003 |
+
if isinstance(padding, int):
|
1004 |
+
return [(padding, padding)] * rank
|
1005 |
+
if isinstance(padding, Sequence) and len(padding) == rank:
|
1006 |
+
new_pad = []
|
1007 |
+
for p in padding:
|
1008 |
+
if isinstance(p, int):
|
1009 |
+
new_pad.append((p, p))
|
1010 |
+
elif isinstance(p, tuple) and len(p) == 2:
|
1011 |
+
new_pad.append(p)
|
1012 |
+
else:
|
1013 |
+
break
|
1014 |
+
if len(new_pad) == rank:
|
1015 |
+
return new_pad
|
1016 |
+
raise ValueError(
|
1017 |
+
f"Invalid padding format: {padding}, should be str, int,"
|
1018 |
+
f" or a sequence of len {rank} where each element is an"
|
1019 |
+
f" int or pair of ints."
|
1020 |
+
)
|
1021 |
+
|
1022 |
+
|
1023 |
+
def _conv_dimension_numbers(input_shape):
|
1024 |
+
"""Computes the dimension numbers based on the input shape."""
|
1025 |
+
ndim = len(input_shape)
|
1026 |
+
lhs_spec = (0, ndim - 1) + tuple(range(1, ndim - 1))
|
1027 |
+
rhs_spec = (ndim - 1, ndim - 2) + tuple(range(0, ndim - 2))
|
1028 |
+
out_spec = lhs_spec
|
1029 |
+
return lax.ConvDimensionNumbers(lhs_spec, rhs_spec, out_spec)
|
1030 |
+
|
1031 |
+
|
1032 |
+
class _Conv(nn.Module):
|
1033 |
+
"""Convolution Module wrapping `lax.conv_general_dilated[_local]`.
|
1034 |
+
|
1035 |
+
Attributes:
|
1036 |
+
features: number of convolution filters.
|
1037 |
+
kernel_size: shape of the convolutional kernel. For 1D convolution,
|
1038 |
+
the kernel size can be passed as an integer. For all other cases, it must
|
1039 |
+
be a sequence of integers.
|
1040 |
+
strides: an integer or a sequence of `n` integers, representing the
|
1041 |
+
inter-window strides (default: 1).
|
1042 |
+
padding: either the string `'SAME'`, the string `'VALID'`, the string
|
1043 |
+
`'CIRCULAR'` (periodic boundary conditions), or a sequence of `n` `(low,
|
1044 |
+
high)` integer pairs that give the padding to apply before and after each
|
1045 |
+
spatial dimension. A single int is interpeted as applying the same padding
|
1046 |
+
in all dims and passign a single int in a sequence causes the same padding
|
1047 |
+
to be used on both sides. `'CAUSAL'` padding for a 1D convolution will
|
1048 |
+
left-pad the convolution axis, resulting in same-sized output.
|
1049 |
+
input_dilation: an integer or a sequence of `n` integers, giving the
|
1050 |
+
dilation factor to apply in each spatial dimension of `inputs`
|
1051 |
+
(default: 1). Convolution with input dilation `d` is equivalent to
|
1052 |
+
transposed convolution with stride `d`.
|
1053 |
+
kernel_dilation: an integer or a sequence of `n` integers, giving the
|
1054 |
+
dilation factor to apply in each spatial dimension of the convolution
|
1055 |
+
kernel (default: 1). Convolution with kernel dilation
|
1056 |
+
is also known as 'atrous convolution'.
|
1057 |
+
feature_group_count: integer, default 1. If specified divides the input
|
1058 |
+
features into groups.
|
1059 |
+
use_bias: whether to add a bias to the output (default: True).
|
1060 |
+
mask: Optional mask for the weights during masked convolution. The mask must
|
1061 |
+
be the same shape as the convolution weight matrix.
|
1062 |
+
dtype: the dtype of the computation (default: infer from input and params).
|
1063 |
+
params_dtype: the dtype passed to parameter initializers (default: float32).
|
1064 |
+
precision: numerical precision of the computation see `jax.lax.Precision`
|
1065 |
+
for details.
|
1066 |
+
kernel_init: initializer for the convolutional kernel.
|
1067 |
+
bias_init: initializer for the bias.
|
1068 |
+
"""
|
1069 |
+
|
1070 |
+
features: int
|
1071 |
+
kernel_size: Sequence[int]
|
1072 |
+
strides: Union[None, int, Sequence[int]] = 1
|
1073 |
+
padding: PaddingLike = "SAME"
|
1074 |
+
input_dilation: Union[None, int, Sequence[int]] = 1
|
1075 |
+
kernel_dilation: Union[None, int, Sequence[int]] = 1
|
1076 |
+
feature_group_count: int = 1
|
1077 |
+
use_bias: bool = True
|
1078 |
+
mask: Optional[Array] = None
|
1079 |
+
dtype: Optional[DType] = None
|
1080 |
+
params_dtype: DType = jnp.float32
|
1081 |
+
precision: PrecisionLike = None
|
1082 |
+
kernel_init: Callable[[PRNGKey, Shape, DType], Array] = nn.initializers.lecun_normal()
|
1083 |
+
bias_init: Callable[[PRNGKey, Shape, DType], Array] = nn.initializers.zeros
|
1084 |
+
conv_general_dilated: ConvGeneralDilatedT = lax.conv_general_dilated
|
1085 |
+
kernel_axes: Tuple[str, ...] = ()
|
1086 |
+
|
1087 |
+
@property
|
1088 |
+
def shared_weights(self) -> bool: # type: ignore
|
1089 |
+
"""Defines whether weights are shared or not between different pixels.
|
1090 |
+
|
1091 |
+
Returns:
|
1092 |
+
`True` to use shared weights in convolution (regular convolution).
|
1093 |
+
`False` to use different weights at different pixels, a.k.a.
|
1094 |
+
"locally connected layer", "unshared convolution", or "local convolution".
|
1095 |
+
|
1096 |
+
"""
|
1097 |
+
...
|
1098 |
+
|
1099 |
+
@nn.compact
|
1100 |
+
def __call__(self, inputs: Array) -> Array:
|
1101 |
+
"""Applies a (potentially unshared) convolution to the inputs.
|
1102 |
+
|
1103 |
+
Args:
|
1104 |
+
inputs: input data with dimensions (*batch_dims, spatial_dims...,
|
1105 |
+
features). This is the channels-last convention, i.e. NHWC for a 2d
|
1106 |
+
convolution and NDHWC for a 3D convolution. Note: this is different from
|
1107 |
+
the input convention used by `lax.conv_general_dilated`, which puts the
|
1108 |
+
spatial dimensions last.
|
1109 |
+
Note: If the input has more than 1 batch dimension, all batch dimensions
|
1110 |
+
are flattened into a single dimension for the convolution and restored
|
1111 |
+
before returning. In some cases directly vmap'ing the layer may yield
|
1112 |
+
better performance than this default flattening approach. If the input
|
1113 |
+
lacks a batch dimension it will be added for the convolution and removed
|
1114 |
+
n return, an allowance made to enable writing single-example code.
|
1115 |
+
|
1116 |
+
Returns:
|
1117 |
+
The convolved data.
|
1118 |
+
"""
|
1119 |
+
|
1120 |
+
if isinstance(self.kernel_size, int):
|
1121 |
+
raise TypeError(
|
1122 |
+
"Expected Conv kernel_size to be a"
|
1123 |
+
" tuple/list of integers (eg.: [3, 3]) but got"
|
1124 |
+
f" {self.kernel_size}."
|
1125 |
+
)
|
1126 |
+
else:
|
1127 |
+
kernel_size = tuple(self.kernel_size)
|
1128 |
+
|
1129 |
+
def maybe_broadcast(x: Optional[Union[int, Sequence[int]]]) -> Tuple[int, ...]:
|
1130 |
+
if x is None:
|
1131 |
+
# backward compatibility with using None as sentinel for
|
1132 |
+
# broadcast 1
|
1133 |
+
x = 1
|
1134 |
+
if isinstance(x, int):
|
1135 |
+
return (x,) * len(kernel_size)
|
1136 |
+
return tuple(x)
|
1137 |
+
|
1138 |
+
# Combine all input batch dimensions into a single leading batch axis.
|
1139 |
+
num_batch_dimensions = inputs.ndim - (len(kernel_size) + 1)
|
1140 |
+
if num_batch_dimensions != 1:
|
1141 |
+
input_batch_shape = inputs.shape[:num_batch_dimensions]
|
1142 |
+
total_batch_size = int(np.prod(input_batch_shape))
|
1143 |
+
flat_input_shape = (total_batch_size,) + inputs.shape[num_batch_dimensions:]
|
1144 |
+
inputs = jnp.reshape(inputs, flat_input_shape)
|
1145 |
+
|
1146 |
+
# self.strides or (1,) * (inputs.ndim - 2)
|
1147 |
+
strides = maybe_broadcast(self.strides)
|
1148 |
+
input_dilation = maybe_broadcast(self.input_dilation)
|
1149 |
+
kernel_dilation = maybe_broadcast(self.kernel_dilation)
|
1150 |
+
|
1151 |
+
padding_lax = canonicalize_padding(self.padding, len(kernel_size))
|
1152 |
+
if padding_lax == "CIRCULAR":
|
1153 |
+
kernel_size_dilated = [(k - 1) * d + 1 for k, d in zip(kernel_size, kernel_dilation)]
|
1154 |
+
zero_pad: List[Tuple[int, int]] = [(0, 0)]
|
1155 |
+
pads = zero_pad + [((k - 1) // 2, k // 2) for k in kernel_size_dilated] + [(0, 0)]
|
1156 |
+
inputs = jnp.pad(inputs, pads, mode="wrap")
|
1157 |
+
padding_lax = "VALID"
|
1158 |
+
elif padding_lax == "CAUSAL":
|
1159 |
+
if len(kernel_size) != 1:
|
1160 |
+
raise ValueError("Causal padding is only implemented for 1D convolutions.")
|
1161 |
+
left_pad = kernel_dilation[0] * (kernel_size[0] - 1)
|
1162 |
+
pads = [(0, 0), (left_pad, 0), (0, 0)]
|
1163 |
+
inputs = jnp.pad(inputs, pads)
|
1164 |
+
padding_lax = "VALID"
|
1165 |
+
|
1166 |
+
dimension_numbers = _conv_dimension_numbers(inputs.shape)
|
1167 |
+
in_features = jnp.shape(inputs)[-1]
|
1168 |
+
|
1169 |
+
if self.shared_weights:
|
1170 |
+
# One shared convolutional kernel for all pixels in the output.
|
1171 |
+
assert in_features % self.feature_group_count == 0
|
1172 |
+
kernel_shape = kernel_size + (in_features // self.feature_group_count, self.features)
|
1173 |
+
|
1174 |
+
else:
|
1175 |
+
if self.feature_group_count != 1:
|
1176 |
+
raise NotImplementedError(
|
1177 |
+
f"`lax.conv_general_dilated_local` does not support "
|
1178 |
+
f"`feature_group_count != 1`, got `{self.feature_group_count}`."
|
1179 |
+
)
|
1180 |
+
|
1181 |
+
# Need to know the spatial output shape of a standard convolution to
|
1182 |
+
# create the unshared convolution kernel.
|
1183 |
+
conv_output_shape = jax.eval_shape(
|
1184 |
+
lambda lhs, rhs: self.conv_general_dilated( # pylint: disable=g-long-lambda
|
1185 |
+
lhs=lhs,
|
1186 |
+
rhs=rhs,
|
1187 |
+
window_strides=strides,
|
1188 |
+
padding=padding_lax,
|
1189 |
+
dimension_numbers=dimension_numbers,
|
1190 |
+
lhs_dilation=input_dilation,
|
1191 |
+
rhs_dilation=kernel_dilation,
|
1192 |
+
),
|
1193 |
+
inputs,
|
1194 |
+
jax.ShapedArray(kernel_size + (in_features, self.features), inputs.dtype),
|
1195 |
+
).shape
|
1196 |
+
|
1197 |
+
# One (unshared) convolutional kernel per each pixel in the output.
|
1198 |
+
kernel_shape = conv_output_shape[1:-1] + (np.prod(kernel_size) * in_features, self.features)
|
1199 |
+
|
1200 |
+
if self.mask is not None and self.mask.shape != kernel_shape:
|
1201 |
+
raise ValueError(
|
1202 |
+
"Mask needs to have the same shape as weights. " f"Shapes are: {self.mask.shape}, {kernel_shape}"
|
1203 |
+
)
|
1204 |
+
|
1205 |
+
kernel = param_with_axes(
|
1206 |
+
"kernel",
|
1207 |
+
self.kernel_init,
|
1208 |
+
kernel_shape,
|
1209 |
+
self.params_dtype,
|
1210 |
+
axes=self.kernel_axes,
|
1211 |
+
)
|
1212 |
+
|
1213 |
+
if self.mask is not None:
|
1214 |
+
kernel *= self.mask
|
1215 |
+
|
1216 |
+
if self.use_bias:
|
1217 |
+
if self.shared_weights:
|
1218 |
+
# One bias weight per output channel, shared between pixels.
|
1219 |
+
bias_shape = (self.features,)
|
1220 |
+
else:
|
1221 |
+
# One bias weight per output entry, unshared betwen pixels.
|
1222 |
+
bias_shape = conv_output_shape[1:]
|
1223 |
+
|
1224 |
+
bias = param_with_axes(
|
1225 |
+
"bias",
|
1226 |
+
self.bias_init,
|
1227 |
+
bias_shape,
|
1228 |
+
self.params_dtype,
|
1229 |
+
axes=(self.kernel_axes[-1],),
|
1230 |
+
)
|
1231 |
+
else:
|
1232 |
+
bias = None
|
1233 |
+
|
1234 |
+
inputs, kernel, bias = promote_dtype(inputs, kernel, bias, dtype=self.dtype)
|
1235 |
+
if self.shared_weights:
|
1236 |
+
y = self.conv_general_dilated(
|
1237 |
+
inputs,
|
1238 |
+
kernel,
|
1239 |
+
strides,
|
1240 |
+
padding_lax,
|
1241 |
+
lhs_dilation=input_dilation,
|
1242 |
+
rhs_dilation=kernel_dilation,
|
1243 |
+
dimension_numbers=dimension_numbers,
|
1244 |
+
feature_group_count=self.feature_group_count,
|
1245 |
+
precision=self.precision,
|
1246 |
+
)
|
1247 |
+
else:
|
1248 |
+
y = lax.conv_general_dilated_local(
|
1249 |
+
lhs=inputs,
|
1250 |
+
rhs=kernel,
|
1251 |
+
window_strides=strides,
|
1252 |
+
padding=padding_lax,
|
1253 |
+
filter_shape=kernel_size,
|
1254 |
+
lhs_dilation=input_dilation,
|
1255 |
+
rhs_dilation=kernel_dilation,
|
1256 |
+
dimension_numbers=dimension_numbers,
|
1257 |
+
precision=self.precision,
|
1258 |
+
)
|
1259 |
+
|
1260 |
+
if self.use_bias:
|
1261 |
+
bias = bias.reshape((1,) * (y.ndim - bias.ndim) + bias.shape)
|
1262 |
+
y += bias
|
1263 |
+
|
1264 |
+
if num_batch_dimensions != 1:
|
1265 |
+
output_shape = input_batch_shape + y.shape[1:]
|
1266 |
+
y = jnp.reshape(y, output_shape)
|
1267 |
+
return y
|
1268 |
+
|
1269 |
+
|
1270 |
+
class Conv(_Conv):
|
1271 |
+
"""Convolution Module wrapping `lax.conv_general_dilated`.
|
1272 |
+
|
1273 |
+
Attributes:
|
1274 |
+
features: number of convolution filters.
|
1275 |
+
kernel_size: shape of the convolutional kernel. For 1D convolution,
|
1276 |
+
the kernel size can be passed as an integer. For all other cases, it must
|
1277 |
+
be a sequence of integers.
|
1278 |
+
strides: an integer or a sequence of `n` integers, representing the
|
1279 |
+
inter-window strides (default: 1).
|
1280 |
+
padding: either the string `'SAME'`, the string `'VALID'`, the string
|
1281 |
+
`'CIRCULAR'` (periodic boundary conditions), or a sequence of `n` `(low,
|
1282 |
+
high)` integer pairs that give the padding to apply before and after each
|
1283 |
+
spatial dimension. A single int is interpeted as applying the same padding
|
1284 |
+
in all dims and passign a single int in a sequence causes the same padding
|
1285 |
+
to be used on both sides. `'CAUSAL'` padding for a 1D convolution will
|
1286 |
+
left-pad the convolution axis, resulting in same-sized output.
|
1287 |
+
input_dilation: an integer or a sequence of `n` integers, giving the
|
1288 |
+
dilation factor to apply in each spatial dimension of `inputs`
|
1289 |
+
(default: 1). Convolution with input dilation `d` is equivalent to
|
1290 |
+
transposed convolution with stride `d`.
|
1291 |
+
kernel_dilation: an integer or a sequence of `n` integers, giving the
|
1292 |
+
dilation factor to apply in each spatial dimension of the convolution
|
1293 |
+
kernel (default: 1). Convolution with kernel dilation
|
1294 |
+
is also known as 'atrous convolution'.
|
1295 |
+
feature_group_count: integer, default 1. If specified divides the input
|
1296 |
+
features into groups.
|
1297 |
+
use_bias: whether to add a bias to the output (default: True).
|
1298 |
+
mask: Optional mask for the weights during masked convolution. The mask must
|
1299 |
+
be the same shape as the convolution weight matrix.
|
1300 |
+
dtype: the dtype of the computation (default: infer from input and params).
|
1301 |
+
params_dtype: the dtype passed to parameter initializers (default: float32).
|
1302 |
+
precision: numerical precision of the computation see `jax.lax.Precision`
|
1303 |
+
for details.
|
1304 |
+
kernel_init: initializer for the convolutional kernel.
|
1305 |
+
bias_init: initializer for the bias.
|
1306 |
+
"""
|
1307 |
+
|
1308 |
+
@property
|
1309 |
+
def shared_weights(self) -> bool:
|
1310 |
+
return True
|
whisper_jax/modeling_flax_whisper.py
ADDED
@@ -0,0 +1,1686 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# coding=utf-8
|
2 |
+
# Copyright 2023 The OpenAI Authors and The HuggingFace Inc. team. All rights reserved.
|
3 |
+
#
|
4 |
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
5 |
+
# you may not use this file except in compliance with the License.
|
6 |
+
# You may obtain a copy of the License at
|
7 |
+
#
|
8 |
+
# http://www.apache.org/licenses/LICENSE-2.0
|
9 |
+
#
|
10 |
+
# Unless required by applicable law or agreed to in writing, software
|
11 |
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
12 |
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
13 |
+
# See the License for the specific language governing permissions and
|
14 |
+
# limitations under the License.
|
15 |
+
""" Flax whisper model."""
|
16 |
+
|
17 |
+
import random
|
18 |
+
from functools import partial
|
19 |
+
from typing import Optional, Tuple
|
20 |
+
|
21 |
+
import flax.linen as nn
|
22 |
+
import jax
|
23 |
+
import jax.numpy as jnp
|
24 |
+
from flax.core.frozen_dict import FrozenDict, freeze, unfreeze
|
25 |
+
from flax.linen import combine_masks, make_causal_mask
|
26 |
+
from flax.linen.attention import dot_product_attention_weights
|
27 |
+
from flax.traverse_util import flatten_dict, unflatten_dict
|
28 |
+
from jax import lax
|
29 |
+
from jax.random import PRNGKey
|
30 |
+
from transformers import WhisperConfig
|
31 |
+
from transformers.generation.flax_logits_process import (
|
32 |
+
FlaxLogitsProcessor,
|
33 |
+
FlaxLogitsProcessorList,
|
34 |
+
FlaxWhisperTimeStampLogitsProcessor,
|
35 |
+
)
|
36 |
+
from transformers.modeling_flax_outputs import (
|
37 |
+
FlaxBaseModelOutput,
|
38 |
+
FlaxBaseModelOutputWithPastAndCrossAttentions,
|
39 |
+
FlaxCausalLMOutputWithCrossAttentions,
|
40 |
+
FlaxSeq2SeqLMOutput,
|
41 |
+
FlaxSeq2SeqModelOutput,
|
42 |
+
)
|
43 |
+
from transformers.modeling_flax_utils import (
|
44 |
+
ACT2FN,
|
45 |
+
FlaxPreTrainedModel,
|
46 |
+
append_call_sample_docstring,
|
47 |
+
append_replace_return_docstrings,
|
48 |
+
overwrite_call_docstring,
|
49 |
+
)
|
50 |
+
from transformers.utils import (
|
51 |
+
add_start_docstrings,
|
52 |
+
add_start_docstrings_to_model_forward,
|
53 |
+
logging,
|
54 |
+
replace_return_docstrings,
|
55 |
+
)
|
56 |
+
|
57 |
+
from whisper_jax import layers
|
58 |
+
from whisper_jax.layers import with_sharding_constraint
|
59 |
+
|
60 |
+
|
61 |
+
logger = logging.get_logger(__name__)
|
62 |
+
|
63 |
+
|
64 |
+
_CHECKPOINT_FOR_DOC = "openai/whisper-tiny"
|
65 |
+
_CONFIG_FOR_DOC = "WhisperConfig"
|
66 |
+
|
67 |
+
|
68 |
+
WHISPER_START_DOCSTRING = r"""
|
69 |
+
This model inherits from [`FlaxPreTrainedModel`]. Check the superclass documentation for the generic methods the
|
70 |
+
library implements for all its models (such as downloading or saving, resizing the input embeddings, pruning heads
|
71 |
+
etc.) This model is also a Flax Linen
|
72 |
+
[flax.nn.Module](https://flax.readthedocs.io/en/latest/_autosummary/flax.nn.module.html) subclass. Use it as a
|
73 |
+
regular Flax Module and refer to the Flax documentation for all matter related to general usage and behavior.
|
74 |
+
Finally, this model supports inherent JAX features such as:
|
75 |
+
- [Just-In-Time (JIT) compilation](https://jax.readthedocs.io/en/latest/jax.html#just-in-time-compilation-jit)
|
76 |
+
- [Automatic Differentiation](https://jax.readthedocs.io/en/latest/jax.html#automatic-differentiation)
|
77 |
+
- [Vectorization](https://jax.readthedocs.io/en/latest/jax.html#vectorization-vmap)
|
78 |
+
- [Parallelization](https://jax.readthedocs.io/en/latest/jax.html#parallelization-pmap)
|
79 |
+
|
80 |
+
Parameters:
|
81 |
+
config ([`WhisperConfig`]): Model configuration class with all the parameters of the model.
|
82 |
+
Initializing with a config file does not load the weights associated with the model, only the
|
83 |
+
configuration. Check out the [`~FlaxPreTrainedModel.from_pretrained`] method to load the model weights.
|
84 |
+
dtype (`jax.numpy.dtype`, *optional*, defaults to `jax.numpy.float32`):
|
85 |
+
The data type of the computation. Can be one of `jax.numpy.float32`, `jax.numpy.float16` (on GPUs) and
|
86 |
+
`jax.numpy.bfloat16` (on TPUs). This can be used to enable mixed-precision training or half-precision
|
87 |
+
inference on GPUs or TPUs. If specified all the computation will be performed with the given `dtype`.
|
88 |
+
**Note that this only specifies the dtype of the computation and does not influence the dtype of model
|
89 |
+
parameters.** If you wish to change the dtype of the model parameters, see [`~FlaxPreTrainedModel.to_fp16`]
|
90 |
+
and [`~FlaxPreTrainedModel.to_bf16`].
|
91 |
+
"""
|
92 |
+
|
93 |
+
WHISPER_INPUTS_DOCSTRING = r"""
|
94 |
+
Args:
|
95 |
+
input_features (`numpy.ndarray` of shape `(batch_size, feature_size, sequence_length)`):
|
96 |
+
Float values mel features extracted from the raw speech waveform. Raw speech waveform can be obtained by
|
97 |
+
loading a `.flac` or `.wav` audio file into an array of type `List[float]` or a `numpy.ndarray`, *e.g.* via
|
98 |
+
the soundfile library (`pip install soundfile`). To prepare the array into `input_features`, the
|
99 |
+
[`WhisperFeatureExtractor`] should be used for extracting the features, padding and conversion into a
|
100 |
+
tensor of type `numpy.ndarray`. See [`~WhisperFeatureExtractor.__call__`]
|
101 |
+
attention_mask (`numpy.ndarray` of shape `(batch_size, sequence_length)`, *optional*):
|
102 |
+
Whisper does not support masking of the `input_features`, this argument is preserved for compatibility, but
|
103 |
+
is not used. By default the silence in the input log mel spectrogram are ignored.
|
104 |
+
decoder_input_ids (`numpy.ndarray` of shape `(batch_size, target_sequence_length)`, *optional*):
|
105 |
+
Indices of decoder input sequence tokens in the vocabulary. Indices can be obtained using
|
106 |
+
[`WhisperTokenizer`]. See [`PreTrainedTokenizer.encode`] and [`PreTrainedTokenizer.__call__`] for details.
|
107 |
+
[What are decoder input IDs?](../glossary#decoder-input-ids) Whisper uses the `decoder_start_token_id` as
|
108 |
+
the starting token for `decoder_input_ids` generation.
|
109 |
+
decoder_attention_mask (`numpy.ndarray` of shape `(batch_size, target_sequence_length)`, *optional*):
|
110 |
+
Default behavior: generate a tensor that ignores pad tokens in `decoder_input_ids`. Causal mask will also
|
111 |
+
be used by default. If you want to change padding behavior, you should modify to your needs. See diagram 1
|
112 |
+
in [the paper](https://arxiv.org/abs/1910.13461) for more information on the default strategy.
|
113 |
+
position_ids (`numpy.ndarray` of shape `(batch_size, sequence_length)`, *optional*):
|
114 |
+
Whisper does not use `position_ids` in the encoder as `input_features` is always the same size and doesn't
|
115 |
+
use masking, but this argument is preserved for compatibility. By default the silence in the input log mel
|
116 |
+
spectrogram are ignored.
|
117 |
+
decoder_position_ids (`numpy.ndarray` of shape `(batch_size, sequence_length)`, *optional*):
|
118 |
+
Indices of positions of each decoder input sequence tokens in the position embeddings. Selected in the
|
119 |
+
range `[0, config.max_position_embeddings - 1]`.
|
120 |
+
output_attentions (`bool`, *optional*):
|
121 |
+
Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
|
122 |
+
tensors for more detail.
|
123 |
+
output_hidden_states (`bool`, *optional*):
|
124 |
+
Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
|
125 |
+
more detail.
|
126 |
+
return_dict (`bool`, *optional*):
|
127 |
+
Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
|
128 |
+
"""
|
129 |
+
|
130 |
+
WHISPER_ENCODE_INPUTS_DOCSTRING = r"""
|
131 |
+
Args:
|
132 |
+
input_features (`numpy.ndarray` of shape `(batch_size, feature_size, sequence_length)`):
|
133 |
+
Float values mel features extracted from the raw speech waveform. Raw speech waveform can be obtained by
|
134 |
+
loading a `.flac` or `.wav` audio file into an array of type `List[float]` or a `numpy.ndarray`, *e.g.* via
|
135 |
+
the soundfile library (`pip install soundfile`). To prepare the array into `input_features`, the
|
136 |
+
[`WhisperFeatureExtractor`] should be used for extracting the mel features, padding and conversion into a
|
137 |
+
tensor of type `numpy.ndarray`. See [`~WhisperFeatureExtractor.__call__`].
|
138 |
+
attention_mask (`numpy.ndarray` of shape `(batch_size, sequence_length)`, *optional*):
|
139 |
+
Whisper does not support masking of the `input_features`, this argument is preserved for compatibility, but
|
140 |
+
is not used. By default the silence in the input log mel spectrogram are ignored.
|
141 |
+
output_attentions (`bool`, *optional*):
|
142 |
+
Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
|
143 |
+
tensors for more detail.
|
144 |
+
output_hidden_states (`bool`, *optional*):
|
145 |
+
Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
|
146 |
+
more detail.
|
147 |
+
return_dict (`bool`, *optional*):
|
148 |
+
Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
|
149 |
+
"""
|
150 |
+
|
151 |
+
WHISPER_DECODE_INPUTS_DOCSTRING = r"""
|
152 |
+
Args:
|
153 |
+
decoder_input_ids (`numpy.ndarray` of shape `(batch_size, target_sequence_length)`):
|
154 |
+
Indices of decoder input sequence tokens in the vocabulary. Indices can be obtained using
|
155 |
+
[`WhisperTokenizer`]. See [`PreTrainedTokenizer.encode`] and [`PreTrainedTokenizer.__call__`] for details.
|
156 |
+
[What are decoder input IDs?](../glossary#decoder-input-ids)
|
157 |
+
encoder_outputs (`tuple(tuple(numpy.ndarray)`):
|
158 |
+
Tuple consists of (`last_hidden_state`, *optional*: `hidden_states`, *optional*: `attentions`)
|
159 |
+
`last_hidden_state` of shape `(batch_size, sequence_length, hidden_size)`, *optional*) is a sequence of
|
160 |
+
hidden-states at the output of the last layer of the encoder. Used in the cross-attention of the decoder.
|
161 |
+
encoder_attention_mask (`numpy.ndarray` of shape `(batch_size, sequence_length)`, *optional*):
|
162 |
+
Whisper does not support masking of the `input_features`, this argument is preserved for compatibility,
|
163 |
+
but it is not used. By default the silence in the input log mel spectrogram are ignored.
|
164 |
+
decoder_attention_mask (`numpy.ndarray` of shape `(batch_size, target_sequence_length)`, *optional*):
|
165 |
+
Default behavior: generate a tensor that ignores pad tokens in `decoder_input_ids`. Causal mask will also
|
166 |
+
be used by default. If you want to change padding behavior, you should modify to your needs. See diagram 1
|
167 |
+
in [the paper](https://arxiv.org/abs/1910.13461) for more information on the default strategy.
|
168 |
+
decoder_position_ids (`numpy.ndarray` of shape `(batch_size, sequence_length)`, *optional*):
|
169 |
+
Indices of positions of each decoder input sequence tokens in the position embeddings. Selected in the
|
170 |
+
range `[0, config.max_position_embeddings - 1]`.
|
171 |
+
past_key_values (`Dict[str, numpy.ndarray]`, *optional*, returned by `init_cache` or when passing previous `past_key_values`):
|
172 |
+
Dictionary of pre-computed hidden-states (key and values in the attention blocks) that can be used for fast
|
173 |
+
auto-regressive decoding. Pre-computed key and value hidden-states are of shape *[batch_size, max_length]*.
|
174 |
+
output_attentions (`bool`, *optional*):
|
175 |
+
Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
|
176 |
+
tensors for more detail.
|
177 |
+
output_hidden_states (`bool`, *optional*):
|
178 |
+
Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
|
179 |
+
more detail.
|
180 |
+
return_dict (`bool`, *optional*):
|
181 |
+
Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
|
182 |
+
"""
|
183 |
+
|
184 |
+
|
185 |
+
class FlaxStaticForceTokensLogitsProcessor(FlaxLogitsProcessor):
|
186 |
+
r"""
|
187 |
+
[`FlaxLogitsProcessor`] that takes a list of pairs of integers which indicates a mapping from generation indices to
|
188 |
+
token indices that will be forced before sampling. The processor will set their log probs to 0 and all other tokens
|
189 |
+
to `-inf` so that they are sampled at their corresponding index. This is a static version of the `transformers` logit
|
190 |
+
processor [`FlaxForceTokensLogitsProcessor`] that is compatible with sharded forced tokens.
|
191 |
+
|
192 |
+
Args:
|
193 |
+
force_token_map (`list`):
|
194 |
+
Map giving token ids and indices where they will be forced to be sampled.
|
195 |
+
"""
|
196 |
+
|
197 |
+
def __init__(self, force_token_map):
|
198 |
+
# The generic `transformers` logit processor builds `force_token_array` as a dictionary - this is not a valid
|
199 |
+
# JAX type, and so we switch to using a JAX array instead
|
200 |
+
force_token_map = jnp.array(force_token_map)
|
201 |
+
# Converts the array of format [[index, token]] containing the tokens to be forced to an array, where the
|
202 |
+
# index of the array corresponds to the index of the token to be forced. For XLA compatibility,
|
203 |
+
# indexes without forced tokens will have a negative value. Note that the last token we ever need to force in
|
204 |
+
# Whisper is at position 3, so we only construct an array up to this index. The native version constructs a tensor
|
205 |
+
# dynamically according to the length of the `force_token_map`. Array shapes need to be concrete for XLA compatibility,
|
206 |
+
# so this is not permitted here.
|
207 |
+
force_token_array = jnp.ones(3, dtype=jnp.int32) * -1
|
208 |
+
for index, token in force_token_map:
|
209 |
+
force_token_array = force_token_array.at[index].set(token)
|
210 |
+
self.force_token_array = jnp.int32(force_token_array)
|
211 |
+
|
212 |
+
def __call__(self, input_ids: jnp.ndarray, scores: jnp.ndarray, cur_len: int) -> jnp.ndarray:
|
213 |
+
def _force_token(generation_idx):
|
214 |
+
batch_size = scores.shape[0]
|
215 |
+
current_token = self.force_token_array[generation_idx]
|
216 |
+
|
217 |
+
new_scores = jnp.ones_like(scores, dtype=scores.dtype) * -float("inf")
|
218 |
+
updates = jnp.zeros((batch_size, 1), dtype=scores.dtype)
|
219 |
+
new_scores = lax.dynamic_update_slice(new_scores, updates, (0, current_token))
|
220 |
+
return new_scores
|
221 |
+
|
222 |
+
scores = lax.cond(
|
223 |
+
cur_len >= self.force_token_array.shape[0],
|
224 |
+
# If the current length is geq than the length of force_token_array, the processor does nothing.
|
225 |
+
lambda: scores,
|
226 |
+
# Otherwise, it may force a certain token.
|
227 |
+
lambda: lax.cond(
|
228 |
+
self.force_token_array[cur_len] >= 0,
|
229 |
+
# Only valid (positive) tokens are forced
|
230 |
+
lambda: _force_token(cur_len),
|
231 |
+
# Otherwise, the processor does nothing.
|
232 |
+
lambda: scores,
|
233 |
+
),
|
234 |
+
)
|
235 |
+
return scores
|
236 |
+
|
237 |
+
|
238 |
+
class FlaxWhisperAttention(nn.Module):
|
239 |
+
config: WhisperConfig
|
240 |
+
embed_dim: int
|
241 |
+
num_heads: int
|
242 |
+
dropout: float = 0.0
|
243 |
+
causal: bool = False
|
244 |
+
bias: bool = True
|
245 |
+
dtype: jnp.dtype = jnp.float32
|
246 |
+
params_dtype: jnp.dtype = jnp.float32
|
247 |
+
|
248 |
+
def setup(self) -> None:
|
249 |
+
self.head_dim = self.embed_dim // self.num_heads
|
250 |
+
if self.head_dim * self.num_heads != self.embed_dim:
|
251 |
+
raise ValueError(
|
252 |
+
f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim}"
|
253 |
+
f" and `num_heads`: {self.num_heads})."
|
254 |
+
)
|
255 |
+
|
256 |
+
dense = partial(
|
257 |
+
layers.DenseGeneral,
|
258 |
+
self.embed_dim,
|
259 |
+
axis=-1,
|
260 |
+
dtype=self.dtype,
|
261 |
+
params_dtype=self.params_dtype,
|
262 |
+
kernel_axes=("embed", "joined_kv"),
|
263 |
+
)
|
264 |
+
|
265 |
+
self.q_proj = dense(use_bias=self.bias)
|
266 |
+
self.k_proj = dense(use_bias=False)
|
267 |
+
self.v_proj = dense(use_bias=self.bias)
|
268 |
+
|
269 |
+
self.out_proj = layers.DenseGeneral(
|
270 |
+
self.embed_dim,
|
271 |
+
axis=-1,
|
272 |
+
dtype=self.dtype,
|
273 |
+
params_dtype=self.params_dtype,
|
274 |
+
kernel_axes=("joined_kv", "embed"),
|
275 |
+
use_bias=self.bias,
|
276 |
+
)
|
277 |
+
|
278 |
+
if self.causal:
|
279 |
+
self.causal_mask = make_causal_mask(
|
280 |
+
jnp.ones((1, self.config.max_target_positions), dtype="bool"), dtype="bool"
|
281 |
+
)
|
282 |
+
|
283 |
+
def __call__(
|
284 |
+
self,
|
285 |
+
hidden_states: jnp.ndarray,
|
286 |
+
key_value_states: Optional[jnp.ndarray] = None,
|
287 |
+
attention_mask: Optional[jnp.ndarray] = None,
|
288 |
+
init_cache: bool = False,
|
289 |
+
deterministic: bool = True,
|
290 |
+
) -> Tuple[jnp.ndarray]:
|
291 |
+
is_cross_attention = key_value_states is not None
|
292 |
+
batch_size = hidden_states.shape[0]
|
293 |
+
|
294 |
+
query_states = self.q_proj(hidden_states)
|
295 |
+
|
296 |
+
if is_cross_attention:
|
297 |
+
key_states = self.k_proj(key_value_states)
|
298 |
+
value_states = self.v_proj(key_value_states)
|
299 |
+
else:
|
300 |
+
key_states = self.k_proj(hidden_states)
|
301 |
+
value_states = self.v_proj(hidden_states)
|
302 |
+
|
303 |
+
query_states = self._split_heads(query_states)
|
304 |
+
key_states = self._split_heads(key_states)
|
305 |
+
value_states = self._split_heads(value_states)
|
306 |
+
|
307 |
+
query_states = with_sharding_constraint(query_states, ("batch", "length", "heads", "kv"))
|
308 |
+
key_states = with_sharding_constraint(key_states, ("batch", "length", "heads", "kv"))
|
309 |
+
value_states = with_sharding_constraint(value_states, ("batch", "length", "heads", "kv"))
|
310 |
+
|
311 |
+
if self.causal:
|
312 |
+
query_length, key_length = query_states.shape[1], key_states.shape[1]
|
313 |
+
if self.has_variable("cache", "cached_key"):
|
314 |
+
mask_shift = self.variables["cache"]["cache_index"]
|
315 |
+
# max_length of cached_key is last dim
|
316 |
+
max_decoder_length = self.variables["cache"]["cached_key"].shape[-1]
|
317 |
+
causal_mask = lax.dynamic_slice(
|
318 |
+
self.causal_mask,
|
319 |
+
(0, 0, mask_shift, 0),
|
320 |
+
(1, 1, query_length, max_decoder_length),
|
321 |
+
)
|
322 |
+
else:
|
323 |
+
causal_mask = self.causal_mask[:, :, :query_length, :key_length]
|
324 |
+
causal_mask = jnp.broadcast_to(causal_mask, (batch_size,) + causal_mask.shape[1:])
|
325 |
+
|
326 |
+
# combine masks if needed
|
327 |
+
if attention_mask is not None and self.causal:
|
328 |
+
attention_mask = jnp.broadcast_to(jnp.expand_dims(attention_mask, axis=(-3, -2)), causal_mask.shape)
|
329 |
+
attention_mask = combine_masks(attention_mask, causal_mask)
|
330 |
+
elif self.causal:
|
331 |
+
attention_mask = causal_mask
|
332 |
+
elif attention_mask is not None:
|
333 |
+
attention_mask = jnp.expand_dims(attention_mask, axis=(-3, -2))
|
334 |
+
|
335 |
+
# During fast autoregressive decoding, we feed one position at a time,
|
336 |
+
# and cache the keys and values step by step.
|
337 |
+
|
338 |
+
if self.causal and (self.has_variable("cache", "cached_key") or init_cache):
|
339 |
+
key_states, value_states, attention_mask = self._concatenate_to_cache(
|
340 |
+
key_states, value_states, query_states, attention_mask
|
341 |
+
)
|
342 |
+
|
343 |
+
# Convert the boolean attention mask to an attention bias.
|
344 |
+
if attention_mask is not None:
|
345 |
+
# attention mask in the form of attention bias
|
346 |
+
attention_bias = lax.select(
|
347 |
+
attention_mask > 0,
|
348 |
+
jnp.full(attention_mask.shape, 0.0).astype(self.dtype),
|
349 |
+
jnp.full(attention_mask.shape, jnp.finfo(self.dtype).min).astype(self.dtype),
|
350 |
+
)
|
351 |
+
else:
|
352 |
+
attention_bias = None
|
353 |
+
|
354 |
+
dropout_rng = None
|
355 |
+
if not deterministic and self.dropout > 0.0:
|
356 |
+
dropout_rng = self.make_rng("dropout")
|
357 |
+
|
358 |
+
attn_weights = dot_product_attention_weights(
|
359 |
+
query_states,
|
360 |
+
key_states,
|
361 |
+
bias=attention_bias,
|
362 |
+
dropout_rng=dropout_rng,
|
363 |
+
dropout_rate=self.dropout,
|
364 |
+
broadcast_dropout=True,
|
365 |
+
deterministic=deterministic,
|
366 |
+
dtype=self.dtype,
|
367 |
+
precision=None,
|
368 |
+
)
|
369 |
+
|
370 |
+
attn_output = jnp.einsum("...hqk,...khd->...qhd", attn_weights, value_states)
|
371 |
+
attn_output = self._merge_heads(attn_output)
|
372 |
+
attn_output = self.out_proj(attn_output)
|
373 |
+
|
374 |
+
return attn_output, attn_weights
|
375 |
+
|
376 |
+
def _split_heads(self, hidden_state) -> jnp.ndarray:
|
377 |
+
return hidden_state.reshape(hidden_state.shape[:2] + (self.num_heads, self.head_dim))
|
378 |
+
|
379 |
+
def _merge_heads(self, hidden_state) -> jnp.ndarray:
|
380 |
+
return hidden_state.reshape(hidden_state.shape[:2] + (self.embed_dim,))
|
381 |
+
|
382 |
+
@nn.compact
|
383 |
+
def _concatenate_to_cache(self, key, value, query, attention_mask):
|
384 |
+
# The following code is largely copied from: https://github.com/google-research/t5x/blob/63d9addf628c6d8c547a407a32095fcb527bb20b/t5x/examples/scalable_t5/layers.py#L280-L284
|
385 |
+
is_initialized = self.has_variable("cache", "cached_key")
|
386 |
+
|
387 |
+
# The key and value have dimension [batch_size, seq_length, num_heads, head_dim],
|
388 |
+
# but we cache them as [batch_size, num_heads, head_dim, seq_length] as a TPU
|
389 |
+
# fusion optimization. This also enables the "scatter via one-hot
|
390 |
+
# broadcast" trick, which means we do a one-hot broadcast instead of a
|
391 |
+
# scatter/gather operations, resulting in a 3-4x speedup in practice.
|
392 |
+
def swap_dims(x):
|
393 |
+
return x[:-3] + tuple(x[i] for i in [-2, -1, -3])
|
394 |
+
|
395 |
+
cached_key = self.variable("cache", "cached_key", jnp.zeros, swap_dims(key.shape), key.dtype)
|
396 |
+
cached_value = self.variable("cache", "cached_value", jnp.zeros, swap_dims(value.shape), value.dtype)
|
397 |
+
cache_index = self.variable("cache", "cache_index", lambda: jnp.array(0, dtype=jnp.int32))
|
398 |
+
|
399 |
+
if is_initialized:
|
400 |
+
batch_size, num_heads, head_dim, seq_length = cached_key.value.shape
|
401 |
+
# During fast autoregressive decoding, we feed one position at a time,
|
402 |
+
# and cache the keys and values step by step.
|
403 |
+
# Sanity shape check of cached key against input query.
|
404 |
+
num_updated_cache_vectors = query.shape[1]
|
405 |
+
expected_shape = (batch_size, 1, num_heads, head_dim)
|
406 |
+
if num_updated_cache_vectors == 1 and expected_shape != query.shape:
|
407 |
+
raise ValueError(
|
408 |
+
f"Autoregressive cache shape error, expected query shape {expected_shape} instead got {query.shape}"
|
409 |
+
)
|
410 |
+
|
411 |
+
# Create a OHE of the current index. NOTE: the index is increased below.
|
412 |
+
cur_index = cache_index.value
|
413 |
+
|
414 |
+
# In order to update the key, value caches with the current key and
|
415 |
+
# value, we move the seq_length axis to the back, similar to what we did for
|
416 |
+
# the cached ones above.
|
417 |
+
# Note these are currently the key and value of a single position, since
|
418 |
+
# we feed one position at a time.
|
419 |
+
one_token_key = jnp.moveaxis(key, -3, -1)
|
420 |
+
one_token_value = jnp.moveaxis(value, -3, -1)
|
421 |
+
|
422 |
+
# Update key, value caches with our new 1d spatial slices.
|
423 |
+
# We implement an efficient scatter into the cache via one-hot
|
424 |
+
# broadcast and addition.
|
425 |
+
if num_updated_cache_vectors > 1:
|
426 |
+
indices = jnp.eye(num_updated_cache_vectors, seq_length)[None, None]
|
427 |
+
key = cached_key.value + jnp.matmul(one_token_key, indices)
|
428 |
+
value = cached_value.value + jnp.matmul(one_token_value, indices)
|
429 |
+
else:
|
430 |
+
one_hot_indices = jax.nn.one_hot(cur_index, seq_length, dtype=key.dtype)
|
431 |
+
key = cached_key.value + one_token_key * one_hot_indices
|
432 |
+
value = cached_value.value + one_token_value * one_hot_indices
|
433 |
+
|
434 |
+
cached_key.value = key
|
435 |
+
cached_value.value = value
|
436 |
+
cache_index.value = cache_index.value + num_updated_cache_vectors
|
437 |
+
|
438 |
+
# Move the keys and values back to their original shapes.
|
439 |
+
key = jnp.moveaxis(key, -1, -3)
|
440 |
+
value = jnp.moveaxis(value, -1, -3)
|
441 |
+
|
442 |
+
# causal mask for cached decoder self-attention: our single query position should only
|
443 |
+
# attend to those key positions that have already been generated and cached, not the
|
444 |
+
# remaining zero elements.
|
445 |
+
pad_mask = jnp.broadcast_to(
|
446 |
+
jnp.arange(seq_length) < cur_index + num_updated_cache_vectors,
|
447 |
+
(batch_size,) + (1, num_updated_cache_vectors, seq_length),
|
448 |
+
)
|
449 |
+
attention_mask = combine_masks(pad_mask, attention_mask)
|
450 |
+
|
451 |
+
return key, value, attention_mask
|
452 |
+
|
453 |
+
|
454 |
+
# Copied from transformers.models.mbart.modeling_flax_mbart.FlaxMBartEncoderLayer with MBart->Whisper
|
455 |
+
class FlaxWhisperEncoderLayer(nn.Module):
|
456 |
+
config: WhisperConfig
|
457 |
+
dtype: jnp.dtype = jnp.float32
|
458 |
+
params_dtype: jnp.dtype = jnp.float32
|
459 |
+
|
460 |
+
def setup(self) -> None:
|
461 |
+
self.embed_dim = self.config.d_model
|
462 |
+
self.self_attn = FlaxWhisperAttention(
|
463 |
+
config=self.config,
|
464 |
+
embed_dim=self.embed_dim,
|
465 |
+
num_heads=self.config.encoder_attention_heads,
|
466 |
+
dropout=self.config.attention_dropout,
|
467 |
+
dtype=self.dtype,
|
468 |
+
params_dtype=self.params_dtype,
|
469 |
+
)
|
470 |
+
self.self_attn_layer_norm = layers.LayerNorm(dtype=self.dtype, epsilon=1e-05, params_dtype=self.params_dtype)
|
471 |
+
self.dropout_layer = nn.Dropout(rate=self.config.dropout)
|
472 |
+
self.activation_fn = ACT2FN[self.config.activation_function]
|
473 |
+
self.activation_dropout_layer = nn.Dropout(rate=self.config.activation_dropout)
|
474 |
+
self.fc1 = layers.DenseGeneral(
|
475 |
+
self.config.encoder_ffn_dim,
|
476 |
+
dtype=self.dtype,
|
477 |
+
params_dtype=self.params_dtype,
|
478 |
+
kernel_axes=("embed", "mlp"),
|
479 |
+
)
|
480 |
+
self.fc2 = layers.DenseGeneral(
|
481 |
+
self.embed_dim,
|
482 |
+
dtype=self.dtype,
|
483 |
+
params_dtype=self.params_dtype,
|
484 |
+
kernel_axes=("mlp", "embed"),
|
485 |
+
)
|
486 |
+
self.final_layer_norm = layers.LayerNorm(dtype=self.dtype, epsilon=1e-05, params_dtype=self.params_dtype)
|
487 |
+
|
488 |
+
def __call__(
|
489 |
+
self,
|
490 |
+
hidden_states: jnp.ndarray,
|
491 |
+
attention_mask: jnp.ndarray,
|
492 |
+
output_attentions: bool = True,
|
493 |
+
deterministic: bool = True,
|
494 |
+
) -> Tuple[jnp.ndarray]:
|
495 |
+
hidden_states = with_sharding_constraint(hidden_states, ("batch", "length", "embed"))
|
496 |
+
|
497 |
+
residual = hidden_states
|
498 |
+
|
499 |
+
layernorm_output = self.self_attn_layer_norm(hidden_states)
|
500 |
+
layernorm_output = with_sharding_constraint(layernorm_output, ("batch", "length", "embed"))
|
501 |
+
|
502 |
+
attn_output, attn_weights = self.self_attn(hidden_states=layernorm_output, attention_mask=attention_mask)
|
503 |
+
attn_output = self.dropout_layer(attn_output, deterministic=deterministic)
|
504 |
+
attn_output = residual + attn_output
|
505 |
+
attn_output = with_sharding_constraint(attn_output, ("batch", "length", "embed"))
|
506 |
+
|
507 |
+
residual = attn_output
|
508 |
+
|
509 |
+
post_layer_norm = self.final_layer_norm(attn_output)
|
510 |
+
post_layer_norm = with_sharding_constraint(post_layer_norm, ("batch", "length", "embed"))
|
511 |
+
|
512 |
+
fc1_output = self.activation_fn(self.fc1(post_layer_norm))
|
513 |
+
fc1_output = self.activation_dropout_layer(fc1_output, deterministic=deterministic)
|
514 |
+
fc1_output = with_sharding_constraint(fc1_output, ("batch", "length", "mlp"))
|
515 |
+
|
516 |
+
hidden_states = self.fc2(fc1_output)
|
517 |
+
hidden_states = self.dropout_layer(hidden_states, deterministic=deterministic)
|
518 |
+
hidden_states = residual + hidden_states
|
519 |
+
hidden_states = with_sharding_constraint(hidden_states, ("batch", "length", "embed"))
|
520 |
+
|
521 |
+
outputs = (hidden_states,)
|
522 |
+
|
523 |
+
if output_attentions:
|
524 |
+
outputs += (attn_weights,)
|
525 |
+
|
526 |
+
return outputs
|
527 |
+
|
528 |
+
|
529 |
+
# Copied from transformers.models.mbart.modeling_flax_mbart.FlaxMBartEncoderLayerCollection with MBart->Whisper
|
530 |
+
class FlaxWhisperEncoderLayerCollection(nn.Module):
|
531 |
+
config: WhisperConfig
|
532 |
+
dtype: jnp.dtype = jnp.float32 # the dtype of the computation
|
533 |
+
params_dtype: jnp.dtype = jnp.float32
|
534 |
+
|
535 |
+
def setup(self):
|
536 |
+
self.layers = [
|
537 |
+
FlaxWhisperEncoderLayer(self.config, name=str(i), dtype=self.dtype, params_dtype=self.params_dtype)
|
538 |
+
for i in range(self.config.encoder_layers)
|
539 |
+
]
|
540 |
+
self.layerdrop = self.config.encoder_layerdrop
|
541 |
+
|
542 |
+
def __call__(
|
543 |
+
self,
|
544 |
+
hidden_states,
|
545 |
+
attention_mask,
|
546 |
+
deterministic: bool = True,
|
547 |
+
output_attentions: bool = False,
|
548 |
+
output_hidden_states: bool = False,
|
549 |
+
return_dict: bool = True,
|
550 |
+
):
|
551 |
+
all_attentions = () if output_attentions else None
|
552 |
+
all_hidden_states = () if output_hidden_states else None
|
553 |
+
|
554 |
+
for encoder_layer in self.layers:
|
555 |
+
if output_hidden_states:
|
556 |
+
all_hidden_states = all_hidden_states + (hidden_states,)
|
557 |
+
# add LayerDrop (see https://arxiv.org/abs/1909.11556 for description)
|
558 |
+
dropout_probability = random.uniform(0, 1)
|
559 |
+
if not deterministic and (dropout_probability < self.layerdrop): # skip the layer
|
560 |
+
layer_outputs = (None, None)
|
561 |
+
else:
|
562 |
+
layer_outputs = encoder_layer(
|
563 |
+
hidden_states,
|
564 |
+
attention_mask,
|
565 |
+
output_attentions,
|
566 |
+
deterministic,
|
567 |
+
)
|
568 |
+
hidden_states = layer_outputs[0]
|
569 |
+
if output_attentions:
|
570 |
+
all_attentions = all_attentions + (layer_outputs[1],)
|
571 |
+
|
572 |
+
if output_hidden_states:
|
573 |
+
all_hidden_states += (hidden_states,)
|
574 |
+
|
575 |
+
outputs = (hidden_states, all_hidden_states, all_attentions)
|
576 |
+
|
577 |
+
if not return_dict:
|
578 |
+
return tuple(v for v in outputs if v is not None)
|
579 |
+
|
580 |
+
return FlaxBaseModelOutput(
|
581 |
+
last_hidden_state=hidden_states, hidden_states=all_hidden_states, attentions=all_attentions
|
582 |
+
)
|
583 |
+
|
584 |
+
|
585 |
+
# Copied from transformers.models.mbart.modeling_flax_mbart.FlaxMBartDecoderLayer with MBart->Whisper
|
586 |
+
class FlaxWhisperDecoderLayer(nn.Module):
|
587 |
+
config: WhisperConfig
|
588 |
+
dtype: jnp.dtype = jnp.float32
|
589 |
+
params_dtype: jnp.dtype = jnp.float32
|
590 |
+
|
591 |
+
def setup(self) -> None:
|
592 |
+
self.embed_dim = self.config.d_model
|
593 |
+
self.self_attn = FlaxWhisperAttention(
|
594 |
+
config=self.config,
|
595 |
+
embed_dim=self.embed_dim,
|
596 |
+
num_heads=self.config.decoder_attention_heads,
|
597 |
+
dropout=self.config.attention_dropout,
|
598 |
+
causal=True,
|
599 |
+
dtype=self.dtype,
|
600 |
+
params_dtype=self.params_dtype,
|
601 |
+
)
|
602 |
+
self.dropout_layer = nn.Dropout(rate=self.config.dropout)
|
603 |
+
self.activation_fn = ACT2FN[self.config.activation_function]
|
604 |
+
self.activation_dropout_layer = nn.Dropout(rate=self.config.activation_dropout)
|
605 |
+
|
606 |
+
self.self_attn_layer_norm = layers.LayerNorm(dtype=self.dtype, epsilon=1e-05, params_dtype=self.params_dtype)
|
607 |
+
self.encoder_attn = FlaxWhisperAttention(
|
608 |
+
config=self.config,
|
609 |
+
embed_dim=self.embed_dim,
|
610 |
+
num_heads=self.config.decoder_attention_heads,
|
611 |
+
dropout=self.config.attention_dropout,
|
612 |
+
dtype=self.dtype,
|
613 |
+
params_dtype=self.params_dtype,
|
614 |
+
)
|
615 |
+
self.encoder_attn_layer_norm = layers.LayerNorm(
|
616 |
+
dtype=self.dtype, epsilon=1e-05, params_dtype=self.params_dtype
|
617 |
+
)
|
618 |
+
self.fc1 = layers.DenseGeneral(
|
619 |
+
self.config.decoder_ffn_dim,
|
620 |
+
dtype=self.dtype,
|
621 |
+
params_dtype=self.params_dtype,
|
622 |
+
kernel_axes=("embed", "mlp"),
|
623 |
+
)
|
624 |
+
self.fc2 = layers.DenseGeneral(
|
625 |
+
self.embed_dim,
|
626 |
+
dtype=self.dtype,
|
627 |
+
params_dtype=self.params_dtype,
|
628 |
+
kernel_axes=("mlp", "embed"),
|
629 |
+
)
|
630 |
+
self.final_layer_norm = layers.LayerNorm(dtype=self.dtype, epsilon=1e-05, params_dtype=self.params_dtype)
|
631 |
+
|
632 |
+
def __call__(
|
633 |
+
self,
|
634 |
+
hidden_states: jnp.ndarray,
|
635 |
+
attention_mask: jnp.ndarray,
|
636 |
+
encoder_hidden_states: Optional[jnp.ndarray] = None,
|
637 |
+
encoder_attention_mask: Optional[jnp.ndarray] = None,
|
638 |
+
init_cache: bool = False,
|
639 |
+
output_attentions: bool = True,
|
640 |
+
deterministic: bool = True,
|
641 |
+
) -> Tuple[jnp.ndarray]:
|
642 |
+
hidden_states = with_sharding_constraint(hidden_states, ("batch", "length", "embed"))
|
643 |
+
|
644 |
+
residual = hidden_states
|
645 |
+
|
646 |
+
layer_norm_output = self.self_attn_layer_norm(hidden_states)
|
647 |
+
layer_norm_output = with_sharding_constraint(layer_norm_output, ("batch", "length", "embed"))
|
648 |
+
|
649 |
+
# Self Attention
|
650 |
+
self_attn_output, self_attn_weights = self.self_attn(
|
651 |
+
hidden_states=layer_norm_output, attention_mask=attention_mask, init_cache=init_cache
|
652 |
+
)
|
653 |
+
self_attn_output = self.dropout_layer(self_attn_output, deterministic=deterministic)
|
654 |
+
self_attn_output = residual + self_attn_output
|
655 |
+
self_attn_output = with_sharding_constraint(self_attn_output, ("batch", "length", "embed"))
|
656 |
+
|
657 |
+
# Cross-Attention Block
|
658 |
+
cross_attn_weights = None
|
659 |
+
if encoder_hidden_states is not None:
|
660 |
+
residual = self_attn_output
|
661 |
+
|
662 |
+
encoder_layer_norm_output = self.encoder_attn_layer_norm(self_attn_output)
|
663 |
+
encoder_layer_norm_output = with_sharding_constraint(
|
664 |
+
encoder_layer_norm_output, ("batch", "length", "embed")
|
665 |
+
)
|
666 |
+
|
667 |
+
cross_attn_output, cross_attn_weights = self.encoder_attn(
|
668 |
+
hidden_states=encoder_layer_norm_output,
|
669 |
+
key_value_states=encoder_hidden_states,
|
670 |
+
attention_mask=encoder_attention_mask,
|
671 |
+
)
|
672 |
+
cross_attn_output = self.dropout_layer(cross_attn_output, deterministic=deterministic)
|
673 |
+
cross_attn_output = residual + cross_attn_output
|
674 |
+
cross_attn_output = with_sharding_constraint(cross_attn_output, ("batch", "length", "embed"))
|
675 |
+
|
676 |
+
# Fully Connected
|
677 |
+
residual = cross_attn_output
|
678 |
+
|
679 |
+
post_layer_norm = self.final_layer_norm(cross_attn_output)
|
680 |
+
post_layer_norm = with_sharding_constraint(post_layer_norm, ("batch", "length", "embed"))
|
681 |
+
|
682 |
+
fc1_output = self.activation_fn(self.fc1(post_layer_norm))
|
683 |
+
fc1_output = self.activation_dropout_layer(fc1_output, deterministic=deterministic)
|
684 |
+
fc1_output = with_sharding_constraint(fc1_output, ("batch", "length", "mlp"))
|
685 |
+
|
686 |
+
hidden_states = self.fc2(fc1_output)
|
687 |
+
hidden_states = self.dropout_layer(hidden_states, deterministic=deterministic)
|
688 |
+
hidden_states = residual + hidden_states
|
689 |
+
hidden_states = with_sharding_constraint(hidden_states, ("batch", "length", "embed"))
|
690 |
+
|
691 |
+
outputs = (hidden_states,)
|
692 |
+
|
693 |
+
if output_attentions:
|
694 |
+
outputs += (self_attn_weights, cross_attn_weights)
|
695 |
+
|
696 |
+
return outputs
|
697 |
+
|
698 |
+
|
699 |
+
# Copied from transformers.models.mbart.modeling_flax_mbart.FlaxMBartDecoderLayerCollection with MBart->Whisper
|
700 |
+
class FlaxWhisperDecoderLayerCollection(nn.Module):
|
701 |
+
config: WhisperConfig
|
702 |
+
dtype: jnp.dtype = jnp.float32 # the dtype of the computation
|
703 |
+
params_dtype: jnp.dtype = jnp.float32
|
704 |
+
|
705 |
+
def setup(self):
|
706 |
+
self.layers = [
|
707 |
+
FlaxWhisperDecoderLayer(self.config, name=str(i), dtype=self.dtype, params_dtype=self.params_dtype)
|
708 |
+
for i in range(self.config.decoder_layers)
|
709 |
+
]
|
710 |
+
self.layerdrop = self.config.decoder_layerdrop
|
711 |
+
|
712 |
+
def __call__(
|
713 |
+
self,
|
714 |
+
hidden_states,
|
715 |
+
attention_mask,
|
716 |
+
encoder_hidden_states: Optional[jnp.ndarray] = None,
|
717 |
+
encoder_attention_mask: Optional[jnp.ndarray] = None,
|
718 |
+
deterministic: bool = True,
|
719 |
+
init_cache: bool = False,
|
720 |
+
output_attentions: bool = False,
|
721 |
+
output_hidden_states: bool = False,
|
722 |
+
return_dict: bool = True,
|
723 |
+
):
|
724 |
+
# decoder layers
|
725 |
+
all_hidden_states = () if output_hidden_states else None
|
726 |
+
all_self_attns = () if output_attentions else None
|
727 |
+
all_cross_attentions = () if (output_attentions and encoder_hidden_states is not None) else None
|
728 |
+
|
729 |
+
for decoder_layer in self.layers:
|
730 |
+
if output_hidden_states:
|
731 |
+
all_hidden_states += (hidden_states,)
|
732 |
+
# add LayerDrop (see https://arxiv.org/abs/1909.11556 for description)
|
733 |
+
dropout_probability = random.uniform(0, 1)
|
734 |
+
if not deterministic and (dropout_probability < self.layerdrop):
|
735 |
+
layer_outputs = (None, None, None)
|
736 |
+
else:
|
737 |
+
layer_outputs = decoder_layer(
|
738 |
+
hidden_states,
|
739 |
+
attention_mask=attention_mask,
|
740 |
+
encoder_hidden_states=encoder_hidden_states,
|
741 |
+
encoder_attention_mask=encoder_attention_mask,
|
742 |
+
init_cache=init_cache,
|
743 |
+
output_attentions=output_attentions,
|
744 |
+
deterministic=deterministic,
|
745 |
+
)
|
746 |
+
|
747 |
+
hidden_states = layer_outputs[0]
|
748 |
+
if output_attentions:
|
749 |
+
all_self_attns += (layer_outputs[1],)
|
750 |
+
|
751 |
+
if encoder_hidden_states is not None:
|
752 |
+
all_cross_attentions += (layer_outputs[2],)
|
753 |
+
|
754 |
+
# add hidden states from the last decoder layer
|
755 |
+
if output_hidden_states:
|
756 |
+
all_hidden_states += (hidden_states,)
|
757 |
+
|
758 |
+
outputs = [hidden_states, all_hidden_states, all_self_attns, all_cross_attentions]
|
759 |
+
|
760 |
+
if not return_dict:
|
761 |
+
return tuple(v for v in outputs if v is not None)
|
762 |
+
|
763 |
+
return FlaxBaseModelOutputWithPastAndCrossAttentions(
|
764 |
+
last_hidden_state=hidden_states,
|
765 |
+
hidden_states=all_hidden_states,
|
766 |
+
attentions=all_self_attns,
|
767 |
+
cross_attentions=all_cross_attentions,
|
768 |
+
)
|
769 |
+
|
770 |
+
|
771 |
+
class FlaxWhisperEncoder(nn.Module):
|
772 |
+
config: WhisperConfig
|
773 |
+
dtype: jnp.dtype = jnp.float32
|
774 |
+
params_dtype: jnp.dtype = jnp.float32
|
775 |
+
|
776 |
+
def setup(self) -> None:
|
777 |
+
self.conv1 = layers.Conv(
|
778 |
+
self.config.d_model,
|
779 |
+
kernel_size=(3,),
|
780 |
+
padding=1,
|
781 |
+
dtype=self.dtype,
|
782 |
+
params_dtype=self.params_dtype,
|
783 |
+
kernel_axes=("channels", "num_mel", "embed"),
|
784 |
+
)
|
785 |
+
self.conv2 = layers.Conv(
|
786 |
+
self.config.d_model,
|
787 |
+
kernel_size=(3,),
|
788 |
+
strides=2,
|
789 |
+
padding=1,
|
790 |
+
dtype=self.dtype,
|
791 |
+
params_dtype=self.params_dtype,
|
792 |
+
kernel_axes=("channels", "embed", "num_mel"),
|
793 |
+
)
|
794 |
+
|
795 |
+
self.dropout_layer = nn.Dropout(rate=self.config.dropout)
|
796 |
+
|
797 |
+
self.layers = FlaxWhisperEncoderLayerCollection(
|
798 |
+
self.config,
|
799 |
+
dtype=self.dtype,
|
800 |
+
params_dtype=self.params_dtype,
|
801 |
+
)
|
802 |
+
self.embed_positions = layers.Embed(
|
803 |
+
self.config.max_source_positions, self.config.d_model, dtype=self.dtype, params_dtype=self.params_dtype
|
804 |
+
)
|
805 |
+
|
806 |
+
self.layer_norm = layers.LayerNorm(dtype=self.dtype, epsilon=1e-05, params_dtype=self.params_dtype)
|
807 |
+
|
808 |
+
def __call__(
|
809 |
+
self,
|
810 |
+
input_features: jnp.ndarray,
|
811 |
+
output_attentions: bool = False,
|
812 |
+
output_hidden_states: bool = False,
|
813 |
+
return_dict: bool = True,
|
814 |
+
deterministic: bool = True,
|
815 |
+
) -> Tuple[jnp.ndarray]:
|
816 |
+
if input_features.shape[1:] != (self.config.num_mel_bins, self.config.max_source_positions * 2):
|
817 |
+
raise ValueError(
|
818 |
+
"input_features.shape[1:], must be equal to (self.config.num_mel_bins,"
|
819 |
+
f" self.config.max_source_positions * 2) (got {input_features.shape[1:]}, but should be"
|
820 |
+
f" ({self.config.num_mel_bins}, {self.config.max_source_positions * 2}))"
|
821 |
+
)
|
822 |
+
|
823 |
+
input_features = input_features.transpose(0, 2, 1)
|
824 |
+
hidden_states = jax.nn.gelu(self.conv1(input_features), approximate=False)
|
825 |
+
hidden_states = with_sharding_constraint(hidden_states, ("batch", "embed", "num_mel"))
|
826 |
+
hidden_states = jax.nn.gelu(self.conv2(hidden_states), approximate=False)
|
827 |
+
hidden_states = with_sharding_constraint(hidden_states, ("batch", "length", "embed"))
|
828 |
+
|
829 |
+
embed_positions = self.embed_positions(jnp.arange(self.config.max_source_positions))
|
830 |
+
hidden_states = hidden_states + embed_positions
|
831 |
+
|
832 |
+
hidden_states = self.dropout_layer(hidden_states, deterministic=deterministic)
|
833 |
+
|
834 |
+
outputs = self.layers(
|
835 |
+
hidden_states,
|
836 |
+
attention_mask=None,
|
837 |
+
deterministic=deterministic,
|
838 |
+
output_attentions=output_attentions,
|
839 |
+
output_hidden_states=output_hidden_states,
|
840 |
+
return_dict=return_dict,
|
841 |
+
)
|
842 |
+
|
843 |
+
last_hidden_states = outputs[0]
|
844 |
+
last_hidden_states = self.layer_norm(last_hidden_states)
|
845 |
+
|
846 |
+
# update the last element in `hidden_states` after applying `layernorm` above
|
847 |
+
hidden_states = None
|
848 |
+
if output_hidden_states:
|
849 |
+
hidden_states = outputs[1]
|
850 |
+
hidden_states = hidden_states[:-1] + (last_hidden_states,)
|
851 |
+
|
852 |
+
if not return_dict:
|
853 |
+
outputs = (last_hidden_states, hidden_states) + (outputs[2:] if output_hidden_states else outputs[1:])
|
854 |
+
return tuple(v for v in outputs if v is not None)
|
855 |
+
|
856 |
+
return FlaxBaseModelOutput(
|
857 |
+
last_hidden_state=last_hidden_states,
|
858 |
+
hidden_states=hidden_states,
|
859 |
+
attentions=outputs.attentions,
|
860 |
+
)
|
861 |
+
|
862 |
+
|
863 |
+
class FlaxWhisperDecoder(nn.Module):
|
864 |
+
config: WhisperConfig
|
865 |
+
dtype: jnp.dtype = jnp.float32
|
866 |
+
params_dtype: jnp.dtype = jnp.float32
|
867 |
+
|
868 |
+
def setup(self) -> None:
|
869 |
+
self.embed_tokens = layers.Embed(
|
870 |
+
self.config.vocab_size, self.config.d_model, dtype=self.dtype, params_dtype=self.params_dtype
|
871 |
+
)
|
872 |
+
self.embed_positions = layers.Embed(
|
873 |
+
self.config.max_target_positions, self.config.d_model, dtype=self.dtype, params_dtype=self.params_dtype
|
874 |
+
)
|
875 |
+
|
876 |
+
self.layers = FlaxWhisperDecoderLayerCollection(self.config, dtype=self.dtype, params_dtype=self.params_dtype)
|
877 |
+
|
878 |
+
self.dropout_layer = nn.Dropout(rate=self.config.dropout)
|
879 |
+
|
880 |
+
self.layer_norm = layers.LayerNorm(dtype=self.dtype, epsilon=1e-5, params_dtype=self.params_dtype)
|
881 |
+
|
882 |
+
def __call__(
|
883 |
+
self,
|
884 |
+
input_ids: jnp.ndarray,
|
885 |
+
attention_mask: jnp.ndarray,
|
886 |
+
position_ids: jnp.ndarray,
|
887 |
+
encoder_hidden_states: Optional[jnp.ndarray] = None,
|
888 |
+
init_cache: bool = False,
|
889 |
+
output_attentions: bool = False,
|
890 |
+
output_hidden_states: bool = False,
|
891 |
+
return_dict: bool = True,
|
892 |
+
deterministic: bool = True,
|
893 |
+
) -> Tuple[jnp.ndarray]:
|
894 |
+
input_embeds = self.embed_tokens(input_ids)
|
895 |
+
position_embeds = self.embed_positions(position_ids)
|
896 |
+
|
897 |
+
hidden_states = input_embeds + position_embeds
|
898 |
+
hidden_states = self.dropout_layer(hidden_states, deterministic=deterministic)
|
899 |
+
|
900 |
+
outputs = self.layers(
|
901 |
+
hidden_states,
|
902 |
+
attention_mask=attention_mask,
|
903 |
+
encoder_hidden_states=encoder_hidden_states,
|
904 |
+
deterministic=deterministic,
|
905 |
+
init_cache=init_cache,
|
906 |
+
output_attentions=output_attentions,
|
907 |
+
output_hidden_states=output_hidden_states,
|
908 |
+
return_dict=return_dict,
|
909 |
+
)
|
910 |
+
|
911 |
+
last_hidden_states = outputs[0]
|
912 |
+
last_hidden_states = self.layer_norm(last_hidden_states)
|
913 |
+
|
914 |
+
# update the last element in `hidden_states` after applying `layernorm` above
|
915 |
+
hidden_states = None
|
916 |
+
if output_hidden_states:
|
917 |
+
hidden_states = outputs[1]
|
918 |
+
hidden_states = hidden_states[:-1] + (last_hidden_states,)
|
919 |
+
|
920 |
+
if not return_dict:
|
921 |
+
outputs = (last_hidden_states, hidden_states) + (outputs[2:] if output_hidden_states else outputs[1:])
|
922 |
+
return tuple(v for v in outputs if v is not None)
|
923 |
+
|
924 |
+
return FlaxBaseModelOutputWithPastAndCrossAttentions(
|
925 |
+
last_hidden_state=last_hidden_states,
|
926 |
+
hidden_states=hidden_states,
|
927 |
+
attentions=outputs.attentions,
|
928 |
+
cross_attentions=outputs.cross_attentions,
|
929 |
+
)
|
930 |
+
|
931 |
+
|
932 |
+
class FlaxWhisperModule(nn.Module):
|
933 |
+
config: WhisperConfig
|
934 |
+
dtype: jnp.dtype = jnp.float32
|
935 |
+
params_dtype: jnp.dtype = jnp.float32
|
936 |
+
|
937 |
+
def setup(self) -> None:
|
938 |
+
self.encoder = FlaxWhisperEncoder(self.config, dtype=self.dtype, params_dtype=self.params_dtype)
|
939 |
+
self.decoder = FlaxWhisperDecoder(self.config, dtype=self.dtype, params_dtype=self.params_dtype)
|
940 |
+
|
941 |
+
def __call__(
|
942 |
+
self,
|
943 |
+
input_features: jnp.ndarray,
|
944 |
+
decoder_input_ids: jnp.ndarray,
|
945 |
+
decoder_attention_mask: jnp.ndarray,
|
946 |
+
decoder_position_ids: jnp.ndarray,
|
947 |
+
output_attentions: bool = False,
|
948 |
+
output_hidden_states: bool = False,
|
949 |
+
return_dict: bool = True,
|
950 |
+
deterministic: bool = True,
|
951 |
+
):
|
952 |
+
encoder_outputs = self.encoder(
|
953 |
+
input_features,
|
954 |
+
output_attentions=output_attentions,
|
955 |
+
output_hidden_states=output_hidden_states,
|
956 |
+
return_dict=return_dict,
|
957 |
+
deterministic=deterministic,
|
958 |
+
)
|
959 |
+
|
960 |
+
decoder_outputs = self.decoder(
|
961 |
+
input_ids=decoder_input_ids,
|
962 |
+
attention_mask=decoder_attention_mask,
|
963 |
+
position_ids=decoder_position_ids,
|
964 |
+
encoder_hidden_states=encoder_outputs[0],
|
965 |
+
output_attentions=output_attentions,
|
966 |
+
output_hidden_states=output_hidden_states,
|
967 |
+
return_dict=return_dict,
|
968 |
+
deterministic=deterministic,
|
969 |
+
)
|
970 |
+
|
971 |
+
if not return_dict:
|
972 |
+
return decoder_outputs + encoder_outputs
|
973 |
+
|
974 |
+
return FlaxSeq2SeqModelOutput(
|
975 |
+
last_hidden_state=decoder_outputs.last_hidden_state,
|
976 |
+
decoder_hidden_states=decoder_outputs.hidden_states,
|
977 |
+
decoder_attentions=decoder_outputs.attentions,
|
978 |
+
cross_attentions=decoder_outputs.cross_attentions,
|
979 |
+
encoder_last_hidden_state=encoder_outputs.last_hidden_state,
|
980 |
+
encoder_hidden_states=encoder_outputs.hidden_states,
|
981 |
+
encoder_attentions=encoder_outputs.attentions,
|
982 |
+
)
|
983 |
+
|
984 |
+
def _get_encoder_module(self):
|
985 |
+
return self.encoder
|
986 |
+
|
987 |
+
def _get_decoder_module(self):
|
988 |
+
return self.decoder
|
989 |
+
|
990 |
+
|
991 |
+
class FlaxWhisperPreTrainedModel(FlaxPreTrainedModel):
|
992 |
+
config_class = WhisperConfig
|
993 |
+
base_model_prefix: str = "model"
|
994 |
+
main_input_name = "input_features"
|
995 |
+
module_class: nn.Module = None
|
996 |
+
|
997 |
+
def __init__(
|
998 |
+
self,
|
999 |
+
config: WhisperConfig,
|
1000 |
+
input_shape: Tuple[int, int, int] = None,
|
1001 |
+
seed: int = 0,
|
1002 |
+
dtype: jnp.dtype = jnp.float32,
|
1003 |
+
params_dtype: jnp.dtype = jnp.float32,
|
1004 |
+
_do_init: bool = True,
|
1005 |
+
**kwargs,
|
1006 |
+
):
|
1007 |
+
if input_shape is None:
|
1008 |
+
input_shape = (1, config.num_mel_bins, 2 * config.max_source_positions)
|
1009 |
+
|
1010 |
+
module = self.module_class(config=config, dtype=dtype, params_dtype=params_dtype, **kwargs)
|
1011 |
+
super().__init__(config, module, input_shape=input_shape, seed=seed, dtype=dtype, _do_init=_do_init)
|
1012 |
+
|
1013 |
+
def init_weights(self, rng: jax.random.PRNGKey, input_shape: Tuple, params: FrozenDict = None) -> FrozenDict:
|
1014 |
+
# init input tensors
|
1015 |
+
input_features = jnp.zeros(input_shape, dtype="f4")
|
1016 |
+
input_features = input_features.at[(..., -1)].set(self.config.eos_token_id)
|
1017 |
+
|
1018 |
+
decoder_input_ids = jnp.zeros((input_shape[0], 1), dtype="i4")
|
1019 |
+
decoder_attention_mask = jnp.ones_like(decoder_input_ids)
|
1020 |
+
|
1021 |
+
batch_size, sequence_length = decoder_input_ids.shape
|
1022 |
+
decoder_position_ids = jnp.broadcast_to(jnp.arange(sequence_length)[None, :], (batch_size, sequence_length))
|
1023 |
+
|
1024 |
+
params_rng, dropout_rng = jax.random.split(rng)
|
1025 |
+
rngs = {"params": params_rng, "dropout": dropout_rng}
|
1026 |
+
|
1027 |
+
random_params = self.module.init(
|
1028 |
+
rngs,
|
1029 |
+
input_features=input_features,
|
1030 |
+
decoder_input_ids=decoder_input_ids,
|
1031 |
+
decoder_attention_mask=decoder_attention_mask,
|
1032 |
+
decoder_position_ids=decoder_position_ids,
|
1033 |
+
)["params"]
|
1034 |
+
|
1035 |
+
if params is not None:
|
1036 |
+
random_params = flatten_dict(unfreeze(random_params))
|
1037 |
+
params = flatten_dict(unfreeze(params))
|
1038 |
+
for missing_key in self._missing_keys:
|
1039 |
+
params[missing_key] = random_params[missing_key]
|
1040 |
+
self._missing_keys = set()
|
1041 |
+
return freeze(unflatten_dict(params))
|
1042 |
+
else:
|
1043 |
+
return random_params
|
1044 |
+
|
1045 |
+
# Copied from transformers.models.bart.modeling_flax_bart.FlaxBartPreTrainedModel.init_cache with Bart->Whisper
|
1046 |
+
def init_cache(self, batch_size, max_length, encoder_outputs):
|
1047 |
+
r"""
|
1048 |
+
Args:
|
1049 |
+
batch_size (`int`):
|
1050 |
+
batch_size used for fast auto-regressive decoding. Defines the batch size of the initialized cache.
|
1051 |
+
max_length (`int`):
|
1052 |
+
maximum possible length for auto-regressive decoding. Defines the sequence length of the initialized
|
1053 |
+
cache.
|
1054 |
+
encoder_outputs (`Union[FlaxBaseModelOutput, tuple(tuple(jnp.ndarray)]`):
|
1055 |
+
`encoder_outputs` consists of (`last_hidden_state`, *optional*: `hidden_states`, *optional*:
|
1056 |
+
`attentions`). `last_hidden_state` of shape `(batch_size, sequence_length, hidden_size)`, *optional*)
|
1057 |
+
is a sequence of hidden-states at the output of the last layer of the encoder. Used in the
|
1058 |
+
cross-attention of the decoder.
|
1059 |
+
"""
|
1060 |
+
# init input variables to retrieve cache
|
1061 |
+
decoder_input_ids = jnp.ones((batch_size, max_length), dtype="i4")
|
1062 |
+
decoder_attention_mask = jnp.ones_like(decoder_input_ids)
|
1063 |
+
decoder_position_ids = jnp.broadcast_to(
|
1064 |
+
jnp.arange(jnp.atleast_2d(decoder_input_ids).shape[-1]), decoder_input_ids.shape
|
1065 |
+
)
|
1066 |
+
|
1067 |
+
def _decoder_forward(module, decoder_input_ids, decoder_attention_mask, decoder_position_ids, **kwargs):
|
1068 |
+
decoder_module = module._get_decoder_module()
|
1069 |
+
return decoder_module(
|
1070 |
+
decoder_input_ids,
|
1071 |
+
decoder_attention_mask,
|
1072 |
+
decoder_position_ids,
|
1073 |
+
**kwargs,
|
1074 |
+
)
|
1075 |
+
|
1076 |
+
init_variables = self.module.init(
|
1077 |
+
jax.random.PRNGKey(0),
|
1078 |
+
decoder_input_ids=decoder_input_ids,
|
1079 |
+
decoder_attention_mask=decoder_attention_mask,
|
1080 |
+
decoder_position_ids=decoder_position_ids,
|
1081 |
+
encoder_hidden_states=encoder_outputs[0],
|
1082 |
+
init_cache=True,
|
1083 |
+
method=_decoder_forward, # we only need to call the decoder to init the cache
|
1084 |
+
)
|
1085 |
+
return unfreeze(init_variables["cache"])
|
1086 |
+
|
1087 |
+
@add_start_docstrings(WHISPER_ENCODE_INPUTS_DOCSTRING)
|
1088 |
+
@replace_return_docstrings(output_type=FlaxBaseModelOutput, config_class=WhisperConfig)
|
1089 |
+
def encode(
|
1090 |
+
self,
|
1091 |
+
input_features: jnp.ndarray,
|
1092 |
+
attention_mask: Optional[jnp.ndarray] = None,
|
1093 |
+
output_attentions: Optional[bool] = None,
|
1094 |
+
output_hidden_states: Optional[bool] = None,
|
1095 |
+
return_dict: Optional[bool] = None,
|
1096 |
+
train: bool = False,
|
1097 |
+
params: dict = None,
|
1098 |
+
dropout_rng: PRNGKey = None,
|
1099 |
+
**kwargs,
|
1100 |
+
):
|
1101 |
+
r"""
|
1102 |
+
Returns:
|
1103 |
+
|
1104 |
+
Example:
|
1105 |
+
|
1106 |
+
```python
|
1107 |
+
>>> from transformers import WhisperProcessor, FlaxWhisperForConditionalGeneration
|
1108 |
+
>>> from datasets import load_dataset
|
1109 |
+
|
1110 |
+
>>> processor = WhisperProcessor.from_pretrained("openai/whisper-tiny.en")
|
1111 |
+
>>> model = FlaxWhisperForConditionalGeneration.from_pretrained("openai/whisper-tiny.en", from_pt=True)
|
1112 |
+
>>> ds = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation")
|
1113 |
+
>>> inputs = processor(ds[0]["audio"]["array"], return_tensors="np")
|
1114 |
+
>>> input_features = inputs.input_features
|
1115 |
+
>>> encoder_outputs = model.encode(input_features=input_features)
|
1116 |
+
```"""
|
1117 |
+
output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
|
1118 |
+
output_hidden_states = (
|
1119 |
+
output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
|
1120 |
+
)
|
1121 |
+
return_dict = return_dict if return_dict is not None else self.config.return_dict
|
1122 |
+
|
1123 |
+
# Handle any PRNG if needed
|
1124 |
+
rngs = {}
|
1125 |
+
if dropout_rng is not None:
|
1126 |
+
rngs["dropout"] = dropout_rng
|
1127 |
+
|
1128 |
+
def _encoder_forward(module, input_features, **kwargs):
|
1129 |
+
encode_module = module._get_encoder_module()
|
1130 |
+
return encode_module(input_features, **kwargs)
|
1131 |
+
|
1132 |
+
return self.module.apply(
|
1133 |
+
{"params": params or self.params},
|
1134 |
+
input_features=jnp.array(input_features, dtype="f4"),
|
1135 |
+
output_attentions=output_attentions,
|
1136 |
+
output_hidden_states=output_hidden_states,
|
1137 |
+
return_dict=return_dict,
|
1138 |
+
deterministic=not train,
|
1139 |
+
rngs=rngs,
|
1140 |
+
method=_encoder_forward,
|
1141 |
+
)
|
1142 |
+
|
1143 |
+
@add_start_docstrings(WHISPER_DECODE_INPUTS_DOCSTRING)
|
1144 |
+
@replace_return_docstrings(output_type=FlaxBaseModelOutputWithPastAndCrossAttentions, config_class=WhisperConfig)
|
1145 |
+
def decode(
|
1146 |
+
self,
|
1147 |
+
decoder_input_ids,
|
1148 |
+
encoder_outputs,
|
1149 |
+
encoder_attention_mask: Optional[jnp.ndarray] = None,
|
1150 |
+
decoder_attention_mask: Optional[jnp.ndarray] = None,
|
1151 |
+
decoder_position_ids: Optional[jnp.ndarray] = None,
|
1152 |
+
past_key_values: dict = None,
|
1153 |
+
output_attentions: Optional[bool] = None,
|
1154 |
+
output_hidden_states: Optional[bool] = None,
|
1155 |
+
return_dict: Optional[bool] = None,
|
1156 |
+
train: bool = False,
|
1157 |
+
params: dict = None,
|
1158 |
+
dropout_rng: PRNGKey = None,
|
1159 |
+
):
|
1160 |
+
r"""
|
1161 |
+
Returns:
|
1162 |
+
|
1163 |
+
Example:
|
1164 |
+
|
1165 |
+
```python
|
1166 |
+
>>> from transformers import WhisperProcessor, FlaxWhisperForConditionalGeneration
|
1167 |
+
>>> from datasets import load_dataset
|
1168 |
+
|
1169 |
+
>>> processor = WhisperProcessor.from_pretrained("openai/whisper-tiny.en")
|
1170 |
+
>>> model = FlaxWhisperForConditionalGeneration.from_pretrained("openai/whisper-tiny.en", from_pt=True)
|
1171 |
+
>>> ds = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation")
|
1172 |
+
>>> inputs = processor(ds[0]["audio"]["array"], return_tensors="np")
|
1173 |
+
>>> input_features = inputs.input_features
|
1174 |
+
>>> encoder_outputs = model.encode(input_features=input_features)
|
1175 |
+
>>> decoder_start_token_id = model.config.decoder_start_token_id
|
1176 |
+
|
1177 |
+
>>> decoder_input_ids = jnp.ones((inputs.input_ids.shape[0], 1), dtype="i4") * decoder_start_token_id
|
1178 |
+
|
1179 |
+
>>> outputs = model.decode(decoder_input_ids, encoder_outputs)
|
1180 |
+
>>> last_decoder_hidden_states = outputs.last_hidden_state
|
1181 |
+
```"""
|
1182 |
+
|
1183 |
+
output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
|
1184 |
+
output_hidden_states = (
|
1185 |
+
output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
|
1186 |
+
)
|
1187 |
+
return_dict = return_dict if return_dict is not None else self.config.return_dict
|
1188 |
+
|
1189 |
+
encoder_hidden_states = encoder_outputs[0]
|
1190 |
+
|
1191 |
+
batch_size, sequence_length = decoder_input_ids.shape
|
1192 |
+
if decoder_position_ids is None:
|
1193 |
+
if past_key_values is not None:
|
1194 |
+
raise ValueError("Make sure to provide `decoder_position_ids` when passing `past_key_values`.")
|
1195 |
+
|
1196 |
+
if decoder_attention_mask is not None:
|
1197 |
+
decoder_position_ids = (decoder_attention_mask.cumsum(-1) * decoder_attention_mask) - 1
|
1198 |
+
else:
|
1199 |
+
decoder_position_ids = jnp.broadcast_to(
|
1200 |
+
jnp.arange(sequence_length)[None, :], (batch_size, sequence_length)
|
1201 |
+
)
|
1202 |
+
|
1203 |
+
if decoder_attention_mask is None:
|
1204 |
+
decoder_attention_mask = jnp.ones((batch_size, sequence_length))
|
1205 |
+
|
1206 |
+
# Handle any PRNG if needed
|
1207 |
+
rngs = {}
|
1208 |
+
if dropout_rng is not None:
|
1209 |
+
rngs["dropout"] = dropout_rng
|
1210 |
+
|
1211 |
+
inputs = {"params": params or self.params}
|
1212 |
+
|
1213 |
+
# if past_key_values are passed then cache is already initialized a private flag init_cache has to be
|
1214 |
+
# passed down to ensure cache is used. It has to be made sure that cache is marked as mutable so that
|
1215 |
+
# it can be changed by FlaxWhisperAttention module
|
1216 |
+
if past_key_values:
|
1217 |
+
inputs["cache"] = past_key_values
|
1218 |
+
mutable = ["cache"]
|
1219 |
+
else:
|
1220 |
+
mutable = False
|
1221 |
+
|
1222 |
+
def _decoder_forward(module, decoder_input_ids, decoder_attention_mask, decoder_position_ids, **kwargs):
|
1223 |
+
decoder_module = module._get_decoder_module()
|
1224 |
+
return decoder_module(
|
1225 |
+
input_ids=decoder_input_ids,
|
1226 |
+
attention_mask=decoder_attention_mask,
|
1227 |
+
position_ids=decoder_position_ids,
|
1228 |
+
**kwargs,
|
1229 |
+
)
|
1230 |
+
|
1231 |
+
outputs = self.module.apply(
|
1232 |
+
inputs,
|
1233 |
+
decoder_input_ids=jnp.array(decoder_input_ids, dtype="i4"),
|
1234 |
+
decoder_attention_mask=jnp.array(decoder_attention_mask, dtype="i4"),
|
1235 |
+
decoder_position_ids=jnp.array(decoder_position_ids, dtype="i4"),
|
1236 |
+
encoder_hidden_states=encoder_hidden_states,
|
1237 |
+
output_attentions=output_attentions,
|
1238 |
+
output_hidden_states=output_hidden_states,
|
1239 |
+
return_dict=return_dict,
|
1240 |
+
deterministic=not train,
|
1241 |
+
rngs=rngs,
|
1242 |
+
mutable=mutable,
|
1243 |
+
method=_decoder_forward,
|
1244 |
+
)
|
1245 |
+
|
1246 |
+
# add updated cache to model output
|
1247 |
+
if past_key_values is not None and return_dict:
|
1248 |
+
outputs, past = outputs
|
1249 |
+
outputs["past_key_values"] = unfreeze(past["cache"])
|
1250 |
+
return outputs
|
1251 |
+
elif past_key_values is not None and not return_dict:
|
1252 |
+
outputs, past = outputs
|
1253 |
+
outputs = outputs[:1] + (unfreeze(past["cache"]),) + outputs[1:]
|
1254 |
+
|
1255 |
+
return outputs
|
1256 |
+
|
1257 |
+
@add_start_docstrings_to_model_forward(WHISPER_INPUTS_DOCSTRING)
|
1258 |
+
def __call__(
|
1259 |
+
self,
|
1260 |
+
input_features: jnp.ndarray,
|
1261 |
+
decoder_input_ids: jnp.ndarray,
|
1262 |
+
attention_mask: Optional[jnp.ndarray] = None,
|
1263 |
+
decoder_attention_mask: Optional[jnp.ndarray] = None,
|
1264 |
+
position_ids: Optional[jnp.ndarray] = None,
|
1265 |
+
decoder_position_ids: Optional[jnp.ndarray] = None,
|
1266 |
+
output_attentions: Optional[bool] = None,
|
1267 |
+
output_hidden_states: Optional[bool] = None,
|
1268 |
+
return_dict: Optional[bool] = None,
|
1269 |
+
train: bool = False,
|
1270 |
+
params: dict = None,
|
1271 |
+
dropout_rng: PRNGKey = None,
|
1272 |
+
):
|
1273 |
+
output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
|
1274 |
+
output_hidden_states = (
|
1275 |
+
output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
|
1276 |
+
)
|
1277 |
+
return_dict = return_dict if return_dict is not None else self.config.return_dict
|
1278 |
+
|
1279 |
+
# prepare decoder inputs
|
1280 |
+
if decoder_position_ids is None:
|
1281 |
+
if decoder_attention_mask is not None:
|
1282 |
+
decoder_position_ids = (decoder_attention_mask.cumsum(-1) * decoder_attention_mask) - 1
|
1283 |
+
else:
|
1284 |
+
batch_size, sequence_length = decoder_input_ids.shape
|
1285 |
+
decoder_position_ids = jnp.broadcast_to(
|
1286 |
+
jnp.arange(sequence_length)[None, :], (batch_size, sequence_length)
|
1287 |
+
)
|
1288 |
+
if decoder_attention_mask is None:
|
1289 |
+
decoder_attention_mask = jnp.ones_like(decoder_input_ids)
|
1290 |
+
|
1291 |
+
# Handle any PRNG if needed
|
1292 |
+
rngs = {"dropout": dropout_rng} if dropout_rng is not None else {}
|
1293 |
+
|
1294 |
+
return self.module.apply(
|
1295 |
+
{"params": params or self.params},
|
1296 |
+
input_features=jnp.array(input_features, dtype="f4"),
|
1297 |
+
decoder_input_ids=jnp.array(decoder_input_ids, dtype="i4"),
|
1298 |
+
decoder_attention_mask=jnp.array(decoder_attention_mask, dtype="i4"),
|
1299 |
+
decoder_position_ids=jnp.array(decoder_position_ids, dtype="i4"),
|
1300 |
+
output_attentions=output_attentions,
|
1301 |
+
output_hidden_states=output_hidden_states,
|
1302 |
+
return_dict=return_dict,
|
1303 |
+
deterministic=not train,
|
1304 |
+
rngs=rngs,
|
1305 |
+
)
|
1306 |
+
|
1307 |
+
|
1308 |
+
@add_start_docstrings(
|
1309 |
+
"The bare Whisper Model transformer outputting raw hidden-states without any specific head on top.",
|
1310 |
+
WHISPER_START_DOCSTRING,
|
1311 |
+
)
|
1312 |
+
class FlaxWhisperModel(FlaxWhisperPreTrainedModel):
|
1313 |
+
config: WhisperConfig
|
1314 |
+
dtype: jnp.dtype = jnp.float32 # the dtype of the computation
|
1315 |
+
params_dtype: jnp.dtype = jnp.float32
|
1316 |
+
module_class = FlaxWhisperModule
|
1317 |
+
|
1318 |
+
|
1319 |
+
append_call_sample_docstring(FlaxWhisperModel, _CHECKPOINT_FOR_DOC, FlaxSeq2SeqModelOutput, _CONFIG_FOR_DOC)
|
1320 |
+
|
1321 |
+
|
1322 |
+
class FlaxWhisperForConditionalGenerationModule(nn.Module):
|
1323 |
+
config: WhisperConfig
|
1324 |
+
dtype: jnp.dtype = jnp.float32
|
1325 |
+
params_dtype: jnp.dtype = jnp.float32
|
1326 |
+
|
1327 |
+
def setup(self) -> None:
|
1328 |
+
self.model = FlaxWhisperModule(config=self.config, dtype=self.dtype, params_dtype=self.params_dtype)
|
1329 |
+
self.lm_head = layers.DenseGeneral(
|
1330 |
+
self.config.vocab_size,
|
1331 |
+
use_bias=False,
|
1332 |
+
dtype=self.dtype,
|
1333 |
+
params_dtype=self.params_dtype,
|
1334 |
+
kernel_axes=("embed", "vocab"),
|
1335 |
+
)
|
1336 |
+
|
1337 |
+
def _get_encoder_module(self):
|
1338 |
+
return self.model.encoder
|
1339 |
+
|
1340 |
+
def _get_decoder_module(self):
|
1341 |
+
return self.model.decoder
|
1342 |
+
|
1343 |
+
def __call__(
|
1344 |
+
self,
|
1345 |
+
input_features,
|
1346 |
+
decoder_input_ids,
|
1347 |
+
decoder_attention_mask: jnp.ndarray = None,
|
1348 |
+
decoder_position_ids: jnp.ndarray = None,
|
1349 |
+
position_ids: jnp.ndarray = None,
|
1350 |
+
attention_mask: jnp.ndarray = None,
|
1351 |
+
output_attentions: bool = False,
|
1352 |
+
output_hidden_states: bool = False,
|
1353 |
+
return_dict: bool = True,
|
1354 |
+
deterministic: bool = True,
|
1355 |
+
):
|
1356 |
+
outputs = self.model(
|
1357 |
+
input_features=input_features,
|
1358 |
+
decoder_input_ids=decoder_input_ids,
|
1359 |
+
decoder_attention_mask=decoder_attention_mask,
|
1360 |
+
decoder_position_ids=decoder_position_ids,
|
1361 |
+
output_attentions=output_attentions,
|
1362 |
+
output_hidden_states=output_hidden_states,
|
1363 |
+
return_dict=return_dict,
|
1364 |
+
deterministic=deterministic,
|
1365 |
+
)
|
1366 |
+
|
1367 |
+
hidden_states = outputs[0]
|
1368 |
+
|
1369 |
+
if self.config.tie_word_embeddings:
|
1370 |
+
shared_embedding = self.model.decoder.embed_tokens.variables["params"]["embedding"]
|
1371 |
+
lm_logits = self.lm_head.apply({"params": {"kernel": shared_embedding.T}}, hidden_states)
|
1372 |
+
else:
|
1373 |
+
lm_logits = self.lm_head(hidden_states)
|
1374 |
+
|
1375 |
+
if not return_dict:
|
1376 |
+
output = (lm_logits,) + outputs[1:]
|
1377 |
+
return output
|
1378 |
+
|
1379 |
+
return FlaxSeq2SeqLMOutput(
|
1380 |
+
logits=lm_logits,
|
1381 |
+
decoder_hidden_states=outputs.decoder_hidden_states,
|
1382 |
+
decoder_attentions=outputs.decoder_attentions,
|
1383 |
+
cross_attentions=outputs.cross_attentions,
|
1384 |
+
encoder_last_hidden_state=outputs.encoder_last_hidden_state,
|
1385 |
+
encoder_hidden_states=outputs.encoder_hidden_states,
|
1386 |
+
encoder_attentions=outputs.encoder_attentions,
|
1387 |
+
)
|
1388 |
+
|
1389 |
+
|
1390 |
+
@add_start_docstrings("The Whisper Model with a language modeling head.", WHISPER_START_DOCSTRING)
|
1391 |
+
class FlaxWhisperForConditionalGeneration(FlaxWhisperPreTrainedModel):
|
1392 |
+
module_class = FlaxWhisperForConditionalGenerationModule
|
1393 |
+
dtype: jnp.dtype = jnp.float32
|
1394 |
+
params_dtype: jnp.dtype = jnp.float32
|
1395 |
+
|
1396 |
+
@add_start_docstrings(WHISPER_DECODE_INPUTS_DOCSTRING)
|
1397 |
+
@replace_return_docstrings(output_type=FlaxCausalLMOutputWithCrossAttentions, config_class=WhisperConfig)
|
1398 |
+
def decode(
|
1399 |
+
self,
|
1400 |
+
decoder_input_ids,
|
1401 |
+
encoder_outputs,
|
1402 |
+
encoder_attention_mask: Optional[jnp.ndarray] = None,
|
1403 |
+
decoder_attention_mask: Optional[jnp.ndarray] = None,
|
1404 |
+
decoder_position_ids: Optional[jnp.ndarray] = None,
|
1405 |
+
past_key_values: dict = None,
|
1406 |
+
output_attentions: Optional[bool] = None,
|
1407 |
+
output_hidden_states: Optional[bool] = None,
|
1408 |
+
return_dict: Optional[bool] = None,
|
1409 |
+
train: bool = False,
|
1410 |
+
params: dict = None,
|
1411 |
+
dropout_rng: PRNGKey = None,
|
1412 |
+
):
|
1413 |
+
r"""
|
1414 |
+
Returns:
|
1415 |
+
|
1416 |
+
Example:
|
1417 |
+
|
1418 |
+
```python
|
1419 |
+
>>> from transformers import WhisperProcessor, FlaxWhisperForConditionalGeneration
|
1420 |
+
>>> from datasets import load_dataset
|
1421 |
+
|
1422 |
+
>>> processor = WhisperProcessor.from_pretrained("openai/whisper-tiny.en")
|
1423 |
+
>>> model = FlaxWhisperForConditionalGeneration.from_pretrained("openai/whisper-tiny.en", from_pt=True)
|
1424 |
+
>>> ds = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation")
|
1425 |
+
>>> inputs = processor(ds[0]["audio"]["array"], return_tensors="np")
|
1426 |
+
>>> input_features = inputs.input_features
|
1427 |
+
>>> encoder_outputs = model.encode(input_features=input_features)
|
1428 |
+
>>> decoder_start_token_id = model.config.decoder_start_token_id
|
1429 |
+
|
1430 |
+
>>> decoder_input_ids = jnp.ones((inputs.input_ids.shape[0], 1), dtype="i4") * decoder_start_token_id
|
1431 |
+
|
1432 |
+
>>> outputs = model.decode(decoder_input_ids, encoder_outputs)
|
1433 |
+
>>> last_decoder_hidden_states = outputs.last_hidden_state
|
1434 |
+
```"""
|
1435 |
+
output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
|
1436 |
+
output_hidden_states = (
|
1437 |
+
output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
|
1438 |
+
)
|
1439 |
+
return_dict = return_dict if return_dict is not None else self.config.return_dict
|
1440 |
+
|
1441 |
+
encoder_hidden_states = encoder_outputs[0]
|
1442 |
+
|
1443 |
+
batch_size, sequence_length = decoder_input_ids.shape
|
1444 |
+
if decoder_position_ids is None:
|
1445 |
+
if past_key_values is not None:
|
1446 |
+
raise ValueError("Make sure to provide `decoder_position_ids` when passing `past_key_values`.")
|
1447 |
+
|
1448 |
+
if decoder_attention_mask is not None:
|
1449 |
+
decoder_position_ids = (decoder_attention_mask.cumsum(-1) * decoder_attention_mask) - 1
|
1450 |
+
else:
|
1451 |
+
decoder_position_ids = jnp.broadcast_to(
|
1452 |
+
jnp.arange(sequence_length)[None, :], (batch_size, sequence_length)
|
1453 |
+
)
|
1454 |
+
if decoder_attention_mask is None:
|
1455 |
+
decoder_attention_mask = jnp.ones((batch_size, sequence_length), dtype="i4")
|
1456 |
+
|
1457 |
+
# Handle any PRNG if needed
|
1458 |
+
rngs = {}
|
1459 |
+
if dropout_rng is not None:
|
1460 |
+
rngs["dropout"] = dropout_rng
|
1461 |
+
|
1462 |
+
inputs = {"params": params or self.params}
|
1463 |
+
|
1464 |
+
# if past_key_values are passed then cache is already initialized a private flag init_cache has to be
|
1465 |
+
# passed down to ensure cache is used. It has to be made sure that cache is marked as mutable so that
|
1466 |
+
# it can be changed by FlaxWhisperAttention module
|
1467 |
+
if past_key_values:
|
1468 |
+
inputs["cache"] = past_key_values
|
1469 |
+
mutable = ["cache"]
|
1470 |
+
else:
|
1471 |
+
mutable = False
|
1472 |
+
|
1473 |
+
def _decoder_forward(module, decoder_input_ids, decoder_attention_mask, decoder_position_ids, **kwargs):
|
1474 |
+
decoder_module = module._get_decoder_module()
|
1475 |
+
outputs = decoder_module(
|
1476 |
+
input_ids=decoder_input_ids,
|
1477 |
+
attention_mask=decoder_attention_mask,
|
1478 |
+
position_ids=decoder_position_ids,
|
1479 |
+
**kwargs,
|
1480 |
+
)
|
1481 |
+
hidden_states = outputs[0]
|
1482 |
+
|
1483 |
+
if self.config.tie_word_embeddings:
|
1484 |
+
shared_embedding = module.model.decoder.embed_tokens.variables["params"]["embedding"]
|
1485 |
+
lm_logits = module.lm_head.apply({"params": {"kernel": shared_embedding.T}}, hidden_states)
|
1486 |
+
else:
|
1487 |
+
lm_logits = module.lm_head(hidden_states)
|
1488 |
+
|
1489 |
+
return lm_logits, outputs
|
1490 |
+
|
1491 |
+
outputs = self.module.apply(
|
1492 |
+
inputs,
|
1493 |
+
decoder_input_ids=jnp.array(decoder_input_ids, dtype="i4"),
|
1494 |
+
decoder_attention_mask=jnp.array(decoder_attention_mask, dtype="i4"),
|
1495 |
+
decoder_position_ids=jnp.array(decoder_position_ids, dtype="i4"),
|
1496 |
+
encoder_hidden_states=encoder_hidden_states,
|
1497 |
+
output_attentions=output_attentions,
|
1498 |
+
output_hidden_states=output_hidden_states,
|
1499 |
+
return_dict=return_dict,
|
1500 |
+
deterministic=not train,
|
1501 |
+
rngs=rngs,
|
1502 |
+
mutable=mutable,
|
1503 |
+
method=_decoder_forward,
|
1504 |
+
)
|
1505 |
+
|
1506 |
+
if past_key_values is None:
|
1507 |
+
lm_logits, decoder_outputs = outputs
|
1508 |
+
else:
|
1509 |
+
(lm_logits, decoder_outputs), past = outputs
|
1510 |
+
|
1511 |
+
if return_dict:
|
1512 |
+
outputs = FlaxCausalLMOutputWithCrossAttentions(
|
1513 |
+
logits=lm_logits,
|
1514 |
+
hidden_states=decoder_outputs.hidden_states,
|
1515 |
+
attentions=decoder_outputs.attentions,
|
1516 |
+
cross_attentions=decoder_outputs.cross_attentions,
|
1517 |
+
)
|
1518 |
+
else:
|
1519 |
+
outputs = (lm_logits,) + decoder_outputs[1:]
|
1520 |
+
|
1521 |
+
# add updated cache to model output
|
1522 |
+
if past_key_values is not None and return_dict:
|
1523 |
+
outputs["past_key_values"] = unfreeze(past["cache"])
|
1524 |
+
return outputs
|
1525 |
+
elif past_key_values is not None and not return_dict:
|
1526 |
+
outputs = outputs[:1] + (unfreeze(past["cache"]),) + outputs[1:]
|
1527 |
+
|
1528 |
+
return outputs
|
1529 |
+
|
1530 |
+
def generate(
|
1531 |
+
self,
|
1532 |
+
input_features,
|
1533 |
+
generation_config=None,
|
1534 |
+
logits_processor=None,
|
1535 |
+
return_timestamps=None,
|
1536 |
+
task=None,
|
1537 |
+
language=None,
|
1538 |
+
is_multilingual=None,
|
1539 |
+
**kwargs,
|
1540 |
+
):
|
1541 |
+
if generation_config is None:
|
1542 |
+
generation_config = self.generation_config
|
1543 |
+
|
1544 |
+
if return_timestamps is not None:
|
1545 |
+
generation_config.return_timestamps = return_timestamps
|
1546 |
+
|
1547 |
+
if task is not None:
|
1548 |
+
generation_config.task = task
|
1549 |
+
|
1550 |
+
if is_multilingual is not None:
|
1551 |
+
generation_config.is_multilingual = is_multilingual
|
1552 |
+
|
1553 |
+
if language is not None:
|
1554 |
+
generation_config.language = language
|
1555 |
+
|
1556 |
+
if kwargs is not None and "decoder_input_ids" in kwargs:
|
1557 |
+
decoder_input_length = len(kwargs["decoder_input_ids"])
|
1558 |
+
else:
|
1559 |
+
decoder_input_length = 1
|
1560 |
+
|
1561 |
+
forced_decoder_ids = []
|
1562 |
+
|
1563 |
+
if hasattr(generation_config, "is_multilingual") and generation_config.is_multilingual:
|
1564 |
+
if hasattr(generation_config, "language"):
|
1565 |
+
forced_decoder_ids.append((1, generation_config.lang_to_id[generation_config.language]))
|
1566 |
+
else:
|
1567 |
+
forced_decoder_ids.append((1, None))
|
1568 |
+
|
1569 |
+
if hasattr(generation_config, "task"):
|
1570 |
+
forced_decoder_ids.append((2, generation_config.task_to_id[generation_config.task]))
|
1571 |
+
else:
|
1572 |
+
forced_decoder_ids.append((2, generation_config.task_to_id["transcribe"]))
|
1573 |
+
|
1574 |
+
if (
|
1575 |
+
hasattr(generation_config, "return_timestamps") and generation_config.return_timestamps
|
1576 |
+
) or return_timestamps:
|
1577 |
+
logits_processor = [
|
1578 |
+
FlaxWhisperTimeStampLogitsProcessor(generation_config, self.config, decoder_input_length)
|
1579 |
+
]
|
1580 |
+
else:
|
1581 |
+
if forced_decoder_ids and forced_decoder_ids[-1][0] != generation_config.no_timestamps_token_id:
|
1582 |
+
idx = forced_decoder_ids[-1][0] + 1 if forced_decoder_ids else 1
|
1583 |
+
forced_decoder_ids.append((idx, generation_config.no_timestamps_token_id))
|
1584 |
+
|
1585 |
+
if len(forced_decoder_ids) > 0:
|
1586 |
+
generation_config.forced_decoder_ids = forced_decoder_ids
|
1587 |
+
|
1588 |
+
return super().generate(
|
1589 |
+
input_features,
|
1590 |
+
generation_config,
|
1591 |
+
logits_processor=logits_processor,
|
1592 |
+
**kwargs,
|
1593 |
+
)
|
1594 |
+
|
1595 |
+
def pipeline_generate(
|
1596 |
+
self,
|
1597 |
+
input_features,
|
1598 |
+
forced_decoder_ids,
|
1599 |
+
return_timestamps=False,
|
1600 |
+
generation_config=None,
|
1601 |
+
**kwargs,
|
1602 |
+
):
|
1603 |
+
if generation_config is None:
|
1604 |
+
generation_config = self.generation_config
|
1605 |
+
|
1606 |
+
# override the generation config forced decoder ids in preference of the ones we have set
|
1607 |
+
generation_config.forced_decoder_ids = None
|
1608 |
+
|
1609 |
+
logits_processor = FlaxLogitsProcessorList()
|
1610 |
+
|
1611 |
+
logits_processor.append(FlaxStaticForceTokensLogitsProcessor(forced_decoder_ids))
|
1612 |
+
|
1613 |
+
if hasattr(generation_config, "return_timestamps") and return_timestamps:
|
1614 |
+
logits_processor.append(FlaxWhisperTimeStampLogitsProcessor(generation_config, self.config, 1))
|
1615 |
+
|
1616 |
+
return super().generate(
|
1617 |
+
input_features,
|
1618 |
+
generation_config,
|
1619 |
+
logits_processor=logits_processor,
|
1620 |
+
**kwargs,
|
1621 |
+
)
|
1622 |
+
|
1623 |
+
def prepare_inputs_for_generation(
|
1624 |
+
self,
|
1625 |
+
decoder_input_ids,
|
1626 |
+
max_length,
|
1627 |
+
attention_mask: Optional[jax.Array] = None,
|
1628 |
+
decoder_attention_mask: Optional[jax.Array] = None,
|
1629 |
+
encoder_outputs=None,
|
1630 |
+
**kwargs,
|
1631 |
+
):
|
1632 |
+
# initializing the cache
|
1633 |
+
batch_size, seq_length = decoder_input_ids.shape
|
1634 |
+
|
1635 |
+
past_key_values = self.init_cache(batch_size, max_length, encoder_outputs)
|
1636 |
+
# Note that usually one would have to put 0's in the attention_mask for x > input_ids.shape[-1] and x < cache_length.
|
1637 |
+
# But since the decoder uses a causal mask, those positions are masked anyways.
|
1638 |
+
# Thus we can create a single static attention_mask here, which is more efficient for compilation
|
1639 |
+
extended_attention_mask = jnp.ones((batch_size, max_length), dtype="i4")
|
1640 |
+
if decoder_attention_mask is not None:
|
1641 |
+
position_ids = decoder_attention_mask.cumsum(-1) - 1
|
1642 |
+
extended_attention_mask = lax.dynamic_update_slice(extended_attention_mask, decoder_attention_mask, (0, 0))
|
1643 |
+
else:
|
1644 |
+
position_ids = jnp.broadcast_to(jnp.arange(seq_length, dtype="i4")[None, :], (batch_size, seq_length))
|
1645 |
+
|
1646 |
+
return {
|
1647 |
+
"past_key_values": past_key_values,
|
1648 |
+
"encoder_outputs": encoder_outputs,
|
1649 |
+
"encoder_attention_mask": attention_mask,
|
1650 |
+
"decoder_attention_mask": extended_attention_mask,
|
1651 |
+
"decoder_position_ids": position_ids,
|
1652 |
+
}
|
1653 |
+
|
1654 |
+
def update_inputs_for_generation(self, model_outputs, model_kwargs):
|
1655 |
+
model_kwargs["past_key_values"] = model_outputs.past_key_values
|
1656 |
+
model_kwargs["decoder_position_ids"] = model_kwargs["decoder_position_ids"][:, -1:] + 1
|
1657 |
+
return model_kwargs
|
1658 |
+
|
1659 |
+
|
1660 |
+
FLAX_WHISPER_CONDITIONAL_GENERATION_DOCSTRING = r"""
|
1661 |
+
Returns:
|
1662 |
+
|
1663 |
+
Transcription example:
|
1664 |
+
|
1665 |
+
```python
|
1666 |
+
>>> from transformers import WhisperProcessor, FlaxWhisperForConditionalGeneration
|
1667 |
+
>>> from datasets import load_dataset
|
1668 |
+
|
1669 |
+
>>> processor = WhisperProcessor.from_pretrained("openai/whisper-tiny.en")
|
1670 |
+
>>> model = FlaxWhisperForConditionalGeneration.from_pretrained("openai/whisper-tiny.en", from_pt=True)
|
1671 |
+
>>> ds = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation")
|
1672 |
+
>>> inputs = processor(ds[0]["audio"]["array"], return_tensors="np")
|
1673 |
+
>>> input_features = inputs.input_features
|
1674 |
+
>>> generated_ids = model.generate(input_ids=input_features)
|
1675 |
+
>>> transcription = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
|
1676 |
+
>>> transcription
|
1677 |
+
' Mr. Quilter is the apostle of the middle classes, and we are glad to welcome his gospel.'
|
1678 |
+
```
|
1679 |
+
"""
|
1680 |
+
|
1681 |
+
overwrite_call_docstring(
|
1682 |
+
FlaxWhisperForConditionalGeneration, WHISPER_INPUTS_DOCSTRING + FLAX_WHISPER_CONDITIONAL_GENERATION_DOCSTRING
|
1683 |
+
)
|
1684 |
+
append_replace_return_docstrings(
|
1685 |
+
FlaxWhisperForConditionalGeneration, output_type=FlaxSeq2SeqLMOutput, config_class=_CONFIG_FOR_DOC
|
1686 |
+
)
|
whisper_jax/partitioner.py
ADDED
@@ -0,0 +1,939 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# coding=utf-8
|
2 |
+
# Copyright 2023 The T5X Authors and The HuggingFace Inc. team.
|
3 |
+
#
|
4 |
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
5 |
+
# you may not use this file except in compliance with the License.
|
6 |
+
# You may obtain a copy of the License at
|
7 |
+
#
|
8 |
+
# http://www.apache.org/licenses/LICENSE-2.0
|
9 |
+
#
|
10 |
+
# Unless required by applicable law or agreed to in writing, software
|
11 |
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
12 |
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
13 |
+
# See the License for the specific language governing permissions and
|
14 |
+
# limitations under the License.
|
15 |
+
|
16 |
+
"""Utilities for partitioning."""
|
17 |
+
|
18 |
+
import abc
|
19 |
+
import collections
|
20 |
+
import dataclasses
|
21 |
+
import typing
|
22 |
+
from typing import Any, Callable, Optional, Sequence, Tuple, Union
|
23 |
+
|
24 |
+
import cached_property
|
25 |
+
import jax
|
26 |
+
import numpy as np
|
27 |
+
from absl import logging
|
28 |
+
from flax import traverse_util
|
29 |
+
from flax.linen import partitioning as flax_partitioning
|
30 |
+
from jax import numpy as jnp
|
31 |
+
from jax import random
|
32 |
+
from jax.experimental import multihost_utils
|
33 |
+
from jax.experimental.mesh_utils import create_hybrid_device_mesh
|
34 |
+
from jax.experimental.pjit import pjit as jax_pjit
|
35 |
+
from jax.sharding import Mesh, PartitionSpec
|
36 |
+
|
37 |
+
|
38 |
+
JaxDevice = Any
|
39 |
+
TpuMesh = Tuple[int, int, int, int] # (x, y, z, num_cores).
|
40 |
+
OtherMesh = Tuple[int, int]
|
41 |
+
HardwareMesh = Union[TpuMesh, OtherMesh]
|
42 |
+
PyTreeDef = type(jax.tree_util.tree_structure(None))
|
43 |
+
TrainState = Any
|
44 |
+
LogicalAxisRules = Sequence[Tuple[str, Optional[str]]]
|
45 |
+
|
46 |
+
if typing.TYPE_CHECKING: # See b/163639353
|
47 |
+
cached_property = property # pylint: disable=invalid-name
|
48 |
+
else:
|
49 |
+
cached_property = cached_property.cached_property
|
50 |
+
|
51 |
+
|
52 |
+
class AxisNames(tuple):
|
53 |
+
"""Tuple of strings specifying name for each axis.
|
54 |
+
|
55 |
+
We create a separate class for this so JAX's pytree utilities can distinguish
|
56 |
+
it from a tuple that should be treated as a pytree, instead treating it as a
|
57 |
+
leaf.
|
58 |
+
"""
|
59 |
+
|
60 |
+
def __new__(cls, *names):
|
61 |
+
return tuple.__new__(AxisNames, names)
|
62 |
+
|
63 |
+
def __repr__(self):
|
64 |
+
return "AxisNames%s" % tuple.__repr__(self)
|
65 |
+
|
66 |
+
|
67 |
+
# pjit wrappers for cpu fallback.
|
68 |
+
# ----------------------------------------------------------------------------
|
69 |
+
# TODO(levskaya): This function is now no different than jax_pjit, but callers
|
70 |
+
# currently depend on `backend` argument
|
71 |
+
def pjit(
|
72 |
+
fun: Callable, # pylint: disable=g-bare-generic
|
73 |
+
in_axis_resources,
|
74 |
+
out_axis_resources,
|
75 |
+
static_argnums: Union[int, Sequence[int]] = (),
|
76 |
+
donate_argnums: Union[int, Sequence[int]] = (),
|
77 |
+
backend: Optional[str] = None,
|
78 |
+
):
|
79 |
+
"""Wrapper for pjit."""
|
80 |
+
del backend
|
81 |
+
return jax_pjit(
|
82 |
+
fun, in_axis_resources, out_axis_resources, static_argnums=static_argnums, donate_argnums=donate_argnums
|
83 |
+
)
|
84 |
+
|
85 |
+
|
86 |
+
# pjit wrappers for cpu fallback.
|
87 |
+
# -----------------------------------------------------------------------------
|
88 |
+
# TODO(levskaya): upstream this fallback behavior to jax pjit.
|
89 |
+
def pjit_with_cpu_fallback(
|
90 |
+
fun: Callable, # pylint: disable=g-bare-generic
|
91 |
+
in_axis_resources,
|
92 |
+
out_axis_resources,
|
93 |
+
static_argnums: Union[int, Sequence[int]] = (),
|
94 |
+
donate_argnums: Union[int, Sequence[int]] = (),
|
95 |
+
backend: Optional[str] = None,
|
96 |
+
):
|
97 |
+
"""Wrapper for pjit that calls normal jit on cpu."""
|
98 |
+
if jax.devices(backend)[0].platform == "cpu":
|
99 |
+
return jax.jit(fun, static_argnums=static_argnums, donate_argnums=donate_argnums)
|
100 |
+
else:
|
101 |
+
return jax_pjit(
|
102 |
+
fun, in_axis_resources, out_axis_resources, static_argnums=static_argnums, donate_argnums=donate_argnums
|
103 |
+
)
|
104 |
+
|
105 |
+
|
106 |
+
def with_sharding_constraint(x, axis_resources):
|
107 |
+
"""Wrapper for pjit with_sharding_constraint, no-op on cpu or outside pjit."""
|
108 |
+
if jax.devices()[0].platform == "cpu" or not global_mesh_defined():
|
109 |
+
return x
|
110 |
+
else:
|
111 |
+
return jax.experimental.pjit.with_sharding_constraint(x, axis_resources)
|
112 |
+
|
113 |
+
|
114 |
+
# pjit Mesh creation functions.
|
115 |
+
# -----------------------------------------------------------------------------
|
116 |
+
def bounds_from_last_device(last_device: JaxDevice) -> HardwareMesh:
|
117 |
+
"""Get the bound from the given last device."""
|
118 |
+
# Must be passed the device at the highest-coordinate corner of the
|
119 |
+
# relevant mesh, which is a requirement we know is satisfied by the last
|
120 |
+
# device in jax.devices().
|
121 |
+
if hasattr(last_device, "coords"):
|
122 |
+
x, y, z = last_device.coords
|
123 |
+
return x + 1, y + 1, z + 1, last_device.core_on_chip + 1
|
124 |
+
else:
|
125 |
+
# On non-TPU platforms, the "mesh" is hosts x devices per host in order
|
126 |
+
# to take advantage of faster within-host interconnect.
|
127 |
+
return jax.host_count(), jax.local_device_count()
|
128 |
+
|
129 |
+
|
130 |
+
def get_coords(device: JaxDevice) -> HardwareMesh:
|
131 |
+
"""Returns the coordinates of the given device."""
|
132 |
+
if hasattr(device, "coords"):
|
133 |
+
return (*device.coords, device.core_on_chip)
|
134 |
+
return (device.process_index, device.id % jax.local_device_count())
|
135 |
+
|
136 |
+
|
137 |
+
def global_mesh_defined():
|
138 |
+
"""Checks if global xmap/pjit mesh resource environment is defined."""
|
139 |
+
maps_env = jax.experimental.maps.thread_resources.env
|
140 |
+
return maps_env.physical_mesh.devices.shape != () # pylint: disable=g-explicit-bool-comparison
|
141 |
+
|
142 |
+
|
143 |
+
def get_mesh(
|
144 |
+
model_parallel_submesh: HardwareMesh,
|
145 |
+
input_devices: Sequence[JaxDevice] = (),
|
146 |
+
input_local_devices: Sequence[JaxDevice] = (),
|
147 |
+
tile_by_host_if_needed: bool = True,
|
148 |
+
backend: Optional[str] = None,
|
149 |
+
) -> Mesh:
|
150 |
+
"""Construct an xmap/pjit Mesh for the given model-parallel submesh.
|
151 |
+
|
152 |
+
The resulting mesh has two resource axes: 'model', with the provided submesh
|
153 |
+
shape, and 'data', which covers the rest of the mesh.
|
154 |
+
|
155 |
+
Args:
|
156 |
+
model_parallel_submesh: a HardwareMesh spec, namely (x,y,z,core) on TPU for
|
157 |
+
a single model-parallel replica's "tile" in the physical device mesh. The
|
158 |
+
first three elements (`x`, `y`, and `z`) should be factors of the pod
|
159 |
+
slice; e.g., if you are using df_4x8, then `x` should be a factor of 4
|
160 |
+
(one of 1, 2, 4), `y` should be a factor of 8 (one of 1, 2, 4, 8), and `z`
|
161 |
+
must be 1, because TPU v3 slices are only 2D. `z` can be >1 for TPU v4
|
162 |
+
(and maybe later TPUs) that allow 3D slices. `core` is the number of cores
|
163 |
+
to use from each TPU node. As communication is usually fastest inside the
|
164 |
+
same node, if you need a tile of more than 1 core, then
|
165 |
+
you should first increase `core`: e.g., for TPU v3, (1,1,1,2) is better
|
166 |
+
than (2,1,1,1). To pick a good spec, try a few possible values until you
|
167 |
+
get high TPU utilization.
|
168 |
+
input_devices: the devices to use, will use jax.devices() if this is not
|
169 |
+
set.
|
170 |
+
input_local_devices: the local devices to use, will use jax.local_devices()
|
171 |
+
if this is not set.
|
172 |
+
tile_by_host_if_needed: JAX currently requires that the parts of any sharded
|
173 |
+
array that are located on one host's local devices form a single
|
174 |
+
contiguous slice. A best effort will be made to achieve this without
|
175 |
+
"tiling" the device assignment over hosts (which can reduce XLA collective
|
176 |
+
performance). If this flag is True, then the device assignment will be
|
177 |
+
tiled over hosts if necessary to satisfy this constraint and create a
|
178 |
+
buildable mesh; if false, mesh construction will fail instead.
|
179 |
+
backend: get devices from the pinned backend, if specified. This is
|
180 |
+
useful for explicitly specifying the devices other than relying on
|
181 |
+
jax_platform_name.
|
182 |
+
|
183 |
+
Returns:
|
184 |
+
A xmap / pjit Mesh containing the virtual device mesh with data, model axes.
|
185 |
+
"""
|
186 |
+
input_devices = input_devices or jax.devices(backend)
|
187 |
+
input_local_devices = input_local_devices or jax.local_devices(0, backend)
|
188 |
+
# Sort input_devices based on coords, as backends might not return devices
|
189 |
+
# in order.
|
190 |
+
last_device = sorted(input_devices, key=get_coords)[-1]
|
191 |
+
last_input_local_devices = sorted(input_local_devices, key=get_coords)[-1]
|
192 |
+
logging.info(
|
193 |
+
"last device coords : %r\nlast local device coords: %r",
|
194 |
+
get_coords(last_device),
|
195 |
+
get_coords(last_input_local_devices),
|
196 |
+
)
|
197 |
+
global_hardware_mesh = bounds_from_last_device(last_device)
|
198 |
+
mesh_ndim = len(global_hardware_mesh)
|
199 |
+
local_hardware_mesh = bounds_from_last_device(last_input_local_devices)
|
200 |
+
mesh_err = (
|
201 |
+
f"each dimension of the model parallel submesh {model_parallel_submesh} "
|
202 |
+
"must be a factor of the corresponding dimension of the global device "
|
203 |
+
f"mesh {global_hardware_mesh}"
|
204 |
+
)
|
205 |
+
assert not any(g % m for g, m in zip(global_hardware_mesh, model_parallel_submesh)), mesh_err
|
206 |
+
assert not any(g % l for g, l in zip(global_hardware_mesh, local_hardware_mesh))
|
207 |
+
devices = np.empty(global_hardware_mesh, dtype=object)
|
208 |
+
for device in input_devices:
|
209 |
+
device_coords = get_coords(device)
|
210 |
+
devices[device_coords] = device
|
211 |
+
tile_by_host = tile_by_host_if_needed
|
212 |
+
if len(global_hardware_mesh) == 4:
|
213 |
+
# enable contiguous local chunks without host tiling by making Z major
|
214 |
+
global_hardware_mesh = typing.cast(Tuple[int, int, int, int], global_hardware_mesh)
|
215 |
+
model_parallel_submesh = typing.cast(Tuple[int, int, int, int], model_parallel_submesh)
|
216 |
+
gx, gy, gz, gc = global_hardware_mesh
|
217 |
+
mx, my, mz, mc = model_parallel_submesh
|
218 |
+
if (mx == gx > 1 and my == mz == 1) or (mx == 1 and my == gy > 1 and mz == gz > 1):
|
219 |
+
logging.info("ensuring YZ plane has a Z-major device order")
|
220 |
+
# YZ should be ZY
|
221 |
+
assert mc == gc, (mc, gc)
|
222 |
+
global_hardware_mesh = gx, gz, gy, gc
|
223 |
+
model_parallel_submesh = mx, mz, my, mc
|
224 |
+
devices = devices.swapaxes(1, 2)
|
225 |
+
tile_by_host = False
|
226 |
+
if (my == gy > 1 and mx == mz == 1) or (my == 1 and mx == gx > 1 and mz == gz > 1):
|
227 |
+
logging.info("ensuring XZ plane has a Z-major device order")
|
228 |
+
# XZ should be ZX
|
229 |
+
assert mc == gc, (mc, gc)
|
230 |
+
global_hardware_mesh = gz, gy, gx, gc
|
231 |
+
model_parallel_submesh = mz, my, mx, mc
|
232 |
+
devices = devices.swapaxes(0, 2)
|
233 |
+
tile_by_host = False
|
234 |
+
if tile_by_host:
|
235 |
+
logging.warning(
|
236 |
+
"Tiling device assignment mesh by hosts, which may lead to "
|
237 |
+
"reduced XLA collective performance. To avoid this, modify "
|
238 |
+
"the model parallel submesh or run with more tasks per host."
|
239 |
+
)
|
240 |
+
tile_err = (
|
241 |
+
"to tile the mesh by hosts, each dimension of the model parallel "
|
242 |
+
"submesh must be either a factor or a multiple of the corresponding "
|
243 |
+
"dimension of the per-host submesh"
|
244 |
+
)
|
245 |
+
|
246 |
+
def dh_dd_mh_md(g: int, m: int, l: int) -> Tuple[int, int, int, int]:
|
247 |
+
"""Split a global mesh dimension into four tiling components.
|
248 |
+
|
249 |
+
Args:
|
250 |
+
g: global mesh bounds dimension size
|
251 |
+
m: model-parallel submesh bounds dimension size
|
252 |
+
l: local submesh bounds dimension size
|
253 |
+
|
254 |
+
Returns:
|
255 |
+
The resulting tuple divides the dimension into the hosts component of
|
256 |
+
the data-parallel submesh, the devices component of the data-parallel
|
257 |
+
submesh, the hosts component of the model-parallel submesh, and the
|
258 |
+
devices component of the model-parallel submesh.
|
259 |
+
"""
|
260 |
+
d = g // m
|
261 |
+
if m >= l:
|
262 |
+
assert not m % l, tile_err
|
263 |
+
return (d, 1, m // l, l)
|
264 |
+
else:
|
265 |
+
assert not l % m, tile_err
|
266 |
+
return (d // (l // m), l // m, 1, m)
|
267 |
+
|
268 |
+
# e.g. [(x_data_hosts, x_data_devs, x_model_hosts, x_model_devs), ...]
|
269 |
+
dh_dd_mh_md_tups = map(dh_dd_mh_md, global_hardware_mesh, model_parallel_submesh, local_hardware_mesh)
|
270 |
+
# reshape to e.g. (x_dh, x_dd, x_mh, x_md, y_dh, ...)
|
271 |
+
devices = devices.reshape(*(s for t in dh_dd_mh_md_tups for s in t)) # pylint: disable=g-complex-comprehension
|
272 |
+
# TODO(jekbradbury): reorder local subgroups for ring locality
|
273 |
+
# Transpose to [data_host], [data_device], [model_host], [model_device]
|
274 |
+
# block ordering e.g. (x_dh, y_dh, ..., x_dd, y_dd, ...)
|
275 |
+
devices = devices.transpose(
|
276 |
+
*(4 * i for i in range(mesh_ndim)),
|
277 |
+
*(4 * i + 1 for i in range(mesh_ndim)),
|
278 |
+
*(4 * i + 2 for i in range(mesh_ndim)),
|
279 |
+
*(4 * i + 3 for i in range(mesh_ndim)),
|
280 |
+
)
|
281 |
+
else:
|
282 |
+
# e.g. [(x_data, x_model), (y_data, y_model), ...]
|
283 |
+
model_data_tups = [(g // m, m) for g, m in zip(global_hardware_mesh, model_parallel_submesh)]
|
284 |
+
# reshape to e.g. (x_data, x_model, y_data, y_model...)
|
285 |
+
devices = devices.reshape(*(s for t in model_data_tups for s in t)) # pylint: disable=g-complex-comprehension
|
286 |
+
# TODO(jekbradbury): reorder small subgroups for ring locality
|
287 |
+
# transpose to e.g. (x_data, y_data, ..., x_model, ...)
|
288 |
+
devices = devices.transpose(*(2 * i for i in range(mesh_ndim)), *(2 * i + 1 for i in range(mesh_ndim)))
|
289 |
+
# reshape to (data, model)
|
290 |
+
devices = devices.reshape(-1, np.prod(model_parallel_submesh))
|
291 |
+
global_mesh = Mesh(devices, ["data", "model"])
|
292 |
+
logging.info("global_mesh axis_names: %s", global_mesh.axis_names)
|
293 |
+
logging.info("global_mesh devices: %s", global_mesh.devices)
|
294 |
+
logging.info("global_mesh devices shape: %s", global_mesh.devices.shape)
|
295 |
+
return global_mesh
|
296 |
+
|
297 |
+
|
298 |
+
def get_cpu_mesh() -> Mesh:
|
299 |
+
"""Trivial mesh for CPU Testing."""
|
300 |
+
devices = np.empty((jax.host_count(), jax.local_device_count()), dtype=object)
|
301 |
+
for device in jax.devices():
|
302 |
+
devices[device.process_index, device.id % jax.local_device_count()] = device
|
303 |
+
return Mesh(devices, ["data", "model"])
|
304 |
+
|
305 |
+
|
306 |
+
def get_gpu_mesh(num_partitions: int) -> Mesh:
|
307 |
+
"""Mesh for GPUs that preferentially places 'model' on NVLink."""
|
308 |
+
nvlink_size = jax.local_device_count()
|
309 |
+
dcn_size = jax.process_count()
|
310 |
+
nvlink_mp = min(num_partitions, nvlink_size)
|
311 |
+
nvlink_dp, extra1 = divmod(nvlink_size, nvlink_mp)
|
312 |
+
dcn_mp, extra2 = divmod(num_partitions, nvlink_mp)
|
313 |
+
assert not (extra1 or extra2), (
|
314 |
+
"number of partitions on GPU must be a factor" " or multiple of the number of local devices"
|
315 |
+
)
|
316 |
+
dcn_dp = dcn_size // dcn_mp
|
317 |
+
|
318 |
+
devices = create_hybrid_device_mesh(
|
319 |
+
mesh_shape=[nvlink_dp, nvlink_mp], dcn_mesh_shape=[dcn_dp, dcn_mp], process_is_granule=True
|
320 |
+
)
|
321 |
+
|
322 |
+
global_mesh = Mesh(devices, ["data", "model"])
|
323 |
+
logging.info("global_mesh axis_names: %s", global_mesh.axis_names)
|
324 |
+
logging.info("global_mesh devices: %s", global_mesh.devices)
|
325 |
+
return global_mesh
|
326 |
+
|
327 |
+
|
328 |
+
def default_mesh(
|
329 |
+
num_partitions: int, model_parallel_submesh: Optional[HardwareMesh] = None, backend: Optional[str] = None
|
330 |
+
) -> Mesh:
|
331 |
+
"""Attempt to return a default mesh for simple cases.
|
332 |
+
|
333 |
+
Args:
|
334 |
+
num_partitions: number of partitions to use, will be ignored if
|
335 |
+
model_parallel_submesh is provided.
|
336 |
+
model_parallel_submesh: 4-tuple that specifies the x,y,z,c submesh to use as
|
337 |
+
the model-parallel device tile.
|
338 |
+
backend: get devices from the pinned backend, if specified. This is useful
|
339 |
+
for explicitly specifying the devices other than relying on
|
340 |
+
jax_platform_name.
|
341 |
+
|
342 |
+
Returns:
|
343 |
+
xmap/pjit 2D Mesh with 'data', 'model' mesh axes.
|
344 |
+
"""
|
345 |
+
last_device = jax.devices(backend)[-1]
|
346 |
+
platform = last_device.platform
|
347 |
+
device_kind = last_device.device_kind
|
348 |
+
bounds = bounds_from_last_device(last_device)
|
349 |
+
|
350 |
+
if model_parallel_submesh:
|
351 |
+
return get_mesh(model_parallel_submesh, backend=backend)
|
352 |
+
|
353 |
+
if platform == "cpu":
|
354 |
+
return get_cpu_mesh()
|
355 |
+
elif platform == "gpu":
|
356 |
+
return get_gpu_mesh(num_partitions)
|
357 |
+
|
358 |
+
mps = None
|
359 |
+
if device_kind in ("TPU v2", "TPU v3"):
|
360 |
+
if num_partitions == 1:
|
361 |
+
mps = (1, 1, 1, 1)
|
362 |
+
elif num_partitions == 2:
|
363 |
+
mps = (1, 1, 1, 2)
|
364 |
+
elif num_partitions == 4:
|
365 |
+
mps = (2, 1, 1, 2)
|
366 |
+
elif num_partitions == 8:
|
367 |
+
mps = (2, 2, 1, 2)
|
368 |
+
elif num_partitions == 16:
|
369 |
+
mps = (4, 2, 1, 2)
|
370 |
+
# assume the use of megacore on TPU v4
|
371 |
+
elif (device_kind == "TPU v4" or device_kind == "TPU v4 lite") and bounds[3] == 1:
|
372 |
+
if num_partitions == 1:
|
373 |
+
mps = (1, 1, 1, 1)
|
374 |
+
elif num_partitions == 2:
|
375 |
+
mps = (1, 2, 1, 1)
|
376 |
+
elif num_partitions == 4:
|
377 |
+
if bounds[0] >= 4:
|
378 |
+
mps = (4, 1, 1, 1)
|
379 |
+
else:
|
380 |
+
mps = (2, 2, 1, 1)
|
381 |
+
elif num_partitions == 8:
|
382 |
+
if bounds[2] >= 8:
|
383 |
+
mps = (1, 1, 8, 1)
|
384 |
+
else:
|
385 |
+
mps = (4, 2, 1, 1)
|
386 |
+
elif num_partitions == 16:
|
387 |
+
if bounds[2] >= 16:
|
388 |
+
mps = (1, 1, 16, 1)
|
389 |
+
elif bounds[0] >= 8:
|
390 |
+
mps = (8, 2, 1, 1)
|
391 |
+
elif bounds[0] >= 4:
|
392 |
+
mps = (4, 4, 1, 1)
|
393 |
+
else:
|
394 |
+
mps = (2, 2, 4, 1)
|
395 |
+
|
396 |
+
if mps is None:
|
397 |
+
raise ValueError(
|
398 |
+
"No default mesh for this configuration: specify " "config.model_parallel_submesh explicitly."
|
399 |
+
)
|
400 |
+
return get_mesh(mps, backend=backend)
|
401 |
+
|
402 |
+
|
403 |
+
# Data chunking helper.
|
404 |
+
# -----------------------------------------------------------------------------
|
405 |
+
@dataclasses.dataclass
|
406 |
+
class LocalChunkInfo:
|
407 |
+
# The logical slice of an array located on this host's local devices.
|
408 |
+
slice: Tuple[slice, ...]
|
409 |
+
# A unique index for this host/local chunk among chunks with the same slice.
|
410 |
+
replica_id: int
|
411 |
+
|
412 |
+
|
413 |
+
class LocalChunker:
|
414 |
+
"""Utility class to aid chunking of sharded arrays in multihost settings."""
|
415 |
+
|
416 |
+
def __init__(self, global_mesh: Mesh):
|
417 |
+
self.global_mesh = global_mesh
|
418 |
+
local_mesh = global_mesh.local_mesh
|
419 |
+
first_local_device = local_mesh.devices.reshape(-1)[0]
|
420 |
+
host_location = collections.OrderedDict(
|
421 |
+
zip(global_mesh.shape.keys(), list(zip(*np.nonzero(global_mesh.devices == first_local_device)))[0])
|
422 |
+
)
|
423 |
+
self.num_chunks = collections.OrderedDict()
|
424 |
+
self.chunk_ids = collections.OrderedDict()
|
425 |
+
self.mesh_axes = list(global_mesh.shape.keys())
|
426 |
+
for mesh_axis in self.mesh_axes:
|
427 |
+
num_devices_per_chunk = local_mesh.shape[mesh_axis]
|
428 |
+
self.num_chunks[mesh_axis] = global_mesh.shape[mesh_axis] // num_devices_per_chunk
|
429 |
+
self.chunk_ids[mesh_axis] = host_location[mesh_axis] // num_devices_per_chunk
|
430 |
+
|
431 |
+
def get_local_chunk_info(
|
432 |
+
self, global_shape: Tuple[int, ...], mesh_axes: Sequence[Optional[str]]
|
433 |
+
) -> LocalChunkInfo:
|
434 |
+
"""Get the local chunk info for a given array shape and sharded axes.
|
435 |
+
|
436 |
+
Args:
|
437 |
+
global_shape: the global, unsharded shape of the array to chunk.
|
438 |
+
mesh_axes: a sequence of names (or None) of equal rank to `global_shape`
|
439 |
+
that specifies which mesh dimensions the array is sharded along.
|
440 |
+
|
441 |
+
Returns:
|
442 |
+
LocalChunkInfo containing the logical slices of the array found on this
|
443 |
+
host's local devices, as well as the replica index for this chunk among
|
444 |
+
chunks with the same slice. The latter is used to determine which
|
445 |
+
host should write this chunk during checkpointing.
|
446 |
+
"""
|
447 |
+
local_slice = [slice(None) for dim in global_shape]
|
448 |
+
sharded_mesh_axes = set()
|
449 |
+
for i, (mesh_axis, size) in enumerate(zip(mesh_axes, global_shape)):
|
450 |
+
if not mesh_axis:
|
451 |
+
continue
|
452 |
+
sharded_mesh_axes.add(mesh_axis)
|
453 |
+
if not isinstance(mesh_axis, str):
|
454 |
+
raise NotImplementedError("TODO(jekbradbury)")
|
455 |
+
chunk_id = self.chunk_ids[mesh_axis]
|
456 |
+
chunk_size = size // self.num_chunks[mesh_axis]
|
457 |
+
local_slice[i] = slice(chunk_id * chunk_size, (chunk_id + 1) * chunk_size)
|
458 |
+
|
459 |
+
replicated_mesh_axes = [mesh_axis for mesh_axis in self.mesh_axes if mesh_axis not in sharded_mesh_axes]
|
460 |
+
replica_id = 0
|
461 |
+
for mesh_axis in replicated_mesh_axes:
|
462 |
+
chunk_id = self.chunk_ids[mesh_axis]
|
463 |
+
replica_id = replica_id * self.num_chunks[mesh_axis] + chunk_id
|
464 |
+
|
465 |
+
return LocalChunkInfo(tuple(local_slice), replica_id)
|
466 |
+
|
467 |
+
|
468 |
+
def standard_logical_axis_rules(
|
469 |
+
activation_partitioning_dims: int = 1,
|
470 |
+
parameter_partitioning_dims: int = 1,
|
471 |
+
additional_rules: Optional[LogicalAxisRules] = None,
|
472 |
+
) -> LogicalAxisRules:
|
473 |
+
"""Default sharding rules for T5X model in terms of logical axis names.
|
474 |
+
|
475 |
+
Args:
|
476 |
+
activation_partitioning_dims: enables 2-D activation sharding when set to 2.
|
477 |
+
parameter_partitioning_dims: enables 2-D parameter sharding when set to 2.
|
478 |
+
additional_rules: additional rules (a sequence of tuples) that will be
|
479 |
+
appended to the standard rules.
|
480 |
+
|
481 |
+
Returns:
|
482 |
+
Sequence of logical axis rules
|
483 |
+
"""
|
484 |
+
logging.info(
|
485 |
+
"`activation_partitioning_dims` = %d, `parameter_partitioning_dims` = %d",
|
486 |
+
activation_partitioning_dims,
|
487 |
+
parameter_partitioning_dims,
|
488 |
+
)
|
489 |
+
|
490 |
+
if activation_partitioning_dims == 1 and parameter_partitioning_dims == 1:
|
491 |
+
rules = [
|
492 |
+
("batch", "data"),
|
493 |
+
("vocab", "model"),
|
494 |
+
("embed", None),
|
495 |
+
("mlp", "model"),
|
496 |
+
("heads", "model"),
|
497 |
+
("kv", None),
|
498 |
+
("joined_kv", "model"), # joined heads+kv dim in 2D attn param layouts
|
499 |
+
]
|
500 |
+
elif activation_partitioning_dims == 2 and parameter_partitioning_dims == 1:
|
501 |
+
rules = [
|
502 |
+
("batch", "data"),
|
503 |
+
("vocab", "model"),
|
504 |
+
("mlp", "model"),
|
505 |
+
("heads", "model"),
|
506 |
+
("kv", None),
|
507 |
+
("joined_kv", "model"),
|
508 |
+
("embed", "model"),
|
509 |
+
]
|
510 |
+
elif activation_partitioning_dims == 1 and parameter_partitioning_dims == 2:
|
511 |
+
rules = [
|
512 |
+
("batch", "data"),
|
513 |
+
("vocab", "model"),
|
514 |
+
("mlp", "model"),
|
515 |
+
("heads", "model"),
|
516 |
+
("kv", None),
|
517 |
+
("joined_kv", "model"),
|
518 |
+
("embed", "data"),
|
519 |
+
]
|
520 |
+
elif activation_partitioning_dims == 2 and parameter_partitioning_dims == 2:
|
521 |
+
rules = [
|
522 |
+
("batch", "data"),
|
523 |
+
("vocab", "model"),
|
524 |
+
("mlp", "model"),
|
525 |
+
("heads", "model"),
|
526 |
+
("kv", None),
|
527 |
+
("joined_kv", "model"),
|
528 |
+
("embed", "model"),
|
529 |
+
("embed", "data"),
|
530 |
+
]
|
531 |
+
else:
|
532 |
+
raise ValueError(
|
533 |
+
f"`activation_partitioning_dims` = {activation_partitioning_dims} "
|
534 |
+
f"`parameter_partitioning_dims` = {parameter_partitioning_dims} "
|
535 |
+
"is not supported."
|
536 |
+
)
|
537 |
+
|
538 |
+
# Add the common rules for the replicated logical axes names.
|
539 |
+
replicated_rules = [
|
540 |
+
("relpos_buckets", None),
|
541 |
+
("abspos_buckets", None),
|
542 |
+
("length", None),
|
543 |
+
("layers", None),
|
544 |
+
("stack", None),
|
545 |
+
("mlp_activations", None),
|
546 |
+
]
|
547 |
+
rules.extend(replicated_rules)
|
548 |
+
|
549 |
+
if additional_rules:
|
550 |
+
rules.extend(additional_rules)
|
551 |
+
|
552 |
+
return rules
|
553 |
+
|
554 |
+
|
555 |
+
# NB: This needs to be top-level for the jax compilation cache.
|
556 |
+
def _id_fn(x, ix):
|
557 |
+
"""Identity function for copying parameters to the devices, sharded."""
|
558 |
+
# A pure identity such as `lambda x, *: x` can get optimized away, so we
|
559 |
+
# include a random.split as a cheap function that cannot be optimized away.
|
560 |
+
y = random.split(random.PRNGKey(jnp.array(ix, dtype=jnp.uint32)))
|
561 |
+
return x, y
|
562 |
+
|
563 |
+
|
564 |
+
@dataclasses.dataclass
|
565 |
+
class DataLayout:
|
566 |
+
"""Represents data layout for the partitioned model."""
|
567 |
+
|
568 |
+
batch_size: int
|
569 |
+
shard_id: int
|
570 |
+
num_shards: int
|
571 |
+
is_first_host_in_replica_set: bool
|
572 |
+
|
573 |
+
|
574 |
+
PartitionedCallable = Callable[..., Any]
|
575 |
+
CompiledPartitionedCallable = Callable[..., Any]
|
576 |
+
|
577 |
+
|
578 |
+
class BasePartitioner(metaclass=abc.ABCMeta):
|
579 |
+
"""Interface for partitioning computations across hardware devices."""
|
580 |
+
|
581 |
+
def __init__(
|
582 |
+
self,
|
583 |
+
num_partitions: Optional[int] = None,
|
584 |
+
model_parallel_submesh: Optional[HardwareMesh] = None,
|
585 |
+
params_on_devices: bool = True,
|
586 |
+
backend: Optional[str] = None,
|
587 |
+
):
|
588 |
+
"""Configures the partitioner.
|
589 |
+
|
590 |
+
Args:
|
591 |
+
num_partitions: the number of partitions to use. Ignored if
|
592 |
+
`model_parallel_submesh` is provided.
|
593 |
+
model_parallel_submesh: 4-tuple that specifies the x,y,z,c submesh to use
|
594 |
+
as the model-parallel device tile. This submesh is used for the larger
|
595 |
+
of the two parameter dimensions, and, if 2-D activation sharding is
|
596 |
+
enabled, for the model dimension of activations. The rest of the mesh is
|
597 |
+
used for data parallelism and, if 2-D parameter sharding is enabled, the
|
598 |
+
other parameter dimension.
|
599 |
+
params_on_devices: whether to keep the params on devices, if False -
|
600 |
+
params stay in the host memory. Note that some partitioners might ignore
|
601 |
+
this setting, for example if they don't support storing all params on
|
602 |
+
device memory.
|
603 |
+
backend: get devices from the pinned backend, if specified. This is useful
|
604 |
+
for explicitly specifying the devices other than relying on
|
605 |
+
jax_platform_name.
|
606 |
+
"""
|
607 |
+
|
608 |
+
if not num_partitions and not model_parallel_submesh:
|
609 |
+
raise ValueError("At least one of `num_partitions` or " "`model_parallel_submesh` must be set.")
|
610 |
+
|
611 |
+
if model_parallel_submesh is not None and len(model_parallel_submesh) != 4:
|
612 |
+
logging.error(
|
613 |
+
(
|
614 |
+
"`model_parallel_submesh` must be either None or a 4-tuple. Got"
|
615 |
+
" `model_parallel_submesh`=%s. A ValueError will be raised"
|
616 |
+
" beginning March 1, 2022."
|
617 |
+
),
|
618 |
+
model_parallel_submesh,
|
619 |
+
)
|
620 |
+
|
621 |
+
if bool(num_partitions) and bool(model_parallel_submesh):
|
622 |
+
logging.error(
|
623 |
+
"At most one of `num_partitions` or `model_parallel_submesh` can be "
|
624 |
+
"set. Got `num_partitions=%s` and `model_parallel_submesh`=%s. A "
|
625 |
+
"ValueError will be raised beginning March 21, 2022.",
|
626 |
+
num_partitions,
|
627 |
+
model_parallel_submesh,
|
628 |
+
)
|
629 |
+
|
630 |
+
self._num_partitions = num_partitions
|
631 |
+
self._model_parallel_submesh = model_parallel_submesh
|
632 |
+
self._params_on_devices = params_on_devices
|
633 |
+
self._data_axis = "data"
|
634 |
+
self._backend = backend
|
635 |
+
|
636 |
+
@property
|
637 |
+
def mesh(self) -> Mesh:
|
638 |
+
raise NotImplementedError
|
639 |
+
|
640 |
+
@property
|
641 |
+
def data_partition_spec(self) -> PartitionSpec:
|
642 |
+
return PartitionSpec(self._data_axis)
|
643 |
+
|
644 |
+
def get_data_layout(self, batch_size: Optional[int] = None, host_index: Optional[int] = None) -> DataLayout:
|
645 |
+
"""Returns filled `DataLayout` based on the partitioned model layout.
|
646 |
+
|
647 |
+
Args:
|
648 |
+
batch_size: if set, indicates the requested batch size. The exception will
|
649 |
+
be raised if this batch size is not compatible with the layout. If not
|
650 |
+
set, the batch size is inferred from the layout.
|
651 |
+
host_index: indicates the host index to use for the calculations, if not
|
652 |
+
set - use JAX-provided one. Should be in [0, num_hosts) interval and the
|
653 |
+
order should match the order of corresponding CPU devices in
|
654 |
+
`jax.devices()`.
|
655 |
+
|
656 |
+
Returns:
|
657 |
+
Filled `DataLayout` structure.
|
658 |
+
"""
|
659 |
+
if host_index is not None:
|
660 |
+
raise NotImplementedError("Explicit host_index is not yet implemented.")
|
661 |
+
if self._data_axis is None:
|
662 |
+
return DataLayout(
|
663 |
+
batch_size=batch_size,
|
664 |
+
shard_id=0,
|
665 |
+
num_shards=1,
|
666 |
+
is_first_host_in_replica_set=(jax.process_index() == 0),
|
667 |
+
)
|
668 |
+
mesh_size = self._local_chunker.global_mesh.shape[self._data_axis]
|
669 |
+
batch_size = batch_size or mesh_size
|
670 |
+
if batch_size % mesh_size:
|
671 |
+
raise ValueError(
|
672 |
+
f"Batch size ({batch_size}) must be divisible by corresponding " f"mesh size ({mesh_size})."
|
673 |
+
)
|
674 |
+
num_shards = self._local_chunker.num_chunks[self._data_axis]
|
675 |
+
if batch_size % num_shards:
|
676 |
+
raise ValueError(f"Batch size ({batch_size}) must be divisible by number of " f"replicas ({num_shards}).")
|
677 |
+
replica_id = self._local_chunker.get_local_chunk_info((batch_size,), [self._data_axis]).replica_id
|
678 |
+
return DataLayout(
|
679 |
+
batch_size=int(batch_size),
|
680 |
+
shard_id=int(self._local_chunker.chunk_ids[self._data_axis]),
|
681 |
+
num_shards=int(num_shards),
|
682 |
+
is_first_host_in_replica_set=(replica_id == 0),
|
683 |
+
)
|
684 |
+
|
685 |
+
def get_local_chunk_info(
|
686 |
+
self, global_shape: Tuple[int, ...], mesh_axes: Sequence[Optional[str]]
|
687 |
+
) -> LocalChunkInfo:
|
688 |
+
"""Returns the local chunk info for a given array shape and sharded axes."""
|
689 |
+
return self._local_chunker.get_local_chunk_info(global_shape, mesh_axes)
|
690 |
+
|
691 |
+
@property
|
692 |
+
def params_on_devices(self):
|
693 |
+
return self._params_on_devices
|
694 |
+
|
695 |
+
def move_params_to_devices(self, train_state: TrainState, train_state_axes: TrainState) -> TrainState:
|
696 |
+
"""Moves the optimizer parameters to devices."""
|
697 |
+
p_id_fn = self.partition(
|
698 |
+
_id_fn,
|
699 |
+
in_axis_resources=(train_state_axes, None),
|
700 |
+
out_axis_resources=(train_state_axes, None),
|
701 |
+
donate_argnums=(0,),
|
702 |
+
)
|
703 |
+
if jax.config.jax_array and jax.process_count() > 1:
|
704 |
+
train_state = multihost_utils.host_local_array_to_global_array(train_state, self.mesh, train_state_axes)
|
705 |
+
train_state, _ = p_id_fn(train_state, jnp.ones((), dtype=jnp.uint32))
|
706 |
+
return train_state
|
707 |
+
|
708 |
+
@property
|
709 |
+
@abc.abstractmethod
|
710 |
+
def _local_chunker(self):
|
711 |
+
"""Returns the chunker that matches the parameters of this partitioner."""
|
712 |
+
raise NotImplementedError
|
713 |
+
|
714 |
+
def get_logical_axes(self, train_state: TrainState) -> TrainState:
|
715 |
+
"""Returns a copy of TrainState with Optional[AxisNames] as leaves."""
|
716 |
+
# By default, return None for the logical axes.
|
717 |
+
return train_state.restore_state(jax.tree_map(lambda x: None, train_state.state_dict()))
|
718 |
+
|
719 |
+
def get_mesh_axes(self, train_state: TrainState) -> TrainState:
|
720 |
+
"""Returns a copy of TrainState with Optional[PartitionSpecs] as leaves."""
|
721 |
+
raise NotImplementedError
|
722 |
+
|
723 |
+
@abc.abstractmethod
|
724 |
+
def partition(
|
725 |
+
self,
|
726 |
+
fn: Callable, # pylint: disable=g-bare-generic
|
727 |
+
in_axis_resources,
|
728 |
+
out_axis_resources,
|
729 |
+
static_argnums: Union[int, Sequence[int]] = (),
|
730 |
+
donate_argnums: Union[int, Sequence[int]] = (),
|
731 |
+
) -> PartitionedCallable:
|
732 |
+
"""Partitions the computation using partitioner-specific implementation.
|
733 |
+
|
734 |
+
Args:
|
735 |
+
fn: the function to partition.
|
736 |
+
in_axis_resources: Pytree of structure matching that of arguments to `fn`,
|
737 |
+
with all actual arguments replaced by resource assignment
|
738 |
+
specifications. It is also valid to specify a pytree prefix (e.g. one
|
739 |
+
value in place of a whole subtree), in which case the leaves get
|
740 |
+
broadcast to all values in that subtree.
|
741 |
+
The valid resource assignment specifications are:
|
742 |
+
`None`: in which case the value will be replicated on all devices
|
743 |
+
`PartitionSpec`: a tuple of length at most equal to the rank of the
|
744 |
+
partitioned value. Each element can be a `None`, a mesh axis or a
|
745 |
+
tuple of mesh axes, and specifies the set of resources assigned to
|
746 |
+
partition the value's dimension matching its position in the spec.
|
747 |
+
out_axis_resources: Like `in_axis_resources`, but specifies resource
|
748 |
+
assignment for function outputs.
|
749 |
+
static_argnums: an optional int or collection of ints that specify which
|
750 |
+
positional arguments to treat as static (compile-time constant) in the
|
751 |
+
partitioned function.
|
752 |
+
donate_argnums: an optional int or collection of ints that specify which
|
753 |
+
argument buffers are "donated" to the computation. It is safe to donate
|
754 |
+
argument buffers if you no longer need them once the computation has
|
755 |
+
finished.
|
756 |
+
|
757 |
+
Returns:
|
758 |
+
A partitioned version of the input function.
|
759 |
+
"""
|
760 |
+
raise NotImplementedError
|
761 |
+
|
762 |
+
@abc.abstractmethod
|
763 |
+
def compile(self, partitioned_fn: PartitionedCallable, *args) -> CompiledPartitionedCallable:
|
764 |
+
"""Compiles and returns the partitioned function, or the original.
|
765 |
+
|
766 |
+
Args:
|
767 |
+
partitioned_fn: The partitioned function.
|
768 |
+
*args: Sample arguments to the partitioned function matching the input
|
769 |
+
shapes that will be passed to the compiled function.
|
770 |
+
|
771 |
+
Returns:
|
772 |
+
The compiled function, or the original if this partitioner does not
|
773 |
+
support compilation.
|
774 |
+
"""
|
775 |
+
raise NotImplementedError
|
776 |
+
|
777 |
+
|
778 |
+
class PjittedFnWithContext(PartitionedCallable):
|
779 |
+
"""Wraps pjitted function to apply the appropriate contexts."""
|
780 |
+
|
781 |
+
def __init__(self, pjitted_fn, partition_mesh: Mesh, logical_axis_rules: flax_partitioning.LogicalRules = ()):
|
782 |
+
self._pjitted_fn = pjitted_fn
|
783 |
+
self._mesh = partition_mesh
|
784 |
+
self._logical_axis_rules = logical_axis_rules
|
785 |
+
|
786 |
+
def __call__(self, *args):
|
787 |
+
with Mesh(self._mesh.devices, self._mesh.axis_names), flax_partitioning.axis_rules(self._logical_axis_rules):
|
788 |
+
return self._pjitted_fn(*args)
|
789 |
+
|
790 |
+
def lower(self, *args):
|
791 |
+
with Mesh(self._mesh.devices, self._mesh.axis_names), flax_partitioning.axis_rules(self._logical_axis_rules):
|
792 |
+
return self._pjitted_fn.lower(*args)
|
793 |
+
|
794 |
+
|
795 |
+
class BasePjitPartitioner(BasePartitioner):
|
796 |
+
"""Partitioner that uses T5X version of jax.pjit."""
|
797 |
+
|
798 |
+
@cached_property
|
799 |
+
def _local_chunker(self) -> LocalChunker:
|
800 |
+
return LocalChunker(self.mesh)
|
801 |
+
|
802 |
+
@cached_property
|
803 |
+
def mesh(self) -> Mesh:
|
804 |
+
return default_mesh(self._num_partitions, self._model_parallel_submesh, self._backend)
|
805 |
+
|
806 |
+
def partition(
|
807 |
+
self,
|
808 |
+
fn: Callable, # pylint: disable=g-bare-generic
|
809 |
+
in_axis_resources,
|
810 |
+
out_axis_resources,
|
811 |
+
static_argnums: Union[int, Sequence[int]] = (),
|
812 |
+
donate_argnums: Union[int, Sequence[int]] = (),
|
813 |
+
) -> PjittedFnWithContext:
|
814 |
+
pjitted = pjit(
|
815 |
+
fn,
|
816 |
+
in_axis_resources=in_axis_resources,
|
817 |
+
out_axis_resources=out_axis_resources,
|
818 |
+
static_argnums=static_argnums,
|
819 |
+
donate_argnums=donate_argnums,
|
820 |
+
backend=self._backend,
|
821 |
+
)
|
822 |
+
|
823 |
+
return PjittedFnWithContext(pjitted, self.mesh)
|
824 |
+
|
825 |
+
def compile(self, partitioned_fn: PjittedFnWithContext, *args) -> CompiledPartitionedCallable:
|
826 |
+
return partitioned_fn.lower(*args).compile()
|
827 |
+
|
828 |
+
|
829 |
+
class PjitPartitioner(BasePjitPartitioner):
|
830 |
+
"""Partitioner that uses named axes and jax.pjit."""
|
831 |
+
|
832 |
+
def __init__(
|
833 |
+
self,
|
834 |
+
num_partitions: Optional[int] = None,
|
835 |
+
model_parallel_submesh: Optional[HardwareMesh] = None,
|
836 |
+
params_on_devices: bool = True,
|
837 |
+
backend: Optional[str] = None,
|
838 |
+
logical_axis_rules: Optional[LogicalAxisRules] = None,
|
839 |
+
use_cpu_pjit: Optional[bool] = False,
|
840 |
+
):
|
841 |
+
"""PjitPartitioner constructor.
|
842 |
+
|
843 |
+
See https://github.com/google-research/text-to-text-transfer-transformer/blob/main/README.mdx/usage/partitioning for details.
|
844 |
+
|
845 |
+
Args:
|
846 |
+
num_partitions: an integer that specifies the size of the model parallel
|
847 |
+
submesh to be automatically selected for the current topology. See
|
848 |
+
`model_parallel_submesh` for details on how this submesh is used.
|
849 |
+
Mutually exlusive with `model_parallel_submesh`.
|
850 |
+
model_parallel_submesh: is a 4-tuple that specifies the `(x, y, z, c)`
|
851 |
+
submesh model-parallel device tile, an axis of accelerator parallelism
|
852 |
+
orthogonal to data parallelism. Array axes in a model's parameters or
|
853 |
+
activations can be sharded over this submesh using axis rules (see
|
854 |
+
`logical_axis_rules`) that map them to 'model'. The effective number of
|
855 |
+
model sub-partitions is equal to `np.prod(model_parallel_submesh)` and
|
856 |
+
must evenly divide the total number of devices (i.e.,
|
857 |
+
`jax.device_count() % np.prod(model_parallel_submesh) == 0`). The rest
|
858 |
+
of the TPU mesh is the data parallel submesh, providing
|
859 |
+
`jax.device_count() // np.prod(model_parallel_submesh)` partitions. It
|
860 |
+
is used for data (batch) parallelism and to shard other array axes that
|
861 |
+
are mapped to 'data'. This argument is mutually exclusive with
|
862 |
+
`num_partitions`.
|
863 |
+
params_on_devices: whether to keep the params on devices, if False -
|
864 |
+
params stay in the host memory. Note that some partitioners might ignore
|
865 |
+
this setting, for example if they don't support storing all params on
|
866 |
+
device memory.
|
867 |
+
backend: get devices from the pinned backend, if specified. This is
|
868 |
+
useful for explicitly specifying the devices other than relying on
|
869 |
+
jax_platform_name.
|
870 |
+
logical_axis_rules: a priority-ordered sequence of KV tuples that maps
|
871 |
+
logical axis names to either `None` (not sharded), 'model' (to shard
|
872 |
+
across the model-parallel submesh), or 'data' (to shard across the
|
873 |
+
data-parallel submesh).
|
874 |
+
use_cpu_pjit: enables wrapper function for pjit which just jits the
|
875 |
+
function if using CPU backend.
|
876 |
+
"""
|
877 |
+
super().__init__(
|
878 |
+
num_partitions=num_partitions,
|
879 |
+
model_parallel_submesh=model_parallel_submesh,
|
880 |
+
params_on_devices=params_on_devices,
|
881 |
+
backend=backend,
|
882 |
+
)
|
883 |
+
if logical_axis_rules is None:
|
884 |
+
logical_axis_rules = standard_logical_axis_rules()
|
885 |
+
self._logical_axis_rules = tuple(logical_axis_rules)
|
886 |
+
(self._data_axis,) = flax_partitioning.logical_to_mesh_axes(["batch"], logical_axis_rules)
|
887 |
+
self._use_cpu_pjit = use_cpu_pjit
|
888 |
+
|
889 |
+
def partition(
|
890 |
+
self,
|
891 |
+
fn: Callable, # pylint: disable=g-bare-generic
|
892 |
+
in_axis_resources,
|
893 |
+
out_axis_resources,
|
894 |
+
static_argnums: Union[int, Sequence[int]] = (),
|
895 |
+
donate_argnums: Union[int, Sequence[int]] = (),
|
896 |
+
) -> PjittedFnWithContext:
|
897 |
+
"""Partitions the function using jax.pjit."""
|
898 |
+
if self._use_cpu_pjit:
|
899 |
+
pjit_fn = pjit_with_cpu_fallback
|
900 |
+
else:
|
901 |
+
pjit_fn = pjit
|
902 |
+
pjitted = pjit_fn(
|
903 |
+
fn,
|
904 |
+
in_axis_resources=in_axis_resources,
|
905 |
+
out_axis_resources=out_axis_resources,
|
906 |
+
static_argnums=static_argnums,
|
907 |
+
donate_argnums=donate_argnums,
|
908 |
+
backend=self._backend,
|
909 |
+
)
|
910 |
+
|
911 |
+
return PjittedFnWithContext(pjitted, self.mesh, self._logical_axis_rules)
|
912 |
+
|
913 |
+
@property
|
914 |
+
def logical_axis_rules(self):
|
915 |
+
"""Returns the logical axis rules."""
|
916 |
+
return self._logical_axis_rules
|
917 |
+
|
918 |
+
def get_logical_axes(self, train_state: TrainState) -> TrainState:
|
919 |
+
"""Returns a copy of TrainState with Optional[AxisNames] as leaves."""
|
920 |
+
return train_state.as_logical_axes()
|
921 |
+
|
922 |
+
def get_mesh_axes(self, train_state: TrainState) -> TrainState:
|
923 |
+
"""Returns a copy of TrainState with Optional[PartitionSpecs] as leaves."""
|
924 |
+
logical_axes = self.get_logical_axes(train_state)
|
925 |
+
|
926 |
+
def _logical_to_mesh_axes(param_name, logical_axes):
|
927 |
+
if logical_axes is None:
|
928 |
+
return None
|
929 |
+
elif logical_axes is traverse_util.empty_node:
|
930 |
+
return traverse_util.empty_node
|
931 |
+
try:
|
932 |
+
return flax_partitioning.logical_to_mesh_axes(logical_axes, self._logical_axis_rules)
|
933 |
+
except ValueError as e:
|
934 |
+
raise ValueError(f"Failed to map logical axes for {param_name}") from e
|
935 |
+
|
936 |
+
flat_logical_axes = traverse_util.flatten_dict(logical_axes.state_dict(), keep_empty_nodes=True, sep="/")
|
937 |
+
flat_mesh_axes = {k: _logical_to_mesh_axes(k, v) for k, v in flat_logical_axes.items()}
|
938 |
+
|
939 |
+
return logical_axes.restore_state(traverse_util.unflatten_dict(flat_mesh_axes, sep="/"))
|
whisper_jax/pipeline.py
ADDED
@@ -0,0 +1,506 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# coding=utf-8
|
2 |
+
# Copyright 2023 The HuggingFace Inc. team.
|
3 |
+
#
|
4 |
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
5 |
+
# you may not use this file except in compliance with the License.
|
6 |
+
# You may obtain a copy of the License at
|
7 |
+
#
|
8 |
+
# http://www.apache.org/licenses/LICENSE-2.0
|
9 |
+
#
|
10 |
+
# Unless required by applicable law or agreed to in writing, software
|
11 |
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
12 |
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
13 |
+
# See the License for the specific language governing permissions and
|
14 |
+
# limitations under the License.
|
15 |
+
|
16 |
+
|
17 |
+
import math
|
18 |
+
|
19 |
+
import jax
|
20 |
+
import jax.numpy as jnp
|
21 |
+
import numpy as np
|
22 |
+
import requests
|
23 |
+
from flax import jax_utils
|
24 |
+
from flax.core.frozen_dict import freeze
|
25 |
+
from flax.training.common_utils import shard
|
26 |
+
from jax.sharding import PartitionSpec as P
|
27 |
+
from transformers import WhisperProcessor, is_tokenizers_available, WhisperFeatureExtractor, WhisperTokenizerFast
|
28 |
+
from transformers.models.whisper.tokenization_whisper import TO_LANGUAGE_CODE, WhisperTokenizer
|
29 |
+
from transformers.pipelines.audio_utils import ffmpeg_read
|
30 |
+
from transformers.utils import logging
|
31 |
+
|
32 |
+
from .modeling_flax_whisper import FlaxWhisperForConditionalGeneration
|
33 |
+
from .partitioner import PjitPartitioner
|
34 |
+
from .train_state import InferenceState
|
35 |
+
|
36 |
+
|
37 |
+
logger = logging.get_logger(__name__)
|
38 |
+
|
39 |
+
# 2D parameter and activation partitioning for DP
|
40 |
+
logical_axis_rules_dp = (
|
41 |
+
("batch", "data"),
|
42 |
+
("mlp", None),
|
43 |
+
("heads", None),
|
44 |
+
("vocab", None),
|
45 |
+
("embed", None),
|
46 |
+
("embed", None),
|
47 |
+
("joined_kv", None),
|
48 |
+
("kv", None),
|
49 |
+
("length", None),
|
50 |
+
("num_mel", None),
|
51 |
+
("channels", None),
|
52 |
+
)
|
53 |
+
|
54 |
+
|
55 |
+
class FlaxWhisperPipline:
|
56 |
+
def __init__(
|
57 |
+
self,
|
58 |
+
checkpoint="openai/whisper-large-v2",
|
59 |
+
dtype=jnp.float32,
|
60 |
+
batch_size=None,
|
61 |
+
max_length=None,
|
62 |
+
):
|
63 |
+
"""
|
64 |
+
Args
|
65 |
+
checkpoint (`str`, *optional*, defaults to `"openai/whisper-large-v2"):
|
66 |
+
The Whisper checkpoint to use with the pipeline. Must be an available checkpoint on the Hugging Face Hub
|
67 |
+
with Flax weights.
|
68 |
+
dtype (`jax.numpy.dtype`, *optional*, defaults to `jax.numpy.float32`):
|
69 |
+
The data type of the computation. Can be one of `jax.numpy.float32`, `jax.numpy.float16` (on GPUs) and
|
70 |
+
`jax.numpy.bfloat16` (on TPUs). This can be used to enable half-precision inference on GPUs or TPUs.
|
71 |
+
If specified all the computation will be performed with the given `dtype`. **Note that this only
|
72 |
+
specifies the dtype of the computation and does not influence the dtype of model parameters.**
|
73 |
+
batch_size (`int`, *optional*, defaults to the minimum per-device batch size, i.e. `jax.local_device_count()`):
|
74 |
+
The batch size to be used in chunking transcription. Beneficial for transcribing long audio files. Passing
|
75 |
+
a batch size in the `__init__` method will be superseded by any batch size passed to the `__call__` method.
|
76 |
+
max_length (`int`, *optional*):
|
77 |
+
The maximum numbers of tokens to generate. Defaults to `model.config.max_length`.
|
78 |
+
"""
|
79 |
+
self.checkpoint = checkpoint
|
80 |
+
self.dtype = dtype
|
81 |
+
|
82 |
+
self.processor = WhisperProcessor.from_pretrained(self.checkpoint)
|
83 |
+
self.feature_extractor = self.processor.feature_extractor
|
84 |
+
# potentially load fast tokenizer if available
|
85 |
+
tokenizer_cls = WhisperTokenizerFast if is_tokenizers_available() else WhisperTokenizer
|
86 |
+
self.tokenizer = tokenizer_cls.from_pretrained(checkpoint)
|
87 |
+
|
88 |
+
self.model, self.params = FlaxWhisperForConditionalGeneration.from_pretrained(
|
89 |
+
self.checkpoint,
|
90 |
+
_do_init=False,
|
91 |
+
dtype=self.dtype,
|
92 |
+
)
|
93 |
+
|
94 |
+
self.max_length = max_length if max_length is not None else self.model.generation_config.max_length
|
95 |
+
self.min_batch_size = jax.local_device_count()
|
96 |
+
self.batch_size = (
|
97 |
+
batch_size if batch_size is not None else self.min_batch_size
|
98 |
+
) # we need a minimum of 1 batch per-device
|
99 |
+
|
100 |
+
def generate(params, input_features, forced_decoder_ids, return_timestamps):
|
101 |
+
output_ids = self.model.pipeline_generate(
|
102 |
+
input_features,
|
103 |
+
params=params,
|
104 |
+
forced_decoder_ids=forced_decoder_ids,
|
105 |
+
return_timestamps=return_timestamps,
|
106 |
+
max_length=self.max_length,
|
107 |
+
)
|
108 |
+
return output_ids
|
109 |
+
|
110 |
+
# use pmap for DP by default - this is compatible on a Colab TPU v2
|
111 |
+
self.params = jax_utils.replicate(self.params)
|
112 |
+
self.p_generate = jax.pmap(
|
113 |
+
generate, "input_features", in_axes=(0, 0, None), out_axes=0, static_broadcasted_argnums=(3,)
|
114 |
+
)
|
115 |
+
self.is_sharded = False
|
116 |
+
|
117 |
+
def shard_params(self, num_mp_partitions=1, logical_axis_rules=logical_axis_rules_dp):
|
118 |
+
def init_fn():
|
119 |
+
input_shape = (1, self.model.config.num_mel_bins, 2 * self.model.config.max_source_positions)
|
120 |
+
|
121 |
+
input_features = jnp.zeros(input_shape, dtype="f4")
|
122 |
+
input_features = input_features.at[(..., -1)].set(self.model.config.eos_token_id)
|
123 |
+
|
124 |
+
decoder_input_ids = jnp.zeros((input_shape[0], 1), dtype="i4")
|
125 |
+
decoder_attention_mask = jnp.ones_like(decoder_input_ids)
|
126 |
+
|
127 |
+
batch_size, sequence_length = decoder_input_ids.shape
|
128 |
+
decoder_position_ids = jnp.broadcast_to(
|
129 |
+
jnp.arange(sequence_length)[None, :], (batch_size, sequence_length)
|
130 |
+
)
|
131 |
+
|
132 |
+
rng = jax.random.PRNGKey(0)
|
133 |
+
init_params = self.model.module.init(
|
134 |
+
rng,
|
135 |
+
input_features=input_features,
|
136 |
+
decoder_input_ids=decoder_input_ids,
|
137 |
+
decoder_attention_mask=decoder_attention_mask,
|
138 |
+
decoder_position_ids=decoder_position_ids,
|
139 |
+
return_dict=False,
|
140 |
+
)
|
141 |
+
return init_params
|
142 |
+
|
143 |
+
# Axis names metadata
|
144 |
+
param_axes = jax.eval_shape(init_fn)["params_axes"]
|
145 |
+
|
146 |
+
# Create InferenceState, since the partitioner expects it
|
147 |
+
state = InferenceState(
|
148 |
+
step=jnp.array(0),
|
149 |
+
params=freeze(self.model.params_shape_tree),
|
150 |
+
params_axes=freeze(param_axes),
|
151 |
+
flax_mutables=None,
|
152 |
+
flax_mutables_axes=param_axes,
|
153 |
+
)
|
154 |
+
|
155 |
+
partitioner = PjitPartitioner(num_partitions=num_mp_partitions, logical_axis_rules=logical_axis_rules)
|
156 |
+
|
157 |
+
mesh_axes = partitioner.get_mesh_axes(state)
|
158 |
+
params_spec = mesh_axes.params
|
159 |
+
|
160 |
+
p_shard_params = partitioner.partition(self.model.to_bf16, (params_spec,), params_spec)
|
161 |
+
|
162 |
+
# This will auto-magically run in mesh context
|
163 |
+
self.params = p_shard_params(freeze(jax_utils.unreplicate(self.params)))
|
164 |
+
self.is_sharded = True
|
165 |
+
|
166 |
+
def generate(params, input_features, forced_decoder_ids, return_timestamps):
|
167 |
+
output_ids = self.model.pipeline_generate(
|
168 |
+
input_features,
|
169 |
+
params=params,
|
170 |
+
forced_decoder_ids=forced_decoder_ids,
|
171 |
+
return_timestamps=return_timestamps,
|
172 |
+
max_length=self.max_length,
|
173 |
+
)
|
174 |
+
return output_ids
|
175 |
+
|
176 |
+
# Use pjit for generate only once we've sharded the params
|
177 |
+
self.p_generate = partitioner.partition(
|
178 |
+
generate,
|
179 |
+
in_axis_resources=(params_spec, P("data"), None),
|
180 |
+
out_axis_resources=P("data"),
|
181 |
+
static_argnums=(3,),
|
182 |
+
)
|
183 |
+
|
184 |
+
def generate(self, input_features, language=None, task=None, return_timestamps=False):
|
185 |
+
forced_decoder_ids = self.get_forced_decoder_ids(
|
186 |
+
language=language, task=task, return_timestamps=return_timestamps
|
187 |
+
)
|
188 |
+
if not self.is_sharded:
|
189 |
+
# if we're using pmap we need to manually replicate the input data across devices and gather the output tokens
|
190 |
+
output_ids = self.p_generate(
|
191 |
+
freeze(self.params), shard(input_features), forced_decoder_ids, return_timestamps
|
192 |
+
).sequences
|
193 |
+
output_ids = jax.device_get(output_ids.reshape(-1, self.max_length))
|
194 |
+
else:
|
195 |
+
# pjit handles replication / gathering for us auto-magically
|
196 |
+
output_ids = self.p_generate(
|
197 |
+
freeze(self.params), input_features, forced_decoder_ids, return_timestamps
|
198 |
+
).sequences
|
199 |
+
return output_ids
|
200 |
+
|
201 |
+
def get_forced_decoder_ids(self, generation_config=None, task=None, language=None, return_timestamps=False):
|
202 |
+
if generation_config is None:
|
203 |
+
generation_config = self.model.generation_config
|
204 |
+
|
205 |
+
if hasattr(generation_config, "is_multilingual"):
|
206 |
+
is_multilingual = generation_config.is_multilingual
|
207 |
+
else:
|
208 |
+
is_multilingual = None
|
209 |
+
|
210 |
+
forced_decoder_ids = []
|
211 |
+
|
212 |
+
if is_multilingual:
|
213 |
+
if language is not None:
|
214 |
+
language = language.lower()
|
215 |
+
if language in generation_config.lang_to_id.keys():
|
216 |
+
language_token = language
|
217 |
+
elif language in TO_LANGUAGE_CODE.values():
|
218 |
+
language_token = f"<|{language}|>"
|
219 |
+
elif language in TO_LANGUAGE_CODE.keys():
|
220 |
+
language_token = f"<|{TO_LANGUAGE_CODE[language]}|>"
|
221 |
+
else:
|
222 |
+
if len(language) == 2:
|
223 |
+
# ISO 639-1 language code
|
224 |
+
acceptable_languages = list(TO_LANGUAGE_CODE.values())
|
225 |
+
elif "<" in language or "|" in language or ">" in language:
|
226 |
+
# generation config language code
|
227 |
+
acceptable_languages = list(generation_config.lang_to_id.keys())
|
228 |
+
else:
|
229 |
+
# language passed as a string
|
230 |
+
acceptable_languages = list(TO_LANGUAGE_CODE.keys())
|
231 |
+
raise ValueError(
|
232 |
+
f"Unsupported language: {language}. Language should be one of:" f" {acceptable_languages}."
|
233 |
+
)
|
234 |
+
forced_decoder_ids.append((1, generation_config.lang_to_id[language_token]))
|
235 |
+
|
236 |
+
if task is not None:
|
237 |
+
forced_decoder_ids.append((2, generation_config.task_to_id[task]))
|
238 |
+
else:
|
239 |
+
forced_decoder_ids.append((2, generation_config.task_to_id["transcribe"]))
|
240 |
+
|
241 |
+
if not return_timestamps:
|
242 |
+
if forced_decoder_ids and forced_decoder_ids[-1][0] != generation_config.no_timestamps_token_id:
|
243 |
+
idx = forced_decoder_ids[-1][0] + 1 if forced_decoder_ids else 1
|
244 |
+
forced_decoder_ids.append((idx, generation_config.no_timestamps_token_id))
|
245 |
+
|
246 |
+
return forced_decoder_ids
|
247 |
+
|
248 |
+
def chunk_iter_with_batch(self, inputs, chunk_len, stride_left, stride_right, batch_size):
|
249 |
+
inputs_len = inputs.shape[0]
|
250 |
+
step = chunk_len - stride_left - stride_right
|
251 |
+
|
252 |
+
all_chunk_start_idx = np.arange(0, inputs_len, step)
|
253 |
+
num_samples = len(all_chunk_start_idx)
|
254 |
+
|
255 |
+
num_batches = math.ceil(num_samples / batch_size)
|
256 |
+
batch_idx = np.array_split(np.arange(num_samples), num_batches)
|
257 |
+
|
258 |
+
for idx in batch_idx:
|
259 |
+
chunk_start_idx = all_chunk_start_idx[idx]
|
260 |
+
chunk_end_idx = chunk_start_idx + chunk_len
|
261 |
+
|
262 |
+
chunks = [inputs[chunk_start:chunk_end] for chunk_start, chunk_end in zip(chunk_start_idx, chunk_end_idx)]
|
263 |
+
processed = self.feature_extractor(
|
264 |
+
chunks, sampling_rate=self.feature_extractor.sampling_rate, return_tensors="np"
|
265 |
+
)
|
266 |
+
|
267 |
+
_stride_left = np.where(chunk_start_idx == 0, 0, stride_left)
|
268 |
+
is_last = np.where(stride_right > 0, chunk_end_idx > inputs_len, chunk_end_idx >= inputs_len)
|
269 |
+
_stride_right = np.where(is_last, 0, stride_right)
|
270 |
+
|
271 |
+
chunk_lens = [chunk.shape[0] for chunk in chunks]
|
272 |
+
strides = [
|
273 |
+
(chunk_l, _stride_l, _stride_r)
|
274 |
+
for chunk_l, _stride_l, _stride_r in zip(chunk_lens, _stride_left, _stride_right)
|
275 |
+
]
|
276 |
+
|
277 |
+
yield {"stride": strides, **processed}
|
278 |
+
|
279 |
+
def preprocess_batch(self, inputs, chunk_length_s=30.0, stride_length_s=None, batch_size=None):
|
280 |
+
if isinstance(inputs, np.ndarray):
|
281 |
+
logger.warning(
|
282 |
+
"Numpy array passed as input - no sampling rate checks will be performed."
|
283 |
+
"It is strongly recommended to pass the input as a dictionary with an 'array' key "
|
284 |
+
"containing the numpy array representing the audio, and a 'sampling_rate' key "
|
285 |
+
"containing the sampling rate associated with the audio array."
|
286 |
+
"Failing to do so can result in silent errors that might be hard to debug."
|
287 |
+
)
|
288 |
+
|
289 |
+
if isinstance(inputs, str):
|
290 |
+
if inputs.startswith("http://") or inputs.startswith("https://"):
|
291 |
+
# We need to actually check for a real protocol, otherwise it's impossible to use a local file
|
292 |
+
# like http_huggingface_co.png
|
293 |
+
inputs = requests.get(inputs).content
|
294 |
+
else:
|
295 |
+
with open(inputs, "rb") as f:
|
296 |
+
inputs = f.read()
|
297 |
+
|
298 |
+
if isinstance(inputs, bytes):
|
299 |
+
inputs = ffmpeg_read(inputs, self.feature_extractor.sampling_rate)
|
300 |
+
|
301 |
+
stride = None
|
302 |
+
if isinstance(inputs, dict):
|
303 |
+
stride = inputs.get("stride", None)
|
304 |
+
# Accepting `"array"` which is the key defined in `datasets` for
|
305 |
+
# better integration
|
306 |
+
if not ("sampling_rate" in inputs and "array" in inputs):
|
307 |
+
raise ValueError(
|
308 |
+
"When passing a dictionary to FlaxWhisperPipline, the dict needs to contain an 'array' key "
|
309 |
+
"containing the numpy array representing the audio, and a 'sampling_rate' key "
|
310 |
+
"containing the sampling rate associated with the audio array."
|
311 |
+
)
|
312 |
+
|
313 |
+
in_sampling_rate = inputs.get("sampling_rate")
|
314 |
+
inputs = inputs.get("array", None)
|
315 |
+
|
316 |
+
if in_sampling_rate != self.feature_extractor.sampling_rate:
|
317 |
+
try:
|
318 |
+
import librosa
|
319 |
+
except ImportError as err:
|
320 |
+
raise ImportError(
|
321 |
+
"To support resampling audio files, please install 'librosa' and 'soundfile'."
|
322 |
+
) from err
|
323 |
+
|
324 |
+
inputs = librosa.resample(
|
325 |
+
inputs, orig_sr=in_sampling_rate, target_sr=self.feature_extractor.sampling_rate
|
326 |
+
)
|
327 |
+
ratio = self.feature_extractor.sampling_rate / in_sampling_rate
|
328 |
+
else:
|
329 |
+
ratio = 1
|
330 |
+
|
331 |
+
if not isinstance(inputs, np.ndarray):
|
332 |
+
raise ValueError(f"We expect a numpy ndarray as input, got `{type(inputs)}`")
|
333 |
+
if len(inputs.shape) != 1:
|
334 |
+
raise ValueError("We expect a single channel audio input for AutomaticSpeechRecognitionPipeline")
|
335 |
+
|
336 |
+
if stride is not None:
|
337 |
+
if stride[0] + stride[1] > inputs.shape[0]:
|
338 |
+
raise ValueError("Stride is too large for input")
|
339 |
+
|
340 |
+
# Stride needs to get the chunk length here, it's going to get
|
341 |
+
# swallowed by the `feature_extractor` later, and then batching
|
342 |
+
# can add extra data in the inputs, so we need to keep track
|
343 |
+
# of the original length in the stride so we can cut properly.
|
344 |
+
stride = (inputs.shape[0], int(round(stride[0] * ratio)), int(round(stride[1] * ratio)))
|
345 |
+
|
346 |
+
if chunk_length_s:
|
347 |
+
if stride_length_s is None:
|
348 |
+
stride_length_s = chunk_length_s / 6
|
349 |
+
|
350 |
+
if isinstance(stride_length_s, (int, float)):
|
351 |
+
stride_length_s = [stride_length_s, stride_length_s]
|
352 |
+
|
353 |
+
chunk_len = round(chunk_length_s * self.feature_extractor.sampling_rate)
|
354 |
+
stride_left = round(stride_length_s[0] * self.feature_extractor.sampling_rate)
|
355 |
+
stride_right = round(stride_length_s[1] * self.feature_extractor.sampling_rate)
|
356 |
+
|
357 |
+
if chunk_len < stride_left + stride_right:
|
358 |
+
raise ValueError("Chunk length must be superior to stride length")
|
359 |
+
|
360 |
+
for item in self.chunk_iter_with_batch(
|
361 |
+
inputs,
|
362 |
+
chunk_len,
|
363 |
+
stride_left,
|
364 |
+
stride_right,
|
365 |
+
batch_size,
|
366 |
+
):
|
367 |
+
yield item
|
368 |
+
else:
|
369 |
+
processed = self.feature_extractor(
|
370 |
+
inputs, sampling_rate=self.feature_extractor.sampling_rate, return_tensors="np"
|
371 |
+
)
|
372 |
+
if stride is not None:
|
373 |
+
processed["stride"] = stride
|
374 |
+
yield processed
|
375 |
+
|
376 |
+
def postprocess(self, model_outputs, return_timestamps=None, return_language=None):
|
377 |
+
# unpack the outputs from list(dict(list)) to list(dict)
|
378 |
+
model_outputs = [dict(zip(output, t)) for output in model_outputs for t in zip(*output.values())]
|
379 |
+
|
380 |
+
time_precision = self.feature_extractor.chunk_length / self.model.config.max_source_positions
|
381 |
+
# Send the chunking back to seconds, it's easier to handle in whisper
|
382 |
+
sampling_rate = self.feature_extractor.sampling_rate
|
383 |
+
for output in model_outputs:
|
384 |
+
if "stride" in output:
|
385 |
+
chunk_len, stride_left, stride_right = output["stride"]
|
386 |
+
# Go back in seconds
|
387 |
+
chunk_len /= sampling_rate
|
388 |
+
stride_left /= sampling_rate
|
389 |
+
stride_right /= sampling_rate
|
390 |
+
output["stride"] = chunk_len, stride_left, stride_right
|
391 |
+
|
392 |
+
text, optional = self.tokenizer._decode_asr(
|
393 |
+
model_outputs,
|
394 |
+
return_timestamps=return_timestamps,
|
395 |
+
return_language=return_language,
|
396 |
+
time_precision=time_precision,
|
397 |
+
)
|
398 |
+
return {"text": text, **optional}
|
399 |
+
|
400 |
+
def forward(self, model_inputs, batch_size=None, language=None, task=None, return_timestamps=False):
|
401 |
+
# We need to keep track of some additional input arguments for post-processing so need to forward these on after running generation
|
402 |
+
input_features = model_inputs.pop("input_features")
|
403 |
+
input_batch_size = input_features.shape[0]
|
404 |
+
|
405 |
+
if input_batch_size != batch_size:
|
406 |
+
padding = np.zeros([batch_size - input_batch_size, *input_features.shape[1:]], input_features.dtype)
|
407 |
+
input_features = np.concatenate([input_features, padding])
|
408 |
+
|
409 |
+
pred_ids = self.generate(input_features, language=language, task=task, return_timestamps=return_timestamps)[
|
410 |
+
:input_batch_size
|
411 |
+
]
|
412 |
+
|
413 |
+
# tokenizer's decode method expects an extra dim - we insert it here for convenience
|
414 |
+
out = {"tokens": pred_ids[:, None, :]}
|
415 |
+
|
416 |
+
stride = model_inputs.pop("stride", None)
|
417 |
+
if stride is not None:
|
418 |
+
out["stride"] = stride
|
419 |
+
|
420 |
+
return out
|
421 |
+
|
422 |
+
def __call__(
|
423 |
+
self,
|
424 |
+
inputs,
|
425 |
+
chunk_length_s=30.0,
|
426 |
+
stride_length_s=None,
|
427 |
+
batch_size=None,
|
428 |
+
language=None,
|
429 |
+
task=None,
|
430 |
+
return_timestamps=None,
|
431 |
+
generate_kwargs=None,
|
432 |
+
):
|
433 |
+
"""
|
434 |
+
Transcribe an audio input sequence to a text transcription, optionally with timestamps.
|
435 |
+
|
436 |
+
Args:
|
437 |
+
inputs (`np.ndarray` or `bytes` or `str` or `dict`):
|
438 |
+
The inputs is either:
|
439 |
+
- `str` that is the filename of the audio file, the file will be read at the correct sampling rate
|
440 |
+
to get the waveform using *ffmpeg*. This requires *ffmpeg* to be installed on the system.
|
441 |
+
- `bytes` is the byte content of an audio file and is interpreted by *ffmpeg* in the
|
442 |
+
same way.
|
443 |
+
- (`np.ndarray` of shape (n, ) of type `np.float32` or `np.float64`)
|
444 |
+
Raw audio assumed to be at the correct sampling rate (16kHz). Note that no further sampling
|
445 |
+
rate check will be done.
|
446 |
+
- `dict` form can be used to pass raw audio sampled at arbitrary `sampling_rate` and let this
|
447 |
+
pipeline do the resampling. The dict must be in the format `{"sampling_rate": int, "array":
|
448 |
+
np.array}`. Optionally an additional argument `"stride": (left: int, right: int)` can be used to
|
449 |
+
ask the pipeline to treat the first `left` samples and last `right` samples to be ignored in
|
450 |
+
decoding (but used at inference to provide more context to the model). In general, this additional
|
451 |
+
stride argument is not required.
|
452 |
+
chunk_length_s (`float`, *optional*, defaults to 30.0):
|
453 |
+
The input length for each chunk. If `chunk_length_s = 0` then chunking is disabled. By default, the chunk
|
454 |
+
length is set 30.0s, equal to Whisper's context window.
|
455 |
+
stride_length_s (`float`, *optional*, defaults to `chunk_length_s / 6`):
|
456 |
+
The length of stride on the left and right of each chunk. Used only with `chunk_length_s > 0`. This enables
|
457 |
+
the model to *see* more context and infer letters better than without this context but the pipeline
|
458 |
+
discards the stride bits at the end to make the final reconstitution as perfect as possible.
|
459 |
+
|
460 |
+
<Tip>
|
461 |
+
|
462 |
+
For more information on how to effectively use `stride_length_s`, refer to the [ASR chunking
|
463 |
+
blog post](https://huggingface.co/blog/asr-chunking).
|
464 |
+
|
465 |
+
</Tip>
|
466 |
+
batch_size (`int`, *optional*, defaults to the minimum per-device batch size, i.e. `jax.local_device_count()`):
|
467 |
+
The batch size to be used in chunking transcription. Beneficial for transcribing long audio files. Passing
|
468 |
+
a batch size in the `__call__` method will supersede any batch size passed to the `__init__`.
|
469 |
+
task (`str`, *optional*):
|
470 |
+
Task to use for generation, either `"transcribe"` or `"translate"`. Defaults to `"transcribe"`.
|
471 |
+
language (`str`, *optional*):
|
472 |
+
Language token to use for generation, can be either in the form of `"<|en|>"`, `"en"` or `"english"`.
|
473 |
+
Defaults to `None`, meaning the language is automatically inferred from the audio input.
|
474 |
+
return_timestamps (*optional*, `bool`):
|
475 |
+
Whether to return timestamps in the prediction. Defaults to False. If set to true, the pipeline
|
476 |
+
will return two keys in the output dictionary: `"text"` containing the text transcription, and `"chunks"`
|
477 |
+
containing the transcription segments chunked by their utterance-level timestamps.
|
478 |
+
|
479 |
+
Return:
|
480 |
+
`Dict`: A dictionary with the following keys:
|
481 |
+
- **text** (`str` ) -- The recognised text.
|
482 |
+
- **chunks** (*optional(, `List[Dict]`)
|
483 |
+
When using `return_timestamps`, the `chunks` will become a list containing all the various text
|
484 |
+
chunks identified by the model, *e.g.* `[{"text": "hi ", "timestamps": (0.5,0.9), {"text":
|
485 |
+
"there", "timestamps": (1.0, 1.5)}]`. The original full text can roughly be recovered by doing
|
486 |
+
`"".join(chunk["text"] for chunk in output["chunks"])`.
|
487 |
+
"""
|
488 |
+
batch_size = batch_size if batch_size is not None else self.batch_size
|
489 |
+
if batch_size % self.min_batch_size != 0:
|
490 |
+
raise ValueError(
|
491 |
+
f"Batch size must be a multiple of the number of JAX devices, but got batch size {batch_size} and num devices {self.min_batch_size}."
|
492 |
+
)
|
493 |
+
|
494 |
+
dataloader = self.preprocess_batch(
|
495 |
+
inputs, chunk_length_s=chunk_length_s, stride_length_s=stride_length_s, batch_size=batch_size
|
496 |
+
)
|
497 |
+
model_outputs = []
|
498 |
+
# iterate over our chunked audio samples
|
499 |
+
for batch in dataloader:
|
500 |
+
model_outputs.append(
|
501 |
+
self.forward(
|
502 |
+
batch, batch_size=batch_size, language=language, task=task, return_timestamps=return_timestamps
|
503 |
+
)
|
504 |
+
)
|
505 |
+
post_processed = self.postprocess(model_outputs, return_timestamps=return_timestamps)
|
506 |
+
return post_processed
|
whisper_jax/train_state.py
ADDED
@@ -0,0 +1,130 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# coding=utf-8
|
2 |
+
# Copyright 2023 The T5X Authors and The HuggingFace Inc. team.
|
3 |
+
#
|
4 |
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
5 |
+
# you may not use this file except in compliance with the License.
|
6 |
+
# You may obtain a copy of the License at
|
7 |
+
#
|
8 |
+
# http://www.apache.org/licenses/LICENSE-2.0
|
9 |
+
#
|
10 |
+
# Unless required by applicable law or agreed to in writing, software
|
11 |
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
12 |
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
13 |
+
# See the License for the specific language governing permissions and
|
14 |
+
# limitations under the License.
|
15 |
+
|
16 |
+
"""Utilities for partitioning."""
|
17 |
+
|
18 |
+
from typing import Any, Mapping, MutableMapping, Optional, Tuple
|
19 |
+
|
20 |
+
import flax.core
|
21 |
+
import flax.serialization
|
22 |
+
import flax.struct
|
23 |
+
import jax.numpy as jnp
|
24 |
+
from flax import traverse_util
|
25 |
+
from flax.core import scope as flax_scope
|
26 |
+
from flax.linen import partitioning as flax_partitioning
|
27 |
+
|
28 |
+
|
29 |
+
EMPTY_DICT = flax.core.freeze({})
|
30 |
+
FrozenDict = flax_scope.FrozenDict
|
31 |
+
FrozenVariableDict = flax_scope.FrozenVariableDict
|
32 |
+
MutableVariableDict = flax_scope.MutableVariableDict
|
33 |
+
VariableDict = flax_scope.VariableDict
|
34 |
+
|
35 |
+
|
36 |
+
def _validate_params_axes(params_axes, params):
|
37 |
+
axis_names = flax_partitioning.get_axis_names(params_axes)
|
38 |
+
missing_params_axes = set(traverse_util.flatten_dict(params, sep="/")) - set(
|
39 |
+
traverse_util.flatten_dict(axis_names, sep="/")
|
40 |
+
)
|
41 |
+
if missing_params_axes:
|
42 |
+
raise ValueError(f"Missing axis names for parameters: {missing_params_axes}")
|
43 |
+
|
44 |
+
|
45 |
+
def _split_variables_and_axes(variables_and_axes: FrozenVariableDict) -> Tuple[FrozenVariableDict, FrozenVariableDict]:
|
46 |
+
"""Splits `variables_and_axes` into two separate dicts with the same keys."""
|
47 |
+
# For each `key`, `key_axes` (if any) are its axes in `variables_and_axes`.
|
48 |
+
variables = {}
|
49 |
+
axes = {}
|
50 |
+
for k, v in variables_and_axes.items():
|
51 |
+
if k.endswith("_axes"):
|
52 |
+
axes[k[:-5]] = v # k without "_axes".
|
53 |
+
_validate_params_axes(v, variables_and_axes[k[:-5]]) # k without "_axes".
|
54 |
+
else:
|
55 |
+
variables[k] = v
|
56 |
+
return flax.core.freeze(variables), flax.core.freeze(axes)
|
57 |
+
|
58 |
+
|
59 |
+
class InferenceState(flax.struct.PyTreeNode):
|
60 |
+
"""State compatible with FlaxOptimTrainState without optimizer state."""
|
61 |
+
|
62 |
+
step: jnp.ndarray
|
63 |
+
params: flax_scope.FrozenVariableDict
|
64 |
+
params_axes: Optional[flax_scope.FrozenVariableDict] = None
|
65 |
+
flax_mutables: flax_scope.FrozenDict = EMPTY_DICT
|
66 |
+
flax_mutables_axes: Optional[flax_scope.FrozenVariableDict] = None
|
67 |
+
|
68 |
+
@classmethod
|
69 |
+
def create(cls, model_variables: FrozenVariableDict) -> "InferenceState":
|
70 |
+
other_variables, params = model_variables.pop("params")
|
71 |
+
if "params_axes" in other_variables:
|
72 |
+
other_variables, params_axes = other_variables.pop("params_axes")
|
73 |
+
_validate_params_axes(params_axes, params)
|
74 |
+
else:
|
75 |
+
params_axes = None
|
76 |
+
|
77 |
+
# Split other_variables into mutables and their corresponding axes.
|
78 |
+
flax_mutables, flax_mutables_axes = _split_variables_and_axes(other_variables)
|
79 |
+
flax_mutables_axes = flax_mutables_axes or None
|
80 |
+
return InferenceState(
|
81 |
+
step=jnp.array(0),
|
82 |
+
params=params,
|
83 |
+
params_axes=params_axes,
|
84 |
+
flax_mutables=flax_mutables,
|
85 |
+
flax_mutables_axes=flax_mutables_axes,
|
86 |
+
)
|
87 |
+
|
88 |
+
@property
|
89 |
+
def param_states(self) -> FrozenVariableDict:
|
90 |
+
"""The optimizer states of the parameters as a PyTree."""
|
91 |
+
raise NotImplementedError("InferenceState has no optimizer states.")
|
92 |
+
|
93 |
+
def apply_gradient(self, *args, **kwargs) -> "InferenceState":
|
94 |
+
raise NotImplementedError("InferenceState does not support `apply_gradient`.")
|
95 |
+
|
96 |
+
def state_dict(self) -> MutableMapping[str, Any]:
|
97 |
+
state_dict = {"target": flax.core.unfreeze(self.params), "state": {"step": self.step}}
|
98 |
+
if self.flax_mutables:
|
99 |
+
state_dict["flax_mutables"] = flax.core.unfreeze(self.flax_mutables)
|
100 |
+
return state_dict
|
101 |
+
|
102 |
+
def replace_step(self, step: jnp.ndarray) -> "InferenceState":
|
103 |
+
return self.replace(step=step)
|
104 |
+
|
105 |
+
def replace_params(self, params: FrozenVariableDict) -> "InferenceState":
|
106 |
+
return self.replace(params=params)
|
107 |
+
|
108 |
+
def replace_flax_mutables(self, flax_mutables: FrozenDict) -> "InferenceState":
|
109 |
+
return self.replace(flax_mutables=flax_mutables)
|
110 |
+
|
111 |
+
def restore_state(self, state_dict: Mapping[str, Any]) -> "InferenceState":
|
112 |
+
return self.replace(
|
113 |
+
params=flax.core.freeze(state_dict["target"]),
|
114 |
+
step=state_dict["state"]["step"],
|
115 |
+
flax_mutables=flax.core.freeze(state_dict["flax_mutables"])
|
116 |
+
if "flax_mutables" in state_dict
|
117 |
+
else EMPTY_DICT,
|
118 |
+
)
|
119 |
+
|
120 |
+
def as_logical_axes(self) -> "InferenceState":
|
121 |
+
# Set step to None so that when the logical axes are processed by the
|
122 |
+
# flax.partitioning.logical_to_mesh_axes function, it will be skipped
|
123 |
+
# because jax.tree_map will short circut and never call the function on the
|
124 |
+
# step.
|
125 |
+
flax_mutables_axes = self.flax_mutables_axes or EMPTY_DICT
|
126 |
+
return InferenceState(
|
127 |
+
step=None,
|
128 |
+
params=flax_partitioning.get_axis_names(self.params_axes),
|
129 |
+
flax_mutables=flax_partitioning.get_axis_names(flax_mutables_axes),
|
130 |
+
)
|