CodyBontecou commited on
Commit
dd0166e
·
1 Parent(s): 034457d
Files changed (1) hide show
  1. handler.py +3 -3
handler.py CHANGED
@@ -3,7 +3,7 @@ from typing import Any, Dict
3
 
4
  class EndpointHandler:
5
  def __init__(self, model_dir: str, **kwargs: Any) -> None:
6
- self.model = AutoModel.from_pretrained(
7
  model_dir,
8
  torch_dtype=torch.bfloat16,
9
  low_cpu_mem_usage=True,
@@ -12,12 +12,12 @@ class EndpointHandler:
12
  device_map=split_model(),
13
  ).eval()
14
 
15
- self.tokenizer = AutoTokenizer.from_pretrained(
16
  model_dir, trust_remote_code=True, use_fast=False
17
  )
18
 
19
  def __call__(self, data: Dict[str, Any]) -> Any:
20
- logger.info(f"Received incoming request with {data=}")
21
 
22
 
23
  if __name__ == "__main__":
 
3
 
4
  class EndpointHandler:
5
  def __init__(self, model_dir: str, **kwargs: Any) -> None:
6
+ self.model = AutoModel.from_pretrained(
7
  model_dir,
8
  torch_dtype=torch.bfloat16,
9
  low_cpu_mem_usage=True,
 
12
  device_map=split_model(),
13
  ).eval()
14
 
15
+ self.tokenizer = AutoTokenizer.from_pretrained(
16
  model_dir, trust_remote_code=True, use_fast=False
17
  )
18
 
19
  def __call__(self, data: Dict[str, Any]) -> Any:
20
+ logger.info(f"Received incoming request with {data=}")
21
 
22
 
23
  if __name__ == "__main__":