IQuestLab/IQuest-Coder-V1-40B-Base-Stage1
29204
1"""2Modified MIT License 3 4Software Copyright© 2025 IQuest Research5 6Our only modification is that, if the Software (or any derivative works7thereof) is used for any of your commercial products or services, you shall8prominently display "IQuest Coder" on the user interface of such product or9service.10Permission is hereby granted, free of charge, to any person obtaining a copy11of this software and associated documentation files (the "Software"), to deal12in the Software without restriction, including without limitation the rights13to use, copy, modify, merge, publish, distribute, sublicense, and/or sell14copies of the Software, and to permit persons to whom the Software is15furnished to do so, subject to the following conditions:16 17The above copyright notice and this permission notice shall be included in all18copies or substantial portions of the Software.19 20THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR21IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,22FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE23AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER24LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,25OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.26"""27 28from typing import Callable, List, Optional, Tuple, Union29 30import torch31import torch.nn as nn32import torch.nn.functional as F33 34from transformers.activations import ACT2FN35from transformers.cache_utils import Cache, DynamicCache, SlidingWindowCache, StaticCache36from transformers.generation import GenerationMixin37from transformers.modeling_attn_mask_utils import AttentionMaskConverter38from transformers.modeling_flash_attention_utils import FlashAttentionKwargs39from transformers.modeling_layers import GradientCheckpointingLayer40from transformers.modeling_outputs import (41 BaseModelOutputWithPast,42 CausalLMOutputWithPast,43 QuestionAnsweringModelOutput,44 SequenceClassifierOutputWithPast,45 TokenClassifierOutput,46)47from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update48from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel49from transformers.processing_utils import Unpack50from transformers.utils import (51 auto_docstring,52 can_return_tuple,53 is_torch_flex_attn_available,54 logging,55)56 57from .configuration_iquestcoder import IQuestCoderConfig58 59 60if is_torch_flex_attn_available():61 from torch.nn.attention.flex_attention import BlockMask62 from transformers.integrations.flex_attention import make_flex_block_causal_mask63 64 65logger = logging.get_logger(__name__)66 67 68# =============================================================================69# Helper Functions70# =============================================================================71 72def rotate_half(x: torch.Tensor) -> torch.Tensor:73 """Rotates half the hidden dims of the input."""74 x1 = x[..., : x.shape[-1] // 2]75 x2 = x[..., x.shape[-1] // 2 :]76 return torch.cat((-x2, x1), dim=-1)77 78 79def apply_rotary_pos_emb(80 q: torch.Tensor,81 k: torch.Tensor,82 cos: torch.Tensor,83 sin: torch.Tensor,84 position_ids: Optional[torch.Tensor] = None,85 unsqueeze_dim: int = 1,86) -> Tuple[torch.Tensor, torch.Tensor]:87 """Applies Rotary Position Embedding to the query and key tensors.88 89 Args:90 q: The query tensor.91 k: The key tensor.92 cos: The cosine part of the rotary embedding.93 sin: The sine part of the rotary embedding.94 position_ids: Deprecated and unused.95 unsqueeze_dim: The dimension along which to unsqueeze cos and sin.96 97 Returns:98 Tuple of query and key tensors rotated using the Rotary Position Embedding.99 """100 # Borrowed from OLMo: preserve original dtypes for numerical stability101 q_dtype, k_dtype = q.dtype, k.dtype102 cos = cos.unsqueeze(unsqueeze_dim)103 sin = sin.unsqueeze(unsqueeze_dim)104 q_embed = (q * cos) + (rotate_half(q) * sin)105 k_embed = (k * cos) + (rotate_half(k) * sin)106 return q_embed.to(q_dtype), k_embed.to(k_dtype)107 108 109def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:110 """111 Expands key/value heads for Grouped Query Attention.112 113 This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep).114 The hidden states go from (batch, num_key_value_heads, seqlen, head_dim) to115 (batch, num_attention_heads, seqlen, head_dim).116 """117 batch, num_key_value_heads, slen, head_dim = hidden_states.shape118 if n_rep == 1:119 return hidden_states120 hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)121 return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)122 123 124def eager_attention_forward(125 module: nn.Module,126 query: torch.Tensor,127 key: torch.Tensor,128 value: torch.Tensor,129 attention_mask: Optional[torch.Tensor],130 scaling: float,131 dropout: float = 0.0,132 **kwargs,133) -> Tuple[torch.Tensor, torch.Tensor]:134 """Standard eager attention implementation."""135 key_states = repeat_kv(key, module.num_key_value_groups)136 value_states = repeat_kv(value, module.num_key_value_groups)137 138 attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling139 if attention_mask is not None:140 causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]141 attn_weights = attn_weights + causal_mask142 143 attn_weights = F.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)144 attn_weights = F.dropout(attn_weights, p=dropout, training=module.training)145 attn_output = torch.matmul(attn_weights, value_states)146 attn_output = attn_output.transpose(1, 2).contiguous()147 148 return attn_output, attn_weights149 150 151# =============================================================================152# Model Components153# =============================================================================154 155class IQuestCoderRMSNorm(nn.Module):156 """Root Mean Square Layer Normalization.157 158 RMSNorm is computationally simpler than LayerNorm while achieving similar159 performance. It normalizes the input by its RMS value.160 """161 162 def __init__(self, hidden_size: int, eps: float = 1e-6):163 super().__init__()164 self.weight = nn.Parameter(torch.ones(hidden_size))165 self.variance_epsilon = eps166 167 def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:168 input_dtype = hidden_states.dtype169 hidden_states = hidden_states.to(torch.float32)170 variance = hidden_states.pow(2).mean(-1, keepdim=True)171 hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)172 return self.weight * hidden_states.to(input_dtype)173 174 def extra_repr(self) -> str:175 return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"176 177 178class IQuestCoderRotaryEmbedding(nn.Module):179 """Rotary Position Embedding (RoPE).180 181 Implements rotary positional embeddings as described in the RoFormer paper.182 Supports various RoPE scaling methods for extended context lengths.183 """184 185 def __init__(self, config: IQuestCoderConfig, device=None):186 super().__init__()187 # BC: "rope_type" was originally "type"188 if hasattr(config, "rope_scaling") and config.rope_scaling is not None:189 self.rope_type = config.rope_scaling.get("rope_type", config.rope_scaling.get("type"))190 else:191 self.rope_type = "default"192 self.max_seq_len_cached = config.max_position_embeddings193 self.original_max_seq_len = config.max_position_embeddings194 195 self.config = config196 self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]197 198 inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device)199 self.register_buffer("inv_freq", inv_freq, persistent=False)200 self.original_inv_freq = self.inv_freq201 202 @torch.no_grad()203 @dynamic_rope_update204 def forward(self, x: torch.Tensor, position_ids: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:205 inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device)206 position_ids_expanded = position_ids[:, None, :].float()207 208 device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu"209 with torch.autocast(device_type=device_type, enabled=False):210 freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)211 emb = torch.cat((freqs, freqs), dim=-1)212 cos = emb.cos() * self.attention_scaling213 sin = emb.sin() * self.attention_scaling214 215 return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)216 217 218class IQuestCoderMLP(nn.Module):219 """Feed-forward network with SwiGLU activation.220 221 Uses the gated linear unit variant with SiLU activation for improved222 performance compared to standard FFN.223 """224 225 def __init__(self, config: IQuestCoderConfig):226 super().__init__()227 self.config = config228 self.hidden_size = config.hidden_size229 self.intermediate_size = config.intermediate_size230 self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=config.mlp_bias)231 self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=config.mlp_bias)232 self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=config.mlp_bias)233 self.act_fn = ACT2FN[config.hidden_act]234 235 def forward(self, x: torch.Tensor) -> torch.Tensor:236 # SwiGLU: down_proj(act_fn(gate_proj(x)) * up_proj(x))237 return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))238 239 240class IQuestCoderAttention(nn.Module):241 """Multi-headed attention with support for Grouped Query Attention (GQA).242 243 Features:244 - Grouped Query Attention for memory efficiency245 - Optional QKV clipping for training stability (from OLMo)246 - Optional sliding window attention (from Qwen2)247 - Rotary Position Embeddings248 """249 250 def __init__(self, config: IQuestCoderConfig, layer_idx: int):251 super().__init__()252 self.config = config253 self.layer_idx = layer_idx254 self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)255 self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads256 self.scaling = self.head_dim ** -0.5257 self.attention_dropout = config.attention_dropout258 self.is_causal = True259 260 # Projection layers261 self.q_proj = nn.Linear(262 config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.attention_bias263 )264 self.k_proj = nn.Linear(265 config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias266 )267 self.v_proj = nn.Linear(268 config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias269 )270 self.o_proj = nn.Linear(271 config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias272 )273 274 def forward(275 self,276 hidden_states: torch.Tensor,277 position_embeddings: Tuple[torch.Tensor, torch.Tensor],278 attention_mask: Optional[torch.Tensor],279 past_key_value: Optional[Cache] = None,280 cache_position: Optional[torch.LongTensor] = None,281 **kwargs: Unpack[FlashAttentionKwargs],282 ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:283 input_shape = hidden_states.shape[:-1]284 hidden_shape = (*input_shape, -1, self.head_dim)285 286 # Compute Q, K, V projections287 query_states = self.q_proj(hidden_states)288 key_states = self.k_proj(hidden_states)289 value_states = self.v_proj(hidden_states)290 291 # [OLMo Feature] Optional QKV clipping for training stability292 if self.config.clip_qkv is not None:293 query_states = query_states.clamp(min=-self.config.clip_qkv, max=self.config.clip_qkv)294 key_states = key_states.clamp(min=-self.config.clip_qkv, max=self.config.clip_qkv)295 value_states = value_states.clamp(min=-self.config.clip_qkv, max=self.config.clip_qkv)296 297 # Reshape to (batch, heads, seq_len, head_dim)298 query_states = query_states.view(hidden_shape).transpose(1, 2)299 key_states = key_states.view(hidden_shape).transpose(1, 2)300 value_states = value_states.view(hidden_shape).transpose(1, 2)301 302 # Apply rotary position embeddings303 cos, sin = position_embeddings304 query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)305 306 # Update KV cache if provided307 if past_key_value is not None:308 cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}309 key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)310 311 # [Qwen2 Feature] Sliding window attention312 sliding_window = None313 if (314 self.config.use_sliding_window315 and getattr(self.config, "sliding_window", None) is not None316 and self.layer_idx >= self.config.max_window_layers317 ):318 sliding_window = self.config.sliding_window319 320 # Select attention implementation321 attention_interface: Callable = eager_attention_forward322 if self.config._attn_implementation != "eager":323 if self.config._attn_implementation == "sdpa" and kwargs.get("output_attentions", False):324 logger.warning_once(325 "`torch.nn.functional.scaled_dot_product_attention` does not support `output_attentions=True`. "326 'Falling back to eager attention. This warning can be removed using the argument '327 '`attn_implementation="eager"` when loading the model.'328 )329 else:330 attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]331 332 # Compute attention333 attn_output, attn_weights = attention_interface(334 self,335 query_states,336 key_states,337 value_states,338 attention_mask,339 dropout=0.0 if not self.training else self.attention_dropout,340 scaling=self.scaling,341 sliding_window=sliding_window,342 **kwargs,343 )344 345 # Reshape and project output346 attn_output = attn_output.reshape(*input_shape, -1).contiguous()347 attn_output = self.o_proj(attn_output)348 349 return attn_output, attn_weights350 351 352class IQuestCoderDecoderLayer(GradientCheckpointingLayer):353 """Transformer decoder layer with pre-normalization.354 355 Architecture: Pre-RMSNorm -> Attention -> Residual -> Pre-RMSNorm -> MLP -> Residual356 """357 358 def __init__(self, config: IQuestCoderConfig, layer_idx: int):359 super().__init__()360 self.hidden_size = config.hidden_size361 self.self_attn = IQuestCoderAttention(config=config, layer_idx=layer_idx)362 self.mlp = IQuestCoderMLP(config)363 self.input_layernorm = IQuestCoderRMSNorm(config.hidden_size, eps=config.rms_norm_eps)364 self.post_attention_layernorm = IQuestCoderRMSNorm(config.hidden_size, eps=config.rms_norm_eps)365 366 # Warn if sliding window is enabled but not properly supported367 if config.use_sliding_window and config._attn_implementation != "flash_attention_2":368 logger.warning_once(369 f"Sliding Window Attention is enabled but not implemented for `{config._attn_implementation}`; "370 "unexpected results may be encountered."371 )372 373 def forward(374 self,375 hidden_states: torch.Tensor,376 attention_mask: Optional[torch.Tensor] = None,377 position_ids: Optional[torch.LongTensor] = None,378 past_key_value: Optional[Cache] = None,379 output_attentions: Optional[bool] = False,380 use_cache: Optional[bool] = False,381 cache_position: Optional[torch.LongTensor] = None,382 position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,383 **kwargs: Unpack[FlashAttentionKwargs],384 ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:385 # Pre-norm + Self Attention386 residual = hidden_states387 hidden_states = self.input_layernorm(hidden_states)388 389 hidden_states, self_attn_weights = self.self_attn(390 hidden_states=hidden_states,391 attention_mask=attention_mask,392 position_ids=position_ids,393 past_key_value=past_key_value,394 output_attentions=output_attentions,395 use_cache=use_cache,396 cache_position=cache_position,397 position_embeddings=position_embeddings,398 **kwargs,399 )400 hidden_states = residual + hidden_states401 402 # Pre-norm + MLP403 residual = hidden_states404 hidden_states = self.post_attention_layernorm(hidden_states)405 hidden_states = self.mlp(hidden_states)406 hidden_states = residual + hidden_states407 408 outputs = (hidden_states,)409 if output_attentions:410 outputs += (self_attn_weights,)411 412 return outputs413 414 415# =============================================================================416# Base Model417# =============================================================================418 419@auto_docstring420class IQuestCoderPreTrainedModel(PreTrainedModel):421 """Base class for IQuestCoder models."""422 423 config_class = IQuestCoderConfig424 base_model_prefix = "model"425 supports_gradient_checkpointing = True426 _no_split_modules = ["IQuestCoderDecoderLayer"]427 _skip_keys_device_placement = ["past_key_values"]428 _supports_flash_attn_2 = True429 _supports_sdpa = True430 _supports_flex_attn = True431 _supports_cache_class = True432 _supports_quantized_cache = True433 _supports_static_cache = True434 _supports_attention_backend = True435 436 def _init_weights(self, module: nn.Module):437 std = self.config.initializer_range438 if isinstance(module, nn.Linear):439 module.weight.data.normal_(mean=0.0, std=std)440 if module.bias is not None:441 module.bias.data.zero_()442 elif isinstance(module, nn.Embedding):443 module.weight.data.normal_(mean=0.0, std=std)444 if module.padding_idx is not None:445 module.weight.data[module.padding_idx].zero_()446 elif isinstance(module, IQuestCoderRMSNorm):447 module.weight.data.fill_(1.0)448 449 450@auto_docstring451class IQuestCoderModel(IQuestCoderPreTrainedModel):452 """453 IQuestCoder Model outputting raw hidden-states without any specific head on top.454 455 This model is compatible with LLaMA weights while incorporating features from OLMo and Qwen2.456 """457 458 def __init__(self, config: IQuestCoderConfig):459 super().__init__(config)460 self.padding_idx = config.pad_token_id461 self.vocab_size = config.vocab_size462 463 self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)464 self.layers = nn.ModuleList(465 [IQuestCoderDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]466 )467 self.norm = IQuestCoderRMSNorm(config.hidden_size, eps=config.rms_norm_eps)468 self.rotary_emb = IQuestCoderRotaryEmbedding(config=config)469 self.gradient_checkpointing = False470 471 # Initialize weights and apply final processing472 self.post_init()473 474 def get_input_embeddings(self) -> nn.Embedding:475 return self.embed_tokens476 477 def set_input_embeddings(self, value: nn.Embedding):478 self.embed_tokens = value479 480 @can_return_tuple481 @auto_docstring482 def forward(483 self,484 input_ids: Optional[torch.LongTensor] = None,485 attention_mask: Optional[torch.Tensor] = None,486 position_ids: Optional[torch.LongTensor] = None,487 past_key_values: Optional[Cache] = None,488 inputs_embeds: Optional[torch.FloatTensor] = None,489 use_cache: Optional[bool] = None,490 output_attentions: Optional[bool] = None,491 output_hidden_states: Optional[bool] = None,492 cache_position: Optional[torch.LongTensor] = None,493 **flash_attn_kwargs: Unpack[FlashAttentionKwargs],494 ) -> BaseModelOutputWithPast:495 output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions496 output_hidden_states = (497 output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states498 )499 use_cache = use_cache if use_cache is not None else self.config.use_cache500 501 if (input_ids is None) ^ (inputs_embeds is not None):502 raise ValueError("You must specify exactly one of input_ids or inputs_embeds")503 504 if self.gradient_checkpointing and self.training and use_cache:505 logger.warning_once(506 "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`."507 )508 use_cache = False509 510 if not isinstance(past_key_values, (type(None), Cache)):511 raise ValueError("The `past_key_values` should be either a `Cache` object or `None`.")512 513 if inputs_embeds is None:514 inputs_embeds = self.embed_tokens(input_ids)515 516 if use_cache and past_key_values is None:517 past_key_values = DynamicCache()518 519 if cache_position is None:520 past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0521 cache_position = torch.arange(522 past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device523 )524 525 if position_ids is None:526 position_ids = cache_position.unsqueeze(0)527 528 causal_mask = self._update_causal_mask(529 attention_mask, inputs_embeds, cache_position, past_key_values, output_attentions530 )531 532 hidden_states = inputs_embeds533 534 # Create position embeddings to be shared across the decoder layers535 position_embeddings = self.rotary_emb(hidden_states, position_ids)536 537 # Decoder layers538 all_hidden_states = () if output_hidden_states else None539 all_self_attns = () if output_attentions else None540 541 for decoder_layer in self.layers[: self.config.num_hidden_layers]:542 if output_hidden_states:543 all_hidden_states += (hidden_states,)544 545 layer_outputs = decoder_layer(546 hidden_states,547 attention_mask=causal_mask,548 position_ids=position_ids,549 past_key_value=past_key_values,550 output_attentions=output_attentions,551 use_cache=use_cache,552 cache_position=cache_position,553 position_embeddings=position_embeddings,554 **flash_attn_kwargs,555 )556 557 hidden_states = layer_outputs[0]558 559 if output_attentions:560 all_self_attns += (layer_outputs[1],)561 562 hidden_states = self.norm(hidden_states)563 564 # Add hidden states from the last decoder layer565 if output_hidden_states:566 all_hidden_states += (hidden_states,)567 568 return BaseModelOutputWithPast(569 last_hidden_state=hidden_states,570 past_key_values=past_key_values if use_cache else None,571 hidden_states=all_hidden_states,572 attentions=all_self_attns,573 )574 575 def _update_causal_mask(576 self,577 attention_mask: Union[torch.Tensor, "BlockMask"],578 input_tensor: torch.Tensor,579 cache_position: torch.Tensor,580 past_key_values: Cache,581 output_attentions: bool = False,582 ):583 if self.config._attn_implementation == "flash_attention_2":584 if attention_mask is not None and past_key_values is not None:585 is_padding_right = attention_mask[:, -1].sum().item() != input_tensor.size()[0]586 if is_padding_right:587 raise ValueError(588 "You are attempting to perform batched generation with padding_side='right'. "589 "This may lead to unexpected behaviour for Flash Attention version of IQuestCoder. "590 "Make sure to call `tokenizer.padding_side = 'left'` before tokenizing the input."591 )592 if attention_mask is not None and 0.0 in attention_mask:593 return attention_mask594 return None595 596 if self.config._attn_implementation == "flex_attention":597 if isinstance(attention_mask, torch.Tensor):598 attention_mask = make_flex_block_causal_mask(attention_mask)599 return attention_mask600 601 past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0602 using_static_cache = isinstance(past_key_values, StaticCache)603 using_sliding_window_cache = isinstance(past_key_values, SlidingWindowCache)604 605 if (606 self.config._attn_implementation == "sdpa"607 and not (using_static_cache or using_sliding_window_cache)608 and not output_attentions609 ):610 if AttentionMaskConverter._ignore_causal_mask_sdpa(611 attention_mask,612 inputs_embeds=input_tensor,613 past_key_values_length=past_seen_tokens,614 sliding_window=self.config.sliding_window if self.config.use_sliding_window else None,615 is_training=self.training,616 ):617 return None618 619 dtype = input_tensor.dtype620 min_dtype = torch.finfo(dtype).min621 sequence_length = input_tensor.shape[1]622 623 if using_sliding_window_cache or using_static_cache:624 target_length = past_key_values.get_max_cache_shape()625 else:626 target_length = (627 attention_mask.shape[-1]628 if isinstance(attention_mask, torch.Tensor)629 else past_seen_tokens + sequence_length + 1630 )631 632 causal_mask = self._prepare_4d_causal_attention_mask_with_cache_position(633 attention_mask,634 sequence_length=sequence_length,635 target_length=target_length,636 dtype=dtype,637 cache_position=cache_position,638 batch_size=input_tensor.shape[0],639 config=self.config,640 past_key_values=past_key_values,641 )642 643 if (644 self.config._attn_implementation == "sdpa"645 and attention_mask is not None646 and attention_mask.device.type in ["cuda", "xpu", "npu"]647 and not output_attentions648 ):649 causal_mask = AttentionMaskConverter._unmask_unattended(causal_mask, min_dtype)650 651 return causal_mask652 653 @staticmethod654 def _prepare_4d_causal_attention_mask_with_cache_position(655 attention_mask: torch.Tensor,656 sequence_length: int,657 target_length: int,658 dtype: torch.dtype,659 cache_position: torch.Tensor,660 batch_size: int,661 config: IQuestCoderConfig,662 past_key_values: Cache,663 ):664 """Creates a causal 4D mask from a 2D mask, or returns the 4D mask if already provided."""665 if attention_mask is not None and attention_mask.dim() == 4:666 causal_mask = attention_mask667 else:668 min_dtype = torch.finfo(dtype).min669 causal_mask = torch.full(670 (sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=cache_position.device671 )672 diagonal_attend_mask = torch.arange(target_length, device=cache_position.device) > cache_position.reshape(673 -1, 1674 )675 676 # [Qwen2 Feature] Handle sliding window mask677 if getattr(config, "use_sliding_window", False) and config.sliding_window is not None:678 if not isinstance(past_key_values, SlidingWindowCache) or sequence_length > target_length:679 sliding_attend_mask = torch.arange(target_length, device=cache_position.device) <= (680 cache_position.reshape(-1, 1) - config.sliding_window681 )682 diagonal_attend_mask.bitwise_or_(sliding_attend_mask)683 684 causal_mask *= diagonal_attend_mask685 causal_mask = causal_mask[None, None, :, :].expand(batch_size, 1, -1, -1)686 687 if attention_mask is not None:688 causal_mask = causal_mask.clone()689 if attention_mask.shape[-1] > target_length:690 attention_mask = attention_mask[:, :target_length]691 mask_length = attention_mask.shape[-1]692 padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :].to(693 causal_mask.device694 )695 padding_mask = padding_mask == 0696 causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill(697 padding_mask, min_dtype698 )699 700 return causal_mask701 702 703# =============================================================================704# Model Heads705# =============================================================================706 707@auto_docstring708class IQuestCoderForCausalLM(IQuestCoderPreTrainedModel, GenerationMixin):709 """IQuestCoder Model with a language modeling head on top for causal LM."""710 711 _tied_weights_keys = ["lm_head.weight"]712 _tp_plan = {"lm_head": "colwise_rep"}713 _pp_plan = {"lm_head": (["hidden_states"], ["logits"])}714 715 def __init__(self, config: IQuestCoderConfig):716 super().__init__(config)717 self.model = IQuestCoderModel(config)718 self.vocab_size = config.vocab_size719 self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)720 721 # Initialize weights and apply final processing722 self.post_init()723 724 def get_input_embeddings(self) -> nn.Embedding:725 return self.model.embed_tokens726 727 def set_input_embeddings(self, value: nn.Embedding):728 self.model.embed_tokens = value729 730 def get_output_embeddings(self) -> nn.Linear:731 return self.lm_head732 733 def set_output_embeddings(self, new_embeddings: nn.Linear):734 self.lm_head = new_embeddings735 736 def set_decoder(self, decoder: IQuestCoderModel):737 self.model = decoder738 739 def get_decoder(self) -> IQuestCoderModel:740 return self.model741 742 @can_return_tuple743 @auto_docstring744 def forward(745 self,746 input_ids: Optional[torch.LongTensor] = None,747 attention_mask: Optional[torch.Tensor] = None,748 position_ids: Optional[torch.LongTensor] = None,749 past_key_values: Optional[Cache] = None,750 inputs_embeds: Optional[torch.FloatTensor] = None,751 labels: Optional[torch.LongTensor] = None,752 use_cache: Optional[bool] = None,753 output_attentions: Optional[bool] = None,754 output_hidden_states: Optional[bool] = None,755 cache_position: Optional[torch.LongTensor] = None,756 logits_to_keep: Union[int, torch.Tensor] = 0,757 **kwargs758 ) -> CausalLMOutputWithPast:759 r"""760 Args:761 labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):762 Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,763 config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored764 (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.765 766 Example:767 ```python768 >>> from transformers import AutoTokenizer769 >>> from modeling_iquestcoder import IQuestCoderForCausalLM770 771 >>> model = IQuestCoderForCausalLM.from_pretrained("path/to/IQuestCoder")772 >>> tokenizer = AutoTokenizer.from_pretrained("path/to/IQuestCoder")773 774 >>> prompt = "Hey, are you conscious? Can you talk to me?"775 >>> inputs = tokenizer(prompt, return_tensors="pt")776 777 >>> # Generate778 >>> generate_ids = model.generate(inputs.input_ids, max_length=30)779 >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]780 "Hey, are you conscious? Can you talk to me?\\nI'm not conscious, but I can talk to you."781 ```782 """783 output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions784 output_hidden_states = (785 output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states786 )787 788 # Decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)789 outputs: BaseModelOutputWithPast = self.model(790 input_ids=input_ids,791 attention_mask=attention_mask,792 position_ids=position_ids,793 past_key_values=past_key_values,794 inputs_embeds=inputs_embeds,795 use_cache=use_cache,796 output_attentions=output_attentions,797 output_hidden_states=output_hidden_states,798 cache_position=cache_position,799 **kwargs,800 )801 802 hidden_states = outputs.last_hidden_state803 # Only compute necessary logits, and do not upcast them to float if we are not computing the loss804 slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep805 logits = self.lm_head(hidden_states[:, slice_indices, :])806 807 loss = None808 if labels is not None:809 loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs)810 811 return CausalLMOutputWithPast(812 loss=loss,813 logits=logits,814 past_key_values=outputs.past_key_values,815 hidden_states=outputs.hidden_states,816 attentions=outputs.attentions,817 )818 819 820@auto_docstring(821 custom_intro="""822 The IQuestCoder Model transformer with a sequence classification head on top (linear layer).823 824 [`IQuestCoderForSequenceClassification`] uses the last token in order to do the classification, as other causal825 models (e.g. GPT-2) do.826 827 Since it does classification on the last token, it requires to know the position of the last token. If a828 `pad_token_id` is defined in the configuration, it finds the last token that is not a padding token in each row.829 If no `pad_token_id` is defined, it simply takes the last value in each row of the batch.830 """831)832class IQuestCoderForSequenceClassification(IQuestCoderPreTrainedModel):833 """IQuestCoder Model with a sequence classification head."""834 835 def __init__(self, config: IQuestCoderConfig):836 super().__init__(config)837 self.num_labels = config.num_labels838 self.model = IQuestCoderModel(config)839 self.score = nn.Linear(config.hidden_size, self.num_labels, bias=False)840 841 # Initialize weights and apply final processing842 self.post_init()843 844 def get_input_embeddings(self) -> nn.Embedding:845 return self.model.embed_tokens846 847 def set_input_embeddings(self, value: nn.Embedding):848 self.model.embed_tokens = value849 850 @can_return_tuple851 @auto_docstring852 def forward(853 self,854 input_ids: Optional[torch.LongTensor] = None,855 attention_mask: Optional[torch.Tensor] = None,856 position_ids: Optional[torch.LongTensor] = None,857 past_key_values: Optional[Cache] = None,858 inputs_embeds: Optional[torch.FloatTensor] = None,859 labels: Optional[torch.LongTensor] = None,860 use_cache: Optional[bool] = None,861 output_attentions: Optional[bool] = None,862 output_hidden_states: Optional[bool] = None,863 ) -> SequenceClassifierOutputWithPast:864 r"""865 labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):866 Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,867 config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss),868 If `config.num_labels > 1` a classification loss is computed (Cross-Entropy).869 """870 transformer_outputs: BaseModelOutputWithPast = self.model(871 input_ids,872 attention_mask=attention_mask,873 position_ids=position_ids,874 past_key_values=past_key_values,875 inputs_embeds=inputs_embeds,876 use_cache=use_cache,877 output_attentions=output_attentions,878 output_hidden_states=output_hidden_states,879 )880 hidden_states = transformer_outputs.last_hidden_state881 logits = self.score(hidden_states)882 883 if input_ids is not None:884 batch_size = input_ids.shape[0]885 else:886 batch_size = inputs_embeds.shape[0]887 888 if self.config.pad_token_id is None and batch_size != 1:889 raise ValueError("Cannot handle batch sizes > 1 if no padding token is defined.")890 if self.config.pad_token_id is None:891 last_non_pad_token = -1892 elif input_ids is not None:893 non_pad_mask = (input_ids != self.config.pad_token_id).to(logits.device, torch.int32)894 token_indices = torch.arange(input_ids.shape[-1], device=logits.device, dtype=torch.int32)895 last_non_pad_token = (token_indices * non_pad_mask).argmax(-1)896 else:897 last_non_pad_token = -1898 logger.warning_once(899 f"{self.__class__.__name__} will not detect padding tokens in `inputs_embeds`. Results may be "900 "unexpected if using padding tokens in conjunction with `inputs_embeds.`"901 )902 903 pooled_logits = logits[torch.arange(batch_size, device=logits.device), last_non_pad_token]904 905 loss = None906 if labels is not None:907 loss = self.loss_function(logits=logits, labels=labels, pooled_logits=pooled_logits, config=self.config)908 909 return SequenceClassifierOutputWithPast(910 loss=loss,911 logits=pooled_logits,912 past_key_values=transformer_outputs.past_key_values,913 hidden_states=transformer_outputs.hidden_states,914 attentions=transformer_outputs.attentions,915 )916 917 918@auto_docstring919class IQuestCoderForTokenClassification(IQuestCoderPreTrainedModel):920 """IQuestCoder Model with a token classification head."""921 922 def __init__(self, config: IQuestCoderConfig):923 super().__init__(config)924 self.num_labels = config.num_labels925 self.model = IQuestCoderModel(config)926 if getattr(config, "classifier_dropout", None) is not None:927 classifier_dropout = config.classifier_dropout928 elif getattr(config, "hidden_dropout", None) is not None:929 classifier_dropout = config.hidden_dropout930 else:931 classifier_dropout = 0.1932 self.dropout = nn.Dropout(classifier_dropout)933 self.score = nn.Linear(config.hidden_size, config.num_labels)934 935 # Initialize weights and apply final processing936 self.post_init()937 938 def get_input_embeddings(self) -> nn.Embedding:939 return self.model.embed_tokens940 941 def set_input_embeddings(self, value: nn.Embedding):942 self.model.embed_tokens = value943 944 @can_return_tuple945 @auto_docstring946 def forward(947 self,948 input_ids: Optional[torch.LongTensor] = None,949 attention_mask: Optional[torch.Tensor] = None,950 position_ids: Optional[torch.LongTensor] = None,951 past_key_values: Optional[Cache] = None,952 inputs_embeds: Optional[torch.FloatTensor] = None,953 labels: Optional[torch.LongTensor] = None,954 use_cache: Optional[bool] = None,955 output_attentions: Optional[bool] = None,956 output_hidden_states: Optional[bool] = None,957 ) -> TokenClassifierOutput:958 r"""959 labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):960 Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,961 config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss),962 If `config.num_labels > 1` a classification loss is computed (Cross-Entropy).963 """964 outputs: BaseModelOutputWithPast = self.model(965 input_ids,966 attention_mask=attention_mask,967 position_ids=position_ids,968 past_key_values=past_key_values,969 inputs_embeds=inputs_embeds,970 use_cache=use_cache,971 output_attentions=output_attentions,972 output_hidden_states=output_hidden_states,973 )974 sequence_output = outputs.last_hidden_state975 sequence_output = self.dropout(sequence_output)976 logits = self.score(sequence_output)977 978 loss = None979 if labels is not None:980 loss = self.loss_function(logits, labels, self.config)981 982 return TokenClassifierOutput(983 loss=loss,984 logits=logits,985 hidden_states=outputs.hidden_states,986 attentions=outputs.attentions,987 )988 989 990@auto_docstring991class IQuestCoderForQuestionAnswering(IQuestCoderPreTrainedModel):992 """IQuestCoder Model with a span classification head for extractive question-answering."""993 994 base_model_prefix = "transformer"995 996 def __init__(self, config: IQuestCoderConfig):997 super().__init__(config)998 self.transformer = IQuestCoderModel(config)999 self.qa_outputs = nn.Linear(config.hidden_size, 2)1000 1001 # Initialize weights and apply final processing1002 self.post_init()1003 1004 def get_input_embeddings(self) -> nn.Embedding:1005 return self.transformer.embed_tokens1006 1007 def set_input_embeddings(self, value: nn.Embedding):1008 self.transformer.embed_tokens = value1009 1010 @can_return_tuple1011 @auto_docstring1012 def forward(1013 self,1014 input_ids: Optional[torch.LongTensor] = None,1015 attention_mask: Optional[torch.Tensor] = None,1016 position_ids: Optional[torch.LongTensor] = None,1017 past_key_values: Optional[Cache] = None,1018 inputs_embeds: Optional[torch.FloatTensor] = None,1019 start_positions: Optional[torch.LongTensor] = None,1020 end_positions: Optional[torch.LongTensor] = None,1021 output_attentions: Optional[bool] = None,1022 output_hidden_states: Optional[bool] = None,1023 **kwargs,1024 ) -> QuestionAnsweringModelOutput:1025 outputs: BaseModelOutputWithPast = self.transformer(1026 input_ids,1027 attention_mask=attention_mask,1028 position_ids=position_ids,1029 past_key_values=past_key_values,1030 inputs_embeds=inputs_embeds,1031 output_attentions=output_attentions,1032 output_hidden_states=output_hidden_states,1033 )1034 1035 sequence_output = outputs.last_hidden_state1036 1037 logits = self.qa_outputs(sequence_output)1038 start_logits, end_logits = logits.split(1, dim=-1)1039 start_logits = start_logits.squeeze(-1).contiguous()1040 end_logits = end_logits.squeeze(-1).contiguous()1041 1042 loss = None1043 if start_positions is not None and end_positions is not None:1044 loss = self.loss_function(start_logits, end_logits, start_positions, end_positions, **kwargs)1045 1046 return QuestionAnsweringModelOutput(1047 loss=loss,1048 start_logits=start_logits,1049 end_logits=end_logits,1050 hidden_states=outputs.hidden_states,1051 attentions=outputs.attentions,1052 )1053 1054 1055__all__ = [1056 "IQuestCoderPreTrainedModel",1057 "IQuestCoderModel",1058 "IQuestCoderForCausalLM",1059 "IQuestCoderForSequenceClassification",1060 "IQuestCoderForTokenClassification",1061 "IQuestCoderForQuestionAnswering",1062]1063 1064