Update modeling_svd_llama.py

modeling_svd_llama.py

@@ -11,7 +11,7 @@ from transformers.pytorch_utils import ALL_LAYERNORM_LAYERS
 from transformers.utils import logging
 from transformers import LlamaForCausalLM
 from transformers.models.llama.modeling_llama import LlamaDecoderLayer, LlamaModel, LlamaRotaryEmbedding, LlamaRMSNorm, repeat_kv, apply_rotary_pos_emb
-from component.configuration_svd_llama import SVDLlamaConfig
+from transformers import LlamaConfig
 
 
 logger = logging.get_logger(__name__)
@@ -21,7 +21,7 @@ _CONFIG_FOR_DOC = "LlamaConfig"
 ALL_LAYERNORM_LAYERS.append(LlamaRMSNorm)
 
 class SVDLlamaMLP(nn.Module):
-    def __init__(self, config: SVDLlamaConfig):
+    def __init__(self, config: LlamaConfig):
         super().__init__()
         self.config = config
         self.hidden_size = config.hidden_size
@@ -48,7 +48,7 @@ class SVDLlamaMLP(nn.Module):
 class SVDLlamaAttention(nn.Module):
     """Multi-headed attention from 'Attention Is All You Need' paper"""
 
-    def __init__(self, config: SVDLlamaConfig, layer_idx: Optional[int] = None):
+    def __init__(self, config: LlamaConfig, layer_idx: Optional[int] = None):
         super().__init__()
         self.config = config
         self.layer_idx = layer_idx
@@ -334,14 +334,14 @@ class SVDLLaMASDPA(SVDLlamaAttention):
 
 
 class SVDLlamaDecoderLayer(LlamaDecoderLayer):
-    def __init__(self, config: SVDLlamaConfig, layer_idx: int):
+    def __init__(self, config: LlamaConfig, layer_idx: int):
         super().__init__(config, layer_idx)
         self.self_attn = SVDLlamaAttention(config=config, layer_idx=layer_idx)
         self.mlp = SVDLlamaMLP(config)
 
 
 class SVDLlamaForCausalLM(LlamaForCausalLM):
-    def __init__(self, config: SVDLlamaConfig):
+    def __init__(self, config: LlamaConfig):
         super().__init__(config)
         self.model = LlamaModel(config)
         self.model.layers = nn.ModuleList(