diff --git a/analysis/sweep/fixed_eval_best_models.py b/analysis/sweep/fixed_eval_best_models.py index 5f5be3c..0844de3 100644 --- a/analysis/sweep/fixed_eval_best_models.py +++ b/analysis/sweep/fixed_eval_best_models.py @@ -160,6 +160,17 @@ def parse_args() -> argparse.Namespace: parser.add_argument("--max_seq_length", type=int, default=128) parser.add_argument("--expected_vocab_size", type=int, default=631) parser.add_argument("--expected_max_position_embeddings", type=int, default=128) + parser.add_argument( + "--representation", + choices=["SELFIES", "SMILES"], + default="SELFIES", + help="Representation of the swept checkpoints (selects the default validation column).", + ) + parser.add_argument( + "--molecule_column", + default=None, + help="Validation parquet column (default: selfies / smiles_canonical_clean by representation).", + ) parser.add_argument( "--limit", type=int, @@ -323,11 +334,13 @@ def load_tokenizer(tokenizer_dir: Path): return AutoTokenizer.from_pretrained(str(tokenizer_dir), trust_remote_code=True) -def load_valid_full(valid_parquet: Path, *, limit: int | None) -> list[str]: +def load_valid_full( + valid_parquet: Path, *, limit: int | None, molecule_column: str = "selfies" +) -> list[str]: if not valid_parquet.exists(): raise FileNotFoundError(f"Missing validation parquet: {valid_parquet}") - frame = pd.read_parquet(valid_parquet, columns=["selfies"]) - seqs = [str(value).strip() for value in frame["selfies"] if str(value).strip()] + frame = pd.read_parquet(valid_parquet, columns=[molecule_column]) + seqs = [str(value).strip() for value in frame[molecule_column] if str(value).strip()] return seqs[:limit] if limit is not None else seqs @@ -337,6 +350,7 @@ def load_valid_train_matched( n_examples: int, seed: int, shuffle_buffer_size: int, + molecule_column: str = "selfies", ) -> list[str]: ds = load_dataset( "parquet", @@ -347,7 +361,7 @@ def load_valid_train_matched( ds = ds.shuffle(seed=seed, buffer_size=shuffle_buffer_size) seqs: list[str] = [] for row in ds: - seq = str(row.get("selfies", "")).strip() + seq = str(row.get(molecule_column, "")).strip() if not seq: continue seqs.append(seq) @@ -605,14 +619,20 @@ def main() -> None: ids_to_tokens = dict(getattr(tokenizer, "ids_to_tokens", {})) log.info(" vocab_size=%d pad=%d mask=%d", vocab_size, pad_token_id, mask_token_id) + molecule_column = args.molecule_column or ( + "smiles_canonical_clean" if args.representation == "SMILES" else "selfies" + ) train_matched_n = min(4096, args.limit) if args.limit is not None else 4096 eval_sets = { - "valid_full": load_valid_full(args.valid_parquet, limit=args.limit), + "valid_full": load_valid_full( + args.valid_parquet, limit=args.limit, molecule_column=molecule_column + ), "valid_4096_train_matched": load_valid_train_matched( args.valid_parquet, n_examples=train_matched_n, seed=242, shuffle_buffer_size=100_000, + molecule_column=molecule_column, ), } for name, seqs in eval_sets.items(): diff --git a/configs/featurizers/modernmolbert_smiles.json b/configs/featurizers/modernmolbert_smiles.json new file mode 100644 index 0000000..57a104e --- /dev/null +++ b/configs/featurizers/modernmolbert_smiles.json @@ -0,0 +1,11 @@ +{ + "batch_size": 32, + "device": "auto", + "max_seq_length": 128, + "model_dir": "runs/best_chembl36_smiles_small/final_model", + "name": "modernmolbert_smiles", + "pooling": "mean", + "representation": "SMILES", + "tokenizer_path": "runs/best_chembl36_smiles_small/final_model", + "type": "modernmolbert_smiles" +} diff --git a/scripts/sweeps/run_sweep.py b/scripts/sweeps/run_sweep.py index 2d79baa..ffbd964 100644 --- a/scripts/sweeps/run_sweep.py +++ b/scripts/sweeps/run_sweep.py @@ -22,12 +22,31 @@ from pathlib import Path # ─── Fixed across all runs ──────────────────────────────────────────────────── +# The prepared chembl36 dataset carries both a `selfies` and a `smiles_canonical_clean` +# column, so both representations reuse the same dataset directory read-only. DATASET_DIR = "data/pretrain/chembl36_selfies" -SELFIES_COLUMN = "selfies" TRAIN_SPLIT = "train" VALIDATION_SPLIT = "valid" -TOKENIZER_PATH = "tokenizer/chembl36_selfies_2m_ape_max2_min3000.json" -TOKENIZER_METADATA_PATH = "tokenizer/chembl36_selfies_2m_ape_max2_min3000.metadata.json" + +# Per-representation tokenizer, dataset column, and masking grid. SELFIES keeps the +# historical defaults; SMILES points at the SMILES APE tokenizer and drops hetero_span +# (its heteroatom bias is SELFIES-bracket-specific and degrades to plain span on SMILES). +REPRESENTATION_DEFAULTS = { + "SELFIES": { + "tokenizer_path": "tokenizer/chembl36_selfies_2m_ape_max2_min3000.json", + "tokenizer_metadata_path": "tokenizer/chembl36_selfies_2m_ape_max2_min3000.metadata.json", + "molecule_column": "selfies", + "masking": ["standard", "span", "hetero_span"], + "run_root_tag": "", + }, + "SMILES": { + "tokenizer_path": "tokenizer/chembl36_smiles_2m_ape_max6_mf3000.json", + "tokenizer_metadata_path": "tokenizer/chembl36_smiles_2m_ape_max6_mf3000.metadata.json", + "molecule_column": "smiles_canonical_clean", + "masking": ["standard", "span"], + "run_root_tag": "smiles_", + }, +} MAX_SEQ_LENGTH = 128 MAX_STEPS = 30000 @@ -69,12 +88,21 @@ def parse_args() -> argparse.Namespace: description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter ) parser.add_argument("--model-size", choices=sorted(PRESETS), required=True) + parser.add_argument( + "--representation", + choices=sorted(REPRESENTATION_DEFAULTS), + default="SELFIES", + help="Molecular string representation (default: SELFIES).", + ) parser.add_argument( "--masking", nargs="+", choices=ALL_MASKING, - default=ALL_MASKING, - help="Masking strategies to sweep (default: all three).", + default=None, + help=( + "Masking strategies to sweep (default: per representation; " + "SELFIES uses all three, SMILES uses standard+span)." + ), ) parser.add_argument( "--mlm-probs", @@ -96,8 +124,21 @@ def parse_args() -> argparse.Namespace: help="Output root (default: runs/chembl36__mask_mlm_lr_sweep).", ) parser.add_argument("--dataset-dir", default=DATASET_DIR) - parser.add_argument("--tokenizer-path", default=TOKENIZER_PATH) - parser.add_argument("--tokenizer-metadata-path", default=TOKENIZER_METADATA_PATH) + parser.add_argument( + "--molecule-column", + default=None, + help="Dataset column (default: per representation).", + ) + parser.add_argument( + "--tokenizer-path", + default=None, + help="Tokenizer vocabulary JSON (default: per representation).", + ) + parser.add_argument( + "--tokenizer-metadata-path", + default=None, + help="Tokenizer metadata JSON (default: per representation).", + ) parser.add_argument( "--dry-run", action="store_true", @@ -144,8 +185,10 @@ def build_command( "modernmolbert.train_selfies_ape_modernbert", "--dataset_name", args.dataset_dir, - "--selfies_column", - SELFIES_COLUMN, + "--representation", + args.representation, + "--molecule_column", + args.molecule_column, "--train_split", TRAIN_SPLIT, "--use_validation_split", @@ -222,7 +265,24 @@ def main() -> None: args = parse_args() preset = PRESETS[args.model_size] learning_rates = args.learning_rates or preset["learning_rates"] - run_root = args.run_root or Path(f"runs/chembl36_{args.model_size}_mask_mlm_lr_sweep") + + # Fill representation-dependent defaults for anything not overridden on the CLI. + rep = REPRESENTATION_DEFAULTS[args.representation] + args.tokenizer_path = args.tokenizer_path or rep["tokenizer_path"] + args.tokenizer_metadata_path = args.tokenizer_metadata_path or rep["tokenizer_metadata_path"] + args.molecule_column = args.molecule_column or rep["molecule_column"] + args.masking = args.masking or rep["masking"] + + invalid = [m for m in args.masking if m not in rep["masking"]] + if invalid: + sys.exit( + f"ERROR: masking {invalid} not valid for representation {args.representation}; " + f"allowed: {rep['masking']}" + ) + + run_root = args.run_root or Path( + f"runs/chembl36_{rep['run_root_tag']}{args.model_size}_mask_mlm_lr_sweep" + ) if not args.dry_run: preflight(args) diff --git a/src/modernmolbert/eval/benchmarking_molecular_models/embed_modernmolbert.py b/src/modernmolbert/eval/benchmarking_molecular_models/embed_modernmolbert.py index 3d53cdf..549ef44 100644 --- a/src/modernmolbert/eval/benchmarking_molecular_models/embed_modernmolbert.py +++ b/src/modernmolbert/eval/benchmarking_molecular_models/embed_modernmolbert.py @@ -44,6 +44,12 @@ def parse_args() -> argparse.Namespace: parser.add_argument("--device", default="auto") parser.add_argument("--max-seq-length", type=int, default=256) parser.add_argument("--pooling", choices=["mean", "cls"], default="mean") + parser.add_argument( + "--representation", + choices=["SELFIES", "SMILES"], + default="SELFIES", + help="Checkpoint's molecular representation (SMILES skips SMILES->SELFIES conversion).", + ) parser.add_argument("--overwrite", action=argparse.BooleanOptionalAction, default=False) return parser.parse_args() @@ -101,6 +107,7 @@ def make_featurizer(args: argparse.Namespace): pooling=args.pooling, device=args.device, batch_size=args.batch_size, + representation=args.representation, ) diff --git a/src/modernmolbert/eval/featurizers/modernmolbert_selfies.py b/src/modernmolbert/eval/featurizers/modernmolbert_selfies.py index 107e4c2..649bda9 100644 --- a/src/modernmolbert/eval/featurizers/modernmolbert_selfies.py +++ b/src/modernmolbert/eval/featurizers/modernmolbert_selfies.py @@ -12,6 +12,7 @@ from modernmolbert.eval.featurizers.base import FeatureBatch from modernmolbert.eval.pooling import mean_pool_excluding_token_ids from modernmolbert.tokenization_ape import APEPreTrainedTokenizer +from modernmolbert.utils import SELFIES_REPRESENTATION, SMILES_REPRESENTATION @dataclass @@ -23,6 +24,7 @@ class ModernMolBERTSelfiesFeaturizer: pooling: Literal["mean", "cls"] = "mean" device: str = "auto" batch_size: int = 32 + representation: Literal["SELFIES", "SMILES"] = SELFIES_REPRESENTATION def __post_init__(self) -> None: self.model_dir = Path(self.model_dir) @@ -35,8 +37,12 @@ def __post_init__(self) -> None: if self.pooling not in {"mean", "cls"}: raise ValueError(f"Unsupported pooling strategy: {self.pooling!r}") + self.representation = str(self.representation).upper() + if self.representation not in {SELFIES_REPRESENTATION, SMILES_REPRESENTATION}: + raise ValueError(f"Unsupported representation: {self.representation!r}") + self._device = self._resolve_device(self.device) - self.tokenizer = _load_ape_tokenizer(self.tokenizer_path) + self.tokenizer = _load_ape_tokenizer(self.tokenizer_path, self.representation) self.model = AutoModel.from_pretrained(self.model_dir) self.model.to(self._device) self.model.eval() @@ -47,37 +53,46 @@ def featurize_smiles( *, batch_size: int | None = None, ) -> FeatureBatch: - import selfies as sf - effective_batch_size = self.batch_size if batch_size is None else batch_size if effective_batch_size <= 0: raise ValueError("batch_size must be positive") - selfies_strings: list[str] = [] + # Build the strings the model actually consumes: for a SMILES checkpoint the + # input SMILES pass through unchanged; for a SELFIES checkpoint each SMILES is + # converted to SELFIES first (rows that fail conversion are marked invalid). + model_strings: list[str] = [] valid_mask = np.zeros(len(smiles), dtype=bool) - for i, smi in enumerate(smiles): - if smi is None: - continue - - text = str(smi).strip() - if not text: - continue - - try: - encoded = sf.encoder(text) - except Exception: - continue - - if not encoded: - continue - - selfies_strings.append(encoded) - valid_mask[i] = True + if self.representation == SMILES_REPRESENTATION: + for i, smi in enumerate(smiles): + if smi is None: + continue + text = str(smi).strip() + if not text: + continue + model_strings.append(text) + valid_mask[i] = True + else: + import selfies as sf + + for i, smi in enumerate(smiles): + if smi is None: + continue + text = str(smi).strip() + if not text: + continue + try: + encoded = sf.encoder(text) + except Exception: + continue + if not encoded: + continue + model_strings.append(encoded) + valid_mask[i] = True hidden_size = int(getattr(self.model.config, "hidden_size", 0)) - if not selfies_strings: + if not model_strings: out = FeatureBatch( X=np.zeros((0, hidden_size), dtype=np.float32), valid_mask=valid_mask, @@ -89,7 +104,7 @@ def featurize_smiles( out.check(n_inputs=len(smiles)) return out - n_valid = len(selfies_strings) + n_valid = len(model_strings) X = np.empty((n_valid, hidden_size), dtype=np.float32) n_batches = math.ceil(n_valid / effective_batch_size) row = 0 @@ -102,9 +117,9 @@ def featurize_smiles( unit="batch", leave=False, ): - batch_strings = selfies_strings[start : start + effective_batch_size] + batch_strings = model_strings[start : start + effective_batch_size] - batch = self._tokenize_selfies_batch(batch_strings) + batch = self._tokenize_batch(batch_strings) batch = {key: value.to(self._device) for key, value in batch.items()} outputs = self.model(**batch) @@ -146,7 +161,8 @@ def featurize( def _metadata(self, *, n_inputs: int, n_valid: int) -> dict[str, object]: return { "featurizer": self.name, - "backend": "modernmolbert_selfies", + "backend": f"modernmolbert_{self.representation.lower()}", + "representation": self.representation, "model_dir": str(self.model_dir), "tokenizer_path": str(self.tokenizer_path), "pooling": self.pooling, @@ -186,23 +202,23 @@ def _special_token_ids(self) -> set[int]: return {int(x) for x in ids if x is not None} - def _tokenize_selfies_batch( + def _tokenize_batch( self, - selfies_strings: list[str], + strings: list[str], ) -> dict[str, torch.Tensor]: - """Tokenize a batch of SELFIES strings with the APE tokenizer.""" + """Tokenize a batch of molecular strings with the APE tokenizer.""" - if isinstance(selfies_strings, str): - raise TypeError("_tokenize_selfies_batch expects list[str], not str") + if isinstance(strings, str): + raise TypeError("_tokenize_batch expects list[str], not str") - if not selfies_strings: - raise ValueError("Cannot tokenize an empty SELFIES batch") + if not strings: + raise ValueError("Cannot tokenize an empty batch") # Batch tokenization: one call for the whole list instead of a per-string # loop + manual padding. On MPS the model forward is fast enough that the # old Python loop became the bottleneck. encoded = self.tokenizer( - selfies_strings, + strings, padding=True, truncation=True, max_length=self.max_seq_length, @@ -215,7 +231,9 @@ def _tokenize_selfies_batch( } -def _load_ape_tokenizer(path: str | Path) -> APEPreTrainedTokenizer: +def _load_ape_tokenizer( + path: str | Path, representation: str = SELFIES_REPRESENTATION +) -> APEPreTrainedTokenizer: """Load APE tokenizer from a file or checkpoint directory. Supported inputs: @@ -228,7 +246,7 @@ def _load_ape_tokenizer(path: str | Path) -> APEPreTrainedTokenizer: path = Path(path) if path.is_file(): - tokenizer = APEPreTrainedTokenizer(representation="SELFIES") + tokenizer = APEPreTrainedTokenizer(representation=representation) tokenizer.load_vocabulary_file(path) return tokenizer @@ -246,7 +264,7 @@ def _load_ape_tokenizer(path: str | Path) -> APEPreTrainedTokenizer: vocab_json = path / "vocab.json" if vocab_json.exists(): - tokenizer = APEPreTrainedTokenizer(representation="SELFIES") + tokenizer = APEPreTrainedTokenizer(representation=representation) tokenizer.load_vocabulary_file(vocab_json) return tokenizer diff --git a/src/modernmolbert/train_selfies_ape_modernbert.py b/src/modernmolbert/train_selfies_ape_modernbert.py index 82e9df9..f31fcb9 100644 --- a/src/modernmolbert/train_selfies_ape_modernbert.py +++ b/src/modernmolbert/train_selfies_ape_modernbert.py @@ -38,25 +38,26 @@ from modernmolbert.utils import ( PUBCHEM10M_DATASET, SELFIES_REPRESENTATION, + SMILES_REPRESENTATION, assert_metadata_representation, assert_representation_compatible, assert_special_ids, compute_tokenization_stats, copy_tokenizer_artifacts, - default_selfies_tokenizer_path, + default_tokenizer_path, eligible_token_ids, encode_sequence, file_sha256, find_local_dataset, get_streaming_dataset, - infer_selfies_column, + infer_molecule_column, infer_validation_split, load_tokenizer_metadata, metadata_path_for_vocab, normalize_sequence, resolve_special_ids, tokenizer_vocab_size, - validate_selfies_sample_shape, + validate_sample_shape, ) DATASET_NAME = PUBCHEM10M_DATASET @@ -74,8 +75,15 @@ def parse_args() -> argparse.Namespace: parser.add_argument( "--tokenizer_vocab_path", type=str, - default=str(default_selfies_tokenizer_path()), - help="SELFIES tokenizer vocabulary JSON.", + default=None, + help="Tokenizer vocabulary JSON. Defaults to the per-representation tokenizer.", + ) + parser.add_argument( + "--representation", + type=str, + choices=[SELFIES_REPRESENTATION, SMILES_REPRESENTATION], + default=SELFIES_REPRESENTATION, + help="Molecular string representation the tokenizer and dataset column use.", ) parser.add_argument( "--tokenizer_metadata_path", @@ -86,11 +94,20 @@ def parse_args() -> argparse.Namespace: # Dataset parser.add_argument("--dataset_name", type=str, default=DATASET_NAME) + parser.add_argument( + "--molecule_column", + type=str, + default=None, + help=( + "Dataset column containing the molecular strings for --representation. " + "Defaults by dataset/representation." + ), + ) parser.add_argument( "--selfies_column", type=str, default=None, - help=("Column containing SELFIES strings. Defaults by dataset."), + help="Deprecated alias for --molecule_column (kept for back-compat).", ) parser.add_argument( "--train_split", @@ -366,7 +383,15 @@ def adjust_args_for_backend(args: argparse.Namespace, backend: str) -> argparse. def resolve_dataset_args(args: argparse.Namespace) -> argparse.Namespace: - args.selfies_column = infer_selfies_column(args.dataset_name, args.selfies_column) + if args.tokenizer_vocab_path is None: + args.tokenizer_vocab_path = str(default_tokenizer_path(args.representation)) + column_override = args.molecule_column or args.selfies_column + # Keep selfies_column as the internal resolved-column name so downstream + # dataset code stays representation-agnostic. + args.selfies_column = infer_molecule_column( + args.dataset_name, args.representation, column_override + ) + args.molecule_column = args.selfies_column args.validation_split = infer_validation_split( args.dataset_name, args.validation_split, @@ -407,7 +432,7 @@ def preview_dataset_and_tokenizer( local = find_local_dataset(args.data_dir, dataset_name=args.dataset_name) log(f"Dataset: {args.dataset_name}") - log(f"SELFIES column: {args.selfies_column}") + log(f"Molecule column: {args.selfies_column}") log(f"Train split: {args.train_split}") log(f"Validation split: {args.validation_split}") log(f"Use validation split: {args.use_validation_split}") @@ -417,7 +442,7 @@ def preview_dataset_and_tokenizer( log(f"Dataset mode: local dataset at {local}") else: log("Dataset mode: streaming from HuggingFace Hub") - log(f"Representation: {SELFIES_REPRESENTATION}") + log(f"Representation: {args.representation}") for i, seq in enumerate(examples, start=1): encoded = encode_sequence(tokenizer, seq, args.max_seq_length) @@ -527,7 +552,7 @@ def load_and_validate_tokenizer( ) metadata = load_tokenizer_metadata(metadata_path) - assert_metadata_representation(metadata, expected_representation=SELFIES_REPRESENTATION) + assert_metadata_representation(metadata, expected_representation=args.representation) recorded_sha = str(metadata.get("tokenizer_sha256", "")) actual_sha = file_sha256(vocab_path) @@ -539,7 +564,7 @@ def load_and_validate_tokenizer( f"metadata={recorded_sha}, file={actual_sha}" ) - tokenizer = APEPreTrainedTokenizer(representation=SELFIES_REPRESENTATION) + tokenizer = APEPreTrainedTokenizer(representation=args.representation) tokenizer.load_vocabulary_file(vocab_path) vocab_size = tokenizer_vocab_size(tokenizer) @@ -552,10 +577,10 @@ def load_and_validate_tokenizer( validation_sequences = _sample_train_partition_sequences( args, n=args.tokenizer_validation_samples ) - validate_selfies_sample_shape(validation_sequences) + validate_sample_shape(validation_sequences, args.representation) assert_representation_compatible( - tokenizer, special_ids, SELFIES_REPRESENTATION, args.max_seq_length + tokenizer, special_ids, args.representation, args.max_seq_length ) stats = compute_tokenization_stats( @@ -832,17 +857,24 @@ def write_run_metadata( tokenizer_metadata = load_tokenizer_metadata(tokenizer_metadata_path) tokenizer_sha256 = str(tokenizer_metadata.get("tokenizer_sha256", "unknown")) + is_smiles = args.representation == SMILES_REPRESENTATION + expected_input = ( + "Canonical SMILES strings." + if is_smiles + else ( + "SELFIES strings only. Convert SMILES before inference using a helper such " + "as smiles_to_selfies()." + ) + ) metadata = { "dataset_name": args.dataset_name, "selfies_column": args.selfies_column, + "molecule_column": args.selfies_column, "train_split": args.train_split, "validation_split": args.validation_split, "use_validation_split": args.use_validation_split, - "representation": SELFIES_REPRESENTATION, - "expected_input": ( - "SELFIES strings only. Convert SMILES before inference using a helper such " - "as smiles_to_selfies()." - ), + "representation": args.representation, + "expected_input": expected_input, "tokenizer_vocab_path": str(tokenizer_vocab_path), "tokenizer_metadata_path": str(tokenizer_metadata_path), "backend": backend, @@ -877,6 +909,13 @@ def write_run_metadata( json.dump(metadata, f, indent=2) final_eval_metrics_text = json.dumps(final_eval_metrics or {}, indent=2, sort_keys=True) + repr_tag = "smiles" if is_smiles else "selfies" + repr_vocab_file = "smiles_vocab.json" if is_smiles else "selfies_vocab.json" + repr_expected = ( + "This checkpoint expects canonical SMILES strings." + if is_smiles + else "This checkpoint expects SELFIES strings only. Convert SMILES before tokenization." + ) model_card = f"""--- license: mit library_name: transformers @@ -884,20 +923,20 @@ def write_run_metadata( tags: - chemistry - molecules -- selfies +- {repr_tag} - modernbert - masked-language-modeling --- -# ModernMolBERT SELFIES Masked Language Model +# ModernMolBERT {args.representation} Masked Language Model -This checkpoint was trained from scratch with ModernBERT for SELFIES masked language modeling. +This checkpoint was trained from scratch with ModernBERT for {args.representation} masked language modeling. ## Representation -`{SELFIES_REPRESENTATION}` +`{args.representation}` -This checkpoint expects SELFIES strings only. Convert SMILES before tokenization. +{repr_expected} ## Tokenizer @@ -909,7 +948,7 @@ def write_run_metadata( Keep these files with the checkpoint: - `vocab.json` -- `selfies_vocab.json` +- `{repr_vocab_file}` - `tokenizer_metadata.json` - `tokenizer_config.json` - `special_tokens_map.json` @@ -919,7 +958,7 @@ def write_run_metadata( `{args.dataset_name}` -SELFIES column: `{args.selfies_column}` +Molecule column: `{args.selfies_column}` ## Model @@ -996,7 +1035,8 @@ def main() -> None: log(f"Backend: {backend}") log(f"bf16={args.bf16}, fp16={args.fp16}") log(f"Dataset: {args.dataset_name}") - log(f"SELFIES column: {args.selfies_column}") + log(f"Representation: {args.representation}") + log(f"Molecule column: {args.selfies_column}") log(f"Train split: {args.train_split}") log(f"Validation split: {args.validation_split}") log(f"Use validation split: {args.use_validation_split}") diff --git a/tests/test_checkpoint_reload.py b/tests/test_checkpoint_reload.py index 848928b..6457115 100644 --- a/tests/test_checkpoint_reload.py +++ b/tests/test_checkpoint_reload.py @@ -1,3 +1,4 @@ +import json from pathlib import Path from argparse import Namespace @@ -142,6 +143,7 @@ def test_write_run_metadata_writes_hub_model_card(tmp_path: Path): output_dir=str(output_dir), dataset_name="data/pretrain/chembl36_selfies", selfies_column="selfies", + representation="SELFIES", train_split="train", validation_split=None, use_validation_split=False, @@ -186,6 +188,59 @@ def test_write_run_metadata_writes_hub_model_card(tmp_path: Path): assert "trust_remote_code=True" in text +def test_write_run_metadata_smiles_model_card(tmp_path: Path): + metadata_path = tmp_path / "smiles_ape_tokenizer.metadata.json" + write_tokenizer_metadata( + metadata_path, + { + "representation": "SMILES", + "tokenizer_sha256": "def456", + "tokenizer_path": "tokenizer/smiles_ape_tokenizer.json", + }, + ) + output_dir = tmp_path / "run" + args = Namespace( + output_dir=str(output_dir), + dataset_name="data/pretrain/chembl36_selfies", + selfies_column="smiles_canonical_clean", + representation="SMILES", + train_split="train", + validation_split=None, + use_validation_split=False, + max_seq_length=128, + mlm_probability=0.3, + masking_strategy="span", + model_size="small", + ) + + write_run_metadata( + args=args, + backend="cpu", + vocab_size=8, + special_ids={ + "pad_token": 1, + "bos_token": 0, + "eos_token": 2, + "unk_token": 3, + "mask_token": 4, + }, + n_params=1234, + tokenizer_stats={"unk_rate": 0.0}, + tokenizer_vocab_path=tmp_path / "smiles_ape_tokenizer.json", + tokenizer_metadata_path=metadata_path, + ) + + text = (output_dir / "final_model" / "README.md").read_text(encoding="utf-8") + assert "SELFIES strings only" not in text + assert "canonical SMILES strings" in text + assert "`SMILES`" in text + assert "smiles_vocab.json" in text + + metadata = json.loads((output_dir / "ape_tokenizer_metadata.json").read_text()) + assert metadata["representation"] == "SMILES" + assert metadata["molecule_column"] == "smiles_canonical_clean" + + def test_copy_tokenizer_metadata_from_anywhere_with_tokenizer_metadata_only(tmp_path: Path) -> None: source = tmp_path / "source" target_root = tmp_path / "final_model" diff --git a/tests/test_eval_modernmolbert_selfies.py b/tests/test_eval_modernmolbert_selfies.py index 7177824..3474708 100644 --- a/tests/test_eval_modernmolbert_selfies.py +++ b/tests/test_eval_modernmolbert_selfies.py @@ -101,7 +101,7 @@ def tiny_modernmolbert_dir(tmp_path: Path, monkeypatch) -> Path: monkeypatch.setattr( mm_selfies, "_load_ape_tokenizer", - lambda path: TinyTokenizer(), + lambda path, representation="SELFIES": TinyTokenizer(), ) monkeypatch.setattr( "modernmolbert.eval.featurizers.modernmolbert_selfies.AutoModel.from_pretrained", diff --git a/tests/test_eval_modernmolbert_smiles.py b/tests/test_eval_modernmolbert_smiles.py new file mode 100644 index 0000000..ece7673 --- /dev/null +++ b/tests/test_eval_modernmolbert_smiles.py @@ -0,0 +1,110 @@ +# tests/test_eval_modernmolbert_smiles.py +"""SMILES-native path of the ModernMolBERT featurizer. + +The SELFIES path of the same featurizer is covered in +test_eval_modernmolbert_selfies.py; here we assert the SMILES checkpoint feeds raw +SMILES to the tokenizer with no SMILES->SELFIES round-trip. +""" + +import sys +from pathlib import Path +from types import SimpleNamespace + +import numpy as np +import torch + +import modernmolbert.eval.featurizers.modernmolbert_selfies as mm_selfies +from modernmolbert.eval.featurizers.modernmolbert_selfies import ( + ModernMolBERTSelfiesFeaturizer, +) + + +class RecordingTokenizer: + """Tokenizer stub that records the exact strings it was asked to encode.""" + + pad_token_id = 1 + bos_token_id = 0 + eos_token_id = 2 + unk_token_id = 3 + mask_token_id = 4 + + def __init__(self) -> None: + self.seen: list[str] = [] + + def __call__(self, texts, *, padding=True, truncation=True, max_length=32, return_tensors="pt"): + self.seen.extend(texts) + ids = [[self.bos_token_id, 5, self.eos_token_id] for _ in texts] + return { + "input_ids": torch.tensor(ids, dtype=torch.long), + "attention_mask": torch.ones((len(texts), 3), dtype=torch.long), + } + + +class TinyModel(torch.nn.Module): + def __init__(self, hidden_size: int = 8): + super().__init__() + self.config = SimpleNamespace(hidden_size=hidden_size) + + def forward(self, input_ids, attention_mask=None): + hidden_size = self.config.hidden_size + values = input_ids.float().unsqueeze(-1).repeat(1, 1, hidden_size) + return SimpleNamespace(last_hidden_state=values) + + +def _make_featurizer(tmp_path: Path, monkeypatch, tokenizer: RecordingTokenizer): + model_dir = tmp_path / "tiny-smiles" + model_dir.mkdir() + monkeypatch.setattr( + mm_selfies, "_load_ape_tokenizer", lambda path, representation="SELFIES": tokenizer + ) + monkeypatch.setattr( + "modernmolbert.eval.featurizers.modernmolbert_selfies.AutoModel.from_pretrained", + lambda path, **kwargs: TinyModel(hidden_size=8), + ) + return ModernMolBERTSelfiesFeaturizer( + model_dir=model_dir, + tokenizer_path=model_dir, + max_seq_length=32, + batch_size=4, + device="cpu", + representation="SMILES", + ) + + +def test_smiles_featurizer_passes_raw_smiles_to_tokenizer(tmp_path, monkeypatch): + tokenizer = RecordingTokenizer() + featurizer = _make_featurizer(tmp_path, monkeypatch, tokenizer) + + batch = featurizer.featurize_smiles(["CCO", "c1ccccc1"]) + + # Raw SMILES reach the tokenizer unchanged (SELFIES encoder would give "[C][C][O]"). + assert tokenizer.seen == ["CCO", "c1ccccc1"] + assert batch.valid_mask.tolist() == [True, True] + assert batch.X.shape == (2, 8) + assert batch.X.dtype == np.float32 + assert batch.metadata["representation"] == "SMILES" + assert batch.metadata["backend"] == "modernmolbert_smiles" + + +def test_smiles_featurizer_never_imports_selfies(tmp_path, monkeypatch): + # Poison the selfies module: if the SMILES path tried to convert, it would raise. + poison = SimpleNamespace(encoder=lambda *_a, **_k: (_ for _ in ()).throw(RuntimeError("no"))) + monkeypatch.setitem(sys.modules, "selfies", poison) + + tokenizer = RecordingTokenizer() + featurizer = _make_featurizer(tmp_path, monkeypatch, tokenizer) + + batch = featurizer.featurize_smiles(["CCO"]) + assert tokenizer.seen == ["CCO"] + assert batch.valid_mask.tolist() == [True] + + +def test_smiles_featurizer_marks_empty_and_none_invalid(tmp_path, monkeypatch): + tokenizer = RecordingTokenizer() + featurizer = _make_featurizer(tmp_path, monkeypatch, tokenizer) + + batch = featurizer.featurize_smiles(["CCO", "", None]) + assert batch.valid_mask.tolist() == [True, False, False] + assert tokenizer.seen == ["CCO"] + assert batch.X.shape == (1, 8) + batch.check(n_inputs=3) diff --git a/tests/test_run_sweep_cli.py b/tests/test_run_sweep_cli.py new file mode 100644 index 0000000..09b8c44 --- /dev/null +++ b/tests/test_run_sweep_cli.py @@ -0,0 +1,51 @@ +# tests/test_run_sweep_cli.py +"""Dry-run coverage for the representation-aware sweep launcher. + +Dry-run skips preflight and only touches stdlib, so the script runs with the bare +interpreter (no uv / heavy deps). +""" + +import subprocess +import sys +from pathlib import Path + +REPO_ROOT = Path(__file__).resolve().parents[1] +SCRIPT = REPO_ROOT / "scripts" / "sweeps" / "run_sweep.py" + + +def _dry_run(*extra: str) -> subprocess.CompletedProcess: + return subprocess.run( + [sys.executable, str(SCRIPT), "--model-size", "small", "--dry-run", *extra], + capture_output=True, + text=True, + cwd=REPO_ROOT, + ) + + +def test_selfies_dry_run_is_unchanged(): + res = _dry_run() + assert res.returncode == 0, res.stderr + out = res.stdout + assert "Total grid: 27" in out + assert "chembl36_selfies_2m_ape_max2_min3000.json" in out + assert "runs/chembl36_small_mask_mlm_lr_sweep/" in out + assert "chembl36_smiles_" not in out + + +def test_smiles_dry_run_uses_smiles_tokenizer_and_column(): + res = _dry_run("--representation", "SMILES") + assert res.returncode == 0, res.stderr + out = res.stdout + # 2 masking x 3 mlm x 3 lr. + assert "Total grid: 18" in out + assert "chembl36_smiles_2m_ape_max6_mf3000.json" in out + assert "--molecule_column smiles_canonical_clean" in out + assert "--representation SMILES" in out + assert "runs/chembl36_smiles_small_mask_mlm_lr_sweep/" in out + assert "hetero_span" not in out + + +def test_smiles_rejects_hetero_span(): + res = _dry_run("--representation", "SMILES", "--masking", "hetero_span") + assert res.returncode != 0 + assert "not valid for representation SMILES" in (res.stderr + res.stdout) diff --git a/tests/test_training_cli.py b/tests/test_training_cli.py index 23daf12..3015136 100644 --- a/tests/test_training_cli.py +++ b/tests/test_training_cli.py @@ -10,6 +10,7 @@ make_eval_dataset, make_train_iterable_dataset, parse_args as parse_train_args, + resolve_dataset_args, sequence_bucket, validate_args, ) @@ -68,6 +69,42 @@ def test_validate_args_rejects_unsupported_cuda_bf16(monkeypatch): validate_args(args, backend="cuda") +def test_representation_defaults_to_selfies(): + with _Argv("--output_dir", "tmp/run"): + args = resolve_dataset_args(parse_train_args()) + + assert args.representation == "SELFIES" + # Column is dataset-inferred; molecule_column mirrors the resolved column. + assert args.molecule_column == args.selfies_column + assert str(args.tokenizer_vocab_path).endswith("selfies_ape_tokenizer.json") + + +def test_representation_smiles_resolves_column_and_tokenizer(): + with _Argv( + "--output_dir", + "tmp/run", + "--representation", + "SMILES", + "--molecule_column", + "smiles_canonical_clean", + "--dataset_name", + "data/pretrain/chembl36_selfies", + ): + args = resolve_dataset_args(parse_train_args()) + + assert args.representation == "SMILES" + assert args.selfies_column == "smiles_canonical_clean" + assert args.molecule_column == "smiles_canonical_clean" + assert str(args.tokenizer_vocab_path).endswith("smiles_ape_tokenizer.json") + + +def test_selfies_column_alias_still_resolves(): + with _Argv("--output_dir", "tmp/run", "--selfies_column", "my_col"): + args = resolve_dataset_args(parse_train_args()) + + assert args.selfies_column == "my_col" + + def test_pretokenized_rows_use_stable_hash_split(monkeypatch): rows = [ {"input_ids": [0, 5, 2]},