llama.cpp: import more defensively

eth-sri · Aug 15, 2023 · f219bd9 · f219bd9
1 parent 7f23adc
commit f219bd9
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 4 deletions.
diff --git a/src/lmql/models/lmtp/backends/llama_cpp_model.py b/src/lmql/models/lmtp/backends/llama_cpp_model.py
@@ -2,14 +2,15 @@
 import sys
 
 import numpy as np
-from llama_cpp import Llama, LlamaTokenizer
 
 import lmql.utils.nputil as nputil
 from lmql.models.lmtp.backends.lmtp_model import (LMTPModel, LMTPModelResult,
                                                   TokenStreamer)
 
 class LlamaCppModel(LMTPModel):
     def __init__(self, model_identifier, **kwargs):
+        from llama_cpp import Llama
+
         self.model_identifier = model_identifier
         self.kwargs = kwargs
 

diff --git a/src/lmql/models/lmtp/backends/lmtp_model.py b/src/lmql/models/lmtp/backends/lmtp_model.py
@@ -107,9 +107,9 @@ def wrapper(loader):
             import importlib
             if module_dependencies is not None:
                 for module in module_dependencies:
-                    try:
-                        importlib.import_module(module)
-                    except ImportError:
+                    # check without importing
+                    spec = importlib.util.find_spec(module)
+                    if spec is None:
                         def error_func(*args, **kwargs):
                             assert False, "To use the {} backend, please install the '{}' package.".format(name, module)
                         LMTPModel.registry[name] = error_func