Fix fused models for tf >= 4.39 (#418)

casper-hansen · Apr 6, 2024 · 5d7b050 · 5d7b050
1 parent 0fa9a2c
commit 5d7b050
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 1 deletion.
diff --git a/awq/modules/fused/model.py b/awq/modules/fused/model.py
@@ -83,6 +83,14 @@ def __init__(self, vocab_size, blocks, embedding, norm):
         self.blocks: List[LlamaLikeBlock] = nn.ModuleList(blocks)
         self.norm = norm
         self.last_forward_num_tokens = 0
+
+    @property
+    def embed_tokens(self):
+        return self.embedding
+
+    @property
+    def layers(self):
+        return self.blocks
 
     @torch.inference_mode()
     def forward(

diff --git a/setup.py b/setup.py
@@ -86,7 +86,7 @@ def get_kernels_whl_url(
 
 requirements = [
     "torch>=2.0.1",
-    "transformers>=4.35.0,<=4.38.2",
+    "transformers>=4.35.0",
     "tokenizers>=0.12.1",
     "typing_extensions>=4.8.0",
     "accelerate",