GradientGuru commited on
Commit
2026f61
1 Parent(s): fc64079

Update modeling_baichuan.py

Browse files
Files changed (1) hide show
  1. modeling_baichuan.py +8 -1
modeling_baichuan.py CHANGED
@@ -4,6 +4,7 @@ import math
4
  from typing import List, Optional, Tuple, Union
5
 
6
  import torch
 
7
  from torch.nn import CrossEntropyLoss
8
  from transformers import PreTrainedModel
9
  from transformers.activations import ACT2FN
@@ -235,7 +236,13 @@ class BaichuanModel(BaichuanPreTrainedModel):
235
  self.post_init()
236
  self.max_cache_pos = config.model_max_length
237
  self.first_run = True
238
-
 
 
 
 
 
 
239
  def get_alibi_mask(self, tensor, seq_length_with_past):
240
  if self.first_run:
241
  self.first_run = False
 
4
  from typing import List, Optional, Tuple, Union
5
 
6
  import torch
7
+ import torch.utils.checkpoint
8
  from torch.nn import CrossEntropyLoss
9
  from transformers import PreTrainedModel
10
  from transformers.activations import ACT2FN
 
236
  self.post_init()
237
  self.max_cache_pos = config.model_max_length
238
  self.first_run = True
239
+
240
+ def get_input_embeddings(self):
241
+ return self.embed_tokens
242
+
243
+ def set_input_embeddings(self, value):
244
+ self.embed_tokens = value
245
+
246
  def get_alibi_mask(self, tensor, seq_length_with_past):
247
  if self.first_run:
248
  self.first_run = False