Simonlob commited on
Commit
f97fc67
·
verified ·
1 Parent(s): a164d52

make hf-token as attribute of model class

Browse files
Files changed (1) hide show
  1. util.py +3 -2
util.py CHANGED
@@ -208,6 +208,7 @@ class KaniModel:
208
  def __init__(self, config, player: NemoAudioPlayer, token: str) -> None:
209
  self.conf = config
210
  self.player = player
 
211
  self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
212
 
213
  print(f"Loading model: {self.conf.model_name}")
@@ -218,13 +219,13 @@ class KaniModel:
218
  self.conf.model_name,
219
  dtype=torch.bfloat16,
220
  device_map=self.conf.device_map,
221
- token=token,
222
  trust_remote_code=True # May be needed for some models
223
  )
224
 
225
  self.tokenizer = AutoTokenizer.from_pretrained(
226
  self.conf.model_name,
227
- token=token,
228
  trust_remote_code=True
229
  )
230
 
 
208
  def __init__(self, config, player: NemoAudioPlayer, token: str) -> None:
209
  self.conf = config
210
  self.player = player
211
+ self.hf_token = token
212
  self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
213
 
214
  print(f"Loading model: {self.conf.model_name}")
 
219
  self.conf.model_name,
220
  dtype=torch.bfloat16,
221
  device_map=self.conf.device_map,
222
+ token=self.hf_token,
223
  trust_remote_code=True # May be needed for some models
224
  )
225
 
226
  self.tokenizer = AutoTokenizer.from_pretrained(
227
  self.conf.model_name,
228
+ token=self.hf_token,
229
  trust_remote_code=True
230
  )
231