Update modeling_phi2_clex.py

modeling_phi2_clex.py

@@ -59,10 +59,7 @@ logger = logging.get_logger(__name__)
 _CHECKPOINT_FOR_DOC = "microsoft/phi-2"
 _CONFIG_FOR_DOC = "CLEXPhiConfig"
 
-PHI_PRETRAINED_MODEL_ARCHIVE_LIST = [
-    "microsoft/phi-2",
-    # See all Phi models at https://huggingface.co/models?filter=phi
-]
+
 
 
 # Copied from transformers.models.llama.modeling_llama._get_unpad_data
@@ -663,7 +660,7 @@ PHI_ATTENTION_CLASSES = {
 class PhiDecoderLayer(nn.Module):
     def __init__(self, config: CLEXPhiConfig, layer_idx: int):
         super().__init__()
-        self.self_attn = PhiFlashAttention2(config, layer_idx=layer_idx)
+        self.self_attn = PHI_ATTENTION_CLASSES[config._attn_implementation](config, layer_idx=layer_idx)
         self.mlp = PhiMLP(config)
         self.input_layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
         self.resid_dropout = nn.Dropout(config.resid_pdrop)