Correção Crítica: Bias Loading e Forward Kwargs (Versão Final)

config.json

@@ -13,7 +13,6 @@
     "intermediate_size": 1024,
     "max_position_embeddings": 1024,
     "rms_norm_eps": 1e-06,
-    "initializer_range": 0.02,
     "model_type": "tinygpt",
     "torch_dtype": "float32"
 }

configuration_tinygpt.py

@@ -1,19 +1,7 @@
 from transformers import PretrainedConfig
-
 class TinyGPTConfig(PretrainedConfig):
     model_type = "tinygpt"
-    def __init__(
-        self,
-        vocab_size=32000,
-        hidden_size=384,
-        num_hidden_layers=12,
-        num_attention_heads=8,
-        intermediate_size=1536,
-        max_position_embeddings=1024,
-        rms_norm_eps=1e-6,
-        initializer_range=0.02,
-        **kwargs,
-    ):
+    def __init__(self, vocab_size=32000, hidden_size=384, num_hidden_layers=12, num_attention_heads=8, intermediate_size=1024, max_position_embeddings=1024, rms_norm_eps=1e-6, initializer_range=0.02, **kwargs):
         self.vocab_size = vocab_size
         self.hidden_size = hidden_size
         self.num_hidden_layers = num_hidden_layers

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:6e30978bab63e4b3d523cb089eeace6028d93af39b418f122899ebdb8275bd91
-size 165985968
+oid sha256:84954ffd48c829d1a6c3a95bc668df9206c2e40562f152d92de821f377afeac0
+size 166134520

modeling_tinygpt.py

@@ -2,13 +2,13 @@ import torch
 import torch.nn as nn
 from transformers import PreTrainedModel
 from .configuration_tinygpt import TinyGPTConfig
+from transformers.modeling_outputs import CausalLMOutputWithPast # Importante para retorno correto
 
 class RMSNorm(nn.Module):
     def __init__(self, dim, eps=1e-6):
         super().__init__()
         self.eps = eps
         self.weight = nn.Parameter(torch.ones(dim))
-
     def forward(self, x):
         var = torch.mean(x ** 2, dim=-1, keepdim=True)
         return x * torch.rsqrt(var + self.eps) * self.weight
@@ -16,10 +16,9 @@ class RMSNorm(nn.Module):
 class MLP(nn.Module):
     def __init__(self, config):
         super().__init__()
-        self.fc_in = nn.Linear(config.hidden_size, config.intermediate_size)
+        self.fc_in = nn.Linear(config.hidden_size, config.intermediate_size, bias=True)
         self.act = nn.GELU()
-        self.fc_out = nn.Linear(config.intermediate_size, config.hidden_size)
-
+        self.fc_out = nn.Linear(config.intermediate_size, config.hidden_size, bias=True)
     def forward(self, x):
         return self.fc_out(self.act(self.fc_in(x)))
 
@@ -29,22 +28,19 @@ class Attention(nn.Module):
         self.n_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)
-
+        self.q_proj = nn.Linear(config.hidden_size, config.hidden_size, bias=True)
+        self.k_proj = nn.Linear(config.hidden_size, config.hidden_size, bias=True)
+        self.v_proj = nn.Linear(config.hidden_size, config.hidden_size, bias=True)
+        self.out_proj = nn.Linear(config.hidden_size, config.hidden_size, bias=True)
     def forward(self, x, mask=None):
         B, T, C = x.shape
         q = self.q_proj(x).view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
         k = self.k_proj(x).view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
         v = self.v_proj(x).view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
-
         att = (q @ k.transpose(-2, -1)) * self.scale
         if mask is not None:
+            if mask.dim() == 2: mask = mask.unsqueeze(0).unsqueeze(0)
             att = att.masked_fill(mask == 0, float('-inf'))
-        
         att = torch.softmax(att, dim=-1)
         out = (att @ v).transpose(1, 2).contiguous().view(B, T, C)
         return self.out_proj(out)
@@ -56,7 +52,6 @@ class Block(nn.Module):
         self.attn = Attention(config)
         self.norm_2 = RMSNorm(config.hidden_size, config.rms_norm_eps)
         self.mlp = MLP(config)
-
     def forward(self, x, mask=None):
         x = x + self.attn(self.norm_1(x), mask)
         x = x + self.mlp(self.norm_2(x))
@@ -79,8 +74,7 @@ class TinyGPTModel(TinyGPTPreTrainedModel):
         self.wpe = nn.Embedding(config.max_position_embeddings, config.hidden_size)
         self.h = nn.ModuleList([Block(config) for _ in range(config.num_hidden_layers)])
         self.ln_f = RMSNorm(config.hidden_size, config.rms_norm_eps)
-
-    def forward(self, input_ids):
+    def forward(self, input_ids, attention_mask=None):
         B, T = input_ids.shape
         pos = torch.arange(0, T, dtype=torch.long, device=input_ids.device)
         x = self.wte(input_ids) + self.wpe(pos)
@@ -94,14 +88,24 @@ class TinyGPTForCausalLM(TinyGPTPreTrainedModel):
         super().__init__(config)
         self.transformer = TinyGPTModel(config)
         self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
-
-    def forward(self, input_ids, labels=None):
-        hidden = self.transformer(input_ids)
+    
+    # AQUI ESTAVA O ERRO! Adicionei **kwargs para engolir return_dict, output_attentions, etc.
+    def forward(self, input_ids, attention_mask=None, labels=None, **kwargs):
+        hidden = self.transformer(input_ids, attention_mask)
         logits = self.lm_head(hidden)
+        
         loss = None
         if labels is not None:
             shift_logits = logits[..., :-1, :].contiguous()
             shift_labels = labels[..., 1:].contiguous()
             loss_fct = nn.CrossEntropyLoss()
             loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))
-        return {"loss": loss, "logits": logits}
+            
+        # Retorna objeto padrão do HF para evitar erros de compatibilidade
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+        )
+
+    def prepare_inputs_for_generation(self, input_ids, **kwargs):
+        return {"input_ids": input_ids}