riaz.somc/transformers
1
0
Fork 0
mirror of https://github.com/huggingface/transformers.git synced 2025-07-03 04:40:06 +06:00
Code Issues Actions 28 Packages Projects Releases Wiki Activity

[flax/mistral] support sliding_window: null in config (#37402)

Browse Source 
flax/mistral: Allow sliding_window to be set to none
This commit is contained in:
Yiding Jia 2025-06-02 07:45:02 -07:00 committed by GitHub
parent 1a25fd2f6d
commit cceab972ba
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 2 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
src/transformers/models/mistral/modeling_flax_mistral.py
View File

@ -240,8 +240,8 @@ class FlaxMistralAttention(nn.Module):
self.v_proj = nn.Dense(self.num_key_value_heads * self.head_dim, use_bias=False, dtype=self.dtype)
self.o_proj = nn.Dense(self.hidden_size, use_bias=False, dtype=self.dtype)
casual_mask = make_causal_mask(jnp.ones((1, config.max_position_embeddings), dtype="bool"), dtype="bool")
self.causal_mask = jnp.triu(casual_mask, k=-config.sliding_window)
self.rotary_emb = FlaxMistralRotaryEmbedding(config, dtype=self.dtype)
self.causal_mask = jnp.triu(casual_mask, k=-(config.sliding_window or 0))
self.rotary_emb = FlaxMistralRotaryEmbedding(self.config, dtype=self.dtype)
def _split_heads(self, hidden_states, num_heads):
return hidden_states.reshape(hidden_states.shape[:2] + (num_heads, self.head_dim))