rixty/KoboldAI-Client
1
1
Fork 0
mirror of https://github.com/KoboldAI/KoboldAI-Client.git synced 2025-06-05 21:59:24 +02:00
Code Issues Projects Releases Wiki Activity

Silently disable MTJ when Jax is not installed

Browse Source
This commit is contained in:
Henk
2023-06-21 17:08:45 +02:00
parent d46663ac0d
commit bbecdaeedb
1 changed files with 15 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

24
modeling/inference_models/hf_mtj/class.py
View File

@@ -29,15 +29,21 @@ class model_backend(HFInferenceModel):
#model_name: str, #model_name: str,
) -> None: ) -> None:
super().__init__() super().__init__()
self.hf_torch = False import importlib
self.model_config = None dependency_exists = importlib.util.find_spec("jax")
self.capabilties = ModelCapabilities( if dependency_exists:
embedding_manipulation=False, self.hf_torch = False
post_token_hooks=False, self.model_config = None
stopper_hooks=False, self.capabilties = ModelCapabilities(
post_token_probs=False, embedding_manipulation=False,
uses_tpu=True, post_token_hooks=False,
) stopper_hooks=False,
post_token_probs=False,
uses_tpu=True,
)
else:
logger.debug("Jax is not installed, hiding TPU backend")
self.disable = True
def is_valid(self, model_name, model_path, menu_path): def is_valid(self, model_name, model_path, menu_path):
# This file shouldn't be imported unless using the TPU # This file shouldn't be imported unless using the TPU