clean up

app_utils/model_utils.py

@@ -1,31 +1,29 @@
-import sys
-# sys.path.insert(0, "/data/yzhouc01/MassSpecGym")
-# sys.path.insert(0, "/data/yzhouc01/FILIP-MS")
+import os
+from pathlib import Path
 
+import yaml
 from rdkit import RDLogger
-from flare.utils.data import get_spec_featurizer, get_mol_featurizer, get_ms_dataset
-from flare.utils.models import get_model
 
-import yaml
+from flare.definitions import REPO_DIR
+from flare.utils.config import resolve_repo_paths
+from flare.utils.data import get_mol_featurizer, get_spec_featurizer
+from flare.utils.models import get_model
 
-# Suppress RDKit warnings and errors
 lg = RDLogger.logger()
 lg.setLevel(RDLogger.CRITICAL)
 
-# Load model and data
 
 def load_model_components():
-    param_pth = 'hparams.yaml'
-    with open(param_pth) as f:
+    param_pth = Path(os.environ.get("FLARE_PARAMS", REPO_DIR / "params.yaml")).expanduser()
+    with open(param_pth, encoding="utf-8") as f:
         params = yaml.load(f, Loader=yaml.FullLoader)
+    resolve_repo_paths(params)
 
-    spec_featurizer = get_spec_featurizer(params['spectra_view'], params)
-    mol_featurizer = get_mol_featurizer(params['molecule_view'], params)
-
-    # load model
+    spec_featurizer = get_spec_featurizer(params["spectra_view"], params)
+    mol_featurizer = get_mol_featurizer(params["molecule_view"], params)
 
-    checkpoint_pth = "pretrained_models/flare.ckpt"
-    params['checkpoint_pth'] = checkpoint_pth
-    model = get_model(params['model'], params)
+    ckpt = os.environ.get("FLARE_CHECKPOINT", str(REPO_DIR / "pretrained_models" / "flare.ckpt"))
+    params["checkpoint_pth"] = ckpt
+    model = get_model(params["model"], params)
 
-    return spec_featurizer, mol_featurizer, model
+    return spec_featurizer, mol_featurizer, model

app_utils/upload_model.py

@@ -1,10 +1,23 @@
+"""Upload a checkpoint to a Hugging Face Space (or model repo). Configure via environment variables."""
+import os
+
 from huggingface_hub import HfApi
 
+# Local checkpoint path (relative to repo root or absolute).
+CKPT_PATH = os.environ.get(
+    "FLARE_UPLOAD_CKPT",
+    str(os.path.join(os.path.dirname(__file__), "..", "pretrained_models", "flare.ckpt")),
+)
+# Destination path inside the repo.
+PATH_IN_REPO = os.environ.get("FLARE_UPLOAD_PATH_IN_REPO", "pretrained_models/flare.ckpt")
+REPO_ID = os.environ.get("HF_REPO_ID", "HassounLab/FLARE")
+REPO_TYPE = os.environ.get("HF_REPO_TYPE", "space")
+
 api = HfApi()
 api.upload_file(
-    path_or_fileobj="experiments/20250913_optimized_filip-model/epoch=1993-train_loss=0.10.ckpt",
-    path_in_repo="epoch=1993-train_loss=0.10.ckpt",
-    repo_id="yzhouchen001/flare",
-    repo_type="space", # automatically uses your saved token
-    use_auth_token=True
-)
+    path_or_fileobj=CKPT_PATH,
+    path_in_repo=PATH_IN_REPO,
+    repo_id=REPO_ID,
+    repo_type=REPO_TYPE,
+    token=os.environ.get("HF_TOKEN"),
+)