Upload handler.py

handler.py

@@ -28,7 +28,7 @@ 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)
+        #bnb_config = BitsAndBytesConfig(load_in_8bit=True, bnb_8bit_compute_dtype=torch.float16)
 
         self.model = AutoModel.from_pretrained(
             model_dir,
@@ -58,7 +58,7 @@ class EndpointHandler():
         return self.tokenizer(text + self.tokenizer.eos_token, max_length=self.max_length, truncation=True, return_tensors='pt').to(self.device)
 
 
-    def extract_attn_vec(self):
+    def extract_attn_vec(self, model):
         return self.model._modules['layers'][-1].self_attn.attn_vec