Schmadge commited on
Commit
cd95ea8
·
1 Parent(s): 9e81bca

Update handler.py

Browse files
Files changed (1) hide show
  1. handler.py +4 -10
handler.py CHANGED
@@ -4,21 +4,15 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
4
  from typing import Any, Dict
5
 
6
  class EndpointHandler:
7
- def __init__(
8
- self,
9
- path='',
10
- torch_dtype=torch.bfloat16,
11
- trust_remote_code=True,
12
- ) -> None:
13
  self.model = AutoModelForCausalLM.from_pretrained(
14
  path,
15
  torch_dtype=torch_dtype,
16
  trust_remote_code=trust_remote_code
17
  )
18
- tokenizer = AutoTokenizer.from_pretrained(
19
- path,
20
- trust_remote_code=trust_remote_code
21
- )
22
  if tokenizer.pad_token_id is None:
23
  warnings.warn(
24
  "pad_token_id is not set for the tokenizer. Using eos_token_id as pad_token_id."
 
4
  from typing import Any, Dict
5
 
6
  class EndpointHandler:
7
+ def __init__(self, path='', torch_dtype=torch.bfloat16, trust_remote_code=True):
8
+
 
 
 
 
9
  self.model = AutoModelForCausalLM.from_pretrained(
10
  path,
11
  torch_dtype=torch_dtype,
12
  trust_remote_code=trust_remote_code
13
  )
14
+ tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-neox-20b")
15
+
 
 
16
  if tokenizer.pad_token_id is None:
17
  warnings.warn(
18
  "pad_token_id is not set for the tokenizer. Using eos_token_id as pad_token_id."