changes

handler.py

@@ -1,14 +1,21 @@
+import torch
 from typing import Dict, List, Any
-from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline, BitsAndBytesConfig
 
 
 class EndpointHandler:
     def __init__(self, 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
+        )
         # load the model
         tokenizer = AutoTokenizer.from_pretrained(path)
-        model = AutoModelForCausalLM.from_pretrained(path, device_map="auto")
+        model = AutoModelForCausalLM.from_pretrained(path, device_map="auto", quantization_config=bnb_config)
         # create inference pipeline
-        self.pipeline = pipeline("text-generation", model=model, tokenizer=tokenizer)
+        self.pipeline = pipeline("text-generation", model=model, tokenizer=tokenizer, torch_dtype=torch.bfloat16)
 
     def __call__(self, data: Any):
         inputs = data.pop("inputs", data)