itlevy commited on
Commit
3209eec
1 Parent(s): e5d0706

DeciLMForCausalLM(DeciLMPreTrainedModel, GenerationMixin) for v4.50 (#16)

Browse files

- DeciLMForCausalLM(DeciLMPreTrainedModel, GenerationMixin) for v4.50 (f3ff1f12f57fe5142bc8902c6a067b870cab022e)

Files changed (1) hide show
  1. modeling_decilm.py +2 -2
modeling_decilm.py CHANGED
@@ -25,7 +25,7 @@ import torch.utils.checkpoint
25
  from torch import nn
26
  from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
27
  from transformers import GenerationConfig
28
- from transformers.generation.utils import NEED_SETUP_CACHE_CLASSES_MAPPING
29
  from transformers.modeling_utils import PreTrainedModel
30
  from transformers.utils import (
31
  add_start_docstrings,
@@ -1131,7 +1131,7 @@ class DeciLMModel(DeciLMPreTrainedModel):
1131
  return causal_mask
1132
 
1133
 
1134
- class DeciLMForCausalLM(DeciLMPreTrainedModel):
1135
  _tied_weights_keys = ["lm_head.weight"]
1136
 
1137
  def __init__(self, config):
 
25
  from torch import nn
26
  from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
27
  from transformers import GenerationConfig
28
+ from transformers.generation.utils import GenerationMixin, NEED_SETUP_CACHE_CLASSES_MAPPING
29
  from transformers.modeling_utils import PreTrainedModel
30
  from transformers.utils import (
31
  add_start_docstrings,
 
1131
  return causal_mask
1132
 
1133
 
1134
+ class DeciLMForCausalLM(DeciLMPreTrainedModel, GenerationMixin):
1135
  _tied_weights_keys = ["lm_head.weight"]
1136
 
1137
  def __init__(self, config):