chore: made optims

Signed-off-by: Suvaditya Mukherjee <[email protected]>

app.py

@@ -10,7 +10,7 @@ import pymupdf
 import gradio as gr
 from qdrant_client import QdrantClient
 from utils import download_pdf_from_gdrive, merge_strings_with_prefix
-from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, QuantoConfig
 
 def rag_query(query: str):
     """
@@ -181,11 +181,14 @@ if __name__ == "__main__":
         model_name,
         torch_dtype=torch.bfloat16,
         device_map="auto",
-        quantization_config = BitsAndBytesConfig(
-            load_in_4bit=True,
-            # bnb_4bit_compute_dtype=torch.float16,
-            # bnb_4bit_quant_type="nf4"
+        quantization_config=QuantoConfig(
+            weights="int8"
         )
+        # quantization_config = BitsAndBytesConfig(
+        #     load_in_4bit=True,
+        #     # bnb_4bit_compute_dtype=torch.float16,
+        #     # bnb_4bit_quant_type="nf4"
+        # )
     )
 
     tokenizer = AutoTokenizer.from_pretrained(model_name)

requirements.txt

@@ -8,7 +8,8 @@ torch
 torchvision
 torchaudio
 accelerate
-bitsandbytes==0.44.1
+# bitsandbytes==0.44.0
 optimum
 wandb
-psutil
+psutil
+optimum-quanto