Upload modeling_perklm.py with huggingface_hub

modeling_perklm.py

@@ -1,56 +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)
+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)