alvarobartt HF staff commited on
Commit
64f317d
1 Parent(s): 06eb103

Update handler.py

Browse files
Files changed (1) hide show
  1. handler.py +2 -2
handler.py CHANGED
@@ -77,7 +77,7 @@ class EndpointHandler:
77
  )
78
 
79
  generation_config = deepcopy(self.generation_config)
80
- generation_config.update(**input.get("generation_parameters", {}))
81
 
82
  messages = [
83
  {
@@ -100,6 +100,6 @@ class EndpointHandler:
100
  generated_ids,
101
  skip_special_tokens=True,
102
  )
103
- predictions.append(generated_texts[0])
104
 
105
  return {"predictions": predictions}
 
77
  )
78
 
79
  generation_config = deepcopy(self.generation_config)
80
+ generation_config.update(**input.get("generation_parameters", {"max_new_tokens": 128}))
81
 
82
  messages = [
83
  {
 
100
  generated_ids,
101
  skip_special_tokens=True,
102
  )
103
+ predictions.append(generated_texts[0][len(prompt):])
104
 
105
  return {"predictions": predictions}