MrD05 commited on
Commit
6815019
·
1 Parent(s): 685e40e

Update handler.py

Browse files
Files changed (1) hide show
  1. handler.py +2 -6
handler.py CHANGED
@@ -1,4 +1,4 @@
1
- from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
2
  from transformers_stream_generator import init_stream_support
3
  import re
4
  init_stream_support()
@@ -22,16 +22,12 @@ Alice Gate: *Alice strides into the room with a smile, her eyes lighting up when
22
 
23
  class EndpointHandler():
24
 
25
- def __init__(self, path = "."):
26
- path = "."
27
- # quantization_config = BitsAndBytesConfig(llm_int8_enable_fp32_cpu_offload = True)
28
  self.tokenizer = AutoTokenizer.from_pretrained(path)
29
  self.model = AutoModelForCausalLM.from_pretrained(
30
  path,
31
  device_map = "auto",
32
  load_in_8bit = True,
33
- torch_dtype = "auto",
34
- low_cpu_mem_usage = True
35
  )
36
 
37
  def __call__(self, data):
 
1
+ from transformers import AutoTokenizer, AutoModelForCausalLM
2
  from transformers_stream_generator import init_stream_support
3
  import re
4
  init_stream_support()
 
22
 
23
  class EndpointHandler():
24
 
25
+ def __init__(self, path = ""):
 
 
26
  self.tokenizer = AutoTokenizer.from_pretrained(path)
27
  self.model = AutoModelForCausalLM.from_pretrained(
28
  path,
29
  device_map = "auto",
30
  load_in_8bit = True,
 
 
31
  )
32
 
33
  def __call__(self, data):