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

@@ -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):

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."
         )