added support for text-generation pipeline

config.json

@@ -25,5 +25,6 @@
   "use_bfloat16": true,
   "transformers_version": "4.35.2",
   "use_cache": true,
-  "vocab_size": 32000
+  "vocab_size": 32000,
+  "tokenizer_class": "LlamaTokenizer"
 }

modeling_decilm.py

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright and license here
+# Copyright and license in the repo.
 """ PyTorch DeciLM model."""
 from .version_check import check_transformers_version
 
@@ -11,18 +11,18 @@ import torch
 import torch.nn.functional as F
 import torch.utils.checkpoint
 from torch import nn
+from transformers.models.auto.modeling_auto import MODEL_FOR_CAUSAL_LM_MAPPING_NAMES
+from transformers.utils import add_start_docstrings, add_start_docstrings_to_model_forward, logging
 
+from .configuration_decilm import DeciLMConfig
+from .transformers_v4_35_2__modeling_attn_mask_utils import _prepare_4d_causal_attention_mask
 from .transformers_v4_35_2__modeling_llama import LlamaMLP, LlamaRMSNorm, LlamaAttention, apply_rotary_pos_emb, \
     repeat_kv, LlamaPreTrainedModel, LLAMA_START_DOCSTRING, LlamaDecoderLayer, LlamaForCausalLM, LlamaModel, \
     BaseModelOutputWithPast, LLAMA_INPUTS_DOCSTRING
-from .transformers_v4_35_2__modeling_attn_mask_utils import _prepare_4d_causal_attention_mask
-from transformers.utils import add_start_docstrings, add_start_docstrings_to_model_forward, logging
-
-from .configuration_decilm import DeciLMConfig
-
-logger = logging.get_logger(__name__)
 
+MODEL_FOR_CAUSAL_LM_MAPPING_NAMES["deci"] = "DeciLMForCausalLM"
 _CONFIG_FOR_DOC = "DeciLMConfig"
+logger = logging.get_logger(__name__)
 
 
 class DeciLMAttention(LlamaAttention):
@@ -61,7 +61,7 @@ class DeciLMAttention(LlamaAttention):
             past_key_value: Optional[Tuple[torch.Tensor]] = None,
             output_attentions: bool = False,
             use_cache: bool = False,
-            # padding_mask: Optional[torch.LongTensor] = None,
+            **kwargs,
     ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
         bsz, q_len, _ = hidden_states.size()
         is_decode = past_key_value is not None
@@ -135,8 +135,7 @@ class DeciLMAttention(LlamaAttention):
         else:
             attn_output = self.o_proj(attn_output)
 
-        if not output_attentions:
-            attn_weights = None
+        attn_weights = None
 
         return attn_output, attn_weights, past_key_value
 

version_check.py

@@ -1,11 +1,11 @@
 import transformers
 from packaging import version
 
-VERSION = "4.35.2"
+MIN_VERSION = "4.35.2"
 
 
 def check_transformers_version():
-    if version.parse(transformers.__version__) < version.parse(VERSION):
+    if version.parse(transformers.__version__) < version.parse(MIN_VERSION):
         raise ImportError(
-            f"You are using transformers=={transformers.__version__}, but transformers>={VERSION} is required to use DeciLM. Please upgrade transformers."
+            f"You are using transformers=={transformers.__version__}, but transformers>={MIN_VERSION} is required to use DeciLM. Please upgrade transformers."
         )