Update cognitivess_model/modeling_cognitivess.py

cognitivess_model/modeling_cognitivess.py

@@ -1,14 +1,116 @@
+# cognitivess_model/modeling_cognitivess.py
+
 import torch
-from torch import nn
+import torch.nn as nn
 from transformers import PreTrainedModel
+from .configuration_cognitivess import CognitivessConfig
+
+class MultiHeadAttention(nn.Module):
+    def __init__(self, hidden_size, num_attention_heads, dropout_prob=0.0):
+        super().__init__()
+        self.num_attention_heads = num_attention_heads
+        self.attention_head_size = hidden_size // num_attention_heads
+        self.all_head_size = self.num_attention_heads * self.attention_head_size
+
+        self.query = nn.Linear(hidden_size, self.all_head_size)
+        self.key = nn.Linear(hidden_size, self.all_head_size)
+        self.value = nn.Linear(hidden_size, self.all_head_size)
+        self.dense = nn.Linear(hidden_size, hidden_size)
+
+        self.dropout = nn.Dropout(dropout_prob)
+
+    def forward(self, hidden_states, attention_mask=None):
+        batch_size, seq_length, hidden_size = hidden_states.size()
+        
+        query_layer = self.query(hidden_states).view(batch_size, seq_length, self.num_attention_heads, self.attention_head_size).transpose(1, 2)
+        key_layer = self.key(hidden_states).view(batch_size, seq_length, self.num_attention_heads, self.attention_head_size).transpose(1, 2)
+        value_layer = self.value(hidden_states).view(batch_size, seq_length, self.num_attention_heads, self.attention_head_size).transpose(1, 2)
+        
+        attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))
+        attention_scores = attention_scores / torch.sqrt(torch.tensor(self.attention_head_size, dtype=torch.float32))
+        
+        if attention_mask is not None:
+            attention_scores = attention_scores + attention_mask
+        
+        attention_probs = torch.nn.functional.softmax(attention_scores, dim=-1)
+        attention_probs = self.dropout(attention_probs)
+        
+        context_layer = torch.matmul(attention_probs, value_layer).transpose(1, 2).contiguous().view(batch_size, seq_length, hidden_size)
+        output_layer = self.dense(context_layer)
+        
+        return output_layer
+
+class FeedForward(nn.Module):
+    def __init__(self, hidden_size, intermediate_size, hidden_act, mlp_bias):
+        super().__init__()
+        self.dense = nn.Linear(hidden_size, intermediate_size, bias=mlp_bias)
+        self.activation = nn.SiLU() if hidden_act == "silu" else nn.ReLU()
+        self.output = nn.Linear(intermediate_size, hidden_size, bias=mlp_bias)
 
-class CognitivessForCausalLM(PreTrainedModel):
-    def __init__(self, config):
+    def forward(self, hidden_states):
+        hidden_states = self.dense(hidden_states)
+        hidden_states = self.activation(hidden_states)
+        hidden_states = self.output(hidden_states)
+        return hidden_states
+
+class TransformerBlock(nn.Module):
+    def __init__(self, hidden_size, num_attention_heads, intermediate_size, hidden_act, layer_norm_eps, mlp_bias, attention_dropout):
+        super().__init__()
+        self.attention = MultiHeadAttention(hidden_size, num_attention_heads, dropout_prob=attention_dropout)
+        self.feed_forward = FeedForward(hidden_size, intermediate_size, hidden_act, mlp_bias)
+        self.layer_norm1 = nn.LayerNorm(hidden_size, eps=layer_norm_eps)
+        self.layer_norm2 = nn.LayerNorm(hidden_size, eps=layer_norm_eps)
+        self.dropout = nn.Dropout(0.1)
+
+    def forward(self, hidden_states, attention_mask=None):
+        # Attention
+        attention_output = self.attention(hidden_states, attention_mask)
+        hidden_states = self.layer_norm1(hidden_states + attention_output)
+
+        # Feed Forward
+        feed_forward_output = self.feed_forward(hidden_states)
+        hidden_states = self.layer_norm2(hidden_states + feed_forward_output)
+        
+        return hidden_states
+
+class CognitivessModel(PreTrainedModel):
+    def __init__(self, config: CognitivessConfig):
         super().__init__(config)
-        self.transformer = nn.Transformer(d_model=config.hidden_size, num_encoder_layers=config.num_hidden_layers)
-        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        self.config = config
+        
+        self.embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id)
+        self.position_embeddings = nn.Embedding(config.max_position_embeddings, config.hidden_size)
+        self.layers = nn.ModuleList([
+            TransformerBlock(
+                config.hidden_size,
+                config.num_attention_heads,
+                config.intermediate_size,
+                config.hidden_act,
+                config.layer_norm_eps,
+                config.mlp_bias,
+                config.attention_dropout
+            )
+            for _ in range(config.num_hidden_layers)
+        ])
+        self.pooler = nn.Linear(config.hidden_size, config.hidden_size)
+        self.layer_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+        self.activation = nn.SiLU() if config.hidden_act == "silu" else nn.ReLU()
+
+    def forward(self, input_ids, attention_mask=None):
+        # Embeddings
+        embeddings = self.embeddings(input_ids)
+        position_ids = torch.arange(input_ids.size(1), dtype=torch.long, device=input_ids.device)
+        position_embeddings = self.position_embeddings(position_ids)
+        embeddings = embeddings + position_embeddings
+
+        # Transformer Layers
+        hidden_states = embeddings
+        for layer in self.layers:
+            hidden_states = layer(hidden_states, attention_mask)
+        
+        # Pooler
+        pooled_output = self.pooler(hidden_states[:, 0])
+        return pooled_output
 
-    def forward(self, input_ids, attention_mask=None, labels=None):
-        outputs = self.transformer(input_ids)
-        logits = self.lm_head(outputs)
-        return logits
+# Define CognitivessForCausalLM as an alias to CognitivessModel if needed
+CognitivessForCausalLM = CognitivessModel