add access_token to model loader
Browse files
app.py
CHANGED
|
@@ -12,6 +12,8 @@ from transformers import AutoTokenizer
|
|
| 12 |
from ddgs import DDGS
|
| 13 |
import spaces # Import spaces early to enable ZeroGPU support
|
| 14 |
|
|
|
|
|
|
|
| 15 |
# Optional: Disable GPU visibility if you wish to force CPU usage
|
| 16 |
# os.environ["CUDA_VISIBLE_DEVICES"] = ""
|
| 17 |
|
|
@@ -96,7 +98,8 @@ def load_pipeline(model_name):
|
|
| 96 |
if model_name in PIPELINES:
|
| 97 |
return PIPELINES[model_name]
|
| 98 |
repo = MODELS[model_name]["repo_id"]
|
| 99 |
-
tokenizer = AutoTokenizer.from_pretrained(repo
|
|
|
|
| 100 |
for dtype in (torch.bfloat16, torch.float16, torch.float32):
|
| 101 |
try:
|
| 102 |
pipe = pipeline(
|
|
@@ -105,8 +108,8 @@ def load_pipeline(model_name):
|
|
| 105 |
tokenizer=tokenizer,
|
| 106 |
trust_remote_code=True,
|
| 107 |
torch_dtype=dtype,
|
| 108 |
-
device_map="auto"
|
| 109 |
-
|
| 110 |
PIPELINES[model_name] = pipe
|
| 111 |
return pipe
|
| 112 |
except Exception:
|
|
|
|
| 12 |
from ddgs import DDGS
|
| 13 |
import spaces # Import spaces early to enable ZeroGPU support
|
| 14 |
|
| 15 |
+
access_token=os.environ['HF_TOKEN']
|
| 16 |
+
|
| 17 |
# Optional: Disable GPU visibility if you wish to force CPU usage
|
| 18 |
# os.environ["CUDA_VISIBLE_DEVICES"] = ""
|
| 19 |
|
|
|
|
| 98 |
if model_name in PIPELINES:
|
| 99 |
return PIPELINES[model_name]
|
| 100 |
repo = MODELS[model_name]["repo_id"]
|
| 101 |
+
tokenizer = AutoTokenizer.from_pretrained(repo,
|
| 102 |
+
token=access_token)
|
| 103 |
for dtype in (torch.bfloat16, torch.float16, torch.float32):
|
| 104 |
try:
|
| 105 |
pipe = pipeline(
|
|
|
|
| 108 |
tokenizer=tokenizer,
|
| 109 |
trust_remote_code=True,
|
| 110 |
torch_dtype=dtype,
|
| 111 |
+
device_map="auto",
|
| 112 |
+
token=access_token)
|
| 113 |
PIPELINES[model_name] = pipe
|
| 114 |
return pipe
|
| 115 |
except Exception:
|