Update modelling_landmark_llama.py

modelling_landmark_llama.py

@@ -32,6 +32,8 @@ from transformers.modeling_utils import PreTrainedModel
 from transformers.utils import add_start_docstrings, add_start_docstrings_to_model_forward, logging, replace_return_docstrings
 from transformers.models.llama.configuration_llama import LlamaConfig
 
+from configuration_landmark_llama import LlamaConfig as LandmarkLlamaConfig
+
 
 logger = logging.get_logger(__name__)
 
@@ -565,7 +567,7 @@ LLAMA_START_DOCSTRING = r"""
     LLAMA_START_DOCSTRING,
 )
 class LlamaPreTrainedModel(PreTrainedModel):
-    config_class = LlamaConfig
+    config_class = LandmarkLlamaConfig
     base_model_prefix = "model"
     supports_gradient_checkpointing = True
     _no_split_modules = ["LlamaDecoderLayer"]
@@ -873,7 +875,7 @@ class LlamaModel(LlamaPreTrainedModel):
 
 
 class LlamaForCausalLM(LlamaPreTrainedModel):
-    def __init__(self, config):
+    def __init__(self, config: LandmarkLlamaConfig):
         super().__init__(config)
         self.model = LlamaModel(config)