A KeyError: 'extra' occurs when I call model.run().

[brilliant199912]

##################################################################################################
from lightning.fabric import seed_everything
from rfd3.engine import RFD3InferenceConfig, RFD3InferenceEngine

Set seed for reproducibility

seed_everything(0)

Configure RFD3 inference

config = RFD3InferenceConfig(
specification={
'length': 80, # Generate 80-residue proteins
},
diffusion_batch_size=2, # Generate 2 structures per batch
)

Initialize engine and run generation

model = RFD3InferenceEngine(**config)
outputs = model.run(
inputs=None, # None for unconditional generation
out_dir=None, # None to return in memory (no file output)
n_batches=1, # Generate 1 batch
)

##################################################################################################

10:24:00 DEBUG transforms: Debug mode is on
Seed set to 0

---

KeyError Traceback (most recent call last)
Cell In[2], line 17
15 # Initialize engine and run generation
16 model = RFD3InferenceEngine(**config)
---> 17 outputs = model.run(
18 inputs=None, # None for unconditional generation
19 out_dir=None, # None to return in memory (no file output)
20 n_batches=1, # Generate 1 batch
21 )

File ~/venv/rc-foundry/lib/python3.12/site-packages/rfd3/engine.py:213, in RFD3InferenceEngine.run(self, inputs, n_batches, out_dir)
211 self._set_out_dir(out_dir)
212 inputs = self._canonicalize_inputs(inputs)
--> 213 design_specifications = self._multiply_specifications(
214 inputs=inputs,
215 n_batches=n_batches,
216 )
217 ensure_inference_sampler_matches_design_spec(
218 design_specifications, self.inference_sampler_overrides
219 )
220 # init before

File ~/venv/rc-foundry/lib/python3.12/site-packages/rfd3/engine.py:394, in RFD3InferenceEngine._multiply_specifications(self, inputs, n_batches)
391 design_specifications = {}
392 for prefix, example_spec in inputs.items():
393 # Record task name in the specification
--> 394 example_spec["extra"]["task_name"] = prefix
396 # ... Create n_batches for example
397 for batch_id in range((n_batches) if exists(n_batches) else 1):
398 # ... Example ID

KeyError: 'extra'
##################################################################################################

Thank you so much for your unparalleled contributions to protein design. I'm currently trying to run foundry/examples/all.ipynb, but I keep encountering a KeyError: 'extra' when calling model.run(). Could you please help me understand what might be causing this issue?

[nwoodall-xyme]

I ran into this issue as well.

Temporary workaround to request the previous version, if you want to try some work in the mean time.

dependencies = [
"rc-foundry[all]==0.1.6",
]

[LorenzoTarricone]

Hi @brilliant199912 and @nwoodall-xyme,

I had the same problem with the IPD design pipeline colab notebook

What worked for me was to simply substitute:

# Configure RFD3 inference
config = RFD3InferenceConfig(
    specification={
        'length': 80  # Generate 80-residue proteins
    },
    diffusion_batch_size=2,  # Generate 2 structures per batch
)

with

# Configure RFD3 inference
config = RFD3InferenceConfig(
    specification={
        'length': 80,  # Generate 80-residue proteins
        'extra': {}
    },
    diffusion_batch_size=2,  # Generate 2 structures per batch
)

If this is the expected use, maybe it should be changed in the notebooks? @rclune . I think I sent a PR fixing this some time ago, but idk why the change is not there 🤷🏻‍♂️

[Ubiquinone-dot]

Thank you Lorenzo; this should've been fixed at the src level but that hasn't gone through to the latest version on pip, will just modify the notebook for now.