Merge pull request #36 from dhruvdcoder/dhruvesh/LLMs

The enormous PR to support larger DLMs and a minimal eval harness

Author: dhruvdcoder

.cursor/rules/conda-env.mdc

@@ -0,0 +1,47 @@
+---
+description: Always use the project conda env for Python / pip / pytest in xlm-core
+alwaysApply: true
+---
+
+# xlm-core conda environment
+
+This workspace has a dedicated conda env. **Always activate it before running
+Python, `pip`, `pytest`, or any tool that imports from the project.**
+
+```bash
+conda activate /scratch3/workspace/dhruveshpate_umass_edu-idlm/xlm-core/.venv_xlm_core
+```
+
+Do **not** fall back to system Python (`/usr/bin/python`) or `pip install
+--break-system-packages`; PEP 668 will block it and packages won't be visible
+to the project. Do **not** create ad-hoc venvs in the workspace.
+
+## Concrete patterns
+
+```bash
+# ❌ BAD — system python, packages won't reach the project env
+python3 -m pip install cairosvg
+pip install --user --break-system-packages cairosvg
+pytest tests/
+
+# ✅ GOOD — activate first, then run normally
+conda activate /scratch3/workspace/dhruveshpate_umass_edu-idlm/xlm-core/.venv_xlm_core
+pip install cairosvg
+pytest tests/
+```
+
+For one-off shell calls where activation state may not persist, prefix
+explicitly:
+
+```bash
+# ✅ GOOD — single-shot invocation without relying on shell state
+/scratch3/workspace/dhruveshpate_umass_edu-idlm/xlm-core/.venv_xlm_core/bin/python -m pytest tests/
+/scratch3/workspace/dhruveshpate_umass_edu-idlm/xlm-core/.venv_xlm_core/bin/pip install cairosvg
+```
+
+## Verifying you're in the right env
+
+```bash
+which python   # should print .../.venv_xlm_core/bin/python
+python -c "import sys; print(sys.prefix)"
+```

.github/workflows/tests.yml

@@ -20,7 +20,7 @@ jobs:
       - name: Install dependencies
         run: |
           python -m pip install --upgrade pip
-          pip install -e .
+          pip install -e ".[all]"
           pip install -e xlm-models/
           pip install -r requirements/test_requirements.txt
 

.gitignore

@@ -47,11 +47,12 @@ site
 .cursorindexingignore
 oracle/fpscores.pkl
 .cursor/plans
+src/xlm/backbones/dream_on
 
 # protein eval — generated artifacts and eval venvs
 vendor/dplm/.venv-esmfold/
 evals/protein_eval/*/esmfold_pdb/
 evals/protein_eval/*/metrics.jsonl
 evals/protein_eval/logs/
 .pytest_cache/
-docs/**/linkedin
+docs/**/linkedin

MANIFEST.in

@@ -0,0 +1,2 @@
+include requirements.txt
+recursive-include requirements *.txt

PR_MESSAGE_SAFE.md

