move flash

webis-de · Dec 16, 2024 · 17d4eb7 · 17d4eb7
1 parent 8b7c25c
commit 17d4eb7
Show file tree

Hide file tree

Showing 3 changed files with 1 addition and 1 deletion.
diff --git a/lightning_ir/flash/__init__.py → lightning_ir/_flash/__init__.py b/lightning_ir/flash/__init__.py → lightning_ir/_flash/__init__.py
diff --git a/lightning_ir/flash/flash_bert.py → lightning_ir/_flash/flash_bert.py b/lightning_ir/flash/flash_bert.py → lightning_ir/_flash/flash_bert.py
diff --git a/lightning_ir/base/model.py b/lightning_ir/base/model.py
@@ -14,7 +14,7 @@
 from transformers import MODEL_MAPPING, BatchEncoding, BertModel
 from transformers.modeling_outputs import ModelOutput
 
-from ..flash import FLASH_ATTENTION_MAP
+from .._flash import FLASH_ATTENTION_MAP
 from .class_factory import LightningIRModelClassFactory
 from .config import LightningIRConfig
 from .external_model_hub import CHECKPOINT_MAPPING, POST_LOAD_CALLBACKS, STATE_DICT_KEY_MAPPING