shim/transformers_shim_runtime.py

"""Re-export names the HF inference toolkit's `utils.py` imports from
`transformers.file_utils` but that transformers 5.x dropped.

The toolkit does, at module-import time:

    from transformers.file_utils import is_tf_available, is_torch_available

In transformers 5.x:
  - `is_torch_available` still lives in `transformers.utils.import_utils`
    and (in this version) is re-exported by `transformers.file_utils`.
  - `is_tf_available` was removed entirely — TensorFlow support was
    dropped from transformers 5.x.

This module is loaded by the sibling `.pth` file at interpreter startup,
before the toolkit's `webservice_starlette` imports run. Best-effort:
silently no-op if transformers is missing (e.g. running during another
package's install).
"""
try:  # pragma: no cover — runs once at interpreter startup
    from transformers import file_utils as _fu  # noqa: I001
    from transformers import utils as _u
    if not hasattr(_fu, "is_torch_available"):
        try:
            from transformers.utils.import_utils import is_torch_available
        except ImportError:
            from transformers.utils import is_torch_available  # type: ignore[no-redef]
        _fu.is_torch_available = is_torch_available
    if not hasattr(_fu, "is_tf_available"):
        # TF was dropped in transformers 5.x. The toolkit only consults
        # this in a top-level guard; returning False short-circuits the
        # `import tensorflow as tf` branch we don't want anyway.
        _fu.is_tf_available = lambda: False
    for _name, _val in (
        ("FLAX_WEIGHTS_NAME", "flax_model.msgpack"),
        ("FLAX_WEIGHTS_INDEX_NAME", "flax_model.msgpack.index.json"),
        ("TF2_WEIGHTS_NAME", "tf_model.h5"),
        ("TF2_WEIGHTS_INDEX_NAME", "tf_model.h5.index.json"),
        ("TF_WEIGHTS_NAME", "model.ckpt"),
    ):
        if not hasattr(_u, _name):
            setattr(_u, _name, _val)
        if not hasattr(_fu, _name):
            setattr(_fu, _name, _val)
except Exception:
    pass

# Disable the HF inference toolkit's diffusers integration. The toolkit
# (in /app/huggingface_inference_toolkit/utils.py) does an unconditional
# `from huggingface_inference_toolkit.diffusers_utils import (...)` which
# in turn does `from diffusers import (...)`. Diffusers' pipeline modules
# transitively reference dozens of transformers 4.x names that are gone
# in 5.x (FLAX_WEIGHTS_NAME, MT5Tokenizer, ...). We don't use diffusers
# at all, so the cheapest fix is to stub out `huggingface_inference_
# toolkit.diffusers_utils` BEFORE the toolkit imports it.
try:
    import sys
    import types
    if "huggingface_inference_toolkit.diffusers_utils" not in sys.modules:
        _stub = types.ModuleType("huggingface_inference_toolkit.diffusers_utils")
        _stub.is_diffusers_available = lambda: False
        # The toolkit imports these names by destructuring; provide
        # placeholders that are never called because is_diffusers_available()
        # short-circuits all real use.
        _stub.get_diffusers_pipeline = lambda *a, **kw: (_ for _ in ()).throw(
            RuntimeError("diffusers stubbed by transformers-shim")
        )
        sys.modules["huggingface_inference_toolkit.diffusers_utils"] = _stub
except Exception:
    pass