diff --git a/src/llms/petals.py b/src/llms/petals.py index 09159a0f..bcb12420 100644 --- a/src/llms/petals.py +++ b/src/llms/petals.py @@ -20,7 +20,7 @@ _ServerInferenceSession_step: None | Callable = None -def _is_batch_size_exception_func(e: BaseException) -> bool: # pragma: no cover +def _is_batch_size_exception_func(e: BaseException) -> bool: from hivemind.p2p.p2p_daemon_bindings.utils import P2PHandlerError return ( @@ -66,7 +66,7 @@ def _patched_on_request_failure( @cache -def _monkey_patch_ServerInferenceSession_step(): # pragma: no cover +def _monkey_patch_ServerInferenceSession_step(): from hivemind.p2p.p2p_daemon_bindings.utils import P2PHandlerError # noqa: F401 try: @@ -80,7 +80,7 @@ def _monkey_patch_ServerInferenceSession_step(): # pragma: no cover _ServerInferenceSession.step = _catch_memory_error__ServerInferenceSession_step -class Petals(HFTransformers): # pragma: no cover +class Petals(HFTransformers): def __init__( self, model_name: str,