Gen ume #199

Sidney-Lisanza · 2025-09-26T15:05:21Z

Description

Brief description of changes made

Type of Change

Bug fix
[ X] New feature
Documentation update
Performance improvement
[X ] Code refactoring

…to gen_ume

src/lobster/callbacks/_structure_decode.py

src/lobster/callbacks/_unconditional_generation.py

src/lobster/data/__init__.py

src/lobster/data/_collate_structure.py

src/lobster/datasets/_structure_dataset.py

src/lobster/model/gen_ume/_checkpoint_utils.py

src/lobster/model/gen_ume/_gen_ume_sequence_structure_encoder.py

src/lobster/model/latent_generator/cmdline/inference.py

src/lobster/transforms/_structure_transforms.py

src/lobster/model/gen_ume/_gen_ume_sequence_structure_encoder.py

src/lobster/model/gen_ume/_gen_ume_sequence_structure_encoder_lightning_module.py

conditional generation to acocund for this

…to gen_ume

src/lobster/hydra_config/experiment/generate_inverse_folding.yaml

examples/generate_inverse_folding.yaml

karinazad · 2025-10-03T22:21:48Z

examples/generate_inverse_folding.yaml

+  # Directory: input_structures: "/path/to/pdb/directory/"
+  # Glob pattern: input_structures: "/path/to/structures/*.pdb"
+  # List of files: input_structures: ["/path/to/file1.pdb", "/path/to/file2.pdb"]
+  input_structures: "/homefs/home/lisanzas/scratch/Develop/lobster/test_data/*.pdb"


could be a relative path to lobster so ./test_data ?

it's a nit though

src/lobster/callbacks/_folding_structure_utils.py

karinazad · 2025-10-03T22:33:10Z

src/lobster/callbacks/_inverse_folding_callback.py

+        # Get validation dataset directly (avoid trainer dependency in val_dataloader)
+        self.val_dataset = self.eval_datamodule._val_dataset
+
+        # Create our own dataloader to avoid trainer state issues


what are the state issues?

this was with trying to use hydra for speciifying the datasets where i wanted something local but hydra wanted to reference the global one used for training

karinazad · 2025-10-03T22:33:55Z

src/lobster/callbacks/_inverse_folding_callback.py

+                    writepdb(filename, x_recon_xyz[i], seq[i])
+                    logger.info(f"Saved {filename}")
+
+            # Only perform ESMFold folding for the first validation batch for speed


why is that?

or rather - do we only check the first batch in general?

we fold with esm for the first 10, but then the rest of the metrics are just for sequence identity calculations

karinazad · 2025-10-03T22:37:57Z

src/lobster/callbacks/_unconditional_generation.py

+            if "vit_decoder" == decoder_name:
+                x_recon_xyz = decoded_x[decoder_name]
+        if generate_sample["sequence_logits"].shape[-1] == 33:
+            seq = convert_lobster_aa_tokenization_to_standard_aa(generate_sample["sequence_logits"], device=device)


this

if generate_sample["sequence_logits"].shape[-1] == 33: seq = convert_lobster_aa_tokenization_to_standard_aa(generate_sample["sequence_logits"], device=device) else: seq = generate_sample["sequence_logits"].argmax(dim=-1) seq[seq > 21] = 20

seems a bit obsure to me

this is for accounting for the tokenizer lobster uses so it has 33 tokens, but i didnt know how else to check for that

src/lobster/cmdline/generate.py

Lisanza and others added 4 commits September 25, 2025 20:16

adding gen ume code

7506105

callbacks for gen

f4de1a7

minor

5adc8e8

ruff fixes

cf3593a

Sidney-Lisanza requested a review from karinazad September 26, 2025 15:05

Lisanza and others added 7 commits September 26, 2025 15:10

Merge branch 'main' of https://github.com/prescient-design/lobster in…

8cab669

…to gen_ume

fix errors

46e11ed

update ruff and submit to slurm

592f8ed

slurm submit

3b1b40d

schedular kwargs

d6fb787

ruff

717c696

fixed multigpu and need to add correct tokenizers for amino acids

29bdffe