Update modeling_phi3.py

modeling_phi3.py

@@ -7,77 +7,15 @@ from transformers.modeling_utils import PreTrainedModel
 from .configuration_phi3 import Phi3Config
 
 
-class Phi3Attention(nn.Module):
-    def __init__(self, config):
-        super().__init__()
-        self.num_heads = config.num_attention_heads
-        self.head_dim = config.hidden_size // config.num_attention_heads
-        self.scale = self.head_dim ** -0.5
-
-        self.q_proj = nn.Linear(config.hidden_size, config.hidden_size)
-        self.k_proj = nn.Linear(config.hidden_size, config.hidden_size)
-        self.v_proj = nn.Linear(config.hidden_size, config.hidden_size)
-        self.out_proj = nn.Linear(config.hidden_size, config.hidden_size)
-
-    def forward(self, x, mask=None):
-        B, T, C = x.size()
-
-        q = self.q_proj(x).view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
-        k = self.k_proj(x).view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
-        v = self.v_proj(x).view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
-
-        attn_weights = torch.matmul(q, k.transpose(-2, -1)) * self.scale
-        if mask is not None:
-            attn_weights = attn_weights.masked_fill(mask == 0, float('-inf'))
-        attn_probs = torch.softmax(attn_weights, dim=-1)
-
-        attn_output = torch.matmul(attn_probs, v).transpose(1, 2).contiguous().view(B, T, C)
-        return self.out_proj(attn_output)
-
-
-class Phi3Block(nn.Module):
-    def __init__(self, config):
-        super().__init__()
-        self.ln1 = nn.LayerNorm(config.hidden_size, eps=config.rms_norm_eps)
-        self.attn = Phi3Attention(config)
-        self.ln2 = nn.LayerNorm(config.hidden_size, eps=config.rms_norm_eps)
-        self.mlp = nn.Sequential(
-            nn.Linear(config.hidden_size, config.intermediate_size),
-            nn.GELU(),
-            nn.Linear(config.intermediate_size, config.hidden_size)
-        )
-
-    def forward(self, x, mask=None):
-        x = x + self.attn(self.ln1(x), mask=mask)
-        x = x + self.mlp(self.ln2(x))
-        return x
-
-
-class Phi3Model(PreTrainedModel):
-    config_class = Phi3Config
-
-    def __init__(self, config):
-        super().__init__(config)
-        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
-        self.blocks = nn.ModuleList([Phi3Block(config) for _ in range(config.num_hidden_layers)])
-        self.ln_f = nn.LayerNorm(config.hidden_size, eps=config.rms_norm_eps)
-
-    def forward(self, input_ids, attention_mask=None):
-        x = self.embed_tokens(input_ids)
-        for block in self.blocks:
-            x = block(x, attention_mask)
-        x = self.ln_f(x)
-        return x
-
-
 class Phi3ForCausalLM(PreTrainedModel):
     config_class = Phi3Config
+    model_type = "phi3"
+    base_model_prefix = "model"
 
     def __init__(self, config):
         super().__init__(config)
         self.model = Phi3Model(config)
         self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
-
         self.post_init()
 
     def forward(self, input_ids, attention_mask=None, labels=None):
@@ -94,7 +32,14 @@ class Phi3ForCausalLM(PreTrainedModel):
         return CausalLMOutputWithPast(
             loss=loss,
             logits=logits,
-            past_key_values=None,
+            past_key_values=None,  # Optional: implement later
             hidden_states=None,
             attentions=None,
         )
+
+    def prepare_inputs_for_generation(self, input_ids, past_key_values=None, attention_mask=None, **kwargs):
+        return {
+            "input_ids": input_ids,
+            "attention_mask": attention_mask,
+            "past_key_values": past_key_values,
+        }