jatingocodeo commited on
Commit
66d83b1
·
verified ·
1 Parent(s): 2862eb9

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +11 -1
app.py CHANGED
@@ -205,10 +205,20 @@ def load_model():
205
  tokenizer.add_special_tokens(special_tokens)
206
 
207
  print("Loading model from Hugging Face Hub...")
 
 
 
 
 
 
 
 
208
  model = AutoModelForCausalLM.from_pretrained(
209
  model_id,
 
210
  torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
211
- trust_remote_code=True
 
212
  )
213
 
214
  # Move model to device manually
 
205
  tokenizer.add_special_tokens(special_tokens)
206
 
207
  print("Loading model from Hugging Face Hub...")
208
+ # Create config first
209
+ config = SmolLM2Config(
210
+ pad_token_id=tokenizer.pad_token_id,
211
+ bos_token_id=tokenizer.bos_token_id,
212
+ eos_token_id=tokenizer.eos_token_id
213
+ )
214
+
215
+ # Load model from Hub
216
  model = AutoModelForCausalLM.from_pretrained(
217
  model_id,
218
+ config=config,
219
  torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
220
+ trust_remote_code=True,
221
+ low_cpu_mem_usage=True
222
  )
223
 
224
  # Move model to device manually