Upload handler.py
Browse files- handler.py +2 -2
handler.py
CHANGED
@@ -28,11 +28,11 @@ class EndpointHandler():
|
|
28 |
self.tokenizer.pad_token = '[PAD]'
|
29 |
self.tokenizer.padding_side = 'left'
|
30 |
|
31 |
-
|
32 |
|
33 |
self.model = AutoModel.from_pretrained(
|
34 |
model_dir,
|
35 |
-
|
36 |
device_map="auto",
|
37 |
trust_remote_code=True,
|
38 |
attn_implementation="eager",
|
|
|
28 |
self.tokenizer.pad_token = '[PAD]'
|
29 |
self.tokenizer.padding_side = 'left'
|
30 |
|
31 |
+
bnb_config = BitsAndBytesConfig(load_in_8bit=True, bnb_8bit_compute_dtype=torch.float16)
|
32 |
|
33 |
self.model = AutoModel.from_pretrained(
|
34 |
model_dir,
|
35 |
+
quantization_config=bnb_config,
|
36 |
device_map="auto",
|
37 |
trust_remote_code=True,
|
38 |
attn_implementation="eager",
|