awq fix

qa_engine/qa_engine.py

@@ -15,6 +15,7 @@ from langchain.llms.base import LLM
 from langchain.embeddings import HuggingFaceEmbeddings, HuggingFaceHubEmbeddings, HuggingFaceInstructEmbeddings
 from langchain.vectorstores import FAISS
 from sentence_transformers import CrossEncoder
+from awq import AutoAWQForCausalLM
 
 from qa_engine import logger
 from qa_engine.response import Response
@@ -63,14 +64,18 @@ class TransformersPipelineModel(LLM):
         self.model_id = model_id
 
         tokenizer = AutoTokenizer.from_pretrained(model_id)
-        model = AutoModelForCausalLM.from_pretrained(
-            model_id,
-            torch_dtype=torch.bfloat16,
-            trust_remote_code=True,
-            load_in_8bit=False,
-            device_map='auto',
-            resume_download=True,
-        )
+        if "AWQ" in model_id:
+            model = AutoAWQForCausalLM.from_quantized(model_id, fuse_layers=True,
+                                          trust_remote_code=False, safetensors=True)
+        else:
+            model = AutoModelForCausalLM.from_pretrained(
+                model_id,
+                torch_dtype=torch.bfloat16,
+                trust_remote_code=True,
+                load_in_8bit=False,
+                device_map='auto',
+                resume_download=True,
+            )
         self.pipeline = transformers.pipeline(
             'text-generation',
             model=model,