hf importable

hf.py

@@ -0,0 +1,21 @@
+import torch
+from transformers import GPT2TokenizerFast
+from modeling_perklm import PerkLMConfig, PerkLM
+
+
+checkpoint = torch.load('perklm.pt', map_location = 'cpu', weights_only = True)
+old_config = checkpoint['config']
+
+hf_config = PerkLMConfig(d_model = old_config['model']['d_model'],
+                        n_heads = old_config['model']['n_heads'],
+                        n_layers = old_config['model']['n_layers'],
+                        d_ff = old_config['model']['d_ff'],
+                        maxt = old_config['model']['maxt'],
+                        dropout = old_config['model']['dropout'],
+                        tokenizer_path = 'tokenizer')
+
+model = PerkLM(hf_config)
+model.transformer.load_state_dict(checkpoint['model_state_dict'])
+
+model.save_pretrained('perklm_hf')
+hf_config.save_pretrained('perklm_hf')

modeling_perklm.py

@@ -0,0 +1,56 @@
+import torch
+from transformers import PreTrainedModel, PretrainedConfig, GPT2TokenizerFast, \
+    AutoConfig, AutoModelForCausalLM
+from transformers.modeling_outputs import CausalLMOutput
+from model import FriendsTransformer
+
+class PerkLMConfig(PretrainedConfig):
+    model_type = "perklm"
+
+    def __init__(self, d_model = 512, n_heads = 8, n_layers = 6, 
+                 d_ff = 2048, maxt = 512, dropout = 0.1, tokenizer_path = None, **kwargs):
+        super().__init__(**kwargs)
+        self.d_model = d_model
+        self.n_heads = n_heads
+        self.n_layers = n_layers
+        self.d_ff = d_ff
+        self.maxt = maxt
+        self.dropout = dropout
+        self.tokenizer_path = tokenizer_path
+
+class PerkLM(PreTrainedModel):
+    config_class = PerkLMConfig
+    _tied_weights_keys = ["transformer.lm_head.weight"]
+
+    def __init__(self, config):
+        super().__init__(config)
+        tokenizer = GPT2TokenizerFast.from_pretrained(config.tokenizer_path)
+        self.transformer = FriendsTransformer(d_model = config.d_model, 
+                                              n_heads = config.n_heads, 
+                                              n_layers = config.n_layers,
+                                              d_ff = config.d_ff,
+                                              dropout = config.dropout,
+                                              maxt = config.maxt,
+                                              tokenizer = tokenizer)
+        self.post_init()
+
+    def forward(self, input_ids, attention_mask = None, 
+                responder = None, labels = None, **kwargs):
+        batch = {'input_ids': input_ids, 
+                'attention_mask': attention_mask if attention_mask is not None \
+                    else torch.ones_like(input_ids), 
+                'responder': responder}
+        logits = self.transformer(batch)
+        
+        loss = None
+        if labels is not None:
+            loss = torch.nn.functional.cross_entropy(logits[:, :-1].reshape(-1, logits.size(-1)), 
+                                                     labels[:, 1:].reshape(-1), ignore_index = -100)
+
+        return CausalLMOutput(loss = loss, logits = logits)
+    
+    def tie_weights(self):
+        self.transformer.lm_head.weight = self.transformer.embedder.embedding.weight
+
+AutoConfig.register("perklm", PerkLMConfig)
+AutoModelForCausalLM.register(PerkLMConfig, PerkLM)

perklm_hf/config.json

@@ -0,0 +1,15 @@
+{
+  "architectures": [
+    "PerkLM"
+  ],
+  "d_ff": 2048,
+  "d_model": 512,
+  "dropout": 0.1,
+  "dtype": "float32",
+  "maxt": 512,
+  "model_type": "perklm",
+  "n_heads": 8,
+  "n_layers": 6,
+  "tokenizer_path": "tokenizer",
+  "transformers_version": "4.57.3"
+}

perklm_hf/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7ae6b37968cc9a69d3231e709f51ea32d5d9ad4d917ff3685126e7ca93c18413
+size 179394920