ESMFold Backend Inference and Tensor Extraction (Issue #43)

[rohan5986]

Related Issue: Contributes to #43

Context & Objective

As discussed, this PR serves as an initial prototype for the backend inference and tracing pipeline required to integrate ESMFold into VizFold. It successfully handles the heavy ML infrastructure tasks—specifically running the forward pass and extracting the required hidden states, attention maps, and predicted structures into a standardized archive format.

Key Changes & Implementations

• ESMFold Inference (extract_esmfold.py): Implemented model loading using the Hugging Face transformers library (EsmForProteinFolding).
  • Configured explicit cache redirection and safetensors enforcement to bypass HPC cluster disk quota limitations and recent PyTorch torch.load security blocks.
• Attention Map Interception: Because the top-level ESMFold wrapper drops attention matrices, I engineered custom PyTorch forward hooks attached to the self-attention modules of all 36 Transformer layers within the ESM-2 stem to manually intercept and extract the raw attention tensors.
• VizFold Schema Formatting (format_vizfold.py): Wrote a formatter to clean the raw PyTorch tensors and convert them to NumPy arrays.
  • Sliced out the implicit <cls> and <eos> padding tokens from the stem's attention maps to ensure sequence length alignment.
  • Isolated the final 3D structural output from the model's 8-step iterative refinement tensor.
  • Packaged the clean tensors (sequence, single representations, attention, positions) into a .pkl archive.

[jayvenn21]

notes for Rohan PR: good stuff, you got the right module path for HF and the / slicing was a key insight. it should already in the backend now. the outputs.s_s extraction is unique to this PR and would be really useful to add on top of our work.

once Jeeva is done with his changes, if you want to add that you can PR against feature/esmfold-backend. the web app could be a separate follow-up too