Prototype branch

.gitignore

@@ -1,20 +1,40 @@
-# Temp files
+# OS and Editor specific
 .DS_Store
 *~
-buck-out/
 
-# Compiled files
+# Virtual Environments (including Hatch environments)
 .venv/
+env/
+venv/
+
+# Python bytecode and cache files
 __pycache__/
+*.py[cod]
+*$py.class
+
+# Cached files created by various build tools
 .*cache/
 
-# Distribution / packaging
+# Distribution / packaging outputs
 /dist/
+*.egg-info/
+.eggs/
 
-# Tests and coverage
+# Tests, coverage, and temporary files
 /data/
-/node_modules/
+node_modules/
+.tox/
+.pytest_cache/
+.coverage
 
-# docs
+# Documentation build outputs
 /docs/generated/
 /docs/_build/
+
+# Buck build output (if applicable)
+buck-out/
+
+# Job outputs & logs
+/output/
+/err/
+

embed_job.sbatch

@@ -0,0 +1,44 @@
+#!/bin/bash
+#SBATCH --job-name=embpy_gen
+#SBATCH --output=err/embpy_gen_output_%A.txt
+#SBATCH --error=err/embpy_gen_error_%A.txt
+#SBATCH --ntasks=2
+#SBATCH --time=3:00:00
+#SBATCH --mem=500Gb
+#SBATCH --gres=gpu:1
+#SBATCH --cpus-per-task=4
+#SBATCH --partition=gpu_p
+#SBATCH --qos=gpu_normal
+#SBATCH --nice=0
+# #SBATCH --constraint="a100_80gb|a100_40gb|h100_80gb"
+# #SBATCH --array=0-2
+
+# If you ever want to run multiple seeds or variants, you could enable an array:
+# #SBATCH --array=0-2
+
+# Ensure log dir exists
+mkdir -p err
+
+# Init shell + env
+source "$HOME/.bashrc"
+mamba activate embpy
+
+echo "▶ Starting MoLFormer embeddings on $(hostname) at $(date)"
+echo "   GPUs: $CUDA_VISIBLE_DEVICES"
+echo "   Python: $(which python)"
+python -c "import torch; print('   PyTorch CUDA:', torch.cuda.is_available())"
+
+# Paths
+PYFILE=/lustre/groups/ml01/workspace/goncalo.pinto/embpy/run_generation_molformer.py
+CSV_IN=/lustre/groups/ml01/workspace/goncalo.pinto/embpy/data/rdkit_200.csv
+OUTDIR=/lustre/groups/ml01/workspace/goncalo.pinto/embpy/output
+
+# Run (choose poolings as you like: cls mean max)
+python -u "$PYFILE" \
+  --csv "$CSV_IN" \
+  --smiles-col smiles \
+  --model molformer_base \
+  --poolings cls mean max \
+  --outdir "$OUTDIR"
+
+echo "✅ Finished at $(date)"

err/embpy_gen_output_26631607.txt

@@ -0,0 +1,4 @@
+▶ Starting embeddings on supergpu08.scidom.de at Thu Jul  3 11:02:29 AM CEST 2025
+   GPUs: 0,1
+   Python: /ictstr01/home/icb/goncalo.pinto/tools/apps/mamba/envs/embpy/bin/python
+Found 18978 genes in CCLE file.

pyproject.toml

