WinterGYC commited on
Commit
25f60a6
1 Parent(s): 9128041

Update handler.py

Browse files
Files changed (1) hide show
  1. handler.py +1 -1
handler.py CHANGED
@@ -12,7 +12,7 @@ class EndpointHandler:
12
  self.tokenizer = AutoTokenizer.from_pretrained(path, trust_remote_code=True)
13
  self.model = AutoModelForCausalLM.from_pretrained(path, device_map="auto", torch_dtype=dtype, trust_remote_code=True)
14
  # create inference pipeline
15
- self.pipeline = pipeline("text-generation", model=model, tokenizer=tokenizer)
16
 
17
  def __call__(self, data: Any) -> List[List[Dict[str, float]]]:
18
  inputs = data.pop("inputs", data)
 
12
  self.tokenizer = AutoTokenizer.from_pretrained(path, trust_remote_code=True)
13
  self.model = AutoModelForCausalLM.from_pretrained(path, device_map="auto", torch_dtype=dtype, trust_remote_code=True)
14
  # create inference pipeline
15
+ self.pipeline = pipeline("text-generation", model=self.model, tokenizer=self.tokenizer)
16
 
17
  def __call__(self, data: Any) -> List[List[Dict[str, float]]]:
18
  inputs = data.pop("inputs", data)