Aluode/PerceptionLabPortable
0
1# coding=utf-82# Copyright 2022 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 at7#8# http://www.apache.org/licenses/LICENSE-2.09#10# Unless required by applicable law or agreed to in writing, software11# 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 and14# limitations under the License.15"""PyTorch Whisper model."""16 17import math18from typing import Callable, Optional, Union19 20import numpy as np21import torch22from torch import nn23from torch.nn import CrossEntropyLoss24 25from ...activations import ACT2FN26from ...cache_utils import Cache, DynamicCache, EncoderDecoderCache27from ...generation import GenerationMixin28from ...masking_utils import create_causal_mask29from ...modeling_flash_attention_utils import (30 FlashAttentionKwargs,31)32from ...modeling_layers import GradientCheckpointingLayer33from ...modeling_outputs import (34 BaseModelOutput,35 BaseModelOutputWithPastAndCrossAttentions,36 CausalLMOutputWithCrossAttentions,37 Seq2SeqLMOutput,38 Seq2SeqModelOutput,39 SequenceClassifierOutput,40)41from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel42from ...processing_utils import Unpack43from ...utils import auto_docstring, logging44from ...utils.deprecation import deprecate_kwarg45from .configuration_whisper import WhisperConfig46from .generation_whisper import WhisperGenerationMixin47 48 49logger = logging.get_logger(__name__)50 51_HIDDEN_STATES_START_POSITION = 152 53 54def sinusoids(length: int, channels: int, max_timescale: float = 10000) -> torch.Tensor:55 """Returns sinusoids for positional embedding"""56 if channels % 2 != 0:57 raise ValueError(58 f"Number of channels has to be divisible by 2 for sinusoidal positional embeddings, got {channels} channels."59 )60 log_timescale_increment = math.log(max_timescale) / (channels // 2 - 1)61 inv_timescales = torch.exp(-log_timescale_increment * torch.arange(channels // 2))62 scaled_time = torch.arange(length).view(-1, 1) * inv_timescales.view(1, -1)63 return torch.cat([scaled_time.sin(), scaled_time.cos()], dim=1)64 65 66# Copied from transformers.models.bart.modeling_bart.shift_tokens_right67def shift_tokens_right(input_ids: torch.Tensor, pad_token_id: int, decoder_start_token_id: int):68 """69 Shift input ids one token to the right.70 """71 shifted_input_ids = input_ids.new_zeros(input_ids.shape)72 shifted_input_ids[:, 1:] = input_ids[:, :-1].clone()73 shifted_input_ids[:, 0] = decoder_start_token_id74 75 if pad_token_id is None:76 raise ValueError("self.model.config.pad_token_id has to be defined.")77 # replace possible -100 values in labels by `pad_token_id`78 shifted_input_ids.masked_fill_(shifted_input_ids == -100, pad_token_id)79 80 return shifted_input_ids81 82 83# Copied from transformers.models.wav2vec2.modeling_wav2vec2._compute_mask_indices84def _compute_mask_indices(85 shape: tuple[int, int],86 mask_prob: float,87 mask_length: int,88 attention_mask: Optional[torch.LongTensor] = None,89 min_masks: int = 0,90) -> np.ndarray:91 """92 Computes random mask spans for a given shape. Used to implement [SpecAugment: A Simple Data Augmentation Method for93 ASR](https://huggingface.co/papers/1904.08779). Note that this method is not optimized to run on TPU and should be run on94 CPU as part of the preprocessing during training.95 96 Args:97 shape: The shape for which to compute masks. This should be of a tuple of size 2 where98 the first element is the batch size and the second element is the length of the axis to span.99 mask_prob: The percentage of the whole axis (between 0 and 1) which will be masked. The number of100 independently generated mask spans of length `mask_length` is computed by101 `mask_prob*shape[1]/mask_length`. Note that due to overlaps, `mask_prob` is an upper bound and the102 actual percentage will be smaller.103 mask_length: size of the mask104 min_masks: minimum number of masked spans105 attention_mask: A (right-padded) attention mask which independently shortens the feature axis of106 each batch dimension.107 """108 batch_size, sequence_length = shape109 110 if mask_length < 1:111 raise ValueError("`mask_length` has to be bigger than 0.")112 113 if mask_length > sequence_length:114 raise ValueError(115 f"`mask_length` has to be smaller than `sequence_length`, but got `mask_length`: {mask_length}"116 f" and `sequence_length`: {sequence_length}`"117 )118 119 # epsilon is used for probabilistic rounding120 epsilon = np.random.rand(1).item()121 122 def compute_num_masked_span(input_length):123 """Given input length, compute how many spans should be masked"""124 num_masked_span = int(mask_prob * input_length / mask_length + epsilon)125 num_masked_span = max(num_masked_span, min_masks)126 127 # make sure num masked span <= sequence_length128 if num_masked_span * mask_length > sequence_length:129 num_masked_span = sequence_length // mask_length130 131 # make sure num_masked span is also <= input_length - (mask_length - 1)132 if input_length - (mask_length - 1) < num_masked_span:133 num_masked_span = max(input_length - (mask_length - 1), 0)134 135 return num_masked_span136 137 # compute number of masked spans in batch138 input_lengths = (139 attention_mask.detach().sum(-1).tolist()140 if attention_mask is not None141 else [sequence_length for _ in range(batch_size)]142 )143 144 # SpecAugment mask to fill145 spec_aug_mask = np.zeros((batch_size, sequence_length), dtype=bool)146 spec_aug_mask_idxs = []147 148 max_num_masked_span = compute_num_masked_span(sequence_length)149 150 if max_num_masked_span == 0:151 return spec_aug_mask152 153 for input_length in input_lengths:154 # compute num of masked spans for this input155 num_masked_span = compute_num_masked_span(input_length)156 157 # get random indices to mask158 spec_aug_mask_idx = np.random.choice(159 np.arange(input_length - (mask_length - 1)), num_masked_span, replace=False160 )161 162 # pick first sampled index that will serve as a dummy index to pad vector163 # to ensure same dimension for all batches due to probabilistic rounding164 # Picking first sample just pads those vectors twice.165 if len(spec_aug_mask_idx) == 0:166 # this case can only happen if `input_length` is strictly smaller then167 # `sequence_length` in which case the last token has to be a padding168 # token which we can use as a dummy mask id169 dummy_mask_idx = sequence_length - 1170 else:171 dummy_mask_idx = spec_aug_mask_idx[0]172 173 spec_aug_mask_idx = np.concatenate(174 [spec_aug_mask_idx, np.ones(max_num_masked_span - num_masked_span, dtype=np.int32) * dummy_mask_idx]175 )176 spec_aug_mask_idxs.append(spec_aug_mask_idx)177 178 spec_aug_mask_idxs = np.array(spec_aug_mask_idxs)179 180 # expand masked indices to masked spans181 spec_aug_mask_idxs = np.broadcast_to(182 spec_aug_mask_idxs[:, :, None], (batch_size, max_num_masked_span, mask_length)183 )184 spec_aug_mask_idxs = spec_aug_mask_idxs.reshape(batch_size, max_num_masked_span * mask_length)185 186 # add offset to the starting indexes so that indexes now create a span187 offsets = np.arange(mask_length)[None, None, :]188 offsets = np.broadcast_to(offsets, (batch_size, max_num_masked_span, mask_length)).reshape(189 batch_size, max_num_masked_span * mask_length190 )191 spec_aug_mask_idxs = spec_aug_mask_idxs + offsets192 193 # ensure that we cannot have indices larger than sequence_length194 if spec_aug_mask_idxs.max() > sequence_length - 1:195 spec_aug_mask_idxs[spec_aug_mask_idxs > sequence_length - 1] = sequence_length - 1196 197 # scatter indices to mask198 np.put_along_axis(spec_aug_mask, spec_aug_mask_idxs, 1, -1)199 200 return spec_aug_mask201 202 203class WhisperPositionalEmbedding(nn.Embedding):204 def __init__(self, num_positions: int, embedding_dim: int, padding_idx: Optional[int] = None):205 super().__init__(num_positions, embedding_dim)206 207 def forward(self, input_ids, past_key_values_length=0, position_ids=None):208 if position_ids is None:209 return self.weight[past_key_values_length : past_key_values_length + input_ids.shape[1]]210 else:211 return self.weight[position_ids]212 213 214def eager_attention_forward(215 module: nn.Module,216 query: torch.Tensor,217 key: torch.Tensor,218 value: torch.Tensor,219 attention_mask: Optional[torch.Tensor],220 scaling: Optional[float] = None,221 dropout: float = 0.0,222 head_mask: Optional[torch.Tensor] = None,223 **kwargs,224):225 if scaling is None:226 scaling = query.size(-1) ** -0.5227 228 attn_weights = torch.matmul(query, key.transpose(2, 3)) * scaling229 if attention_mask is not None and attention_mask.ndim == 4:230 attn_weights = attn_weights + attention_mask[:, :, :, : key.shape[-2]]231 232 attn_weights = nn.functional.softmax(attn_weights, dim=-1)233 234 if head_mask is not None:235 attn_weights = attn_weights * head_mask.view(1, -1, 1, 1)236 237 attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)238 attn_output = torch.matmul(attn_weights, value)239 attn_output = attn_output.transpose(1, 2).contiguous()240 241 return attn_output, attn_weights242 243 244class WhisperAttention(nn.Module):245 """Multi-headed attention from 'Attention Is All You Need' paper"""246 247 def __init__(248 self,249 embed_dim: int,250 num_heads: int,251 dropout: float = 0.0,252 is_decoder: bool = False,253 bias: bool = True,254 is_causal: bool = False,255 layer_idx: Optional[int] = None,256 config: Optional[WhisperConfig] = None,257 ):258 super().__init__()259 self.embed_dim = embed_dim260 self.num_heads = num_heads261 self.dropout = dropout262 self.head_dim = embed_dim // num_heads263 self.config = config264 265 if (self.head_dim * num_heads) != self.embed_dim:266 raise ValueError(267 f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim}"268 f" and `num_heads`: {num_heads})."269 )270 self.scaling = self.head_dim**-0.5271 self.is_decoder = is_decoder272 self.is_causal = is_causal273 274 if layer_idx is None and is_decoder:275 logger.warning_once(276 f"Instantiating a decoder {self.__class__.__name__} without passing `layer_idx` is not recommended and "277 "will to errors during the forward call, if caching is used. Please make sure to provide a `layer_idx` "278 "when creating this class."279 )280 self.layer_idx = layer_idx281 282 self.k_proj = nn.Linear(embed_dim, embed_dim, bias=False)283 self.v_proj = nn.Linear(embed_dim, embed_dim, bias=bias)284 self.q_proj = nn.Linear(embed_dim, embed_dim, bias=bias)285 self.out_proj = nn.Linear(embed_dim, embed_dim, bias=bias)286 287 @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")288 def forward(289 self,290 hidden_states: torch.Tensor,291 key_value_states: Optional[torch.Tensor] = None,292 past_key_values: Optional[Cache] = None,293 attention_mask: Optional[torch.Tensor] = None,294 layer_head_mask: Optional[torch.Tensor] = None,295 output_attentions: bool = False,296 cache_position: Optional[torch.Tensor] = None,297 # TODO: we need a refactor so that the different attention modules can get their specific kwargs298 # ATM, we have mixed things encoder, decoder, and encoder-decoder attn299 **kwargs: Unpack[FlashAttentionKwargs],300 ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[tuple[torch.Tensor]]]:301 """Input shape: Batch x Time x Channel"""302 303 # if key_value_states are provided this layer is used as a cross-attention layer304 # for the decoder305 is_cross_attention = key_value_states is not None306 307 # determine input shapes308 bsz, tgt_len = hidden_states.shape[:-1]309 q_input_shape = (bsz, tgt_len, -1, self.head_dim)310 311 # Scaling is susceptible to floating point arithmetics' inprecisions312 # which can lead to different results (this is dependent from model313 # to model, e.g. whisper is one such case). We therefore keep the314 # original order of scaling to follow the original implementation315 # and enforce no scaling (1.0) in the attention call below.316 query_states = self.q_proj(hidden_states) * self.scaling317 query_states = query_states.view(*q_input_shape)318 query_states = query_states.transpose(1, 2).contiguous()319 320 # Check is encoder-decoder model is being used. Otherwise we'll get `DynamicCache`321 if past_key_values is not None and isinstance(past_key_values, EncoderDecoderCache):322 is_updated = past_key_values.is_updated.get(self.layer_idx)323 if is_cross_attention:324 # after the first generated id, we can subsequently re-use all key/value_states from cache325 past_key_values.is_updated[self.layer_idx] = True326 past_key_values = past_key_values.cross_attention_cache327 else:328 past_key_values = past_key_values.self_attention_cache329 330 # use key_value_states if cross attention331 current_states = key_value_states if key_value_states is not None else hidden_states332 if is_cross_attention and past_key_values and is_updated:333 # reuse k,v, cross_attentions334 key_states = past_key_values.layers[self.layer_idx].keys335 value_states = past_key_values.layers[self.layer_idx].values336 else:337 key_states = self.k_proj(current_states).view(bsz, -1, self.num_heads, self.head_dim)338 value_states = self.v_proj(current_states).view(bsz, -1, self.num_heads, self.head_dim)339 key_states = key_states.transpose(1, 2).contiguous()340 value_states = value_states.transpose(1, 2).contiguous()341 if past_key_values is not None:342 # save all key/value_states to cache to be re-used for fast auto-regressive generation343 cache_position = cache_position if not is_cross_attention else None344 key_states, value_states = past_key_values.update(345 key_states, value_states, self.layer_idx, {"cache_position": cache_position}346 )347 348 attention_interface: Callable = eager_attention_forward349 if self.config._attn_implementation != "eager":350 attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]351 352 attn_output, attn_weights = attention_interface(353 self,354 query_states,355 key_states,356 value_states,357 attention_mask,358 dropout=0.0 if not self.training else self.dropout,359 scaling=1.0,360 output_attentions=output_attentions,361 head_mask=layer_head_mask,362 **kwargs,363 )364 365 attn_output = attn_output.reshape(bsz, tgt_len, -1).contiguous()366 attn_output = self.out_proj(attn_output)367 368 return attn_output, attn_weights369 370 371# Copied from transformers.models.mbart.modeling_mbart.MBartEncoderLayer with MBart->Whisper, MBART->WHISPER372class WhisperEncoderLayer(GradientCheckpointingLayer):373 def __init__(self, config: WhisperConfig):374 super().__init__()375 self.embed_dim = config.d_model376 377 self.self_attn = WhisperAttention(378 embed_dim=self.embed_dim,379 num_heads=config.encoder_attention_heads,380 dropout=config.attention_dropout,381 config=config,382 )383 self.self_attn_layer_norm = nn.LayerNorm(self.embed_dim)384 self.dropout = config.dropout385 self.activation_fn = ACT2FN[config.activation_function]386 self.activation_dropout = config.activation_dropout387 self.fc1 = nn.Linear(self.embed_dim, config.encoder_ffn_dim)388 self.fc2 = nn.Linear(config.encoder_ffn_dim, self.embed_dim)389 self.final_layer_norm = nn.LayerNorm(self.embed_dim)390 391 def forward(392 self,393 hidden_states: torch.Tensor,394 attention_mask: torch.Tensor,395 layer_head_mask: torch.Tensor,396 output_attentions: bool = False,397 ) -> torch.Tensor:398 """399 Args:400 hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`401 attention_mask (`torch.FloatTensor`): attention mask of size402 `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values.403 layer_head_mask (`torch.FloatTensor`): mask for attention heads in a given layer of size404 `(encoder_attention_heads,)`.405 output_attentions (`bool`, *optional*):406 Whether or not to return the attentions tensors of all attention layers. See `attentions` under407 returned tensors for more detail.408 """409 residual = hidden_states410 hidden_states = self.self_attn_layer_norm(hidden_states)411 hidden_states, attn_weights = self.self_attn(412 hidden_states=hidden_states,413 attention_mask=attention_mask,414 layer_head_mask=layer_head_mask,415 output_attentions=output_attentions,416 )417 hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)418 hidden_states = residual + hidden_states419 420 residual = hidden_states421 hidden_states = self.final_layer_norm(hidden_states)422 hidden_states = self.activation_fn(self.fc1(hidden_states))423 hidden_states = nn.functional.dropout(hidden_states, p=self.activation_dropout, training=self.training)424 hidden_states = self.fc2(hidden_states)425 hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)426 hidden_states = residual + hidden_states427 428 if hidden_states.dtype == torch.float16:429 clamp_value = torch.finfo(hidden_states.dtype).max - 1000430 hidden_states = torch.clamp(hidden_states, min=-clamp_value, max=clamp_value)431 432 return hidden_states, attn_weights433 434 435class WhisperDecoderLayer(GradientCheckpointingLayer):436 def __init__(self, config: WhisperConfig, layer_idx: Optional[int] = None):437 super().__init__()438 self.embed_dim = config.d_model439 440 self.self_attn = WhisperAttention(441 embed_dim=self.embed_dim,442 num_heads=config.decoder_attention_heads,443 dropout=config.attention_dropout,444 is_decoder=True,445 is_causal=True,446 layer_idx=layer_idx,447 config=config,448 )449 self.dropout = config.dropout450 self.activation_fn = ACT2FN[config.activation_function]451 self.activation_dropout = config.activation_dropout452 453 self.self_attn_layer_norm = nn.LayerNorm(self.embed_dim)454 self.encoder_attn = WhisperAttention(455 self.embed_dim,456 config.decoder_attention_heads,457 dropout=config.attention_dropout,458 is_decoder=True,459 layer_idx=layer_idx,460 config=config,461 )462 self.encoder_attn_layer_norm = nn.LayerNorm(self.embed_dim)463 self.fc1 = nn.Linear(self.embed_dim, config.decoder_ffn_dim)464 self.fc2 = nn.Linear(config.decoder_ffn_dim, self.embed_dim)465 self.final_layer_norm = nn.LayerNorm(self.embed_dim)466 467 @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")468 def forward(469 self,470 hidden_states: torch.Tensor,471 attention_mask: Optional[torch.Tensor] = None,472 encoder_hidden_states: Optional[torch.Tensor] = None,473 encoder_attention_mask: Optional[torch.Tensor] = None,474 layer_head_mask: Optional[torch.Tensor] = None,475 cross_attn_layer_head_mask: Optional[torch.Tensor] = None,476 past_key_values: Optional[EncoderDecoderCache] = None,477 output_attentions: Optional[bool] = False,478 use_cache: Optional[bool] = True,479 cache_position: Optional[torch.LongTensor] = None,480 ) -> torch.Tensor:481 """482 Args:483 hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`484 attention_mask (`torch.FloatTensor`): attention mask of size485 `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values.486 encoder_hidden_states (`torch.FloatTensor`):487 cross attention input to the layer of shape `(batch, seq_len, embed_dim)`488 encoder_attention_mask (`torch.FloatTensor`): encoder attention mask of size489 `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values.490 layer_head_mask (`torch.FloatTensor`): mask for attention heads in a given layer of size491 `(encoder_attention_heads,)`.492 cross_attn_layer_head_mask (`torch.FloatTensor`): mask for cross-attention heads in a given layer of493 size `(decoder_attention_heads,)`.494 past_key_values (`Cache`): cached past key and value projection states495 output_attentions (`bool`, *optional*):496 Whether or not to return the attentions tensors of all attention layers. See `attentions` under497 returned tensors for more detail.498 """499 residual = hidden_states500 hidden_states = self.self_attn_layer_norm(hidden_states)501 502 # Self Attention503 hidden_states, self_attn_weights = self.self_attn(504 hidden_states=hidden_states,505 past_key_values=past_key_values,506 attention_mask=attention_mask,507 layer_head_mask=layer_head_mask,508 output_attentions=output_attentions,509 cache_position=cache_position,510 )511 hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)512 hidden_states = residual + hidden_states513 514 # Cross-Attention Block515 cross_attn_weights = None516 if encoder_hidden_states is not None:517 residual = hidden_states518 hidden_states = self.encoder_attn_layer_norm(hidden_states)519 hidden_states, cross_attn_weights = self.encoder_attn(520 hidden_states=hidden_states,521 key_value_states=encoder_hidden_states,522 attention_mask=encoder_attention_mask,523 layer_head_mask=cross_attn_layer_head_mask,524 past_key_values=past_key_values,525 output_attentions=output_attentions,526 )527 hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)528 hidden_states = residual + hidden_states529 530 # Fully Connected531 residual = hidden_states532 hidden_states = self.final_layer_norm(hidden_states)533 hidden_states = self.activation_fn(self.fc1(hidden_states))534 hidden_states = nn.functional.dropout(hidden_states, p=self.activation_dropout, training=self.training)535 hidden_states = self.fc2(hidden_states)536 hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)537 hidden_states = residual + hidden_states538 539 outputs = (hidden_states,)540 541 if output_attentions:542 outputs += (self_attn_weights, cross_attn_weights)543 544 return outputs545 546 547@auto_docstring548class WhisperPreTrainedModel(PreTrainedModel):549 config: WhisperConfig550 base_model_prefix = "model"551 main_input_name = "input_features"552 supports_gradient_checkpointing = True553 _no_split_modules = ["WhisperEncoderLayer", "WhisperDecoderLayer"]554 _supports_flash_attn = True555 _supports_sdpa = True556 _supports_flex_attn = True557 558 _can_compile_fullgraph = True559 560 def _init_weights(self, module):561 std = self.config.init_std562 if isinstance(module, (nn.Linear, nn.Conv1d)):563 module.weight.data.normal_(mean=0.0, std=std)564 if module.bias is not None:565 module.bias.data.zero_()566 elif isinstance(module, nn.Embedding):567 module.weight.data.normal_(mean=0.0, std=std)568 if module.padding_idx is not None:569 module.weight.data[module.padding_idx].zero_()570 elif isinstance(module, nn.LayerNorm):571 module.weight.data.fill_(1.0)572 module.bias.data.zero_()573 elif isinstance(module, WhisperEncoder):574 module.embed_positions.weight.copy_(sinusoids(*module.embed_positions.weight.shape))575 elif isinstance(module, WhisperForAudioClassification):576 if self.config.use_weighted_layer_sum:577 module.layer_weights.data.fill_(1.0 / (self.config.num_hidden_layers + 1))578 579 def _get_feat_extract_output_lengths(self, input_lengths: torch.LongTensor):580 """581 Computes the output length of the convolutional layers582 """583 input_lengths = (input_lengths - 1) // 2 + 1584 585 return input_lengths586 587 588class WhisperEncoder(WhisperPreTrainedModel):589 """590 Transformer encoder consisting of *config.encoder_layers* self attention layers. Each layer is a591 [`WhisperEncoderLayer`].592 593 Args:594 config: WhisperConfig595 """596 597 def __init__(self, config: WhisperConfig):598 super().__init__(config)599 self.dropout = config.dropout600 self.layerdrop = config.encoder_layerdrop601 602 embed_dim = config.d_model603 self.num_mel_bins = config.num_mel_bins604 self.padding_idx = config.pad_token_id605 self.max_source_positions = config.max_source_positions606 self.embed_scale = math.sqrt(embed_dim) if config.scale_embedding else 1.0607 608 self.conv1 = nn.Conv1d(self.num_mel_bins, embed_dim, kernel_size=3, padding=1)609 self.conv2 = nn.Conv1d(embed_dim, embed_dim, kernel_size=3, stride=2, padding=1)610 611 self.embed_positions = nn.Embedding(self.max_source_positions, embed_dim)612 self.embed_positions.requires_grad_(False)613 614 self.layers = nn.ModuleList([WhisperEncoderLayer(config) for _ in range(config.encoder_layers)])615 self.layer_norm = nn.LayerNorm(config.d_model)616 617 self.gradient_checkpointing = False618 # Initialize weights and apply final processing619 self.post_init()620 621 def _freeze_parameters(self):622 for param in self.parameters():623 param.requires_grad = False624 self._requires_grad = False625 626 def get_input_embeddings(self) -> nn.Module:627 return self.conv1628 629 def set_input_embeddings(self, value: nn.Module):630 self.conv1 = value631 632 def forward(633 self,634 input_features,635 attention_mask=None,636 head_mask=None,637 output_attentions=None,638 output_hidden_states=None,639 return_dict=None,640 ):641 r"""642 Args:643 input_features (`torch.LongTensor` of shape `(batch_size, feature_size, sequence_length)`):644 Float values of mel features extracted from the raw speech waveform. Raw speech waveform can be645 obtained by loading a `.flac` or `.wav` audio file into an array of type `list[float]`, a646 `numpy.ndarray` or a `torch.Tensor`, *e.g.* via the torchcodec library (`pip install torchcodec`) or647 the soundfile library (`pip install soundfile`). To prepare the array into648 `input_features`, the [`AutoFeatureExtractor`] should be used for extracting the mel features, padding649 and conversion into a tensor of type `torch.FloatTensor`. See [`~WhisperFeatureExtractor.__call__`]650 attention_mask (`torch.Tensor`)`, *optional*):651 Whisper does not support masking of the `input_features`, this argument is preserved for compatibility,652 but it is not used. By default the silence in the input log mel spectrogram are ignored.653 head_mask (`torch.Tensor` of shape `(encoder_layers, encoder_attention_heads)`, *optional*):654 Mask to nullify selected heads of the attention modules. Mask values selected in `[0, 1]`:655 656 - 1 indicates the head is **not masked**,657 - 0 indicates the head is **masked**.658 output_attentions (`bool`, *optional*):659 Whether or not to return the attentions tensors of all attention layers. See `attentions` under660 returned tensors for more detail.661 output_hidden_states (`bool`, *optional*):662 Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors663 for more detail.664 return_dict (`bool`, *optional*):665 Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.666 """667 668 expected_seq_length = self.config.max_source_positions * self.conv1.stride[0] * self.conv2.stride[0]669 if input_features.shape[-1] != expected_seq_length:670 raise ValueError(671 f"Whisper expects the mel input features to be of length {expected_seq_length}, but found {input_features.shape[-1]}. Make sure to pad the input mel features to {expected_seq_length}."672 )673 674 output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions675 output_hidden_states = (676 output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states677 )678 return_dict = return_dict if return_dict is not None else self.config.use_return_dict679 inputs_embeds = nn.functional.gelu(self.conv1(input_features))680 inputs_embeds = nn.functional.gelu(self.conv2(inputs_embeds))681 682 inputs_embeds = inputs_embeds.permute(0, 2, 1)683 all_positions = torch.arange(self.embed_positions.num_embeddings, device=inputs_embeds.device)684 685 hidden_states = inputs_embeds + self.embed_positions(all_positions)686 hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)687 688 encoder_states = () if output_hidden_states else None689 all_attentions = () if output_attentions else None690 691 # check if head_mask has a correct number of layers specified if desired692 if head_mask is not None:693 assert head_mask.size()[0] == (len(self.layers)), (694 f"The head_mask should be specified for {len(self.layers)} layers, but it is for {head_mask.size()[0]}."695 )696 697 for idx, encoder_layer in enumerate(self.layers):698 if output_hidden_states:699 encoder_states = encoder_states + (hidden_states,)700 # add LayerDrop (see https://huggingface.co/papers/1909.11556 for description)701 to_drop = False702 if self.training:703 dropout_probability = torch.rand([])704 if dropout_probability < self.layerdrop: # skip the layer705 to_drop = True706 707 if to_drop:708 layer_outputs = (None, None)709 else:710 layer_outputs = encoder_layer(711 hidden_states,712 None,713 layer_head_mask=(head_mask[idx] if head_mask is not None else None),714 output_attentions=output_attentions,715 )716 717 hidden_states = layer_outputs[0]718 719 if output_attentions:720 all_attentions = all_attentions + (layer_outputs[1],)721 722 hidden_states = self.layer_norm(hidden_states)723 if output_hidden_states:724 encoder_states = encoder_states + (hidden_states,)725 726 if not return_dict:727 return tuple(v for v in [hidden_states, encoder_states, all_attentions] if v is not None)728 return BaseModelOutput(729 last_hidden_state=hidden_states, hidden_states=encoder_states, attentions=all_attentions730 )731 732 733class WhisperDecoder(WhisperPreTrainedModel):734 """735 Transformer decoder consisting of *config.decoder_layers* layers. Each layer is a [`WhisperDecoderLayer`]736 737 Args:738 config: WhisperConfig739 """740 741 main_input_name = "input_ids"742 743 def __init__(self, config: WhisperConfig):744 super().__init__(config)745 self.dropout = config.dropout746 self.layerdrop = config.decoder_layerdrop747 self.padding_idx = config.pad_token_id748 self.max_target_positions = config.max_target_positions749 self.max_source_positions = config.max_source_positions750 self.embed_scale = math.sqrt(config.d_model) if config.scale_embedding else 1.0751 752 self.embed_tokens = nn.Embedding(config.vocab_size, config.d_model, self.padding_idx)753 self.embed_positions = WhisperPositionalEmbedding(self.max_target_positions, config.d_model)754 755 self.layers = nn.ModuleList(756 [WhisperDecoderLayer(config, layer_idx) for layer_idx in range(config.decoder_layers)]757 )758 759 self.layer_norm = nn.LayerNorm(config.d_model)760 761 self.gradient_checkpointing = False762 # Initialize weights and apply final processing763 self.post_init()764 765 def forward(766 self,767 input_ids=None,768 attention_mask=None,769 encoder_hidden_states=None,770 head_mask=None,771 cross_attn_head_mask=None,772 past_key_values=None,773 inputs_embeds=None,774 position_ids=None,775 use_cache=None,776 output_attentions=None,777 output_hidden_states=None,778 return_dict=None,779 cache_position=None,780 ):781 r"""782 Args:783 input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):784 Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you785 provide it.786 787 Indices can be obtained using [`WhisperTokenizer`]. See [`PreTrainedTokenizer.encode`] and788 [`PreTrainedTokenizer.__call__`] for details.789 790 [What are input IDs?](../glossary#input-ids)791 attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):792 Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:793 794 - 1 for tokens that are **not masked**,795 - 0 for tokens that are **masked**.796 797 [What are attention masks?](../glossary#attention-mask)798 encoder_hidden_states (`torch.FloatTensor` of shape `(batch_size, encoder_sequence_length, hidden_size)`, *optional*):799 Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention800 of the decoder.801 head_mask (`torch.Tensor` of shape `(decoder_layers, decoder_attention_heads)`, *optional*):802 Mask to nullify selected heads of the attention modules. Mask values selected in `[0, 1]`:803 804 - 1 indicates the head is **not masked**,805 - 0 indicates the head is **masked**.806 807 cross_attn_head_mask (`torch.Tensor` of shape `(decoder_layers, decoder_attention_heads)`, *optional*):808 Mask to nullify selected heads of the attention modules in encoder to avoid performing cross-attention809 on hidden heads. Mask values selected in `[0, 1]`:810 811 - 1 indicates the head is **not masked**,812 - 0 indicates the head is **masked**.813 814 past_key_values (`EncoderDecoderCache` or `tuple(tuple(torch.FloatTensor))`, *optional*):815 It is a [`~cache_utils.Cache`] instance. For more details, see our [kv cache guide](https://huggingface.co/docs/transformers/en/kv_cache).816 817 If `past_key_values` are used, the user can optionally input only the last `decoder_input_ids` (those818 that don't have their past key value states given to this model) of shape `(batch_size, 1)` instead of819 all `decoder_input_ids` of shape `(batch_size, sequence_length)`.820 inputs_embeds (`torch.FloatTensor` of821 shape `(batch_size, sequence_length, hidden_size)`, *optional*): Optionally, instead of passing822 `input_ids` you can choose to directly pass an embedded representation. This is useful if you want more823 control over how to convert `input_ids` indices into associated vectors than the model's internal824 embedding lookup matrix.825 output_attentions (`bool`, *optional*):826 Whether or not to return the attentions tensors of all attention layers. See `attentions` under827 returned tensors for more detail.828 output_hidden_states (`bool`, *optional*):829 Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors830 for more detail.831 return_dict (`bool`, *optional*):832 Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.833 cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*):834 Indices depicting the position of the input sequence tokens in the sequence. It is used to update the835 cache in the correct position and to infer the complete sequence length.836 """837 output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions838 output_hidden_states = (839 output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states840 )841 use_cache = use_cache if use_cache is not None else self.config.use_cache842 return_dict = return_dict if return_dict is not None else self.config.use_return_dict843 844 # retrieve input_ids and inputs_embeds845 if input_ids is not None and inputs_embeds is not None:846 raise ValueError("You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time")847 elif input_ids is not None:848 input_shape = input_ids.size()849 input_ids = input_ids.view(-1, input_shape[-1])850 elif inputs_embeds is not None:851 input_shape = inputs_embeds.size()[:-1]852 else:853 raise ValueError("You have to specify either decoder_input_ids or decoder_inputs_embeds")854 855 if inputs_embeds is None:856 inputs_embeds = self.embed_tokens(input_ids)857 858 if use_cache and past_key_values is None:859 if self.config.is_encoder_decoder:860 past_key_values = EncoderDecoderCache(861 DynamicCache(config=self.config), DynamicCache(config=self.config)862 )863 else:864 past_key_values = DynamicCache(config=self.config)865 866 past_key_values_length = 0867 if cache_position is not None:868 past_key_values_length = cache_position[0]869 elif past_key_values is not None:870 past_key_values_length = past_key_values.get_seq_length()871 872 if cache_position is None:873 cache_position = torch.arange(874 past_key_values_length, past_key_values_length + input_shape[1], device=inputs_embeds.device875 )876 877 if position_ids is None:878 position_ids = cache_position.unsqueeze(0).repeat(input_shape[0], 1)879 880 # embed positions881 if input_ids is not None:882 positions = self.embed_positions(883 input_ids, past_key_values_length=past_key_values_length, position_ids=position_ids884 )885 else:886 positions = self.embed_positions(887 inputs_embeds, past_key_values_length=past_key_values_length, position_ids=position_ids888 )889 890 hidden_states = inputs_embeds + positions.to(inputs_embeds.device)891 hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)892 893 causal_mask = create_causal_mask(894 config=self.config,895 input_embeds=inputs_embeds,896 attention_mask=attention_mask,897 cache_position=cache_position,898 past_key_values=past_key_values,899 position_ids=position_ids,900 )901 902 if self.gradient_checkpointing and self.training:903 if use_cache:904 logger.warning_once(905 "`use_cache = True` is incompatible with gradient checkpointing. Setting `use_cache = False`..."906 )907 use_cache = False908 # decoder layers909 all_hidden_states = () if output_hidden_states else None910 all_self_attns = () if output_attentions else None911 all_cross_attentions = () if (output_attentions and encoder_hidden_states is not None) else None912 913 # check if head_mask/cross_attn_head_mask has a correct number of layers specified if desired914 for attn_mask, mask_name in zip([head_mask, cross_attn_head_mask], ["head_mask", "cross_attn_head_mask"]):915 if attn_mask is not None:916 assert attn_mask.size()[0] == (len(self.layers)), (917 f"The `{mask_name}` should be specified for {len(self.layers)} layers, but it is for"918 f" {head_mask.size()[0]}."919 )920 for idx, decoder_layer in enumerate(self.layers):921 # add LayerDrop (see https://huggingface.co/papers/1909.11556 for description)922 if output_hidden_states:923 all_hidden_states += (hidden_states,)924 if self.training:925 dropout_probability = torch.rand([])926 if dropout_probability < self.layerdrop:927 continue928 929 layer_outputs = decoder_layer(930 hidden_states,931 attention_mask=causal_mask,932 encoder_hidden_states=encoder_hidden_states,933 layer_head_mask=(head_mask[idx] if head_mask is not None else None),934 cross_attn_layer_head_mask=(cross_attn_head_mask[idx] if cross_attn_head_mask is not None else None),935 past_key_values=past_key_values if use_cache else None,936 output_attentions=output_attentions,937 use_cache=use_cache,938 cache_position=cache_position,939 )940 hidden_states = layer_outputs[0]941 942 if output_attentions:943 all_self_attns += (layer_outputs[1],)944 945 if encoder_hidden_states is not None:946 all_cross_attentions += (layer_outputs[2],)947 948 hidden_states = self.layer_norm(hidden_states)949 # add hidden states from the last decoder layer950 if output_hidden_states:951 all_hidden_states += (hidden_states,)952 953 next_cache = past_key_values if use_cache else None954 if not return_dict:955 return tuple(956 v957 for v in [hidden_states, next_cache, all_hidden_states, all_self_attns, all_cross_attentions]958 if v is not None959 )960 return BaseModelOutputWithPastAndCrossAttentions(961 last_hidden_state=hidden_states,962 past_key_values=next_cache,963 hidden_states=all_hidden_states,964 attentions=all_self_attns,965 cross_attentions=all_cross_attentions,966 )967 968 969@auto_docstring970class WhisperModel(WhisperPreTrainedModel):971 def __init__(self, config: WhisperConfig):972 super().__init__(config)973 974 self.encoder = WhisperEncoder(config)975 self.decoder = WhisperDecoder(config)976 # Initialize weights and apply final processing977 self.post_init()978 979 def get_input_embeddings(self):980 return self.decoder.embed_tokens981 982 def set_input_embeddings(self, value):983 self.decoder.embed_tokens = value984 985 def get_encoder(self):986 return self.encoder987 988 def freeze_encoder(self):989 """990 Calling this function will disable the gradient computation for the Whisper encoder so that its parameters will991 not be updated during training.992 """993 self.encoder._freeze_parameters()994 995 def _mask_input_features(996 self,997 input_features: torch.FloatTensor,998 attention_mask: Optional[torch.LongTensor] = None,999 ):1000 """1001 Masks extracted features along time axis and/or along feature axis according to1002 [SpecAugment](https://huggingface.co/papers/1904.08779).1003 """1004 1005 # `config.apply_spec_augment` can set masking to False1006 if not getattr(self.config, "apply_spec_augment", True):1007 return input_features1008 1009 # generate indices & apply SpecAugment along time axis1010 batch_size, hidden_size, sequence_length = input_features.size()1011 1012 if self.config.mask_time_prob > 0 and self.training:1013 # generate indices & apply SpecAugment along time axis1014 mask_time_indices = _compute_mask_indices(1015 (batch_size, sequence_length),1016 mask_prob=self.config.mask_time_prob,1017 mask_length=self.config.mask_time_length,1018 attention_mask=attention_mask,1019 min_masks=self.config.mask_time_min_masks,1020 )1021 mask_time_indices = torch.tensor(mask_time_indices, device=input_features.device, dtype=torch.bool)1022 mask_time_indices = mask_time_indices[:, None].expand(-1, hidden_size, -1)1023 input_features[mask_time_indices] = 01024 1025 if self.config.mask_feature_prob > 0 and self.training:1026 # generate indices & apply SpecAugment along feature axis1027 mask_feature_indices = _compute_mask_indices(1028 (batch_size, hidden_size),1029 mask_prob=self.config.mask_feature_prob,1030 mask_length=self.config.mask_feature_length,1031 min_masks=self.config.mask_feature_min_masks,1032 )1033 mask_feature_indices = torch.tensor(mask_feature_indices, device=input_features.device, dtype=torch.bool)1034 input_features[mask_feature_indices] = 01035 1036 return input_features1037 1038 @auto_docstring1039 def forward(1040 self,1041 input_features: Optional[torch.FloatTensor] = None,1042 attention_mask: Optional[torch.LongTensor] = None,1043 decoder_input_ids: Optional[torch.LongTensor] = None,1044 decoder_attention_mask: Optional[torch.LongTensor] = None,1045 head_mask: Optional[torch.Tensor] = None,1046 decoder_head_mask: Optional[torch.Tensor] = None,1047 cross_attn_head_mask: Optional[torch.Tensor] = None,1048 encoder_outputs: Optional[tuple[tuple[torch.FloatTensor]]] = None,1049 past_key_values: Optional[Cache] = None,1050 decoder_inputs_embeds: Optional[tuple[torch.FloatTensor]] = None,1051 decoder_position_ids: Optional[tuple[torch.LongTensor]] = None,1052 use_cache: Optional[bool] = None,1053 output_attentions: Optional[bool] = None,1054 output_hidden_states: Optional[bool] = None,1055 return_dict: Optional[bool] = None,1056 cache_position: Optional[torch.LongTensor] = None,1057 ) -> Union[tuple[torch.Tensor], Seq2SeqModelOutput]:1058 r"""1059 decoder_input_ids (`torch.LongTensor` of shape `(batch_size, target_sequence_length)`, *optional*):1060 Indices of decoder input sequence tokens in the vocabulary.1061 1062 Indices can be obtained using [`WhisperTokenizer`]. See [`PreTrainedTokenizer.encode`] and1063 [`PreTrainedTokenizer.__call__`] for details.1064 1065 [What are decoder input IDs?](../glossary#decoder-input-ids)1066 1067 Whisper uses the `decoder_start_token_id` as the starting token for `decoder_input_ids` generation. If1068 `past_key_values` is used, optionally only the last `decoder_input_ids` have to be input (see1069 `past_key_values`).1070 decoder_attention_mask (`torch.LongTensor` of shape `(batch_size, target_sequence_length)`, *optional*):1071 Default behavior: generate a tensor that ignores pad tokens in `decoder_input_ids`. Causal mask will also1072 be used by default.1073 1074 If you want to change padding behavior, you should read1075 [`modeling_whisper._prepare_decoder_attention_mask`] and modify to your needs. See diagram 1 in [the BART1076 paper](https://huggingface.co/papers/1910.13461) for more information on the default strategy.1077 cross_attn_head_mask (`torch.Tensor` of shape `(decoder_layers, decoder_attention_heads)`, *optional*):1078 Mask to nullify selected heads of the cross-attention modules. Mask values selected in `[0, 1]`:1079 1080 - 1 indicates the head is **not masked**,1081 - 0 indicates the head is **masked**.1082 decoder_position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):1083 Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,1084 config.n_positions - 1]`.1085 1086 [What are position IDs?](../glossary#position-ids)1087 1088 Example:1089 ```python1090 >>> import torch1091 >>> from transformers import AutoFeatureExtractor, WhisperModel1092 >>> from datasets import load_dataset1093 1094 >>> model = WhisperModel.from_pretrained("openai/whisper-base")1095 >>> feature_extractor = AutoFeatureExtractor.from_pretrained("openai/whisper-base")1096 >>> ds = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation")1097 >>> inputs = feature_extractor(ds[0]["audio"]["array"], return_tensors="pt")1098 >>> input_features = inputs.input_features1099 >>> decoder_input_ids = torch.tensor([[1, 1]]) * model.config.decoder_start_token_id1100 >>> last_hidden_state = model(input_features, decoder_input_ids=decoder_input_ids).last_hidden_state1101 >>> list(last_hidden_state.shape)1102 [1, 2, 512]1103 ```"""1104 output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions1105 output_hidden_states = (1106 output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states1107 )1108 use_cache = use_cache if use_cache is not None else self.config.use_cache1109 return_dict = return_dict if return_dict is not None else self.config.use_return_dict1110 1111 if encoder_outputs is None:1112 input_features = self._mask_input_features(input_features, attention_mask=attention_mask)1113 1114 encoder_outputs = self.encoder(1115 input_features,1116 head_mask=head_mask,1117 output_attentions=output_attentions,1118 output_hidden_states=output_hidden_states,1119 return_dict=return_dict,1120 )1121 # If the user passed a tuple for encoder_outputs, we wrap it in a BaseModelOutput when return_dict=True1122 elif return_dict and not isinstance(encoder_outputs, BaseModelOutput):1123 encoder_outputs = BaseModelOutput(1124 last_hidden_state=encoder_outputs[0],1125 hidden_states=encoder_outputs[1] if len(encoder_outputs) > 1 else None,1126 attentions=encoder_outputs[2] if len(encoder_outputs) > 2 else None,1127 )1128 1129 # decoder outputs consists of (dec_features, past_key_values, dec_hidden, dec_attn)1130 decoder_outputs = self.decoder(1131 input_ids=decoder_input_ids,1132 attention_mask=decoder_attention_mask,1133 encoder_hidden_states=encoder_outputs[0],1134 head_mask=decoder_head_mask,1135 cross_attn_head_mask=cross_attn_head_mask,1136 past_key_values=past_key_values,1137 inputs_embeds=decoder_inputs_embeds,1138 position_ids=decoder_position_ids,1139 use_cache=use_cache,1140 output_attentions=output_attentions,1141 output_hidden_states=output_hidden_states,1142 return_dict=return_dict,1143 cache_position=cache_position,1144 )1145 1146 if not return_dict:1147 return decoder_outputs + encoder_outputs1148 1149 return Seq2SeqModelOutput(1150 last_hidden_state=decoder_outputs.last_hidden_state,1151 past_key_values=decoder_outputs.past_key_values,1152 decoder_hidden_states=decoder_outputs.hidden_states,1153 decoder_attentions=decoder_outputs.attentions,1154 cross_attentions=decoder_outputs.cross_attentions,1155 encoder_last_hidden_state=encoder_outputs.last_hidden_state,1156 encoder_hidden_states=encoder_outputs.hidden_states,1157 encoder_attentions=encoder_outputs.attentions,1158 )1159 1160 1161@auto_docstring(1162 custom_intro="""1163 The Whisper Model with a language modeling head. Can be used for automatic speech recognition.1164 """1165)1166class WhisperForConditionalGeneration(WhisperGenerationMixin, WhisperPreTrainedModel):1167 base_model_prefix = "model"1168 _tied_weights_keys = ["proj_out.weight"]1169 1170 def __init__(self, config: WhisperConfig):1171 super().__init__(config)1172 self.model = WhisperModel(config)1173 self.proj_out = nn.Linear(config.d_model, config.vocab_size, bias=False)1174 self.max_target_positions = config.max_target_positions1175 1176 # Initialize weights and apply final processing1177 self.post_init()1178 1179 def get_encoder(self):1180 return self.model.get_encoder()1181 1182 def get_decoder(self):1183 return self.model.get_decoder()1184 1185 def get_output_embeddings(self):1186 return self.proj_out1187 1188 def set_output_embeddings(self, new_embeddings):1189 self.proj_out = new_embeddings1190 1191 def get_input_embeddings(self) -> nn.Module:1192 return self.model.get_input_embeddings()1193 1194 def freeze_encoder(self):1195 """1196 Calling this function will disable the gradient computation for the Whisper encoder so that its parameters will1197 not be updated during training.1198 """1199 self.model.encoder._freeze_parameters()1200