from_pretrained: revision pinning
Browse files- inference.py +5 -2
inference.py
CHANGED
|
@@ -72,12 +72,15 @@ class FusionEmbedder:
|
|
| 72 |
|
| 73 |
# ------------------------------------------------------------------ loading
|
| 74 |
@classmethod
|
| 75 |
-
def from_pretrained(cls, repo_or_path: str, device: str = "cuda",
|
|
|
|
|
|
|
|
|
|
| 76 |
if os.path.exists(repo_or_path):
|
| 77 |
path = repo_or_path
|
| 78 |
else:
|
| 79 |
from huggingface_hub import hf_hub_download
|
| 80 |
-
path = hf_hub_download(repo_or_path, CKPT_FILE)
|
| 81 |
return cls(path, device=device, **kw)
|
| 82 |
|
| 83 |
# ------------------------------------------------------------------ helpers
|
|
|
|
| 72 |
|
| 73 |
# ------------------------------------------------------------------ loading
|
| 74 |
@classmethod
|
| 75 |
+
def from_pretrained(cls, repo_or_path: str, device: str = "cuda",
|
| 76 |
+
revision: Optional[str] = None, **kw) -> "FusionEmbedder":
|
| 77 |
+
"""Load from a local checkpoint path or an HF repo. ``revision`` pins a repo
|
| 78 |
+
tag/commit (e.g. ``"v0.1-preview"``, ``"v0.2-preview"``); default is latest."""
|
| 79 |
if os.path.exists(repo_or_path):
|
| 80 |
path = repo_or_path
|
| 81 |
else:
|
| 82 |
from huggingface_hub import hf_hub_download
|
| 83 |
+
path = hf_hub_download(repo_or_path, CKPT_FILE, revision=revision)
|
| 84 |
return cls(path, device=device, **kw)
|
| 85 |
|
| 86 |
# ------------------------------------------------------------------ helpers
|