Update modeling_cohere.py

modeling_cohere.py

@@ -52,6 +52,9 @@ from transformers.utils import (
 )
 from .configuration_cohere import CohereConfig
 
+if is_flash_attn_2_available():
+    from flash_attn import flash_attn_func, flash_attn_varlen_func
+    from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input  # noqa
 
 logger = logging.get_logger(__name__)
 

requirements.txt

@@ -0,0 +1 @@
+flash_attn