Upload 6 files

config.json

@@ -1,10 +1,10 @@
 {
   "architectures": [
-    "EnigmaModel"
+    "EnigmaForCausalLM"
   ],
   "auto_map": {
     "AutoConfig": "enigma_module.EnigmaConfig",
-    "AutoModel": "enigma_module.EnigmaModel"
+    "AutoModelForCausalLM": "enigma_module.EnigmaForCausalLM"
   },
   "dtype": "float32",
   "hidden_size": 128,

enigma_module.py

@@ -1,6 +1,7 @@
 import torch
 import torch.nn as nn
 from transformers import PreTrainedModel, PretrainedConfig
+from transformers.modeling_outputs import CausalLMOutputWithPast
 
 class EnigmaConfig(PretrainedConfig):
     model_type = "enigma"
@@ -19,10 +20,32 @@ class EnigmaModel(PreTrainedModel):
         self.linear = nn.Linear(config.hidden_size, config.hidden_size)
         self.post_init()
         
-    def forward(self, input_ids):
+    def forward(self, input_ids, **kwargs):
         x = self.embedding(input_ids)
         return self.linear(x)
 
-# Registrando para permitir AutoModel e AutoConfig (trust_remote_code=True)
+class EnigmaForCausalLM(PreTrainedModel):
+    config_class = EnigmaConfig
+    
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = EnigmaModel(config)
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        self.post_init()
+        
+    def forward(self, input_ids, labels=None, **kwargs):
+        hidden_states = self.model(input_ids)
+        logits = self.lm_head(hidden_states)
+        loss = None
+        if labels is not None:
+            loss_fct = nn.CrossEntropyLoss()
+            loss = loss_fct(logits.view(-1, self.config.vocab_size), labels.view(-1))
+        return CausalLMOutputWithPast(loss=loss, logits=logits)
+        
+    def prepare_inputs_for_generation(self, input_ids, **kwargs):
+        return {"input_ids": input_ids}
+
+# Registrando para permitir AutoModel, AutoConfig e AutoModelForCausalLM
 EnigmaConfig.register_for_auto_class()
 EnigmaModel.register_for_auto_class("AutoModel")
+EnigmaForCausalLM.register_for_auto_class("AutoModelForCausalLM")

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:78e404d4dc8e88bc893ccc0b188a5e095da978c676a4cde1c1a7bce429eb70ed
-size 2626336
+oid sha256:72ba2e2671bd22265a6aa09c43682e31209e85aac91be3d872a4bb04e142d6df
+size 5186432

tokenizer_config.json

@@ -0,0 +1,13 @@
+{
+  "add_prefix_space": false,
+  "backend": "tokenizers",
+  "bos_token": "<|endoftext|>",
+  "eos_token": "<|endoftext|>",
+  "errors": "replace",
+  "is_local": false,
+  "local_files_only": false,
+  "model_max_length": 1024,
+  "pad_token": null,
+  "tokenizer_class": "GPT2Tokenizer",
+  "unk_token": "<|endoftext|>"
+}