srishtirai commited on
Commit
df181e8
Β·
verified Β·
1 Parent(s): 55215ba

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +11 -1
app.py CHANGED
@@ -6,6 +6,15 @@ import os
6
  from transformers import AutoTokenizer, AutoModelForCausalLM
7
  from peft import PeftModel, PeftConfig
8
 
 
 
 
 
 
 
 
 
 
9
  # βœ… Ensure offload directory exists
10
  os.makedirs("offload", exist_ok=True)
11
 
@@ -26,7 +35,8 @@ def load_model(model_path):
26
  base_model_name,
27
  torch_dtype=torch.float16, # Use FP16 to reduce memory usage
28
  device_map="auto", # Automatically distribute across CPU/GPU
29
- offload_folder="offload" # βœ… Prevents memory crashes
 
30
  )
31
 
32
  # βœ… Load LoRA adapter with `is_trainable=False`
 
6
  from transformers import AutoTokenizer, AutoModelForCausalLM
7
  from peft import PeftModel, PeftConfig
8
 
9
+ # βœ… Retrieve Hugging Face token from environment variable (Secret)
10
+ HF_TOKEN = os.getenv("HF_TOKEN") # βœ… Now using stored secret, not hardcoded!
11
+
12
+ if HF_TOKEN is None:
13
+ raise ValueError("❌ ERROR: Hugging Face token is missing! Please set HF_TOKEN in Hugging Face Secrets.")
14
+
15
+ # βœ… Authenticate with Hugging Face
16
+ login(token=HF_TOKEN)
17
+
18
  # βœ… Ensure offload directory exists
19
  os.makedirs("offload", exist_ok=True)
20
 
 
35
  base_model_name,
36
  torch_dtype=torch.float16, # Use FP16 to reduce memory usage
37
  device_map="auto", # Automatically distribute across CPU/GPU
38
+ offload_folder="offload", # βœ… Prevents memory crashes
39
+ use_auth_token=HF_TOKEN # βœ… Authenticate model loading
40
  )
41
 
42
  # βœ… Load LoRA adapter with `is_trainable=False`