abtonmoy commited on
Commit
a815bcd
·
verified ·
1 Parent(s): 37f999f

from_pretrained: revision pinning

Browse files
Files changed (1) hide show
  1. 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", **kw) -> "FusionEmbedder":
 
 
 
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