ksee commited on
Commit
57d7832
·
1 Parent(s): 57a0d4c

Update handler.py

Browse files
Files changed (1) hide show
  1. handler.py +1 -2
handler.py CHANGED
@@ -8,9 +8,8 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
8
  class EndpointHandler():
9
  def __init__(self, path=""):
10
  config = PeftConfig.from_pretrained(path)
11
- model = AutoModelForCausalLM.from_pretrained(config.base_model_name_or_path, load_in_8bit=True, device_map='auto')
12
  self.tokenizer = AutoTokenizer.from_pretrained(config.base_model_name_or_path)
13
- # Load the Lora model
14
  self.model = PeftModel.from_pretrained(model, path)
15
 
16
  def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
 
8
  class EndpointHandler():
9
  def __init__(self, path=""):
10
  config = PeftConfig.from_pretrained(path)
11
+ model = AutoModelForCausalLM.from_pretrained(config.base_model_name_or_path, device_map="auto")#, load_in_8bit=True, device_map='auto')
12
  self.tokenizer = AutoTokenizer.from_pretrained(config.base_model_name_or_path)
 
13
  self.model = PeftModel.from_pretrained(model, path)
14
 
15
  def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]: