[QUESTION] Document indices extended in add_document function in IndexedDatasetBuilder but not in add_item

[clarence-lee-sheng]

Dear Megatron-LM team

In the IndexedDatasetBuilder class add_item function (

Megatron-LM/megatron/core/datasets/indexed_dataset.py

Lines 767 to 779 in 912be78

	def add_item(self, tensor: torch.Tensor, mode: int = 0) -> None:
	"""Add a single item to the dataset
	Args:
	tensor (torch.Tensor): The item to add to the data file
	mode (int, optional): The mode for the item. Defaults to 0.
	"""
	np_array = numpy.array(tensor.numpy(), dtype=self.dtype)
	self.data_file.write(np_array.tobytes(order="C"))
	self.sequence_lengths.append(np_array.size)
	if self.multimodal:
	self.sequence_modes.append(mode)

), the document indices is not extended, however for the add_document function, the document indices are extended (

Megatron-LM/megatron/core/datasets/indexed_dataset.py

Lines 781 to 798 in 912be78

	def add_document(
	self, tensor: torch.Tensor, lengths: List[int], modes: Optional[List[int]] = None
	) -> None:
	"""Add an entire document to the dataset
	Args:
	tensor (torch.Tensor): The document to add
	lengths (List[int]): The lengths of each item in the document
	modes (Optional[List[int]], optional): The modes for each item in the document. Defaults to None.
	"""
	np_array = numpy.array(tensor, dtype=self.dtype)
	self.data_file.write(np_array.tobytes(order="C"))
	self.sequence_lengths.extend(lengths)
	self.document_indices.append(len(self.sequence_lengths))
	if self.multimodal:
	self.sequence_modes.extend(modes if modes is not None else [0] * lengths)

). May I ask why is that the case and when do we use add_item?

Thank you!