Model: Prompt users to install extras if dependencies don't exist

Ex: tokenizers, lmfe, outlines. Signed-off-by: kingbri <[email protected]>
theroyallab · Mar 23, 2024 · 1755f28 · 1755f28
1 parent f952b81
commit 1755f28
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 4 deletions.
diff --git a/backends/exllamav2/grammar.py b/backends/exllamav2/grammar.py
@@ -71,8 +71,9 @@ def add_json_schema_filter(
             logger.error(
                 "Skipping JSON schema parsing because "
                 "lm-format-enforcer is not installed.\n"
-                "Please run the following command: "
-                "pip install lm-format-enforcer"
+                "Please run the following command in your environment "
+                "to install extra packages:\n"
+                "pip install -U .[extras]"
             )
 
             return
@@ -113,7 +114,9 @@ def add_ebnf_filter(
         except ImportError:
             logger.error(
                 "Skipping EBNF parsing because Outlines is not installed.\n"
-                "Please run the following command: pip install outlines"
+                "Please run the following command in your environment "
+                "to install extra packages:\n"
+                "pip install -U .[extras]"
             )
 
             return

diff --git a/backends/exllamav2/model.py b/backends/exllamav2/model.py
@@ -399,7 +399,19 @@ def progress(loaded_modules: int, total_modules: int)
         ExLlamaV2Tokenizer.extended_id_to_piece = {}
         ExLlamaV2Tokenizer.extended_piece_to_id = {}
 
-        self.tokenizer = ExLlamaV2Tokenizer(self.config)
+        try:
+            self.tokenizer = ExLlamaV2Tokenizer(self.config)
+        except AssertionError as exc:
+            if "HF tokenizer" in str(exc):
+                raise ImportError(
+                    "Could not create ExllamaV2's tokenizer for this model "
+                    "because tokenizers is not installed.\n"
+                    "Please run the following command in your environment "
+                    "to install extra packages:\n"
+                    "pip install -U .[extras]"
+                ) from exc
+            else:
+                raise exc
 
         # Calculate autosplit reserve for all GPUs
         gpu_count = torch.cuda.device_count()