fix exaone import (#659)

casper-hansen · Nov 27, 2024 · 0d1906f · 0d1906f
1 parent 4b10b31
commit 0d1906f
Showing 1 changed file with 8 additions and 4 deletions.
diff --git a/awq/models/exaone.py b/awq/models/exaone.py
@@ -4,10 +4,14 @@
 from awq.utils.fused_utils import fuse_qkv
 from awq.modules.fused.block import LlamaLikeBlock
 from awq.modules.fused.model import LlamaLikeModel
-from transformers.models.exaone.modeling_exaone import (
-    ExaoneBlock as OldExaoneBlock,
-    ExaoneForCausalLM as OldExaoneForCausalLM,
-)
+try:
+    from transformers.models.exaone.modeling_exaone import (
+        ExaoneBlock as OldExaoneBlock,
+        ExaoneForCausalLM as OldExaoneForCausalLM,
+    )
+except:
+    OldExaoneBlock = None
+    OldExaoneForCausalLM = None
 from awq.modules.fused.norm import FasterTransformerRMSNorm