Improvement: Support other formats than .db

[tgmaxson]

It is unclear why data must be given in a .db format in the fine tuning script. Is it possible to get support for generic ase-compatible files? I am unsure if there is some specific to a .db file for training (although I will admit it would allow for asynchronous updates of the training set dynamically).

As a related note, is there a pragmatic way to supply my own dataset batch by batch? I guess I would do this by defining my own dataloader?

[vsimkus]

Yes, if you want to supply your own data batch by batch or use other data formats, you can modify build_train_loader in the fine tuning script to return a custom data loader. Any coding agent should be able to do this!

To implement a custom data loader you should see AseSqliteDataset for inspiration on how the ase.Atoms should be pre-processed and converted to our internal AtomGraphs object, see

orb-models/orb_models/common/dataset/ase_sqlite_dataset.py

Lines 57 to 101 in 1b7b785

	def __getitem__(self, idx: int) -> AbstractAtomBatch:
	"""Fetch an item from the db.
	Args:
	idx: An index to fetch from the db file and convert to an AtomGraphs.
	Returns:
	A AtomGraphs object containing everything the model needs as input,
	positions and atom types and other auxillary information, such as
	fine tuning targets, or global graph features.
	"""
	# Validate index bounds
	if idx < 0 or idx >= len(self):
	raise IndexError(f"Index {idx} out of range for dataset of length {len(self)}")
	# Sqlite db is 1 indexed.
	row = self.db.get(idx + 1)
	atoms = row.toatoms()
	# Features and targets are stored in a mutable atoms.info dict.
	# This dict may be modified as part of an augmentation e.g.
	# Force and stress targets are transformed when rotating a system.
	# These will add {node,graph,edge}_{features,targets} keys to the atoms.info
	atoms.info = {}
	atoms.info.update(self.feature_config.extract(row, self.name, "features"))
	atoms.info.update(self.target_config.extract(row, self.name, "targets"))
	for augmentation in self.augmentations:
	augmentation(atoms)
	# Remove preexisting constraints and then apply any of our own constraints.
	# Constraints have modelling implications e.g. for partial-diffusion / force-prediction.
	atoms.set_constraint(None)
	for constraint_fn in self.constraints:
	constraint_fn(atoms, row.data, self.name)
	return self._from_ase_atoms(
	atoms=atoms,
	device=torch.device("cpu"),
	output_dtype=self.dtype,
	system_id=idx,
	edge_method="knn_alchemi",
	half_supercell=False,
	graph_construction_dtype=self.dtype,
	)

for details.