Update modeling_gad2.py

modeling_gad2.py

@@ -16,6 +16,11 @@ class GAD2Config(PretrainedConfig):
         self.num_intents = num_intents
         self.memory_size = memory_size
         self.seq_len = seq_len
+        
+        # --- תיקון השגיאה: הוספת שמות סטנדרטיים שהספרייה מחפשת ---
+        self.num_hidden_layers = num_layers
+        self.hidden_size = d_model
+        self.num_attention_heads = nhead
 
 class RMSNorm(nn.Module):
     def __init__(self, dim, eps=1e-6):
@@ -47,7 +52,6 @@ class AdaptiveMemory(nn.Module):
         q = self.q_proj(x)
         kv = self.kv_proj(self.memory).view(-1, 2, d)
         k, v = kv[:, 0].unsqueeze(0).expand(b, -1, -1), kv[:, 1].unsqueeze(0).expand(b, -1, -1)
-        # Scaled Dot Product Attention
         attn_out = F.scaled_dot_product_attention(q.unsqueeze(1), k.unsqueeze(1), v.unsqueeze(1))
         return self.out_proj(attn_out.squeeze(1))
 
@@ -63,18 +67,13 @@ class GAD2_Block(nn.Module):
         h = self.norm1(x)
         b, t, d = h.size()
         qkv = self.qkv(h).view(b, t, 3, self.nhead, self.head_dim).permute(2, 0, 3, 1, 4)
-        q, k, v = qkv[0], qkv[1], qkv[2]
-        
-        attn_out = F.scaled_dot_product_attention(q, k, v, is_causal=True)
+        attn_out = F.scaled_dot_product_attention(qkv[0], qkv[1], qkv[2], is_causal=True)
         x = x + self.proj(attn_out.transpose(1, 2).reshape(b, t, d))
-        
         h = self.norm2(x)
-        x = x + self.w3(F.silu(self.w1(h)) * self.w2(h))
-        return x
+        return x + self.w3(F.silu(self.w1(h)) * self.w2(h))
 
 class GAD2ForAgenticModeling(PreTrainedModel):
     config_class = GAD2Config
-    
     def __init__(self, config):
         super().__init__(config)
         self.token_emb = nn.Embedding(config.vocab_size, config.d_model)
@@ -84,33 +83,14 @@ class GAD2ForAgenticModeling(PreTrainedModel):
         self.post_init()
 
     def forward(self, input_ids, labels=None, attention_mask=None, **kwargs):
-        # input_ids עשוי להגיע כחלק מ-kwargs ב-generate
-        if input_ids is None and "inputs_embeds" in kwargs:
-             return None # טיפול במקרים חריגים של inputs_embeds
-             
         x = self.token_emb(input_ids)
-        
-        # Agentic Integration
         x = x + self.intent_evolver(x) + self.memory_module(x)
-        
-        # Transformer Blocks
-        for block in self.blocks:
+        for block in self.blocks: 
             x = block(x)
-            
         logits = self.output_head(self.final_norm(x))
-        
-        loss = None
-        if labels is not None:
-            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), labels.view(-1))
-            
-        return CausalLMOutputWithPast(
-            loss=loss,
-            logits=logits,
-            past_key_values=None,
-            hidden_states=None,
-            attentions=None,
-        )
+        loss = F.cross_entropy(logits.view(-1, logits.size(-1)), labels.view(-1)) if labels is not None else None
+        return CausalLMOutputWithPast(loss=loss, logits=logits)
 
     def prepare_inputs_for_generation(self, input_ids, **kwargs):
-        # פונקציה חיונית עבור .generate()
-        return {"input_ids": input_ids}
+        # חשוב מאוד: הוספת תמיכה ב-kwargs כדי ש-generate לא יתבלבל
+        return {"input_ids": input_ids, **kwargs}