poiccard commited on
Commit
0cffbf7
1 Parent(s): 624d54c

Update handler.py

Browse files
Files changed (1) hide show
  1. handler.py +1 -2
handler.py CHANGED
@@ -3,7 +3,6 @@
3
  from typing import Dict, List, Any
4
  import torch
5
  from transformers import AutoTokenizer, AutoModelForCausalLM
6
- import json
7
 
8
  class EndpointHandler():
9
  def __init__(self, path=""):
@@ -36,7 +35,7 @@ class EndpointHandler():
36
  response = self.tokenizer.batch_decode(
37
  generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=True
38
  )[0]
39
- final_response = response.split("### Response: [|AI|]").json()
40
 
41
 
42
  return [{"generated_text": final_response[-1]}]
 
3
  from typing import Dict, List, Any
4
  import torch
5
  from transformers import AutoTokenizer, AutoModelForCausalLM
 
6
 
7
  class EndpointHandler():
8
  def __init__(self, path=""):
 
35
  response = self.tokenizer.batch_decode(
36
  generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=True
37
  )[0]
38
+ final_response = response.split("### Response: [|AI|]")
39
 
40
 
41
  return [{"generated_text": final_response[-1]}]