quantization trial

handler.py

@@ -1,6 +1,8 @@
 import logging
 from typing import Dict, List, Any
-from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline, StoppingCriteria, StoppingCriteriaList
+import torch
+
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline, StoppingCriteria, StoppingCriteriaList, BitsAndBytesConfig
 
 # Configure logging
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
@@ -10,7 +12,13 @@ class EndpointHandler():
         logging.info("Initializing EndpointHandler with model path: %s", path)
         tokenizer = AutoTokenizer.from_pretrained(path)
         tokenizer.pad_token = tokenizer.eos_token
-        self.model = AutoModelForCausalLM.from_pretrained(path)
+        bnb_config = BitsAndBytesConfig(
+            load_in_4bit=True,
+            bnb_4bit_use_double_quant=True,
+            bnb_4bit_quant_type="nf4",
+            bnb_4bit_compute_dtype=torch.bfloat16
+        )
+        self.model = AutoModelForCausalLM.from_pretrained(path, quantization_config=bnb_config)
         self.tokenizer = tokenizer
         self.stopping_criteria = StoppingCriteriaList([StopAtPeriodCriteria(tokenizer)])