Kowsher commited on
Commit
cd690fb
·
1 Parent(s): 231705d

Update handler.py

Browse files
Files changed (1) hide show
  1. handler.py +5 -3
handler.py CHANGED
@@ -1,16 +1,18 @@
1
  import torch
2
 
3
  from typing import Any, Dict
4
- from transformers import AutoModelForCausalLM, AutoTokenizer
5
  from transformers.models.auto import modeling_auto
6
 
7
 
8
  class EndpointHandler:
9
  def __init__(self, path=""):
 
 
10
  # load model and tokenizer from path
11
- self.tokenizer = AutoTokenizer.from_pretrained(path)
12
  self.model = AutoModelForCausalLM.from_pretrained(
13
- path, device_map="auto", torch_dtype=torch.float16, trust_remote_code=True
14
  )
15
  self.device = "cuda" if torch.cuda.is_available() else "cpu"
16
 
 
1
  import torch
2
 
3
  from typing import Any, Dict
4
+ from transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig
5
  from transformers.models.auto import modeling_auto
6
 
7
 
8
  class EndpointHandler:
9
  def __init__(self, path=""):
10
+ print('starting machine')
11
+ config = AutoConfig.from_pretrained("Kowsher/Egol_model", trust_remote_code=True)
12
  # load model and tokenizer from path
13
+ self.tokenizer = AutoTokenizer.from_pretrained(path, trust_remote_code=True)
14
  self.model = AutoModelForCausalLM.from_pretrained(
15
+ path, device_map="auto", torch_dtype=torch.float16, config = config, trust_remote_code=True
16
  )
17
  self.device = "cuda" if torch.cuda.is_available() else "cpu"
18