Upload handler.py
Browse files- handler.py +2 -2
handler.py
CHANGED
@@ -28,7 +28,7 @@ class EndpointHandler():
|
|
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,
|
@@ -58,7 +58,7 @@ class EndpointHandler():
|
|
58 |
return self.tokenizer(text + self.tokenizer.eos_token, max_length=self.max_length, truncation=True, return_tensors='pt').to(self.device)
|
59 |
|
60 |
|
61 |
-
def extract_attn_vec(self):
|
62 |
return self.model._modules['layers'][-1].self_attn.attn_vec
|
63 |
|
64 |
|
|
|
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,
|
|
|
58 |
return self.tokenizer(text + self.tokenizer.eos_token, max_length=self.max_length, truncation=True, return_tensors='pt').to(self.device)
|
59 |
|
60 |
|
61 |
+
def extract_attn_vec(self, model):
|
62 |
return self.model._modules['layers'][-1].self_attn.attn_vec
|
63 |
|
64 |
|