Are DOM features (node_ids, parent_node_ids, etc.) supposed to be passed to the model during pre-training?

[LahadMbacke]

Hi, I have a question about DataCollatorForDOMNodeMask.torch_call in data_collator.py.

The collator only returns input_ids and labels to the model:

return {"input_ids": inputs, "labels": labels}
But the model's forward() also expects node_ids, parent_node_ids, sibling_node_ids, depth_ids and tag_ids (the P0–P4 tree position embeddings from the paper). Since they are not returned by the collator, the model receives None for all of them and falls back to padding values.

Is this intentional, or is it a bug? I looked through the full codebase and couldn't find anywhere else where these features are passed to the model during training.

Thanks!

[sueun-dev]

yeah, your read matches current main.

I checked a temp clone at 9813261: DataCollatorForDOMNodeMask.torch_call reads node_ids, but only to build the whole-node mask labels. The batch it returns is just input_ids and labels. src/train.py passes that collator into Trainer, while DOMLMForMaskedLM.forward still accepts node_ids, parent_node_ids, sibling_node_ids, depth_ids, and tag_ids; DOMLMEmbeddings.forward fills any missing ones with the configured pad ids.

So I wouldn't call the P0-P4 DOM features active in this pretraining path as-is. Can't speak for intent, but if they're meant to be used, the collator needs to pad and return those feature tensors too.