Adds support for AWQ to use quantization_layer_structure hooks (keras-team#2511)

Author: JyotinderSingh

keras_hub/src/models/causal_lm.py

@@ -431,7 +431,7 @@ def _post_quantize(self, mode, **kwargs):
         self.generate_function = None
 
     def get_quantization_layer_structure(self, mode):
-        if mode != "gptq":
+        if mode not in ["gptq", "awq"]:
             return None
 
         backbone = self.backbone

keras_hub/src/models/gemma/gemma_causal_lm.py

@@ -433,7 +433,7 @@ def default_layer_intercept_fn(x, unused_i):
         return per_token_loss
 
     def get_quantization_layer_structure(self, mode):
-        if mode != "gptq":
+        if mode not in ["gptq", "awq"]:
             return None
 
         # Wrap embedding + scaling

keras_hub/src/models/gpt2/gpt2_causal_lm.py

@@ -422,7 +422,7 @@ def default_layer_intercept_fn(x, unused_i):
         return per_token_loss
 
     def get_quantization_layer_structure(self, mode):
-        if mode != "gptq":
+        if mode not in ["gptq", "awq"]:
             return None
 
         backbone = self.backbone

keras_hub/src/models/masked_lm.py

@@ -86,7 +86,7 @@ def compile(
         )
 
     def get_quantization_layer_structure(self, mode):
-        if mode != "gptq":
+        if mode not in ["gptq", "awq"]:
             return None
 
         backbone = self.backbone