krilecy
/

e5-mistral-7b-instruct

Feature Extraction

Inference Endpoints

text-generation-inference

text-embeddings-inference

Model card Files Files and versions Community

krilecy commited on Mar 28

Commit

4f3226b

•

1 Parent(s): 93bd2da

Upload handler.py

Files changed (1) hide show

handler.py +2 -2

handler.py CHANGED Viewed

@@ -28,11 +28,11 @@ class EndpointHandler():
         self.tokenizer.pad_token = '[PAD]'
         self.tokenizer.padding_side = 'left'
-        #bnb_config = BitsAndBytesConfig(load_in_8bit=True, bnb_8bit_compute_dtype=torch.float16)
         self.model = AutoModel.from_pretrained(
             model_dir,
-            #quantization_config=bnb_config,
             device_map="auto",
             trust_remote_code=True,
             attn_implementation="eager",

         self.tokenizer.pad_token = '[PAD]'
         self.tokenizer.padding_side = 'left'
+        bnb_config = BitsAndBytesConfig(load_in_8bit=True, bnb_8bit_compute_dtype=torch.float16)
         self.model = AutoModel.from_pretrained(
             model_dir,
+            quantization_config=bnb_config,
             device_map="auto",
             trust_remote_code=True,
             attn_implementation="eager",