Upload MyLLaMa

config.json

@@ -2,6 +2,10 @@
   "architectures": [
     "MyLLaMa"
   ],
+  "auto_map": {
+    "AutoConfig": "configure_for_hf.MyLLaMaConfig",
+    "AutoModelForCausalLM": "configure_for_hf.MyLLaMa"
+  },
   "embed_dim": 1536,
   "model_type": "LLaMa",
   "n_chckpnt_segments": 24,

configure_for_hf.py

@@ -7,6 +7,7 @@ from transformers import (
     PretrainedConfig,
     PreTrainedModel,
 )
+from collections import namedtuple
 
 from .llama import CustomAttentionLLaMa
 
@@ -70,11 +71,16 @@ class MyLLaMa(PreTrainedModel):
         pad_mask = torch.where(
             tensor == self.model.tokenizer.pad_token_id, False, True
         ).cuda()
+        outs = namedtuple("output", ["logits", "loss"])
         logits = self.model(tensor, att_mask, pad_mask)["logits"]
+
+        outs.logits = logits
+
         if labels is not None:
             loss = nn.functional.cross_entropy(logits, labels)
-            return {"loss": loss, "logits": logits}
-        return {"logits": logits}
+            outs.loss = loss
+
+        return outs
 
 
 AutoConfig.register("LLaMa", MyLLaMaConfig)