Fix hf_hub_download import

worker.py

@@ -1,79 +1,3 @@
-"""
-CASCADE worker process
-"""
-
-import os
-import sys
-import warnings
-warnings.filterwarnings("ignore", category=UserWarning)
-warnings.filterwarnings("ignore", category=DeprecationWarning)
-warnings.filterwarnings("ignore", category=FutureWarning)
-
-import json
-import math
-import pickle
-import datetime
-from io import StringIO
-
-import redis
-import numpy as np
-
-os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
-
-import keras
-from keras.models import load_model
-from NMR_Prediction.apply import (
-    preprocess_C, preprocess_H,
-    evaluate_C, evaluate_H,
-    RBFSequence,
-)
-from nfp.layers import (
-    MessageLayer, GRUStep, Squeeze, EdgeNetwork,
-    ReduceBondToPro, ReduceBondToAtom,
-    GatherAtomToBond, ReduceAtomToPro,
-)
-from nfp.models import GraphModel
-
-import pandas as pd
-from rdkit import Chem
-from rdkit.Chem import AllChem
-from rdkit.Chem import SDWriter
-from NMR_Prediction.genConf import genConf
-
-MODEL_PATH_C = os.path.join("NMR_Prediction", "schnet_edgeupdate",   "best_model.hdf5")
-MODEL_PATH_H = os.path.join("NMR_Prediction", "schnet_edgeupdate_H", "best_model.hdf5")
-PREPROCESSOR_PATH = os.path.join("NMR_Prediction", "preprocessor.p")
-
-custom_objects = {
-    "MessageLayer": MessageLayer,
-    "GRUStep": GRUStep,
-    "Squeeze": Squeeze,
-    "EdgeNetwork": EdgeNetwork,
-    "ReduceBondToPro": ReduceBondToPro,
-    "ReduceBondToAtom": ReduceBondToAtom,
-    "GatherAtomToBond": GatherAtomToBond,
-    "ReduceAtomToPro": ReduceAtomToPro,
-    "GraphModel": GraphModel,
-}
-
-print("Loading 13C model...", flush=True)
-model_C = load_model(MODEL_PATH_C, custom_objects=custom_objects)
-print("Loading 1H model...", flush=True)
-model_H = load_model(MODEL_PATH_H, custom_objects=custom_objects)
-print("Both models loaded.", flush=True)
-
-with open(PREPROCESSOR_PATH, "rb") as f:
-    preprocessor = pickle.load(f)["preprocessor"]
-
-redis_client = redis.StrictRedis(
-    host="localhost", port=6379, db=0, decode_responses=True
-)
-
-# ── Analytics logging to HF Dataset ──────────────────────────────────────────
-_HF_TOKEN       = os.environ.get("HF_TOKEN", "")
-_ANALYTICS_REPO = "patonlab/analytics"
-_ANALYTICS_FILE = "data.csv"
-
 def _log_prediction():
     """Append one row to the existing patonlab/analytics data.csv.
     Format matches the alfabet log: space,timestamp
@@ -81,19 +5,19 @@ def _log_prediction():
     if not _HF_TOKEN:
         return
     try:
-        from huggingface_hub import HfApi
+        from huggingface_hub import HfApi, hf_hub_download
         import tempfile
 
         api = HfApi(token=_HF_TOKEN)
         timestamp = datetime.datetime.utcnow().isoformat()
 
-        # Download the current CSV, append a row, re-upload
+        # Download current CSV, append a row, re-upload
         with tempfile.TemporaryDirectory() as tmpdir:
-            local_path = os.path.join(tmpdir, "data.csv")
-            api.hf_hub_download(
+            local_path = hf_hub_download(
                 repo_id=_ANALYTICS_REPO,
                 filename=_ANALYTICS_FILE,
                 repo_type="dataset",
+                token=_HF_TOKEN,
                 local_dir=tmpdir,
             )
             with open(local_path, "a") as f:
@@ -107,167 +31,4 @@ def _log_prediction():
                 commit_message=f"log: cascade prediction {timestamp[:10]}",
             )
     except Exception as e:
-        print(f"Analytics logging failed (non-fatal): {e}", flush=True)
-
-
-def _mol_to_sdf(mol, conf_id=0):
-    sio = StringIO()
-    w = SDWriter(sio)
-    w.write(mol, confId=conf_id)
-    w.close()
-    return sio.getvalue()
-
-
-def _build_sdfs_from_genconf(mol_with_confs, ids):
-    """
-    Build SDF strings directly from the genConf mol using real conformer IDs.
-    ids is a list of (energy, conf_id) tuples sorted by energy (lowest first).
-    Returns (sdfs, energy_order).
-    """
-    sdfs = []
-    energy_order = []
-    for energy, conf_id in ids:
-        try:
-            sdf = _mol_to_sdf(mol_with_confs, conf_id=int(conf_id))
-            if sdf.strip():
-                sdfs.append(sdf)
-                energy_order.append(int(conf_id))
-        except Exception as e:
-            print(f"SDF error for conf_id={conf_id}: {e}", flush=True)
-    return sdfs, energy_order
-
-
-def _boltzmann_average(spread_df):
-    spread_df["b_weight"] = spread_df["relative_E"].apply(
-        lambda x: math.exp(-x / (0.001987 * 298.15))
-    )
-    df_group = spread_df.set_index(["mol_id", "atom_index", "cf_id"]).groupby(level=[0, 1])
-    final = []
-    for (m_id, a_id), df in df_group:
-        ws = (df["b_weight"] * df["predicted"]).sum() / df["b_weight"].sum()
-        final.append([m_id, a_id, ws])
-    final = pd.DataFrame(final, columns=["mol_id", "atom_index", "Shift"])
-    final["atom_index"] = final["atom_index"].apply(lambda x: x + 1)
-    return final.round(2).astype(dtype={"atom_index": "int"})
-
-
-def _fmt_weighted(final_df):
-    return "".join(f"{int(r['atom_index'])},{r['Shift']:.2f};" for _, r in final_df.iterrows())
-
-
-def _fmt_conf_shifts(spread_df, energy_order):
-    parts = []
-    for cf_id in energy_order:
-        sub = spread_df[spread_df["cf_id"] == cf_id]
-        if len(sub) == 0:
-            continue
-        parts.append("".join(f"{int(r['atom_index'])},{r['predicted']:.2f};" for _, r in sub.iterrows()))
-    return "!".join(parts)
-
-
-def _fmt_relative_E(spread_df, energy_order):
-    total_bw = spread_df.groupby("cf_id")["b_weight"].first().sum()
-    parts = []
-    for cf_id in energy_order:
-        sub = spread_df[spread_df["cf_id"] == cf_id]
-        if len(sub) == 0:
-            continue
-        e = round(sub["relative_E"].iloc[0], 2)
-        bw = round(sub["b_weight"].iloc[0] / total_bw, 4)
-        parts.append(f"{e},{bw},")
-    return "!".join(parts)
-
-
-def run_job(task_id, smiles, type_):
-    result_key = f"task_result_{task_id}"
-    try:
-        mol = Chem.MolFromSmiles(smiles)
-        AllChem.EmbedMolecule(mol, useRandomCoords=True)
-        mol_with_h = Chem.AddHs(mol, addCoords=True)
-
-        # Single conformer search
-        mol_with_confs, ids, nr = genConf(mol_with_h, rms=-1, nc=200, efilter=10.0, rmspost=0.5)
-        print(f"genConf: {len(ids)} conformers", flush=True)
-
-        conf_sdfs, energy_order = _build_sdfs_from_genconf(mol_with_confs, ids)
-
-        mols = [Chem.MolFromSmiles(smiles)]
-        for m in mols:
-            AllChem.EmbedMolecule(m, useRandomCoords=True)
-        mols = [Chem.AddHs(m, addCoords=True) for m in mols]
-
-        # Suppress duplicate genConf stdout during preprocess
-        _stdout, _stderr = sys.stdout, sys.stderr
-        sys.stdout = sys.stderr = open(os.devnull, 'w')
-        try:
-            if type_ == "C":
-                inputs, df, conf_mols = preprocess_C(mols, preprocessor, keep_all_cf=True)
-            else:
-                inputs, df, conf_mols = preprocess_H(mols, preprocessor, keep_all_cf=True)
-        finally:
-            sys.stdout.close()
-            sys.stdout, sys.stderr = _stdout, _stderr
-
-        if type_ == "C":
-            predicted = evaluate_C(inputs, preprocessor, model_C)
-        else:
-            predicted = evaluate_H(inputs, preprocessor, model_H)
-
-        if len(inputs) == 0:
-            raise RuntimeError("No conformers generated")
-
-        spread_df = pd.DataFrame(columns=["mol_id", "atom_index", "relative_E", "cf_id"])
-        for _, r in df.iterrows():
-            n = len(r["atom_index"])
-            tmp = pd.DataFrame({
-                "mol_id": [r["mol_id"]] * n,
-                "atom_index": r["atom_index"],
-                "relative_E": [r["relative_E"]] * n,
-                "cf_id": [r["cf_id"]] * n,
-            })
-            spread_df = pd.concat([spread_df, tmp], sort=True)
-
-        spread_df["predicted"] = predicted
-        spread_df["b_weight"] = spread_df["relative_E"].apply(
-            lambda x: math.exp(-x / (0.001987 * 298.15))
-        )
-        spread_df["atom_index"] = spread_df["atom_index"].apply(lambda x: x + 1)
-        spread_df = spread_df.round(2)
-
-        final_df = _boltzmann_average(
-            spread_df.copy().assign(
-                atom_index=spread_df["atom_index"].apply(lambda x: x - 1)
-            )
-        )
-
-        result = {
-            "smiles": smiles,
-            "type_": type_,
-            "conf_sdfs": conf_sdfs,
-            "weightedShiftTxt": _fmt_weighted(final_df),
-            "confShiftTxt": _fmt_conf_shifts(spread_df, energy_order),
-            "relative_E": _fmt_relative_E(spread_df, energy_order),
-        }
-        redis_client.set(result_key, json.dumps(result), ex=3600)
-        print(f"Task {task_id} complete — {len(conf_sdfs)} conformers", flush=True)
-
-        # Log to analytics dataset (non-blocking)
-        _log_prediction()
-
-    except Exception as e:
-        import traceback; traceback.print_exc()
-        redis_client.set(result_key, json.dumps({"errMessage": str(e)}), ex=3600)
-
-
-print("Worker ready, waiting for jobs...", flush=True)
-while True:
-    item = redis_client.blpop("task_queue", timeout=5)
-    if item is None:
-        continue
-    _, task_id = item
-    detail = redis_client.get(f"task_detail_{task_id}")
-    if not detail:
-        continue
-    detail = json.loads(detail)
-    print(f"Processing task {task_id}  smiles={detail['smiles']}  type={detail['type_']}", flush=True)
-    run_job(task_id, detail["smiles"], detail["type_"])
+        print(f"Analytics logging failed (non-fatal): {e}", flush=True)