@@ -0,0 +1,58 @@
+# Consolidate SAFE / molgen task code in `xlm.tasks.safe_molgen`
+
+## Motivation
+
+- Single canonical module for SAFE bracket encoding, fragment preprocessing, and post-hoc evaluators (`DeNovoEval`, `FragmentEval`). Deprecated TorchMetrics molgen training metrics were **not** carried forward.
+- Align FlexMDM configs and lightning_train SAFE datasets with the learned-noise pipeline (fragment infill, logging fields, checkpoint cadence where copied).
+- Remove `src/xlm/tasks/molgen.py` so there is no duplicate or shim layer.
+
+## What changed
+
+### Python
+
+- **Added** [`src/xlm/tasks/safe_molgen.py`](src/xlm/tasks/safe_molgen.py): based on learned-noise `molgen_new.py` (GenMol-aligned `BracketSAFEConverter`, FCD hooks, grouped `FragmentEval`, etc.). Module name avoids clashing with the PyPI **`safe`** package.
+- **Deleted** [`src/xlm/tasks/molgen.py`](src/xlm/tasks/molgen.py).
+- **Removed** deprecated TorchMetrics path: no `MolGenMetric` / `molgen_update_fn` in `safe_molgen.py`; deleted `src/xlm/configs/lightning_train/metrics/molgen.yaml`. Use **`post_hoc_evaluator`** + `DeNovoEval` / `FragmentEval` for molecule metrics.
+
+### Configs (Hydra)
+
+- **Copied / updated** from learned-noise:
+  - `xlm-models/flexmdm/configs/experiment/safe_flexmdm.yaml`, `safe_flexmdm_fragment.yaml`
+  - `xlm-models/flexmdm/configs/datamodule/safe_flexmdm.yaml`, `safe_flexmdm_fragment.yaml`
+  - `src/xlm/configs/lightning_train/datamodule/safe_fragment.yaml`
+  - `src/xlm/configs/lightning_train/datasets/safe_fragment_test.yaml`, `safe_fragment_val.yaml`
+  - `datasets/safe.yaml`, `safe_train.yaml`, `safe_val.yaml`, `safe_test.yaml`
+  - `tokenizer/safe.yaml`, `post_hoc_evaluator/denovo.yaml`
+- **Retargeted** all `xlm.tasks.molgen` / `xlm.tasks.molgen_new` references to **`xlm.tasks.safe_molgen`** (including `qm9_flexmdm` post-hoc evaluator).
+- **Fixed** `safe_flexmdm_fragment` datamodule `print_batch_fn`: `lflexmdm...` → **`flexmdm.datamodule_flexmdm.print_batch_flexmdm`**.
+
+### Requirements
+
+- Comment update in [`requirements/molgen_requirements.txt`](requirements/molgen_requirements.txt) (datamol note refers to `xlm.tasks.safe_molgen`). Dependency set already matched learned-noise (`datamol`, `torch_fcd`, `fcd`, etc.); `datasets` remains on core install.
+
+## Config migration cheat sheet
+
+| Old | New |
+|-----|-----|
+| `xlm.tasks.molgen.*` | `xlm.tasks.safe_molgen.*` |
+| `xlm.tasks.molgen_new.*` | `xlm.tasks.safe_molgen.*` |
+| `xlm.tasks.safe.*` (internal task module; clashes with PyPI `safe`) | `xlm.tasks.safe_molgen.*` |
+
+Unchanged: YAML keys such as **`molgen_task_name`** (task selector, not a Python module).
+
+## Verification
+
+- `python -m py_compile src/xlm/tasks/safe_molgen.py` (passes).
+- Full import test requires a molgen-capable env (`torch`, `rdkit`, `safe-mol`, `datamol`, `pytdc`, …). Example:
+
+  ```bash
+  conda activate /scratch3/workspace/dhruveshpate_umass_edu-idlm/xlm-core/.venv_xlm_core  # or your env
+  cd xlm-core && PYTHONPATH=src python -c "from xlm.tasks.safe_molgen import DeNovoEval, FragmentEval"
+  ```
+
+- Recommended: Hydra compose training config with `experiment=safe_flexmdm` and `experiment=safe_flexmdm_fragment` once `flexmdm` is on the Hydra search path (`xlm-models` discovery).
+
+## Follow-ups
+
+- Delete this file before or after merge if you do not want it in the repo long term.
+- Watch for downstream repos that imported `xlm.tasks.molgen` in Python code (xlm-core YAMLs are updated; external callers must switch to `xlm.tasks.safe_molgen`).

requirements.txt

@@ -9,6 +9,7 @@ hydra-submitit-launcher
 datasets >= 3.3.2, < 4.0.0
 wandb
 transformers
+regex
 more-itertools
 torchdata>=0.9.0
 rich
@@ -19,4 +20,5 @@ torch-ema
 pydot
 tabulate
 pandas
-simple_slurm
+simple_slurm
+math-verify

requirements/llm_eval.txt

@@ -0,0 +1 @@
+math-verify[antlr4_9_3]

requirements/molgen_requirements.txt

@@ -0,0 +1,26 @@
+# Requirements for Molecule Generation. Largely based on GenMol: https://github.com/NVIDIA-Digital-Bio/genmol/blob/main/env/requirements.txt
+# Optional heavyweight install via setuptools extra: pip install "xlm-core[molgen]"
+# Lightweight SAFE stack without Biomemo: pip install "xlm-core[safe]"
+# bionemo sub-packages
+bionemo-moco>=0.0.2.1
+
+# external dependencies
+#torch==2.6.0
+#numpy==1.26.4
+jaxtyping>=0.2.34
+pot>=0.9.5
+# scikit-learn==1.2.2 install this after due to soft bionemo-moco conflict
+einops
+pandas
+safe-mol==0.1.14
+datamol  # BracketSAFEConverter in xlm.tasks.safe_molgen uses datamol with safe.converter.SAFEConverter
+pytdc==0.4.1
+easydict
+openbabel-wheel==3.1.1.22
+rdkit
+# PyTorch FCD backend (import path: fcd_torch); PyPI distribution is hyphenated:
+fcd-torch
+fcd
+transformers==4.57.1 # other versions may not work.
+setuptools<82.0.0
+#numpy==1.26.4

