Deepakvictor commited on
Commit
f8269e3
1 Parent(s): d3f59c0

Update handler.py

Browse files
Files changed (1) hide show
  1. handler.py +2 -1
handler.py CHANGED
@@ -9,7 +9,8 @@ class EndpointHandler():
9
  self.model = AutoModelForSeq2SeqLM.from_pretrained(path)
10
 
11
  def __call__(self, data: str) -> str:
12
- inp = self.tokenizer(data, return_tensors="pt")
 
13
  with torch.inference_mode():
14
  out= self.model.generate(**inp)
15
  final_output = self.tokenizer.batch_decode(out,skip_special_tokens=True)
 
9
  self.model = AutoModelForSeq2SeqLM.from_pretrained(path)
10
 
11
  def __call__(self, data: str) -> str:
12
+ inputs = data.pop("inputs", data)
13
+ inp = self.tokenizer(inputs, return_tensors="pt")
14
  with torch.inference_mode():
15
  out= self.model.generate(**inp)
16
  final_output = self.tokenizer.batch_decode(out,skip_special_tokens=True)