Update modeling_phi3.py

modeling_phi3.py

@@ -7,6 +7,69 @@ 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"
@@ -32,7 +95,7 @@ class Phi3ForCausalLM(PreTrainedModel):
         return CausalLMOutputWithPast(
             loss=loss,
             logits=logits,
-            past_key_values=None,  # Optional: implement later
+            past_key_values=None,  # Future: return actual cache if implemented
             hidden_states=None,
             attentions=None,
         )