Add exception handling

1ac3f976 · Casper Hansen · bd56c301 · 1ac3f976
Commit 1ac3f976 authored Sep 27, 2023 by Casper Hansen
Hide whitespace changes
Inline Side-by-side

Showing with 9 additions and 1 deletion

awq/models/mistral.py awq/models/mistral.py +9 -1

No files found.
--- a/awq/models/mistral.py
+++ b/awq/models/mistral.py
+import logging
 from .base import BaseAWQForCausalLM
-from transformers.models.mistral.modeling_mistral import MistralDecoderLayer, MistralForCausalLM
+
+try:
+    from transformers.models.mistral.modeling_mistral import MistralDecoderLayer, MistralForCausalLM
+except:
+    # TODO: Remove once released on PyPi
+    logging.warning("You need the latest transformers 4.34.0.dev0: pip install git+https://github.com/huggingface/transformers.git")
+    MistralForCausalLM = None
+    MistralDecoderLayer = None

 class MistralAWQForCausalLM(BaseAWQForCausalLM):
    layer_type = "MistralDecoderLayer"