Spaces:
Sleeping
Sleeping
Update initializer.py to use explicit Hub filenames
Browse files- models/initializer.py +5 -5
models/initializer.py
CHANGED
|
@@ -29,14 +29,14 @@ def download_llm() -> str:
|
|
| 29 |
"""
|
| 30 |
local_model_path = hf_hub_download(
|
| 31 |
repo_id=config.HF_MODEL_HUB,
|
| 32 |
-
filename=
|
| 33 |
token=config.HF_TOKEN,
|
| 34 |
local_dir=config.LLM_LOCAL_DIR
|
| 35 |
)
|
| 36 |
|
| 37 |
hf_hub_download(
|
| 38 |
repo_id=config.HF_MODEL_HUB,
|
| 39 |
-
filename=
|
| 40 |
token=config.HF_TOKEN,
|
| 41 |
local_dir=config.CONFIG_LOCAL_DIR
|
| 42 |
)
|
|
@@ -49,9 +49,9 @@ def download_reranker() -> str:
|
|
| 49 |
"""
|
| 50 |
return hf_hub_download(
|
| 51 |
repo_id=config.HF_MODEL_HUB,
|
| 52 |
-
filename=
|
| 53 |
token=config.HF_TOKEN,
|
| 54 |
-
local_dir=
|
| 55 |
)
|
| 56 |
|
| 57 |
def load_llm(local_model_path: str) -> CustomModel:
|
|
@@ -68,7 +68,7 @@ def load_llm(local_model_path: str) -> CustomModel:
|
|
| 68 |
)
|
| 69 |
except Exception:
|
| 70 |
# Fallback: manual load if file naming is non-standard (e.g., .pth)
|
| 71 |
-
_config = AutoConfig.from_pretrained(config.
|
| 72 |
model = CustomModel(_config)
|
| 73 |
state_dict = torch.load(local_model_path, map_location="cpu")
|
| 74 |
model.load_state_dict(state_dict)
|
|
|
|
| 29 |
"""
|
| 30 |
local_model_path = hf_hub_download(
|
| 31 |
repo_id=config.HF_MODEL_HUB,
|
| 32 |
+
filename=config.HF_LLM_FILENAME,
|
| 33 |
token=config.HF_TOKEN,
|
| 34 |
local_dir=config.LLM_LOCAL_DIR
|
| 35 |
)
|
| 36 |
|
| 37 |
hf_hub_download(
|
| 38 |
repo_id=config.HF_MODEL_HUB,
|
| 39 |
+
filename=config.HF_CONFIG_FILENAME,
|
| 40 |
token=config.HF_TOKEN,
|
| 41 |
local_dir=config.CONFIG_LOCAL_DIR
|
| 42 |
)
|
|
|
|
| 49 |
"""
|
| 50 |
return hf_hub_download(
|
| 51 |
repo_id=config.HF_MODEL_HUB,
|
| 52 |
+
filename=config.HF_RERANKER_FILENAME,
|
| 53 |
token=config.HF_TOKEN,
|
| 54 |
+
local_dir=config.RERANKER_LOCAL_DIR
|
| 55 |
)
|
| 56 |
|
| 57 |
def load_llm(local_model_path: str) -> CustomModel:
|
|
|
|
| 68 |
)
|
| 69 |
except Exception:
|
| 70 |
# Fallback: manual load if file naming is non-standard (e.g., .pth)
|
| 71 |
+
_config = AutoConfig.from_pretrained(config.CONFIG_LOCAL_DIR)
|
| 72 |
model = CustomModel(_config)
|
| 73 |
state_dict = torch.load(local_model_path, map_location="cpu")
|
| 74 |
model.load_state_dict(state_dict)
|