requirements/safe_extra.txt

@@ -0,0 +1,6 @@
+# Pinned SAFE / cheminformatics stack pulled from requirements/molgen_requirements.txt
+# Install: pip install "xlm-core[safe]"
+safe-mol==0.1.14
+datamol
+rdkit
+pytdc==0.4.1

setup.py

@@ -9,15 +9,44 @@
 
 PATH_ROOT = os.path.dirname(__file__)
 
+
 def load_requirements(
     path_dir: str = PATH_ROOT, comment_char: str = "#"
 ) -> List:
     with open(os.path.join(path_dir, "requirements.txt"), "r") as file:
         reqs = [ln.strip() for ln in file.readlines()]
     return reqs
 
+
+def load_requirements_optional(relative_filename: str) -> List[str]:
+    """Extras list for setuptools (strip blanks and inline ``#`` comments)."""
+    path = os.path.join(PATH_ROOT, relative_filename)
+    with open(path, "r") as file:
+        lines = file.readlines()
+    reqs = []
+    for ln in lines:
+        s = ln.split("#")[0].strip()
+        if s:
+            reqs.append(s)
+    return reqs
+
+
 install_requires = load_requirements()
 
+_safe_reqs = load_requirements_optional("requirements/safe_extra.txt")
+_molgen_reqs = load_requirements_optional("requirements/molgen_requirements.txt")
+_llm_eval_reqs = load_requirements_optional("requirements/llm_eval.txt")
+
+extras_require = {
+    "safe": _safe_reqs,
+    "molgen": _molgen_reqs,
+    "llm_eval": _llm_eval_reqs,
+    "all": list(
+        dict.fromkeys(_safe_reqs + _molgen_reqs + _llm_eval_reqs),
+    ),
+}
+
+
 setup(
     name="xlm-core",
     version=VERSION["VERSION"],
@@ -31,25 +60,30 @@ def load_requirements(
             "tests",
         ],
     ),
-    description='XLM Framework',
+    description="XLM Framework",
     long_description="""
     XLM is a unified framework for developing and comparing small non-autoregressive language models. It uses PyTorch as the deep learning framework, PyTorch Lightning for training utilities, and Hydra for configuration management. XLM provides core components for flexible data handling and training, useful architectural implementations for non-autoregressive workflows, and support for arbitrary runtime code injection. Custom model implementations that leverage the core components of xlm can be found in the xlm-models package. The package also includes a few preconfigured synthetic planning and language-modeling datasets.
 
     Usage:
         pip install xlm-core
-
-    Command usage:
+        pip install "xlm-core[safe]"     # optional: SAFE molecule preprocessing / evaluators
+        pip install "xlm-core[molgen]"   # optional: fuller GenMol / Biomemo stack (molgen_requirements.txt)
+        pip install "xlm-core[llm_eval]" # optional: ANTLR build of math-verify (LLM benchmarks)
+        pip install "xlm-core[all]"      # union of safe + molgen + llm_eval (used in CI)
         xlm job_type=[JOB_TYPE] job_name=[JOB_NAME] experiment=[CONFIG_PATH]
        
         The job_type argument can be one of train ,eval and generate. The experiment argument should point to the root hydra config file.
 """,
-    install_requires = install_requires,
-    project_urls={
-        "Source Code": "https://github.com/dhruvdcoder/xlm-core"
-    },
+    install_requires=install_requires,
+    extras_require=extras_require,
+    project_urls={"Source Code": "https://github.com/dhruvdcoder/xlm-core"},
     package_dir={"": "src"},
     package_data={
-        "xlm": ["configs/**/*.yaml", "configs/**/*.yml"],
+        "xlm": [
+            "configs/**/*.yaml",
+            "configs/**/*.yml",
+            "src/xlm/tasks/zinc_len.pkl",
+        ],
     },
     keywords=[
         "AI",
@@ -63,7 +97,7 @@ def load_requirements(
         "console_scripts": [
             "xlm=xlm.__main__:main",
             "xlm-scaffold=xlm.commands.scaffold_model:main",
-            "xlm-push-to-hub=xlm.commands.push_to_hub:main"
+            "xlm-push-to-hub=xlm.commands.push_to_hub:main",
         ],
     },
     python_requires=">=3.11",