| """ |
| Download model checkpoints and install dependencies at startup. |
| """ |
| import os |
| import subprocess |
| import sys |
| from pathlib import Path |
|
|
| BASE_DIR = Path(__file__).parent |
|
|
|
|
| def install_pyg_deps(): |
| """Install PyG dependencies with correct torch version.""" |
| try: |
| import torch |
| torch_version = torch.__version__.split("+")[0] |
| cuda_version = "cu" + torch.version.cuda.replace(".", "") if torch.cuda.is_available() else "cpu" |
| pyg_url = f"https://data.pyg.org/whl/torch-{torch_version}+{cuda_version}.html" |
| print(f"Installing PyG deps from: {pyg_url}") |
| subprocess.run([ |
| sys.executable, "-m", "pip", "install", |
| "torch-scatter", "torch-sparse", "torch-cluster", "pyg-lib", |
| "-f", pyg_url, "--quiet", "--no-deps" |
| ], check=False) |
| print("PyG deps installed") |
| except Exception as e: |
| print(f"PyG deps install warning: {e}") |
|
|
|
|
| def download_evodiffmol_dataset(): |
| """Download EvoDiffMol moses dataset.""" |
| dataset_path = BASE_DIR / "datasets/moses/without_h/raw/valid.npz" |
| if dataset_path.exists(): |
| print("EvoDiffMol dataset already exists") |
| return |
| |
| print("Downloading EvoDiffMol dataset...") |
| dataset_path.parent.mkdir(parents=True, exist_ok=True) |
| |
| try: |
| from huggingface_hub import hf_hub_download |
| hf_hub_download( |
| repo_id="scofieldlinlin/EvoDiffMol", |
| filename="datasets/moses/without_h/raw/valid.npz", |
| local_dir=str(BASE_DIR), |
| ) |
| print("EvoDiffMol dataset downloaded") |
| except Exception as e: |
| print(f"Dataset download warning: {e}") |
| |
| try: |
| script = BASE_DIR / "EvoDiffMol/datasets/process_datasets.py" |
| if script.exists(): |
| subprocess.run([sys.executable, str(script), "--dataset", "moses"], |
| cwd=str(BASE_DIR / "EvoDiffMol"), check=False) |
| except Exception as e2: |
| print(f"Dataset process warning: {e2}") |
|
|
|
|
| def setup(): |
| |
| if not (BASE_DIR / "RLMolLM").exists(): |
| print("Cloning RLMolLM...") |
| subprocess.run(["git", "clone", "https://github.com/XiaoboLinlin/RLMolLM.git"], |
| cwd=BASE_DIR, check=True) |
|
|
| |
| if not (BASE_DIR / "EvoDiffMol").exists(): |
| print("Cloning EvoDiffMol...") |
| subprocess.run(["git", "clone", "https://github.com/XiaoboLinlin/EvoDiffMol.git"], |
| cwd=BASE_DIR, check=True) |
|
|
| |
| if not (BASE_DIR / "RLMolLM" / "rlmollm.egg-info").exists(): |
| print("Installing RLMolLM...") |
| subprocess.run([sys.executable, "-m", "pip", "install", "-e", |
| str(BASE_DIR / "RLMolLM"), "--no-deps", "--quiet"], check=False) |
|
|
| |
| if not (BASE_DIR / "EvoDiffMol" / "evodiffmol.egg-info").exists(): |
| print("Installing EvoDiffMol...") |
| subprocess.run([sys.executable, "-m", "pip", "install", "-e", |
| str(BASE_DIR / "EvoDiffMol"), "--no-deps", "--quiet"], check=False) |
|
|
| |
| install_pyg_deps() |
|
|
| |
| rlmollm_checkpoint = BASE_DIR / "RLMolLM/assets/models/moses_pretrained.pt" |
| if not rlmollm_checkpoint.exists(): |
| print("Downloading RLMolLM checkpoint...") |
| os.chdir(BASE_DIR / "RLMolLM") |
| subprocess.run([sys.executable, "download_assets.py", "--dataset", "moses"], check=False) |
| os.chdir(BASE_DIR) |
|
|
| |
| evodiff_checkpoint = BASE_DIR / "EvoDiffMol/assets/checkpoints/moses_without_h_80.pt" |
| if not evodiff_checkpoint.exists(): |
| print("Downloading EvoDiffMol checkpoint...") |
| os.chdir(BASE_DIR / "EvoDiffMol") |
| subprocess.run([sys.executable, "assets/download_checkpoint.py"], check=False) |
| os.chdir(BASE_DIR) |
|
|
| |
| download_evodiffmol_dataset() |
|
|
| print("Setup complete!") |
|
|
|
|
| if __name__ == "__main__": |
| setup() |
