adamluc commited on
Commit
8846580
·
1 Parent(s): 6948a7c

Upload handler.py

Browse files
Files changed (1) hide show
  1. handler.py +1 -1
handler.py CHANGED
@@ -8,7 +8,7 @@ class EndpointHandler:
8
  tokenizer = AutoTokenizer.from_pretrained(path)
9
  model = AutoModelForCausalLM.from_pretrained(path, low_cpu_mem_usage=True)
10
  # create inference pipeline
11
- self.pipeline = pipeline("text-generation", model=model, tokenizer=tokenizer, device=device)
12
 
13
  def __call__(self, data: Any) -> List[List[Dict[str, float]]]:
14
  inputs = data.pop("inputs", data)
 
8
  tokenizer = AutoTokenizer.from_pretrained(path)
9
  model = AutoModelForCausalLM.from_pretrained(path, low_cpu_mem_usage=True)
10
  # create inference pipeline
11
+ self.pipeline = pipeline("text-generation", model=model, tokenizer=tokenizer)
12
 
13
  def __call__(self, data: Any) -> List[List[Dict[str, float]]]:
14
  inputs = data.pop("inputs", data)