XBFDI commited on
Commit
3b6d28f
·
verified ·
1 Parent(s): 7bacf93

Update modeling_easyformer.py

Browse files
Files changed (1) hide show
  1. modeling_easyformer.py +11 -2
modeling_easyformer.py CHANGED
@@ -1,7 +1,11 @@
1
-
2
  import math, torch, torch.nn as nn, torch.nn.functional as F
3
  from transformers import PreTrainedModel, GenerationMixin
4
- from .configuration_easyformer import EasyFormerConfig
 
 
 
 
5
 
6
  class RMSNorm(nn.Module):
7
  def __init__(self, d, eps=1e-5):
@@ -50,6 +54,11 @@ class EasyFormerPreTrainedModel(PreTrainedModel):
50
  supports_gradient_checkpointing = False
51
 
52
  class EasyFormerLMHeadModel(EasyFormerPreTrainedModel, GenerationMixin):
 
 
 
 
 
53
  def __init__(self, config):
54
  super().__init__(config)
55
  self.cfg = config
 
1
+ # modeling_easyformer.py
2
  import math, torch, torch.nn as nn, torch.nn.functional as F
3
  from transformers import PreTrainedModel, GenerationMixin
4
+
5
+ try:
6
+ from .configuration_easyformer import EasyFormerConfig
7
+ except ImportError:
8
+ from configuration_easyformer import EasyFormerConfig
9
 
10
  class RMSNorm(nn.Module):
11
  def __init__(self, d, eps=1e-5):
 
54
  supports_gradient_checkpointing = False
55
 
56
  class EasyFormerLMHeadModel(EasyFormerPreTrainedModel, GenerationMixin):
57
+ config_class = EasyFormerConfig
58
+ base_model_prefix = "easyformer"
59
+ _tied_weights_keys = ["lm_head.weight"]
60
+ all_tied_weights_keys = ["lm_head.weight"]
61
+
62
  def __init__(self, config):
63
  super().__init__(config)
64
  self.cfg = config