@@ -24,18 +24,30 @@ classifiers = [
 ]
 dependencies = [
   "anndata",
+  "biopython",
+  "borzoi-pytorch",
   "enformer-pytorch",
+  "esm",
+  "ipywidgets",
+  "matplotlib",
   "numpy",
   "pandas",
+  "pyarrow",
+  "pyensembl",
+  "pytest",
+  "rdkit",            # For molecule embeddings
+  "seaborn",
   "session-info",
-  "torch>=2.0",
+  "torch",
+  "transformers",
 ]
 optional-dependencies.dev = [
   "pre-commit",
   "twine>=6.1",
 ]
 optional-dependencies.doc = [
   "docutils>=0.8,!=0.18.*,!=0.19.*",
+  "huggingface-hub",
   "ipykernel",
   "ipython",
   "myst-nb>=1.1",
@@ -48,17 +60,17 @@ optional-dependencies.doc = [
   "sphinxcontrib-bibtex>=1",
   "sphinxext-opengraph",
 ]
+optional-dependencies.molecule = [ "transformers>=4" ] # If ChemBERTa uses transformers
+optional-dependencies.protein = [ "transformers>=4" ] # If ESM uses transformers
+optional-dependencies.sentence_transformers = [ "sentence-transformers>=2" ]
 optional-dependencies.test = [
   "coverage",
   "pytest",
 ]
-# Add optional dependencies for specific model types
-optional-dependencies.transformers = ["transformers>=4.0"]
-optional-dependencies.sentence_transformers = ["sentence-transformers>=2.0"]
 # You could create combined extras
-optional-dependencies.text = ["transformers>=4.0", "sentence-transformers>=2.0"]
-optional-dependencies.protein = ["transformers>=4.0"] # If ESM uses transformers
-optional-dependencies.molecule = ["transformers>=4.0"] # If ChemBERTa uses transformers
+optional-dependencies.text = [ "sentence-transformers>=2", "transformers>=4" ]
+# Add optional dependencies for specific model types
+optional-dependencies.transformers = [ "transformers>=4" ]
 # https://docs.pypi.org/project_metadata/#project-urls
 urls.Documentation = "https://embpy.readthedocs.io/"
 urls.Homepage = "https://github.com/grpinto/embpy"

run_generation.py

@@ -0,0 +1,153 @@
+#!/usr/bin/env python3
+import argparse
+import os
+import time
+from collections.abc import Iterable
+
+import numpy as np
+import pandas as pd
+
+from embpy.embedder import BioEmbedder
+
+VALID_POOLINGS = ["mean", "cls", "max"]
+
+
+def chunked(seq: list[str], size: int) -> Iterable[list[str]]:
+    for i in range(0, len(seq), size):
+        yield seq[i : i + size]
+
+
+def write_chunk(path: str, df: pd.DataFrame, wrote_header: bool):
+    df.to_csv(path, mode="a", header=not wrote_header, index=True)
+
+
+def main(
+    in_csv: str,
+    out_prefix: str,
+    smiles_col: str | None,
+    id_col: str | None,
+    model_key: str,
+    batch_size: int,
+):
+    # 1) Read input (first row is header)
+    df = pd.read_csv(in_csv)
+    if smiles_col is None:
+        smiles_col = df.columns[0]
+    if smiles_col not in df.columns:
+        raise SystemExit(f"SMILES column '{smiles_col}' not found. Columns: {list(df.columns)}")
+
+    if id_col is None:
+        id_col = smiles_col
+    if id_col not in df.columns:
+        raise SystemExit(f"ID column '{id_col}' not found. Columns: {list(df.columns)}")
+
+    smiles_all = df[smiles_col].astype(str).tolist()
+    ids_all = df[id_col].astype(str).tolist()
+    print(f"Found {len(smiles_all)} molecules in input file.")
+
+    # 2) Initialize embedder (auto: CUDA/MPS/CPU)
+    embedder = BioEmbedder(device="auto")
+
+    # 3) Prepare outputs (make dir if prefix includes folders)
+    out_dir = os.path.dirname(out_prefix)
+    if out_dir:
+        os.makedirs(out_dir, exist_ok=True)
+
+    tmp_paths: dict[str, str] = {p: f"{out_prefix}_{p}.csv.part" for p in VALID_POOLINGS}
+    final_paths: dict[str, str] = {p: f"{out_prefix}_{p}.csv" for p in VALID_POOLINGS}
+    wrote_header = dict.fromkeys(VALID_POOLINGS, False)
+    invalid_ids: list[str] = []
+
+    # 4) Embed (do all three poolings; keep rows valid for ALL poolings)
+    t0 = time.time()
+    offset = 0
+    for sm_chunk in chunked(smiles_all, batch_size):
+        id_chunk = ids_all[offset : offset + len(sm_chunk)]
+        print(f"→ Embedding rows {offset}..{offset + len(sm_chunk) - 1}")
+
+        # Run each pooling; this returns list[np.ndarray|None] preserving order
+        by_pool = {}
+        for p in VALID_POOLINGS:
+            by_pool[p] = embedder.embed_molecules_batch(
+                identifiers=sm_chunk,
+                model=model_key,  # "chemberta_zinc_v1"
+                pooling_strategy=p,  # "mean" | "cls" | "max"
+            )
+
+        # Determine rows valid across ALL poolings (mostly guards invalid SMILES)
+        valid_mask = []
+        for i in range(len(sm_chunk)):
+            ok = all(by_pool[p][i] is not None for p in VALID_POOLINGS)
+            valid_mask.append(ok)
+            if not ok:
+                invalid_ids.append(id_chunk[i])
+
+        # Append one chunk per pooling
+        for p in VALID_POOLINGS:
+            # find dim from first valid
+            dim = None
+            for i, ok in enumerate(valid_mask):
+                if ok:
+                    dim = by_pool[p][i].shape[0]
+                    break
+            if dim is None:
+                continue  # nothing valid in this chunk for this pooling
+
+            cols = [f"e{i}" for i in range(dim)]
+            rows = []
+            idxs = []
+            for rid, emb, ok in zip(id_chunk, by_pool[p], valid_mask, strict=False):
+                if not ok:
+                    continue
+                if emb.shape[0] != dim:
+                    raise RuntimeError(f"Inconsistent dim for ID {rid} in pooling '{p}': {emb.shape[0]} vs {dim}")
+                rows.append(emb)
+                idxs.append(rid)
+
+            if rows:
+                chunk_df = pd.DataFrame(np.vstack(rows), columns=cols, index=idxs)
+                write_chunk(tmp_paths[p], chunk_df, wrote_header[p])
+                wrote_header[p] = True
+
+        offset += len(sm_chunk)
+
+    # 5) Finalize files
+    for p in VALID_POOLINGS:
+        if wrote_header[p]:
+            os.replace(tmp_paths[p], final_paths[p])
+            print(f"Wrote: {final_paths[p]}")
+        else:
+            if os.path.exists(tmp_paths[p]):
+                os.remove(tmp_paths[p])
+            print(f"No valid rows for pooling '{p}'")
+
+    # Log invalids (if any)
+    if invalid_ids:
+        bad_path = f"{out_prefix}_invalid_smiles.txt"
+        with open(bad_path, "w") as fh:
+            for rid in invalid_ids:
+                fh.write(str(rid) + "\n")
+        print(f"Logged {len(invalid_ids)} invalid SMILES to: {bad_path}")
+
+    t1 = time.time()
+    print(f"✓ Finished ChemBERTa embeddings in {t1 - t0:.1f}s")
+
+
+if __name__ == "__main__":
+    ap = argparse.ArgumentParser(description="Generate ChemBERTa embeddings (mean/cls/max) via BioEmbedder.")
+    ap.add_argument("--in_csv", required=True, help="Input CSV with header.")
+    ap.add_argument("--out_prefix", required=True, help="Output prefix. Writes <prefix>_{mean,cls,max}.csv")
+    ap.add_argument("--smiles_col", default=None, help="Column with SMILES (default: first column).")
+    ap.add_argument("--id_col", default=None, help="Column to use as row index (default: SMILES column).")
+    ap.add_argument("--model_key", default="chemberta_zinc_v1", help="Registry key in BioEmbedder.")
+    ap.add_argument("--batch_size", type=int, default=1024, help="Batch size for processing.")
+    args = ap.parse_args()
+
+    main(
+        in_csv=args.in_csv,
+        out_prefix=args.out_prefix,
+        smiles_col=args.smiles_col,
+        id_col=args.id_col,
+        model_key=args.model_key,
+        batch_size=args.batch_size,
+    )