Upload 7 files

config.json

@@ -4,6 +4,7 @@
   ],
   "auto_map": {
     "AutoConfig": "enigma_module.EnigmaConfig",
+    "AutoModel": "enigma_module.EnigmaModel",
     "AutoModelForCausalLM": "enigma_module.EnigmaForCausalLM"
   },
   "dtype": "float32",

enigma_module.py

@@ -24,7 +24,9 @@ class EnigmaModel(PreTrainedModel):
         x = self.embedding(input_ids)
         return self.linear(x)
 
-class EnigmaForCausalLM(PreTrainedModel):
+from transformers.generation import GenerationMixin
+
+class EnigmaForCausalLM(PreTrainedModel, GenerationMixin):
     config_class = EnigmaConfig
     
     def __init__(self, config):

generation_config.json

@@ -0,0 +1,6 @@
+{
+  "_from_model_config": true,
+  "output_attentions": false,
+  "output_hidden_states": false,
+  "transformers_version": "5.8.0.dev0"
+}

model.safetensors

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