feat: Add mmCIF file support for macromolecular structures

[behroozazarkhalili]

Summary

This PR adds support for loading mmCIF (macromolecular Crystallographic Information File) files with load_dataset(), following the ImageFolder pattern where one row = one structure.

Based on feedback from @lhoestq in #7930, this approach makes datasets more practical for ML workflows:

• Each row is independent, enabling train/test splits and shuffling
• Easy to add labels (folder-based) and metadata (metadata.jsonl)
• Compatible with Dataset Viewer (one 3D render per row)

Architecture

Uses FolderBasedBuilder pattern (like ImageFolder, AudioFolder):

class MmcifFolder(FolderBasedBuilder):
    BASE_FEATURE = ProteinStructure
    BASE_COLUMN_NAME = "structure"
    EXTENSIONS = [".cif", ".mmcif"]

New ProteinStructure Feature Type

# Arrow schema for lazy loading
pa.struct({"bytes": pa.binary(), "path": pa.string()})

# Decoded: returns structure file content as string
dataset = load_dataset("mmcif", data_dir="structures/")
print(dataset[0]["structure"])  # Full mmCIF file content

Supported Extensions

.cif, .mmcif

Usage

from datasets import load_dataset

# Load from directory
dataset = load_dataset("mmcif", data_dir="protein_structures/")

# Load with folder-based labels
# structures/
#   enzymes/
#     1abc.cif
#   receptors/
#     2def.cif
dataset = load_dataset("mmcif", data_dir="structures/")
print(dataset[0])  # {"structure": "data_...", "label": "enzymes"}

# Load with metadata
# structures/
#   1abc.cif
#   metadata.jsonl  # {"file_name": "1abc.cif", "resolution": 2.5}
dataset = load_dataset("mmcif", data_dir="structures/")
print(dataset[0])  # {"structure": "data_...", "resolution": 2.5}

# Drop labels or metadata
dataset = load_dataset("mmcif", data_dir="structures/", drop_labels=True)
dataset = load_dataset("mmcif", data_dir="structures/", drop_metadata=True)

Test Results

All 24 mmCIF tests + 15 ProteinStructure feature tests pass.

Related PRs

• Add lightweight PDB (Protein Data Bank) file support #7926 - PDB support (same pattern)
• Proposal: Protein 3D Structure Visualization for Dataset Viewer #7930 - Protein 3D visualization proposal

References

• mmCIF specification: https://mmcif.wwpdb.org/
• PDB archive: https://www.rcsb.org/

cc @lhoestq @georgia-hf

[lhoestq]

Hi ! is there a lib in python that parses the text of the files ? Maybe it would be more useful to return a parsed object rather than just the data as string

(and sorry for the delay reveiwing this !)

[behroozazarkhalili]

Hi @lhoestq, thanks — and no worries about the delay!

Good questions. Quick status + a proposal:

On a parsing library: there are mature options — gemmi, biotite, and Biopython all parse mmCIF/PDB. The catch is they're all compiled/heavyweight dependencies (Biopython is ~150 MB installed; gemmi and biotite ship C/C++ extensions), which is a lot to pull into datasets for a single loader. I kept this consistent with the other lightweight bio loaders here (FASTA/FASTQ, and the GenBank PR), which use small pure-Python parsers and add zero dependencies.

On returning a parsed object instead of a string: agreed — and that's actually a regression I introduced. The original version of this PR parsed _atom_site into typed columns (type_symbol, label_atom_id, Cartn_x/y/z, occupancy, B-factor, …). When I moved to the one-row-per-structure layout you suggested in #7930, I replaced that with a ProteinStructure feature that currently just returns the raw file text — so it lost the parsing (the include_hetatm filter went with it).

Proposal — combine both so each row is one structure and a parsed object, with no new dependency:

• one row = one structure (keeps independence for splits/shuffling/metadata, per Proposal: Protein 3D Structure Visualization for Dataset Viewer #7930)
• the structure column becomes a parsed struct of atom-level arrays from the existing lightweight parser, e.g.

structure: {
  atom_name: List[str], element: List[str],
  residue:   List[str], chain_id: List[str], seq_id: List[int],
  x: List[float], y: List[float], z: List[float],
  occupancy: List[float], b_factor: List[float],
}

• reintroduce include_hetatm (and similar filters) as config options on top of the parsed schema
• optionally keep the raw text available too (e.g. ProteinStructure(decode=False) for the bytes, or a raw field) for users who want to run their own parser

That keeps it queryable and ML-friendly while staying dependency-free. Does that direction work for you? Happy to expose the raw bytes alongside the parsed fields if you think both are useful.

[behroozazarkhalili]

Hi @lhoestq — following up on the parsed-object direction we discussed. I've implemented it.

Each row is still one structure (one row = one structure, per #7930), but the structure column is now parsed on decode into a struct-of-arrays of atom-level columns instead of returning raw text. The columns use the PDBx/mmCIF _atom_site dictionary names so PDB- and mmCIF-derived datasets share one vocabulary:

>>> ds = load_dataset("mmcif", data_dir="structures")
>>> ds[0]["structure"]
{'type_symbol': ['N', 'C', ...], 'label_atom_id': ['N', 'CA', ...],
 'label_comp_id': ['ALA', 'ALA', ...], 'label_asym_id': ['A', 'A', ...],
 'label_seq_id': [1, 1, ...], 'Cartn_x': [0.0, 1.458, ...],
 'Cartn_y': [...], 'Cartn_z': [...], 'occupancy': [...], 'B_iso_or_equiv': [...]}

Design notes:

• Dependency-free. Parsing is a small pure-Python module (features/protein_parsing.py); no Biopython/gemmi/biotite. A single canonical column→dtype table is shared by both the PDB and mmCIF parsers and by the feature, so there's no per-format duplication.
• Consistent with the existing feature pattern. ProteinStructure parses lazily in decode_example, exactly like Image/Audio/Mesh decode their stored bytes. Storage stays {bytes, path}, so ProteinStructure(decode=False) still gives the raw file for anyone who wants their own parser.
• Config options: include_hetatm (default True) to drop ligands/water, and columns=[...] to select a subset — e.g. load_dataset("mmcif", data_dir=..., include_hetatm=False, columns=["label_atom_id", "Cartn_x", "Cartn_y", "Cartn_z"]).

All existing tests pass and I added coverage for the parsed schema, the HETATM filter, and column selection. Docs in loading.mdx are updated to match. The PDB PR (#7926) shares the same parser and feature so they stay consistent. Happy to adjust the column set or naming if you'd prefer something different.

[lhoestq]

I feel like it would be better for people to extend the parsing and to add useful processing methods if the decoding returns an object like a biopython structure (esp. since the community has been contributing to biopython mostly so they can contribute there). IMO It's ok if the dependency is not that lightweight, what matters the most is to get an object that users can easily process/convert/prepare to train models