utils.py

import os
from unittest.mock import patch
from transformers.dynamic_module_utils import get_imports

def fixed_get_imports(filename: str | os.PathLike) -> list[str]:
    """Workaround for flash_attn import issue."""
    if not str(filename).endswith(("modeling_phi.py", "configuration_florence2.py")):
        return get_imports(filename)
    imports = get_imports(filename)
    if "flash_attn" in imports:
        imports.remove("flash_attn")
    return imports

def load_model_without_flash_attn(model_loader):
    """Load a model using the flash_attn workaround."""
    with patch("transformers.dynamic_module_utils.get_imports", fixed_get_imports):
        return model_loader()